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 { 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 { 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 { 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 { } if (std::optional capsule = mlirApiObjectToCapsule(src)) { value = mlirPythonCapsuleToContext(capsule->ptr()); - return !mlirContextIsNull(value); + return pyErrClearIfFalse(!mlirContextIsNull(value)); } return false; } @@ -173,7 +184,7 @@ struct type_caster { 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 { } 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 { 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 { 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 { 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 { 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 { 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 { 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 { 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);