working

Author: pranavm-nvidia

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

@@ -212,7 +212,9 @@ mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
     accum += std::string(str.data, str.length);
   };
 
-  cppOpts->layerMetadataCallback = [&](Operation *op) {
+  // Capturing by reference here will cause `callback` to point to the wrong
+  // place at the time this callback is invoked.
+  cppOpts->layerMetadataCallback = [=](Operation *op) {
     std::string accum;
     void *appendCtx = reinterpret_cast<void *>(&accum);
     callback(wrap(op), appendFunc, appendCtx, userData);

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

@@ -68,9 +68,7 @@ class PyStableHLOToExecutableOptions
           mtrtPythonCapsuleToStableHLOToExecutableOptions,
           mtrtPythonStableHLOToExecutableOptionsToCapsule};
 
-  std::function<void(MlirOperation op, MlirStringCallback append,
-                     void *appendCtx, void *userData)>
-      callback;
+  std::function<std::string(MlirOperation)> callback;
 };
 } // namespace
 
@@ -228,6 +226,25 @@ static void bindTensorRTPluginAdaptorObjects(py::module m) {
 #endif
 #endif
 
+void layerMetadataCallback(MlirOperation op, MlirStringCallback append,
+                           void *appendCtx, void *userDataVoid) {
+  // auto pyCallback =
+  //     *static_cast<std::function<std::string(MlirOperation)>
+  //     *>(userDataVoid);
+
+  // std::string result;
+  // try {
+  //   result = pyCallback(op);
+  // } catch (const std::exception &e) {
+  //   // TODO: What to do here? Change callback to return a status
+  //   // instead of void?
+  //   // return mtrtStatusCreate(
+  //   //     MTRT_StatusCode::MTRT_StatusCode_Unknown, e.what());
+  // }
+
+  // append(MlirStringRef{result.data(), result.size()}, appendCtx);
+}
+
 PYBIND11_MODULE(_api, m) {
 
   populateCommonBindingsInModule(m);
@@ -285,7 +302,7 @@ PYBIND11_MODULE(_api, m) {
             // Since we're constructing a C callback, our closures must not
             // capture. We can pass in the Python callback via the userData
             // argument.
-            self.callback = [](MlirOperation op, MlirStringCallback append,
+            auto callback = [](MlirOperation op, MlirStringCallback append,
                                void *appendCtx, void *userDataVoid) {
               auto pyCallback =
                   *static_cast<std::function<std::string(MlirOperation)> *>(
@@ -304,13 +321,10 @@ PYBIND11_MODULE(_api, m) {
               append(MlirStringRef{result.data(), result.size()}, appendCtx);
             };
 
+            self.callback = pyCallback;
             THROW_IF_MTRT_ERROR(
                 mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
-                    self,
-                    self.callback.target<void(
-                        MlirOperation op, MlirStringCallback append,
-                        void *appendCtx, void *userData)>(),
-                    reinterpret_cast<void *>(&pyCallback)));
+                    self, callback, reinterpret_cast<void *>(&self.callback)));
           },
           py::arg("callback"), py::keep_alive<1, 2>{})
 #endif