diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h index 2dd35c097c796..8dcf91e5806bd 100644 --- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h +++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h @@ -20,6 +20,7 @@ #define MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H #include +#include #include "mlir-c/Diagnostics.h" #include "mlir-c/IR.h" @@ -43,18 +44,14 @@ namespace detail { /// with a raw handle (unowned). The returned object's lifetime may not extend /// beyond the apiObject handle without explicitly having its refcount increased /// (i.e. on return). -static nanobind::object mlirApiObjectToCapsule(nanobind::handle apiObject) { +static std::optional +mlirApiObjectToCapsule(nanobind::handle apiObject) { if (PyCapsule_CheckExact(apiObject.ptr())) return nanobind::borrow(apiObject); nanobind::object api = nanobind::getattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR, nanobind::none()); - if (api.is_none()) { - std::string repr = nanobind::cast(nanobind::repr(apiObject)); - throw nanobind::type_error( - (llvm::Twine("Expected an MLIR object (got ") + repr + ").") - .str() - .c_str()); - } + if (api.is_none()) + return {}; return api; } @@ -67,12 +64,9 @@ static nanobind::object mlirApiObjectToCapsule(nanobind::handle apiObject) { template <> struct type_caster { NB_TYPE_CASTER(MlirAffineMap, const_name("MlirAffineMap")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToAffineMap(capsule.ptr()); - if (mlirAffineMapIsNull(value)) { - return false; - } + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToAffineMap(capsule->ptr()); return !mlirAffineMapIsNull(value); } static handle from_cpp(MlirAffineMap v, rv_policy, @@ -90,9 +84,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirAttribute, const_name("MlirAttribute")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToAttribute(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToAttribute(capsule->ptr()); return !mlirAttributeIsNull(value); } static handle from_cpp(MlirAttribute v, rv_policy, @@ -111,9 +105,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirBlock, const_name("MlirBlock")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToBlock(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToBlock(capsule->ptr()); return !mlirBlockIsNull(value); } }; @@ -122,7 +116,7 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirContext, const_name("MlirContext")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { if (src.is_none()) { // Gets the current thread-bound context. // TODO: This raises an error of "No current context" currently. @@ -132,8 +126,8 @@ struct type_caster { .attr("Context") .attr("current"); } - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToContext(capsule.ptr()); + std::optional capsule = mlirApiObjectToCapsule(src); + value = mlirPythonCapsuleToContext(capsule->ptr()); return !mlirContextIsNull(value); } }; @@ -142,9 +136,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirDialectRegistry, const_name("MlirDialectRegistry")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToDialectRegistry(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToDialectRegistry(capsule->ptr()); return !mlirDialectRegistryIsNull(value); } static handle from_cpp(MlirDialectRegistry v, rv_policy, @@ -162,15 +156,15 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirLocation, const_name("MlirLocation")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { if (src.is_none()) { // Gets the current thread-bound context. src = nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir")) .attr("Location") .attr("current"); } - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToLocation(capsule.ptr()); + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToLocation(capsule->ptr()); return !mlirLocationIsNull(value); } static handle from_cpp(MlirLocation v, rv_policy, @@ -188,9 +182,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirModule, const_name("MlirModule")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToModule(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToModule(capsule->ptr()); return !mlirModuleIsNull(value); } static handle from_cpp(MlirModule v, rv_policy, @@ -209,12 +203,13 @@ template <> struct type_caster { NB_TYPE_CASTER(MlirFrozenRewritePatternSet, const_name("MlirFrozenRewritePatternSet")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule.ptr()); + 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; } - static handle from_cpp(MlirFrozenRewritePatternSet v, rv_policy, handle) { + static handle from_cpp(MlirFrozenRewritePatternSet v, rv_policy, + handle) noexcept { nanobind::object capsule = nanobind::steal( mlirPythonFrozenRewritePatternSetToCapsule(v)); return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("rewrite")) @@ -228,9 +223,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirOperation, const_name("MlirOperation")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToOperation(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToOperation(capsule->ptr()); return !mlirOperationIsNull(value); } static handle from_cpp(MlirOperation v, rv_policy, @@ -250,9 +245,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirValue, const_name("MlirValue")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToValue(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToValue(capsule->ptr()); return !mlirValueIsNull(value); } static handle from_cpp(MlirValue v, rv_policy, @@ -273,9 +268,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirPassManager, const_name("MlirPassManager")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToPassManager(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToPassManager(capsule->ptr()); return !mlirPassManagerIsNull(value); } }; @@ -284,9 +279,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirTypeID, const_name("MlirTypeID")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToTypeID(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToTypeID(capsule->ptr()); return !mlirTypeIDIsNull(value); } static handle from_cpp(MlirTypeID v, rv_policy, @@ -306,9 +301,9 @@ struct type_caster { template <> struct type_caster { NB_TYPE_CASTER(MlirType, const_name("MlirType")) - bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) { - nanobind::object capsule = mlirApiObjectToCapsule(src); - value = mlirPythonCapsuleToType(capsule.ptr()); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if (auto capsule = mlirApiObjectToCapsule(src)) + value = mlirPythonCapsuleToType(capsule->ptr()); return !mlirTypeIsNull(value); } static handle from_cpp(MlirType t, rv_policy, @@ -462,9 +457,10 @@ class mlir_attribute_subclass : public pure_subclass { nanobind::object newCf = nanobind::cpp_function( [superCls, isaFunction, captureTypeName]( nanobind::object cls, nanobind::object otherAttribute) { - MlirAttribute rawAttribute = - nanobind::cast(otherAttribute); - if (!isaFunction(rawAttribute)) { + MlirAttribute rawAttribute; + if (!nanobind::try_cast(otherAttribute, + rawAttribute) || + !isaFunction(rawAttribute)) { auto origRepr = nanobind::cast(nanobind::repr(otherAttribute)); throw std::invalid_argument( @@ -543,8 +539,9 @@ class mlir_type_subclass : public pure_subclass { nanobind::object newCf = nanobind::cpp_function( [superCls, isaFunction, captureTypeName](nanobind::object cls, nanobind::object otherType) { - MlirType rawType = nanobind::cast(otherType); - if (!isaFunction(rawType)) { + MlirType rawType; + if (!nanobind::try_cast(otherType, rawType) || + !isaFunction(rawType)) { auto origRepr = nanobind::cast(nanobind::repr(otherType)); throw std::invalid_argument((llvm::Twine("Cannot cast type to ") + @@ -625,8 +622,9 @@ class mlir_value_subclass : public pure_subclass { nanobind::object newCf = nanobind::cpp_function( [superCls, isaFunction, captureValueName](nanobind::object cls, nanobind::object otherValue) { - MlirValue rawValue = nanobind::cast(otherValue); - if (!isaFunction(rawValue)) { + MlirValue rawValue; + if (!nanobind::try_cast(otherValue, rawValue) || + !isaFunction(rawValue)) { auto origRepr = nanobind::cast(nanobind::repr(otherValue)); throw std::invalid_argument((llvm::Twine("Cannot cast value to ") + diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py index fd678f8321fd9..694616696a9e2 100644 --- a/mlir/test/python/dialects/python_test.py +++ b/mlir/test/python/dialects/python_test.py @@ -361,7 +361,9 @@ def testCustomAttribute(): try: TestAttr(42) except TypeError as e: - assert "Expected an MLIR object" in str(e) + assert "Expected an MLIR object (got 42)" in str(e) + except ValueError as e: + assert "Cannot cast attribute to TestAttr (from 42)" in str(e) else: raise @@ -406,7 +408,9 @@ def testCustomType(): try: TestType(42) except TypeError as e: - assert "Expected an MLIR object" in str(e) + assert "Expected an MLIR object (got 42)" in str(e) + except ValueError as e: + assert "Cannot cast type to TestType (from 42)" in str(e) else: raise diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp index dd0a9f2b66ea6..108df8da86a7e 100644 --- a/mlir/test/python/lib/PythonTestModuleNanobind.cpp +++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp @@ -113,8 +113,10 @@ NB_MODULE(_mlirPythonTestNanobind, m) { .attr(MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR)( mlirRankedTensorTypeID)( nanobind::cpp_function([valueCls](const nb::object &valueObj) { - nb::object capsule = mlirApiObjectToCapsule(valueObj); - MlirValue v = mlirPythonCapsuleToValue(capsule.ptr()); + std::optional capsule = + mlirApiObjectToCapsule(valueObj); + assert(capsule.has_value() && "capsule is not null"); + MlirValue v = mlirPythonCapsuleToValue(capsule.value().ptr()); MlirType t = mlirValueGetType(v); // This is hyper-specific in order to exercise/test registering a // value caster from cpp (but only for a single test case; see