[mlir:python] Fix crash in from_python in type casters.

[ingomueller-net]

This PR fixes a crash due to a failed assertion in the from_python implementations of the type casters. The assertion obviously only triggers if assertions are enabled, which isn't the case for many Python installations, and if a Python capsule of the wrong type is attempted to be used, so this this isn't triggered easily. The problem is that the conversion from Python capsules may set the Python error indicator but the callers of the type casters do not expect that. In fact, if there are several operloads of a function, the first may cause the error indicator to be set and the second runs into the assertion. The fix is to unset the error indicator after a failed capsule conversion, which is indicated with the return value of the function anyways.

In alternative fix would be to unset the error indicator inside the mlirPythonCapsuleTo* functions; however, their documentations does say that the Python error indicator is set, so I assume that some callers may want to see the indicator and that the responsibility to handle it is on them.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[github-actions]

✅ With the latest revision this PR passed the Python code formatter.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ingo Müller (ingomueller-net)

Changes

This PR fixes a crash due to a failed assertion in the from_python implementations of the type casters. The assertion obviously only triggers if assertions are enabled, which isn't the case for many Python installations, and if a Python capsule of the wrong type is attempted to be used, so this this isn't triggered easily. The problem is that the conversion from Python capsules may set the Python error indicator but the callers of the type casters do not expect that. In fact, if there are several operloads of a function, the first may cause the error indicator to be set and the second runs into the assertion. The fix is to unset the error indicator after a failed capsule conversion, which is indicated with the return value of the function anyways.

In alternative fix would be to unset the error indicator inside the mlirPythonCapsuleTo* functions; however, their documentations does say that the Python error indicator is set, so I assume that some callers may want to see the indicator and that the responsibility to handle it is on them.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (+24-13)
• (modified) mlir/test/python/dialects/python_test.py (+35)
• (modified) mlir/test/python/lib/PythonTestModuleNanobind.cpp (+27)

diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 6669433550a00..68329e580975c 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -72,6 +72,17 @@ mlirApiObjectToCapsule(nanobind::handle apiObject) {
   return api;
 }
 
+/// Clears the Python error indicator if the given condition `val` is false and
+/// returns the condition. This is needed in `from_python` of the type casters
+/// below, where a failed conversion from a Python capsule sets the Python error
+/// indicator but the caller of the caster expects or may even need a clean
+/// error indicator.
+inline bool pyErrClearIfFalse(bool val) {
+  if (!val)
+    PyErr_Clear();
+  return val;
+}
+
 // Note: Currently all of the following support cast from nanobind::object to
 // the Mlir* C-API type, but only a few light-weight, context-bound ones
 // implicitly cast the other way because the use case has not yet emerged and
@@ -85,7 +96,7 @@ struct type_caster<MlirAffineMap> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToAffineMap(capsule->ptr());
-      return !mlirAffineMapIsNull(value);
+      return pyErrClearIfFalse(!mlirAffineMapIsNull(value));
     }
     return false;
   }
@@ -108,7 +119,7 @@ struct type_caster<MlirAttribute> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToAttribute(capsule->ptr());
-      return !mlirAttributeIsNull(value);
+      return pyErrClearIfFalse(!mlirAttributeIsNull(value));
     }
     return false;
   }
@@ -131,7 +142,7 @@ struct type_caster<MlirBlock> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToBlock(capsule->ptr());
-      return !mlirBlockIsNull(value);
+      return pyErrClearIfFalse(!mlirBlockIsNull(value));
     }
     return false;
   }
@@ -159,7 +170,7 @@ struct type_caster<MlirContext> {
     }
     if (std::optional<nanobind::object> capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToContext(capsule->ptr());
-      return !mlirContextIsNull(value);
+      return pyErrClearIfFalse(!mlirContextIsNull(value));
     }
     return false;
   }
@@ -173,7 +184,7 @@ struct type_caster<MlirDialectRegistry> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToDialectRegistry(capsule->ptr());
-      return !mlirDialectRegistryIsNull(value);
+      return pyErrClearIfFalse(!mlirDialectRegistryIsNull(value));
     }
     return false;
   }
@@ -200,7 +211,7 @@ struct type_caster<MlirLocation> {
     }
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToLocation(capsule->ptr());
-      return !mlirLocationIsNull(value);
+      return pyErrClearIfFalse(!mlirLocationIsNull(value));
     }
     return false;
   }
@@ -222,7 +233,7 @@ struct type_caster<MlirModule> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToModule(capsule->ptr());
-      return !mlirModuleIsNull(value);
+      return pyErrClearIfFalse(!mlirModuleIsNull(value));
     }
     return false;
   }
@@ -246,7 +257,7 @@ struct type_caster<MlirFrozenRewritePatternSet> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule->ptr());
-      return value.ptr != nullptr;
+      return pyErrClearIfFalse(value.ptr != nullptr);
     }
     return false;
   }
@@ -269,7 +280,7 @@ struct type_caster<MlirOperation> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToOperation(capsule->ptr());
-      return !mlirOperationIsNull(value);
+      return pyErrClearIfFalse(!mlirOperationIsNull(value));
     }
     return false;
   }
@@ -293,7 +304,7 @@ struct type_caster<MlirValue> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToValue(capsule->ptr());
-      return !mlirValueIsNull(value);
+      return pyErrClearIfFalse(!mlirValueIsNull(value));
     }
     return false;
   }
@@ -319,7 +330,7 @@ struct type_caster<MlirPassManager> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToPassManager(capsule->ptr());
-      return !mlirPassManagerIsNull(value);
+      return pyErrClearIfFalse(!mlirPassManagerIsNull(value));
     }
     return false;
   }
@@ -332,7 +343,7 @@ struct type_caster<MlirTypeID> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToTypeID(capsule->ptr());
-      return !mlirTypeIDIsNull(value);
+      return pyErrClearIfFalse(!mlirTypeIDIsNull(value));
     }
     return false;
   }
@@ -356,7 +367,7 @@ struct type_caster<MlirType> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (auto capsule = mlirApiObjectToCapsule(src)) {
       value = mlirPythonCapsuleToType(capsule->ptr());
-      return !mlirTypeIsNull(value);
+      return pyErrClearIfFalse(!mlirTypeIsNull(value));
     }
     return false;
   }
diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index d4b394b8c7e06..655862f957771 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -14,6 +14,7 @@
     TestType,
     TestTensorValue,
     TestIntegerRankedTensorType,
+    take_module_or_operation,
 )
 
 test.register_python_test_dialect(get_dialect_registry())
@@ -1033,3 +1034,37 @@ def testVariadicAndNormalRegionOp():
 
             assert isinstance(region_op.opview, OpView)
             assert isinstance(region_op.operation.opview, OpView)
+
+
+# Regression test for the dirty-error-state crash in `NanobindAdaptors.h`
+# `from_python` type casters (#191764).
+#
+# !!! This only fails with a debug version of Python. !!!
+#
+# Uses an overloaded function: overload 1 takes `MlirOperation`, overload 2
+# takes `MlirModule`. When called with an `ir.Module`:
+#
+#   1. `nanobind` tries overload 1 (`MlirOperation`). `from_python` gets the
+#      `Module`'s `_CAPIPtr` capsule, then `mlirPythonCapsuleToOperation` calls
+#      `PyCapsule_GetPointer` with `"mlir.ir.Operation._CAPIPtr"` — but the
+#      capsule is named `"mlir.ir.Module._CAPIPtr"`. `PyCapsule_GetPointer`
+#      returns `NULL` and sets `PyErr_Occurred()`. `from_python` returns `false`.
+#
+#   2. `nanobind` tries overload 2 (`MlirModule`). `from_python` calls
+#      `mlirApiObjectToCapsule` --> `nanobind::getattr(obj, "_CAPIPtr")` -->
+#      `_PyType_LookupRef`.
+#
+# Without the fix:
+#   `_PyType_LookupRef` asserts `!PyErr_Occurred()` --> `SIGABRT`.
+#
+# With the fix (`PyErr_Clear` in `from_python` after failed capsule conversion):
+#   Overload 2 succeeds and returns `"module"`.
+# CHECK-LABEL: testOverloadWithWrongPythonCapsule
+@run
+def testOverloadWithWrongPythonCapsule():
+    with Context():
+        module = Module.parse("module {}")
+
+    # CHECK: result = module
+    result = take_module_or_operation(module)
+    print(f"result = {result}")
diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
index e9754749352b1..36fc7b433752d 100644
--- a/mlir/test/python/lib/PythonTestModuleNanobind.cpp
+++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
@@ -147,6 +147,33 @@ NB_MODULE(_mlirPythonTestNanobind, m) {
       },
       nb::arg("context").none() = nb::none());
 
+  // Reproducer for the failed assertion `_PyType_LookupRef` triggered by
+  // `NanobindAdaptors.h::from_python` type casters.
+  //
+  // Two overloads of the same function: one takes `MlirOperation`, the other
+  // takes `MlirModule`. When called with an `ir.Module`:
+  //
+  //   1. nanobind tries overload 1 (`MlirOperation`). `from_python` calls
+  //      `mlirApiObjectToCapsule` (succeeds — `Module` has `_CAPIPtr`), then
+  //      `mlirPythonCapsuleToOperation`, whose `PyCapsule_GetPointer` fails on
+  //      the capsule-name mismatch and sets `PyErr_Occurred()`.
+  //      `from_python` returns false.
+  //
+  //   If `PyErr` is still set and assertions are enabled:
+  //   2. nanobind tries overload 2 (`MlirModule`).  `from_python` calls
+  //      `mlirApiObjectToCapsule` --> `nanobind::getattr(obj, "_CAPIPtr")` -->
+  //      CPython's `_PyType_LookupRef` --> `assert(!PyErr_Occurred())` -->
+  //      `SIGABRT`.
+  //
+  //   If `PyErr_Clear` is called after failed capsule conversion:
+  //   2. `PyErr` is clear --> overload 2 succeeds --> returns "module".
+  m.def(
+      "take_module_or_operation",
+      [](MlirOperation) { return std::string("operation"); }, nb::arg("arg"));
+  m.def(
+      "take_module_or_operation",
+      [](MlirModule) { return std::string("module"); }, nb::arg("arg"));
+
   using namespace python_test;
   PyTestAttr::bind(m);
   PyTestType::bind(m);

[makslevental]

Man this is a subtle bug - good job...

[ingomueller-net]

I had help from Claude. It came up with five BS fixes before this one; I had to nudge it quite a bit into the right direction... 😉

[joker-eph]

I use Claude all the time, it's great (even though requires some hand holding like you report here). One note for the future: the project policy here https://llvm.org/docs/AIToolPolicy.html requires the PR description (/ future commit message) to include an "Assisted-by: " mention of the AI help.