drop the setattr design

Author: PragmaTwice

mlir/lib/Bindings/Python/Pass.cpp

@@ -56,7 +56,12 @@ class PyPassManager {
 
 /// Create the `mlir.passmanager` here.
 void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
-  constexpr const char *mlirExternalPassAttr = "__mlir_external_pass__";
+  //----------------------------------------------------------------------------
+  // Mapping of MlirExternalPass
+  //----------------------------------------------------------------------------
+  nb::class_<MlirExternalPass>(m, "ExternalPass")
+      .def("signal_failure",
+           [](MlirExternalPass pass) { mlirExternalPassSignalFailure(pass); });
 
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
@@ -186,27 +191,8 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
             };
             callbacks.run = [](MlirOperation op, MlirExternalPass pass,
                                void *userData) {
-              auto callable =
-                  nb::borrow<nb::callable>(static_cast<PyObject *>(userData));
-              nb::setattr(callable, mlirExternalPassAttr,
-                          nb::capsule(pass.ptr));
-              callable(op);
-              // delete it to avoid that it is used after
-              // the external pass is freed by the pass manager
-              nb::delattr(callable, mlirExternalPassAttr);
+              nb::handle(static_cast<PyObject *>(userData))(op, pass);
             };
-            nb::setattr(run, "signal_pass_failure", nb::cpp_function([run]() {
-                          nb::capsule cap;
-                          try {
-                            cap = run.attr(mlirExternalPassAttr);
-                          } catch (nb::python_error &e) {
-                            throw std::runtime_error(
-                                "signal_pass_failure() should always be called "
-                                "from the __call__ method");
-                          }
-                          mlirExternalPassSignalFailure(
-                              MlirExternalPass{cap.data()});
-                        }));
             auto externalPass = mlirCreateExternalPass(
                 passID, mlirStringRefCreate(name->data(), name->length()),
                 mlirStringRefCreate(argument.data(), argument.length()),

mlir/test/python/python_pass.py

@@ -64,12 +64,12 @@ def testCustomPass():
         """
         )
 
-        def custom_pass_1(op):
+        def custom_pass_1(op, pass_):
             print("hello from pass 1!!!", file=sys.stderr)
 
         class CustomPass2:
-            def __call__(self, m):
-                apply_patterns_and_fold_greedily(m, frozen)
+            def __call__(self, op, pass_):
+                apply_patterns_and_fold_greedily(op, frozen)
 
         custom_pass_2 = CustomPass2()
 
@@ -89,9 +89,9 @@ def __call__(self, m):
 
         # test signal_pass_failure
         class CustomPassThatFails:
-            def __call__(self, m):
+            def __call__(self, op, pass_):
                 print("hello from pass that fails")
-                self.signal_pass_failure()
+                pass_.signal_failure()
 
         custom_pass_that_fails = CustomPassThatFails()
 
@@ -103,9 +103,3 @@ def __call__(self, m):
             pm.run(module)
         except Exception as e:
             print(f"caught exception: {e}")
-
-        # CHECK: caught exception: signal_pass_failure() should always be called from the __call__ method
-        try:
-            custom_pass_that_fails.signal_pass_failure()
-        except Exception as e:
-            print(f"caught exception: {e}")