adds py bindings

Author: pranavm-nvidia

mlir-tensorrt/compiler/include/mlir-tensorrt-c/Compiler/Compiler.h

@@ -79,8 +79,7 @@ MLIR_CAPI_EXPORTED MTRT_Status mtrtStableHloToExecutableOptionsSetDebugOptions(
 
 MLIR_CAPI_EXPORTED MTRT_Status
 mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
-    MTRT_StableHLOToExecutableOptions options,
-    const char *(*callback)(MlirOperation));
+    MTRT_StableHLOToExecutableOptions options, void *callback);
 
 MLIR_CAPI_EXPORTED MTRT_Status mtrtStableHloToExecutableOptionsDestroy(
     MTRT_StableHLOToExecutableOptions options);

mlir-tensorrt/compiler/lib/CAPI/Compiler/Compiler.cpp

@@ -201,10 +201,10 @@ MTRT_Status mtrtStableHloToExecutableOptionsSetDebugOptions(
 
 MTRT_Status
 mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
-    MTRT_StableHLOToExecutableOptions options,
-    const char *(*callback)(MlirOperation)) {
+    MTRT_StableHLOToExecutableOptions options, void *callback) {
   StableHLOToExecutableOptions *cppOpts = unwrap(options);
-  cppOpts->layerMetadataCallback = callback;
+  cppOpts->layerMetadataCallback =
+      reinterpret_cast<std::string (*)(MlirOperation)>(callback);
   return mtrtStatusGetOk();
 }
 

mlir-tensorrt/compiler/lib/Compiler/TensorRTExtension/TensorRTExtension.cpp

@@ -65,7 +65,7 @@ void StableHLOToExecutableTensorRTExtension::populatePasses(
     tensorrt::buildTensorRTModuleTransformationPipeline(
         trtPM, translationOptions.enableStronglyTyped);
     trtPM.addPass(tensorrt::createTranslateTensorRTPass(
-        nullptr, translationOptions, options.layerMetadataCallback));
+        nullptr, options.layerMetadataCallback, translationOptions));
     return;
   }
 

mlir-tensorrt/python/bindings/Compiler/CompilerPyBind.cpp

@@ -20,9 +20,12 @@
 #include "pybind11/pybind11.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include <pybind11/attr.h>
+#include <pybind11/functional.h>
+#include <stdexcept>
 
 #ifdef MLIR_TRT_TARGET_TENSORRT
 #include "mlir-tensorrt-dialect/Utils/NvInferAdaptor.h"
+#include "mlir-tensorrt-dialect/Utils/Types.h"
 #endif
 
 namespace py = pybind11;
@@ -270,7 +273,27 @@ PYBIND11_MODULE(_api, m) {
           py::arg("enabled"),
           py::arg("debug_types") = std::vector<std::string>{},
           py::arg("dump_ir_tree_dir") = py::none(),
-          py::arg("dump_tensorrt_dir") = py::none());
+          py::arg("dump_tensorrt_dir") = py::none())
+
+#ifdef MLIR_TRT_TARGET_TENSORRT
+      .def(
+          "set_tensorrt_translation_metadata_callback",
+          [](PyStableHLOToExecutableOptions &self, MetadataCallbackT callback) {
+            auto *ptr = callback.target<MetadataCallbackT>();
+
+            if (!ptr)
+              throw std::runtime_error{
+                  "Metadata callback has incorrect signature. Expected a "
+                  "function that accepts an MLIR operation and returns a "
+                  "string."};
+
+            THROW_IF_MTRT_ERROR(
+                mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
+                    self, reinterpret_cast<void *>(ptr)));
+          },
+          py::arg("callback"))
+#endif
+      ;
 
   m.def(
       "compiler_stablehlo_to_executable",
@@ -308,4 +331,4 @@ PYBIND11_MODULE(_api, m) {
   bindTensorRTPluginAdaptorObjects(m);
 #endif
 #endif
-}
+}

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Target/TranslateToTensorRT.h

@@ -223,12 +223,9 @@ FailureOr<TensorRTEngineResult> buildFunction(
 /// TensorRTBuilderContext.
 std::unique_ptr<mlir::Pass> createTranslateTensorRTPass(
     std::shared_ptr<tensorrt::TensorRTBuilderContext> context,
+    mlirtrt::MetadataCallbackT layerMetadataCallback,
     TensorRTTranslationOptions options =
-        TensorRTTranslationOptions::fromCLFlags(),
-    // TODO: Add a sane default here:
-    mlirtrt::MetadataCallbackT layerMetadataCallback = [](MlirOperation op) {
-      return "";
-    });
+        TensorRTTranslationOptions::fromCLFlags());
 
 /// Register llvm::cl opts related to TensorRT translation. This should be
 /// called before having LLVM parse CL options.

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/Types.h

@@ -3,9 +3,10 @@
 
 #include "mlir/IR/Operation.h"
 #include <functional>
+#include <string>
 
 namespace mlirtrt {
-using MetadataCallbackT = std::function<const char *(MlirOperation)>;
+using MetadataCallbackT = std::function<std::string(MlirOperation)>;
 } // namespace mlirtrt
 
 #endif // MLIR_TENSORRT_UTILS_TYPES_H

mlir-tensorrt/tensorrt/lib/Target/TranslateToTensorRT.cpp

@@ -830,8 +830,8 @@ class TranslateToTensorRTEnginePass
 
 std::unique_ptr<mlir::Pass> tensorrt::createTranslateTensorRTPass(
     std::shared_ptr<tensorrt::TensorRTBuilderContext> context,
-    TensorRTTranslationOptions options,
-    mlirtrt::MetadataCallbackT layerMetadataCallback) {
+    mlirtrt::MetadataCallbackT layerMetadataCallback,
+    TensorRTTranslationOptions options) {
   return std::make_unique<TranslateToTensorRTEnginePass>(context, options,
                                                          layerMetadataCallback);
 }