[MLIR][Python] Add the ability to signal pass failures in python-defined passes (#157613)

This is a follow-up PR for #156000.

In this PR we add the ability to signal pass failures
(`signal_pass_failure()`) in python-defined passes.

To achieve this, we expose `MlirExternalPass` via `nb::class_` with a
method `signal_pass_failure()`, and the callable passed to `pm.add(..)`
now accepts two arguments (`op: MlirOperation, pass_:
MlirExternalPass`).

For example:
```python
def custom_pass_that_fails(op, pass_):
    if some_condition:
        pass_.signal_pass_failure()
    # do something
```

Author: PragmaTwice

mlir/lib/Bindings/Python/Pass.cpp

@@ -56,6 +56,13 @@ class PyPassManager {
 
 /// Create the `mlir.passmanager` here.
 void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
+  //----------------------------------------------------------------------------
+  // Mapping of MlirExternalPass
+  //----------------------------------------------------------------------------
+  nb::class_<MlirExternalPass>(m, "ExternalPass")
+      .def("signal_pass_failure",
+           [](MlirExternalPass pass) { mlirExternalPassSignalFailure(pass); });
+
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
@@ -182,9 +189,9 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
             callbacks.clone = [](void *) -> void * {
               throw std::runtime_error("Cloning Python passes not supported");
             };
-            callbacks.run = [](MlirOperation op, MlirExternalPass,
+            callbacks.run = [](MlirOperation op, MlirExternalPass pass,
                                void *userData) {
-              nb::borrow<nb::callable>(static_cast<PyObject *>(userData))(op);
+              nb::handle(static_cast<PyObject *>(userData))(op, pass);
             };
             auto externalPass = mlirCreateExternalPass(
                 passID, mlirStringRefCreate(name->data(), name->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()
 
@@ -86,3 +86,17 @@ def __call__(self, m):
         # CHECK: llvm.mul
         pm.add("convert-arith-to-llvm")
         pm.run(module)
+
+        # test signal_pass_failure
+        def custom_pass_that_fails(op, pass_):
+            print("hello from pass that fails")
+            pass_.signal_pass_failure()
+
+        pm = PassManager("any")
+        pm.add(custom_pass_that_fails, "CustomPassThatFails")
+        # CHECK: hello from pass that fails
+        # CHECK: caught exception: Failure while executing pass pipeline
+        try:
+            pm.run(module)
+        except Exception as e:
+            print(f"caught exception: {e}")