Merge changes from llvm#157369

Co-authored-by: Maksim Levental <[email protected]>

Author: PragmaTwice

mlir/lib/Bindings/Python/MainModule.cpp

@@ -139,7 +139,4 @@ NB_MODULE(_mlir, m) {
   auto passManagerModule =
       m.def_submodule("passmanager", "MLIR Pass Management Bindings");
   populatePassManagerSubmodule(passManagerModule);
-  auto passesModule =
-      m.def_submodule("passes", "MLIR Pass Infrastructure Bindings");
-  populatePassSubmodule(passesModule);
 }

mlir/lib/Bindings/Python/Pass.cpp

@@ -14,7 +14,6 @@
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
 // clang-format on
-#include "nanobind/trampoline.h"
 
 namespace nb = nanobind;
 using namespace nb::literals;
@@ -23,81 +22,6 @@ using namespace mlir::python;
 
 namespace {
 
-// A base class for defining passes in Python
-// Users are expected to subclass this and implement the `run` method, e.g.
-// ```
-// class MyPass(Pass):
-//   def __init__(self):
-//     super().__init__("MyPass", ..)
-//     # other init stuff..
-//   def run(self, operation):
-//     # do something with operation..
-//     pass
-// ```
-class PyPassBase {
-public:
-  PyPassBase(std::string name, std::string argument, std::string description,
-             std::string opName)
-      : name(std::move(name)), argument(std::move(argument)),
-        description(std::move(description)), opName(std::move(opName)) {
-    callbacks.construct = [](void *obj) {};
-    callbacks.destruct = [](void *obj) {
-      nb::handle(static_cast<PyObject *>(obj)).dec_ref();
-    };
-    callbacks.run = [](MlirOperation op, MlirExternalPass, void *obj) {
-      auto handle = nb::handle(static_cast<PyObject *>(obj));
-      nb::cast<PyPassBase *>(handle)->run(op);
-    };
-    callbacks.clone = [](void *obj) -> void * {
-      nb::object copy = nb::module_::import_("copy");
-      nb::object deepcopy = copy.attr("deepcopy");
-      return deepcopy(obj).release().ptr();
-    };
-    callbacks.initialize = nullptr;
-  }
-
-  // this method should be overridden by subclasses in Python.
-  virtual void run(MlirOperation op) = 0;
-
-  virtual ~PyPassBase() = default;
-
-  // Make an MlirPass instance on-the-fly that wraps this object.
-  // Note that passmanager will take the ownership of the returned
-  // object and release it when appropriate.
-  MlirPass make() {
-    auto *obj = nb::find(this).release().ptr();
-    return mlirCreateExternalPass(
-        mlirTypeIDCreate(this), mlirStringRefCreate(name.data(), name.length()),
-        mlirStringRefCreate(argument.data(), argument.length()),
-        mlirStringRefCreate(description.data(), description.length()),
-        mlirStringRefCreate(opName.data(), opName.size()), 0, nullptr,
-        callbacks, obj);
-  }
-
-  const std::string &getName() const { return name; }
-  const std::string &getArgument() const { return argument; }
-  const std::string &getDescription() const { return description; }
-  const std::string &getOpName() const { return opName; }
-
-private:
-  MlirExternalPassCallbacks callbacks;
-
-  std::string name;
-  std::string argument;
-  std::string description;
-  std::string opName;
-};
-
-// A trampoline class upon PyPassBase.
-// Refer to
-// https://nanobind.readthedocs.io/en/latest/classes.html#overriding-virtual-functions-in-python
-class PyPass : PyPassBase {
-public:
-  NB_TRAMPOLINE(PyPassBase, 1);
-
-  void run(MlirOperation op) override { NB_OVERRIDE_PURE(run, op); }
-};
-
 /// Owning Wrapper around a PassManager.
 class PyPassManager {
 public:
@@ -130,26 +54,6 @@ class PyPassManager {
 
 } // namespace
 
-void mlir::python::populatePassSubmodule(nanobind::module_ &m) {
-  //----------------------------------------------------------------------------
-  // Mapping of the Python-defined Pass interface
-  //----------------------------------------------------------------------------
-  nb::class_<PyPassBase, PyPass>(m, "Pass")
-      .def(nb::init<std::string, std::string, std::string, std::string>(),
-           "name"_a, nb::kw_only(), "argument"_a = "", "description"_a = "",
-           "op_name"_a = "", "Create a new Pass.")
-      .def("run", &PyPassBase::run, "operation"_a,
-           "Run the pass on the provided operation.")
-      .def_prop_ro("name",
-                   [](const PyPassBase &self) { return self.getName(); })
-      .def_prop_ro("argument",
-                   [](const PyPassBase &self) { return self.getArgument(); })
-      .def_prop_ro("description",
-                   [](const PyPassBase &self) { return self.getDescription(); })
-      .def_prop_ro("op_name",
-                   [](const PyPassBase &self) { return self.getOpName(); });
-}
-
 /// Create the `mlir.passmanager` here.
 void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
   //----------------------------------------------------------------------------
@@ -256,11 +160,44 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
           "Add textual pipeline elements to the pass manager. Throws a "
           "ValueError if the pipeline can't be parsed.")
       .def(
-          "add",
-          [](PyPassManager &passManager, PyPassBase &pass) {
-            mlirPassManagerAddOwnedPass(passManager.get(), pass.make());
+          "add_python_pass",
+          [](PyPassManager &passManager, const nb::callable &run,
+             std::optional<std::string> &name, const std::string &argument,
+             const std::string &description, const std::string &opName) {
+            if (!name.has_value()) {
+              name = nb::cast<std::string>(
+                  nb::borrow<nb::str>(run.attr("__name__")));
+            }
+            MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
+            MlirTypeID passID =
+                mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+            MlirExternalPassCallbacks callbacks;
+            callbacks.construct = [](void *obj) {
+              (void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();
+            };
+            callbacks.destruct = [](void *obj) {
+              (void)nb::handle(static_cast<PyObject *>(obj)).dec_ref();
+            };
+            callbacks.initialize = nullptr;
+            callbacks.clone = [](void *) -> void * {
+              throw std::runtime_error("Cloning Python passes not supported");
+            };
+            callbacks.run = [](MlirOperation op, MlirExternalPass,
+                               void *userData) {
+              nb::borrow<nb::callable>(static_cast<PyObject *>(userData))(op);
+            };
+            auto externalPass = mlirCreateExternalPass(
+                passID, mlirStringRefCreate(name->data(), name->length()),
+                mlirStringRefCreate(argument.data(), argument.length()),
+                mlirStringRefCreate(description.data(), description.length()),
+                mlirStringRefCreate(opName.data(), opName.size()),
+                /*nDependentDialects*/ 0, /*dependentDialects*/ nullptr,
+                callbacks, /*userData*/ run.ptr());
+            mlirPassManagerAddOwnedPass(passManager.get(), externalPass);
           },
-          "pass"_a, "Add a python-defined pass to the pass manager.")
+          "run"_a, "name"_a.none() = nb::none(), "argument"_a.none() = "",
+          "description"_a.none() = "", "op_name"_a.none() = "",
+          "Add a python-defined pass to the pass manager.")
       .def(
           "run",
           [](PyPassManager &passManager, PyOperationBase &op,

mlir/lib/Bindings/Python/Pass.h

@@ -15,7 +15,6 @@ namespace mlir {
 namespace python {
 
 void populatePassManagerSubmodule(nanobind::module_ &m);
-void populatePassSubmodule(nanobind::module_ &m);
 
 } // namespace python
 } // namespace mlir

mlir/lib/CAPI/IR/Pass.cpp

@@ -145,10 +145,14 @@ class ExternalPass : public Pass {
       : Pass(passID, opName), id(passID), name(name), argument(argument),
         description(description), dependentDialects(dependentDialects),
         callbacks(callbacks), userData(userData) {
-    callbacks.construct(userData);
+    if (callbacks.construct)
+      callbacks.construct(userData);
   }
 
-  ~ExternalPass() override { callbacks.destruct(userData); }
+  ~ExternalPass() override {
+    if (callbacks.destruct)
+      callbacks.destruct(userData);
+  }
 
   StringRef getName() const override { return name; }
   StringRef getArgument() const override { return argument; }

mlir/python/CMakeLists.txt

@@ -20,7 +20,6 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python
   SOURCES
     _mlir_libs/__init__.py
     ir.py
-    passes.py
     passmanager.py
     rewrite.py
     dialects/_ods_common.py

mlir/python/mlir/passes.py

@@ -3,7 +3,6 @@
 import gc, sys
 from mlir.ir import *
 from mlir.passmanager import *
-from mlir.passes import *
 from mlir.dialects.builtin import ModuleOp
 from mlir.dialects import pdl
 from mlir.rewrite import *
@@ -54,13 +53,6 @@ def testCustomPass():
         pdl_module = make_pdl_module()
         frozen = PDLModule(pdl_module).freeze()
 
-        class CustomPass(Pass):
-            def __init__(self):
-                super().__init__("CustomPass", op_name="builtin.module")
-
-            def run(self, m):
-                apply_patterns_and_fold_greedily_with_op(m, frozen)
-
         module = ModuleOp.parse(
             r"""
             module {
@@ -72,12 +64,24 @@ def run(self, m):
         """
         )
 
+        def custom_pass_1(op):
+            print("hello from pass 1!!!", file=sys.stderr)
+
+        class CustomPass2:
+            def __call__(self, m):
+                apply_patterns_and_fold_greedily_with_op(m, frozen)
+
+        custom_pass_2 = CustomPass2()
+
         pm = PassManager("any")
         pm.enable_ir_printing()
 
-        # CHECK-LABEL: Dump After CustomPass
+        # CHECK: hello from pass 1!!!
+        # CHECK-LABEL: Dump After custom_pass_1
+        pm.add_python_pass(custom_pass_1)
+        # CHECK-LABEL: Dump After CustomPass2
         # CHECK: arith.muli
-        pm.add(CustomPass())
+        pm.add_python_pass(custom_pass_2, "CustomPass2")
         # CHECK-LABEL: Dump After ArithToLLVMConversionPass
         # CHECK: llvm.mul
         pm.add("convert-arith-to-llvm")