Skip to content

Instantly share code, notes, and snippets.

@makslevental
Created July 10, 2025 15:08
Show Gist options
  • Save makslevental/b224ffca7f15e273a4897975cda28b4c to your computer and use it in GitHub Desktop.
Save makslevental/b224ffca7f15e273a4897975cda28b4c to your computer and use it in GitHub Desktop.
Index: mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (revision 468275dc491e88d68def668f1a23cde9ba952e5e)
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (date 1752159288603)
@@ -43,18 +43,14 @@
/// 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<nanobind::object>
+mlirApiObjectToCapsule(nanobind::handle apiObject) {
if (PyCapsule_CheckExact(apiObject.ptr()))
return nanobind::borrow<nanobind::object>(apiObject);
nanobind::object api =
nanobind::getattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR, nanobind::none());
- if (api.is_none()) {
- std::string repr = nanobind::cast<std::string>(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 +63,9 @@
template <>
struct type_caster<MlirAffineMap> {
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 +83,9 @@
template <>
struct type_caster<MlirAttribute> {
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 +104,9 @@
template <>
struct type_caster<MlirBlock> {
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 +115,7 @@
template <>
struct type_caster<MlirContext> {
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 +125,8 @@
.attr("Context")
.attr("current");
}
- nanobind::object capsule = mlirApiObjectToCapsule(src);
- value = mlirPythonCapsuleToContext(capsule.ptr());
+ if (auto capsule = mlirApiObjectToCapsule(src))
+ value = mlirPythonCapsuleToContext(capsule->ptr());
return !mlirContextIsNull(value);
}
};
@@ -142,9 +135,9 @@
template <>
struct type_caster<MlirDialectRegistry> {
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 +155,15 @@
template <>
struct type_caster<MlirLocation> {
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 +181,9 @@
template <>
struct type_caster<MlirModule> {
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 +202,13 @@
struct type_caster<MlirFrozenRewritePatternSet> {
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<nanobind::object>(
mlirPythonFrozenRewritePatternSetToCapsule(v));
return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("rewrite"))
@@ -228,9 +222,9 @@
template <>
struct type_caster<MlirOperation> {
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 +244,9 @@
template <>
struct type_caster<MlirValue> {
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 +267,9 @@
template <>
struct type_caster<MlirPassManager> {
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 +278,9 @@
template <>
struct type_caster<MlirTypeID> {
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 +300,9 @@
template <>
struct type_caster<MlirType> {
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,
Index: mlir/test/python/lib/PythonTestModuleNanobind.cpp
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
--- a/mlir/test/python/lib/PythonTestModuleNanobind.cpp (revision 468275dc491e88d68def668f1a23cde9ba952e5e)
+++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp (date 1752159288608)
@@ -113,8 +113,10 @@
.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<nb::object> capsule =
+ mlirApiObjectToCapsule(valueObj);
+ assert(capsule.has_value() && "capsule is not null");
+ MlirValue v = mlirPythonCapsuleToValue(capsule->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
Index: mlir/test/python/dialects/python_test.py
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
--- a/mlir/test/python/dialects/python_test.py (revision 468275dc491e88d68def668f1a23cde9ba952e5e)
+++ b/mlir/test/python/dialects/python_test.py (date 1752160097302)
@@ -361,7 +361,9 @@
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 @@
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment