Skip to content

[MLIR][Python] fixup Context and Location stubs and NanobindAdaptors #161433

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 2, 2025
Merged

[MLIR][Python] fixup Context and Location stubs and NanobindAdaptors #161433

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 23 additions & 15 deletions mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,8 @@ mlirApiObjectToCapsule(nanobind::handle apiObject) {
/// Casts object <-> MlirAffineMap.
template <>
struct type_caster<MlirAffineMap> {
NB_TYPE_CASTER(MlirAffineMap, const_name("MlirAffineMap"))
NB_TYPE_CASTER(MlirAffineMap,
const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.AffineMap")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToAffineMap(capsule->ptr());
Expand All @@ -138,7 +139,8 @@ struct type_caster<MlirAffineMap> {
/// Casts object <-> MlirAttribute.
template <>
struct type_caster<MlirAttribute> {
NB_TYPE_CASTER(MlirAttribute, const_name("MlirAttribute"))
NB_TYPE_CASTER(MlirAttribute,
const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Attribute")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToAttribute(capsule->ptr());
Expand All @@ -161,7 +163,7 @@ struct type_caster<MlirAttribute> {
/// Casts object -> MlirBlock.
template <>
struct type_caster<MlirBlock> {
NB_TYPE_CASTER(MlirBlock, const_name("MlirBlock"))
NB_TYPE_CASTER(MlirBlock, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Block")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToBlock(capsule->ptr());
Expand All @@ -174,7 +176,8 @@ struct type_caster<MlirBlock> {
/// Casts object -> MlirContext.
template <>
struct type_caster<MlirContext> {
NB_TYPE_CASTER(MlirContext, const_name("MlirContext"))
NB_TYPE_CASTER(MlirContext,
const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Context")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (src.is_none()) {
// Gets the current thread-bound context.
Expand All @@ -192,7 +195,8 @@ struct type_caster<MlirContext> {
/// Casts object <-> MlirDialectRegistry.
template <>
struct type_caster<MlirDialectRegistry> {
NB_TYPE_CASTER(MlirDialectRegistry, const_name("MlirDialectRegistry"))
NB_TYPE_CASTER(MlirDialectRegistry,
const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.DialectRegistry")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToDialectRegistry(capsule->ptr());
Expand All @@ -214,7 +218,8 @@ struct type_caster<MlirDialectRegistry> {
/// Casts object <-> MlirLocation.
template <>
struct type_caster<MlirLocation> {
NB_TYPE_CASTER(MlirLocation, const_name("MlirLocation"))
NB_TYPE_CASTER(MlirLocation,
const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Location")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (src.is_none()) {
// Gets the current thread-bound context.
Expand All @@ -240,7 +245,7 @@ struct type_caster<MlirLocation> {
/// Casts object <-> MlirModule.
template <>
struct type_caster<MlirModule> {
NB_TYPE_CASTER(MlirModule, const_name("MlirModule"))
NB_TYPE_CASTER(MlirModule, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Module")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToModule(capsule->ptr());
Expand All @@ -262,8 +267,9 @@ struct type_caster<MlirModule> {
/// Casts object <-> MlirFrozenRewritePatternSet.
template <>
struct type_caster<MlirFrozenRewritePatternSet> {
NB_TYPE_CASTER(MlirFrozenRewritePatternSet,
const_name("MlirFrozenRewritePatternSet"))
NB_TYPE_CASTER(
MlirFrozenRewritePatternSet,
const_name(MAKE_MLIR_PYTHON_QUALNAME("rewrite.FrozenRewritePatternSet")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule->ptr());
Expand All @@ -285,7 +291,8 @@ struct type_caster<MlirFrozenRewritePatternSet> {
/// Casts object <-> MlirOperation.
template <>
struct type_caster<MlirOperation> {
NB_TYPE_CASTER(MlirOperation, const_name("MlirOperation"))
NB_TYPE_CASTER(MlirOperation,
const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Operation")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToOperation(capsule->ptr());
Expand All @@ -309,7 +316,7 @@ struct type_caster<MlirOperation> {
/// Casts object <-> MlirValue.
template <>
struct type_caster<MlirValue> {
NB_TYPE_CASTER(MlirValue, const_name("MlirValue"))
NB_TYPE_CASTER(MlirValue, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Value")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToValue(capsule->ptr());
Expand All @@ -334,7 +341,8 @@ struct type_caster<MlirValue> {
/// Casts object -> MlirPassManager.
template <>
struct type_caster<MlirPassManager> {
NB_TYPE_CASTER(MlirPassManager, const_name("MlirPassManager"))
NB_TYPE_CASTER(MlirPassManager, const_name(MAKE_MLIR_PYTHON_QUALNAME(
"passmanager.PassManager")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToPassManager(capsule->ptr());
Expand All @@ -347,7 +355,7 @@ struct type_caster<MlirPassManager> {
/// Casts object <-> MlirTypeID.
template <>
struct type_caster<MlirTypeID> {
NB_TYPE_CASTER(MlirTypeID, const_name("MlirTypeID"))
NB_TYPE_CASTER(MlirTypeID, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToTypeID(capsule->ptr());
Expand All @@ -371,7 +379,7 @@ struct type_caster<MlirTypeID> {
/// Casts object <-> MlirType.
template <>
struct type_caster<MlirType> {
NB_TYPE_CASTER(MlirType, const_name("MlirType"))
NB_TYPE_CASTER(MlirType, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Type")))
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
if (auto capsule = mlirApiObjectToCapsule(src)) {
value = mlirPythonCapsuleToType(capsule->ptr());
Expand All @@ -394,7 +402,7 @@ struct type_caster<MlirType> {
/// Casts MlirStringRef -> object.
template <>
struct type_caster<MlirStringRef> {
NB_TYPE_CASTER(MlirStringRef, const_name("MlirStringRef"))
NB_TYPE_CASTER(MlirStringRef, const_name("str"))
static handle from_cpp(MlirStringRef s, rv_policy,
cleanup_list *cleanup) noexcept {
return nanobind::str(s.data, s.length).release();
Expand Down
58 changes: 23 additions & 35 deletions mlir/lib/Bindings/Python/IRCore.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3219,13 +3219,11 @@ void mlir::python::populateIRCore(nb::module_ &m) {
nb::arg("end_line"), nb::arg("end_col"),
nb::arg("context") = nb::none(), kContextGetFileRangeDocstring)
.def("is_a_file", mlirLocationIsAFileLineColRange)
.def_prop_ro(
"filename",
[](MlirLocation loc) {
return mlirIdentifierStr(
mlirLocationFileLineColRangeGetFilename(loc));
},
nb::sig("def filename(self) -> str"))
.def_prop_ro("filename",
[](MlirLocation loc) {
return mlirIdentifierStr(
mlirLocationFileLineColRangeGetFilename(loc));
})
.def_prop_ro("start_line", mlirLocationFileLineColRangeGetStartLine)
.def_prop_ro("start_col", mlirLocationFileLineColRangeGetStartColumn)
.def_prop_ro("end_line", mlirLocationFileLineColRangeGetEndLine)
Expand Down Expand Up @@ -3274,12 +3272,10 @@ void mlir::python::populateIRCore(nb::module_ &m) {
nb::arg("name"), nb::arg("childLoc") = nb::none(),
nb::arg("context") = nb::none(), kContextGetNameLocationDocString)
.def("is_a_name", mlirLocationIsAName)
.def_prop_ro(
"name_str",
[](MlirLocation loc) {
return mlirIdentifierStr(mlirLocationNameGetName(loc));
},
nb::sig("def name_str(self) -> str"))
.def_prop_ro("name_str",
[](MlirLocation loc) {
return mlirIdentifierStr(mlirLocationNameGetName(loc));
})
.def_prop_ro("child_loc",
[](PyLocation &self) {
return PyLocation(self.getContext(),
Expand Down Expand Up @@ -3453,15 +3449,13 @@ void mlir::python::populateIRCore(nb::module_ &m) {
return concreteOperation.getContext().getObject();
},
"Context that owns the Operation")
.def_prop_ro(
"name",
[](PyOperationBase &self) {
auto &concreteOperation = self.getOperation();
concreteOperation.checkValid();
MlirOperation operation = concreteOperation.get();
return mlirIdentifierStr(mlirOperationGetName(operation));
},
nb::sig("def name(self) -> str"))
.def_prop_ro("name",
[](PyOperationBase &self) {
auto &concreteOperation = self.getOperation();
concreteOperation.checkValid();
MlirOperation operation = concreteOperation.get();
return mlirIdentifierStr(mlirOperationGetName(operation));
})
.def_prop_ro("operands",
[](PyOperationBase &self) {
return PyOpOperandList(self.getOperation().getRef());
Expand Down Expand Up @@ -3597,12 +3591,11 @@ void mlir::python::populateIRCore(nb::module_ &m) {
},
"Reports if the operation is attached to its parent block.")
.def("erase", [](PyOperationBase &self) { self.getOperation().erase(); })
.def(
"walk", &PyOperationBase::walk, nb::arg("callback"),
nb::arg("walk_order") = MlirWalkPostOrder,
// clang-format off
nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder = " MAKE_MLIR_PYTHON_QUALNAME("ir.WalkOrder.POST_ORDER") ") -> None")
// clang-format on
.def("walk", &PyOperationBase::walk, nb::arg("callback"),
nb::arg("walk_order") = MlirWalkPostOrder,
// clang-format off
nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder) -> None")
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i removed the default val here because stubgen gets confused (for whatever reason) and tries to do import WalkOrder even though WalkOrder is defined at this point in the stub........... 🤷

// clang-format on
);

nb::class_<PyOperation, PyOperationBase>(m, "Operation")
Expand Down Expand Up @@ -4118,7 +4111,6 @@ void mlir::python::populateIRCore(nb::module_ &m) {
[](PyNamedAttribute &self) {
return mlirIdentifierStr(self.namedAttr.name);
},
nb::sig("def name(self) -> str"),
"The name of the NamedAttribute binding")
.def_prop_ro(
"attr",
Expand Down Expand Up @@ -4336,17 +4328,15 @@ void mlir::python::populateIRCore(nb::module_ &m) {
kValueReplaceAllUsesWithDocstring)
.def(
"replace_all_uses_except",
[](MlirValue self, MlirValue with, PyOperation &exception) {
[](PyValue &self, PyValue &with, PyOperation &exception) {
MlirOperation exceptedUser = exception.get();
mlirValueReplaceAllUsesExcept(self, with, 1, &exceptedUser);
},
nb::arg("with_"), nb::arg("exceptions"),
nb::sig("def replace_all_uses_except(self, with_: Value, exceptions: "
"Operation) -> None"),
kValueReplaceAllUsesExceptDocstring)
.def(
"replace_all_uses_except",
[](MlirValue self, MlirValue with, nb::list exceptions) {
[](PyValue &self, PyValue &with, const nb::list &exceptions) {
// Convert Python list to a SmallVector of MlirOperations
llvm::SmallVector<MlirOperation> exceptionOps;
for (nb::handle exception : exceptions) {
Expand All @@ -4358,8 +4348,6 @@ void mlir::python::populateIRCore(nb::module_ &m) {
exceptionOps.data());
},
nb::arg("with_"), nb::arg("exceptions"),
nb::sig("def replace_all_uses_except(self, with_: Value, exceptions: "
"Sequence[Operation]) -> None"),
kValueReplaceAllUsesExceptDocstring)
.def(
"replace_all_uses_except",
Expand Down
6 changes: 2 additions & 4 deletions mlir/lib/Bindings/Python/IRModule.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -273,8 +273,7 @@ class DefaultingPyMlirContext
: public Defaulting<DefaultingPyMlirContext, PyMlirContext> {
public:
using Defaulting::Defaulting;
static constexpr const char kTypeDescription[] =
MAKE_MLIR_PYTHON_QUALNAME("ir.Context");
static constexpr const char kTypeDescription[] = "Context";
static PyMlirContext &resolve();
};

Expand Down Expand Up @@ -500,8 +499,7 @@ class DefaultingPyLocation
: public Defaulting<DefaultingPyLocation, PyLocation> {
public:
using Defaulting::Defaulting;
static constexpr const char kTypeDescription[] =
MAKE_MLIR_PYTHON_QUALNAME("ir.Location");
static constexpr const char kTypeDescription[] = "Location";
static PyLocation &resolve();

operator MlirLocation() const { return *get(); }
Expand Down
4 changes: 2 additions & 2 deletions mlir/lib/Bindings/Python/IRTypes.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1010,7 +1010,7 @@ class PyTupleType : public PyConcreteType<PyTupleType> {
},
nb::arg("elements"), nb::arg("context") = nb::none(),
// clang-format off
nb::sig("def get_tuple(elements: Sequence[Type], context: mlir.ir.Context | None = None) -> TupleType"),
nb::sig("def get_tuple(elements: Sequence[Type], context: Context | None = None) -> TupleType"),
// clang-format on
"Create a tuple type");
c.def(
Expand Down Expand Up @@ -1070,7 +1070,7 @@ class PyFunctionType : public PyConcreteType<PyFunctionType> {
},
nb::arg("inputs"), nb::arg("results"), nb::arg("context") = nb::none(),
// clang-format off
nb::sig("def get(inputs: Sequence[Type], results: Sequence[Type], context: mlir.ir.Context | None = None) -> FunctionType"),
nb::sig("def get(inputs: Sequence[Type], results: Sequence[Type], context: Context | None = None) -> FunctionType"),
// clang-format on
"Gets a FunctionType from a list of input and result types");
c.def_prop_ro(
Expand Down
6 changes: 0 additions & 6 deletions mlir/lib/Bindings/Python/MainModule.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,9 +115,6 @@ NB_MODULE(_mlir, m) {
});
},
"typeid"_a, nb::kw_only(), "replace"_a = false,
// clang-format off
nb::sig("def register_type_caster(typeid: " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID") ", *, replace: bool = False) -> object"),
// clang-format on
"Register a type caster for casting MLIR types to custom user types.");
m.def(
MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR,
Expand All @@ -130,9 +127,6 @@ NB_MODULE(_mlir, m) {
});
},
"typeid"_a, nb::kw_only(), "replace"_a = false,
// clang-format off
nb::sig("def register_value_caster(typeid: " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID") ", *, replace: bool = False) -> object"),
// clang-format on
"Register a value caster for casting MLIR values to custom user values.");

// Define and populate IR submodule.
Expand Down