[Mlir-commits] [mlir] [MLIR][Python] Handle errors in dialect conversion properly (PR #183320)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Twice (PragmaTwice)

<details>
<summary>Changes</summary>

Before this, MLIR error capture in `apply_partial_conversion` and `apply_full_conversion` wasn’t handled, which meant any `emitError` would crash the entire program. This PR adds the handling.


---
Full diff: https://github.com/llvm/llvm-project/pull/183320.diff


2 Files Affected:

- (modified) mlir/lib/Bindings/Python/Rewrite.cpp (+4-2) 
- (modified) mlir/test/python/rewrite.py (+19) 


``````````diff
diff --git a/mlir/lib/Bindings/Python/Rewrite.cpp b/mlir/lib/Bindings/Python/Rewrite.cpp
index 256dd0a78f809..e370552c00a9a 100644
--- a/mlir/lib/Bindings/Python/Rewrite.cpp
+++ b/mlir/lib/Bindings/Python/Rewrite.cpp
@@ -769,10 +769,11 @@ void populateRewriteSubmodule(nb::module_ &m) {
              std::optional<PyConversionConfig> config) {
             if (!config)
               config.emplace(PyConversionConfig());
+            PyMlirContext::ErrorCapture errors(op.getOperation().getContext());
             MlirLogicalResult status = mlirApplyPartialConversion(
                 op.getOperation(), target.get(), set.get(), config->get());
             if (mlirLogicalResultIsFailure(status))
-              throw std::runtime_error("partial conversion failed");
+              throw MLIRError("partial conversion failed", errors.take());
           },
           "op"_a, "target"_a, "set"_a, "config"_a = nb::none(),
           "Applies a partial conversion on the given operation.")
@@ -783,10 +784,11 @@ void populateRewriteSubmodule(nb::module_ &m) {
              std::optional<PyConversionConfig> config) {
             if (!config)
               config.emplace(PyConversionConfig());
+            PyMlirContext::ErrorCapture errors(op.getOperation().getContext());
             MlirLogicalResult status = mlirApplyFullConversion(
                 op.getOperation(), target.get(), set.get(), config->get());
             if (mlirLogicalResultIsFailure(status))
-              throw std::runtime_error("full conversion failed");
+              throw MLIRError("full conversion failed", errors.take());
           },
           "op"_a, "target"_a, "set"_a, "config"_a = nb::none(),
           "Applies a full conversion on the given operation.");
diff --git a/mlir/test/python/rewrite.py b/mlir/test/python/rewrite.py
index d805831369ce0..40dff039425a9 100644
--- a/mlir/test/python/rewrite.py
+++ b/mlir/test/python/rewrite.py
@@ -299,3 +299,22 @@ def convert_muli(op, adaptor, type_converter, rewriter):
         # CHECK:     return %3 : i64
         # CHECK: }
         print(module)
+
+        module = ModuleOp.parse(
+            r"""
+            module {
+                func.func @f(%0: i64) -> i64 {
+                    %1 = arith.constant 3 : i64
+                    %2 = arith.addi %0, %1 : i64
+                    %3 = arith.muli %2, %1 : i64
+                    return %3 : i64
+                }
+            }
+            """
+        )
+        try:
+            apply_partial_conversion(module, target, frozen)
+        except MLIRError as e:
+            # CHECK: caught exception: partial conversion failed
+            # CHECK: failed to legalize unresolved materialization from ('!smt.int') to ('i64') that remained live after conversion
+            print("caught exception:", e)

``````````

</details>


https://github.com/llvm/llvm-project/pull/183320