sketch

Author: makslevental

mlir/include/mlir-c/Pass.h

@@ -153,15 +153,15 @@ struct MlirExternalPassCallbacks {
   /// The callback is called before the pass is run, allowing a chance to
   /// initialize any complex state necessary for running the pass.
   /// See Pass::initialize(MLIRContext *).
-  MlirLogicalResult (*initialize)(MlirContext ctx, void *userData);
+  MlirLogicalResult (*initialize)(void *userData, MlirContext ctx);
 
   /// This callback is called when the pass is cloned.
   /// See Pass::clonePass().
   void *(*clone)(void *userData);
 
   /// This callback is called when the pass is run.
   /// See Pass::runOnOperation().
-  void (*run)(MlirOperation op, MlirExternalPass pass, void *userData);
+  void (*run)(void *userData, MlirOperation op, MlirExternalPass pass);
 };
 typedef struct MlirExternalPassCallbacks MlirExternalPassCallbacks;
 

mlir/lib/Bindings/Python/Pass.cpp

@@ -51,6 +51,55 @@ class PyPassManager {
 
 } // namespace
 
+class PythonPass {
+public:
+  explicit PythonPass(py::object passObj) : passObj(std::move(passObj)) {}
+
+  void *construct() {}
+  void *destruct() {}
+
+  MlirLogicalResult *initialize(MlirContext ctx) {}
+  void *clone() {}
+
+  void run(MlirOperation op, MlirExternalPass pass) {}
+
+  py::object passObj;
+};
+
+template <typename T, typename R>
+void *void_cast(R (T::*f)()) {
+  union {
+    R (T::*pf)();
+    void *p;
+  };
+  pf = f;
+  return p;
+}
+
+template <typename classT, typename memberT>
+union u_ptm_cast {
+  memberT pmember;
+  void *pvoid;
+};
+
+MlirExternalPassCallbacks makeTestExternalPassCallbacks() {
+  return (MlirExternalPassCallbacks){
+      reinterpret_cast<decltype(MlirExternalPassCallbacks::construct)>(
+          void_cast(&PythonPass::construct)),
+      reinterpret_cast<decltype(MlirExternalPassCallbacks::destruct)>(
+          void_cast(&PythonPass::destruct)),
+      nullptr,
+      reinterpret_cast<decltype(MlirExternalPassCallbacks::clone)>(
+          void_cast(&PythonPass::clone)),
+      reinterpret_cast<decltype(MlirExternalPassCallbacks::run)>(
+          u_ptm_cast<PythonPass,
+                     void (PythonPass::*)(MlirOperation, MlirExternalPass)>{
+              &PythonPass::run}
+              .pvoid),
+
+  };
+}
+
 /// Create the `mlir.passmanager` here.
 void mlir::python::populatePassManagerSubmodule(py::module &m) {
   //----------------------------------------------------------------------------
@@ -114,6 +163,27 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
           "pipeline"_a,
           "Add textual pipeline elements to the pass manager. Throws a "
           "ValueError if the pipeline can't be parsed.")
+      .def_static(
+          "create_external_pass",
+          [](py::object &passObj) {
+            PythonPass pass = PythonPass(passObj);
+
+            MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
+            MlirTypeID passID =
+                mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+            MlirStringRef name =
+                mlirStringRefCreateFromCString("TestExternalPass");
+            MlirStringRef description = mlirStringRefCreateFromCString("");
+            MlirStringRef emptyOpName = mlirStringRefCreateFromCString("");
+            MlirStringRef argument =
+                mlirStringRefCreateFromCString("test-external-pass");
+
+            auto cbs = makeTestExternalPassCallbacks();
+
+            MlirPass externalPass =
+                mlirCreateExternalPass(passID, name, argument, description,
+                                       emptyOpName, 0, NULL, cbs, &pass);
+          })
       .def(
           "run",
           [](PyPassManager &passManager, PyOperationBase &op,
@@ -144,4 +214,8 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
           },
           "Print the textual representation for this PassManager, suitable to "
           "be passed to `parse` for round-tripping.");
+
+  py::class_<PythonPass>(m, "PythonPass", py::module_local())
+      .def(py::init<>(
+          [](py::object pass) { return PythonPass(std::move(pass)); }));
 }

mlir/lib/CAPI/IR/Pass.cpp

@@ -138,7 +138,7 @@ class ExternalPass : public Pass {
 protected:
   LogicalResult initialize(MLIRContext *ctx) override {
     if (callbacks.initialize)
-      return unwrap(callbacks.initialize(wrap(ctx), userData));
+      return unwrap(callbacks.initialize(userData, wrap(ctx)));
     return success();
   }
 
@@ -149,7 +149,7 @@ class ExternalPass : public Pass {
   }
 
   void runOnOperation() override {
-    callbacks.run(wrap(getOperation()), wrap(this), userData);
+    callbacks.run(userData, wrap(getOperation()), wrap(this));
   }
 
   std::unique_ptr<Pass> clonePass() const override {