finishes implementation, adds test

Author: pranavm-nvidia

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

@@ -28,6 +28,8 @@
 #include "mlir-c/Support.h"
 #include "mlir-executor-c/Common/Common.h"
 #include "mlir-executor-c/Support/Status.h"
+#include "mlir-tensorrt-dialect/Utils/Types.h"
+#include "mlir/CAPI/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -60,6 +62,12 @@ typedef struct MTRT_StableHLOToExecutableOptions {
   void *ptr;
 } MTRT_StableHLOToExecutableOptions;
 
+typedef struct MTRT_MetadataCallback {
+  void *ptr;
+} MTRT_MetadataCallback;
+
+DEFINE_C_API_PTR_METHODS(MTRT_MetadataCallback, mlirtrt::MetadataCallbackT)
+
 MLIR_CAPI_EXPORTED MTRT_Status mtrtStableHloToExecutableOptionsCreate(
     MTRT_CompilerClient client, MTRT_StableHLOToExecutableOptions *options,
     int32_t tensorRTBuilderOptLevel, bool tensorRTStronglyTyped);
@@ -79,7 +87,7 @@ MLIR_CAPI_EXPORTED MTRT_Status mtrtStableHloToExecutableOptionsSetDebugOptions(
 
 MLIR_CAPI_EXPORTED MTRT_Status
 mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
-    MTRT_StableHLOToExecutableOptions options, void *callback);
+    MTRT_StableHLOToExecutableOptions options, MTRT_MetadataCallback callback);
 
 MLIR_CAPI_EXPORTED MTRT_Status mtrtStableHloToExecutableOptionsDestroy(
     MTRT_StableHLOToExecutableOptions options);

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

@@ -132,7 +132,9 @@ struct StableHLOToExecutableOptions : public mlir::OptionsContext {
   DebugOptions debugOptions;
 
   // TODO: Add a sane default here:
-  MetadataCallbackT layerMetadataCallback = [](MlirOperation op) { return ""; };
+  MetadataCallbackT layerMetadataCallback = [](const MlirOperation &op) {
+    return "";
+  };
 
   /// Base class for extensions associated with StableHloToExecutableTask.
   class ExtensionBase : public TaskExtensionBase {

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

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

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

@@ -279,17 +279,20 @@ PYBIND11_MODULE(_api, m) {
       .def(
           "set_tensorrt_translation_metadata_callback",
           [](PyStableHLOToExecutableOptions &self, MetadataCallbackT callback) {
-            auto *ptr = callback.target<MetadataCallbackT>();
+            // 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."};
+            // if (!ptr)
+            //   throw std::runtime_error{
+            //       "Metadata callback has incorrect signature. Expected a "
+            //       "function that accepts an MLIR operation and returns a "
+            //       "string."};
+            // [](PyStableHLOToExecutableOptions &self, py::object callback) {
+
+            //   auto ptr = callback.ptr();
 
             THROW_IF_MTRT_ERROR(
                 mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
-                    self, reinterpret_cast<void *>(ptr)));
+                    self, wrap(&callback)));
           },
           py::arg("callback"))
 #endif

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Target/TensorRTEncodingOpInterface/NetworkEncoder.h

@@ -145,7 +145,7 @@ class NvInferNetworkEncoder {
 
   /// Set the name of the `trtLayer` to a unique string that contains the op
   /// name and location information from `sourceOp`.
-  void setName(nvinfer1::ILayer *layer, Operation *sourceOp);
+  void setMetadata(nvinfer1::ILayer *layer, Operation *sourceOp);
 
   // Check if network uses fp16 types.
   bool hasFp16Usage() const { return usesFp16; }
@@ -245,7 +245,7 @@ class NvInferNetworkEncoder {
 
   // TODO: Where to use this? encodeOp doesn't have a way for us to access the
   // layers.
-  std::function<std::string(MlirOperation)> layerMetadataCallback;
+  mlirtrt::MetadataCallbackT layerMetadataCallback;
 };
 
 //===----------------------------------------------------------------------===//

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

@@ -215,9 +215,8 @@ FailureOr<TensorRTEngineResult> buildFunction(
     const TensorRTTranslationOptions &options =
         TensorRTTranslationOptions::fromCLFlags(),
     // TODO: Add a sane default here:
-    mlirtrt::MetadataCallbackT layerMetadataCallback = [](MlirOperation op) {
-      return "";
-    });
+    mlirtrt::MetadataCallbackT layerMetadataCallback =
+        [](const MlirOperation &op) { return ""; });
 
 /// Create an instance of a translate-to-tensorrt pass using an existing
 /// TensorRTBuilderContext.