[Mlir-commits] [mlir] [MLIR][Python] Make Python-defined dialect loading context-aware (PR #210501)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 23 08:56:06 PDT 2026
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/210501
>From cd6ecf25efeece948522106a74e83d711c973027 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 18 Jul 2026 18:37:27 +0800
Subject: [PATCH 1/3] [MLIR][Python] Make Python-defined dialect loading
context-aware
---
mlir/include/mlir-c/IR.h | 7 ++
mlir/lib/Bindings/Python/IRCore.cpp | 8 +++
mlir/lib/CAPI/IR/IR.cpp | 5 ++
mlir/python/mlir/dialects/ext.py | 29 ++++----
mlir/test/python/dialects/ext.py | 66 +++++++++++++++++++
.../python/dialects/transform_op_interface.py | 2 +-
...ansform_pattern_descriptor_op_interface.py | 2 +-
7 files changed, 100 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 311234e42df60..98ea0dfc00a1b 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -140,6 +140,13 @@ mlirContextGetNumLoadedDialects(MlirContext context);
MLIR_CAPI_EXPORTED MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
MlirStringRef name);
+/// Gets the dialect instance owned by the given context using the dialect
+/// namespace to identify it. If the dialect is not loaded by the context,
+/// returns null. Use mlirContextGetOrLoadDialect to load a dialect if it is
+/// registered with the context.
+MLIR_CAPI_EXPORTED MlirDialect mlirContextGetLoadedDialect(MlirContext context,
+ MlirStringRef name);
+
/// Set threading mode (must be set to false to mlir-print-ir-after-all).
MLIR_CAPI_EXPORTED void mlirContextEnableMultithreading(MlirContext context,
bool enable);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index ab75bc20e123e..9ec6d048dbadf 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3408,6 +3408,14 @@ void populateIRCore(nb::module_ &m) {
},
"dialect_name"_a,
"Gets or loads a dialect by name, returning its descriptor object.")
+ .def(
+ "is_dialect_loaded",
+ [=](PyMlirContext &self, std::string &name) {
+ MlirDialect dialect = mlirContextGetLoadedDialect(
+ self.get(), {name.data(), name.size()});
+ return !mlirDialectIsNull(dialect);
+ },
+ "dialect_name"_a, "Checks if a dialect is loaded in the context.")
.def_prop_rw(
"allow_unregistered_dialects",
[](PyMlirContext &self) -> bool {
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 94442e2be19a4..d5d545a4e1a71 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -97,6 +97,11 @@ MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
return wrap(unwrap(context)->getOrLoadDialect(unwrap(name)));
}
+MlirDialect mlirContextGetLoadedDialect(MlirContext context,
+ MlirStringRef name) {
+ return wrap(unwrap(context)->getLoadedDialect(unwrap(name)));
+}
+
bool mlirContextIsRegisteredOperation(MlirContext context, MlirStringRef name) {
return unwrap(context)->isOperationRegistered(unwrap(name));
}
diff --git a/mlir/python/mlir/dialects/ext.py b/mlir/python/mlir/dialects/ext.py
index 10ac002c75d96..b4a392ff5cb36 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -957,18 +957,11 @@ def _emit_module(cls) -> ir.Module:
return m
@classmethod
- def load(
- cls,
- *,
- reload: bool = False,
- ) -> None:
- if hasattr(cls, "_mlir_module") and not reload:
- if cls._mlir_module.context is not ir.Context.current:
- raise RuntimeError(
- "This dialect was loaded in a different context. "
- "Please set reload=True to reload the dialect in the current context."
- )
- return
+ def load(cls) -> None:
+ if ir.Context.current.is_dialect_loaded(cls.DIALECT_NAMESPACE):
+ raise RuntimeError(
+ f"Dialect '{cls.DIALECT_NAMESPACE}' has already been loaded in the current context."
+ )
cls._mlir_module = cls._emit_module()
pm = PassManager()
@@ -980,19 +973,21 @@ def load(
for op in cls.operations:
op._attach_traits()
- _cext.globals._register_dialect_impl(cls.DIALECT_NAMESPACE, cls, replace=reload)
+ _cext.globals._register_dialect_impl(cls.DIALECT_NAMESPACE, cls, replace=True)
+ # typeids for dynamic types and attributes are context-dependent,
+ # so we need to register them for every MLIR context.
for type_ in cls.types:
typeid = ir.DynamicType.lookup_typeid(type_.type_name)
- _cext.register_type_caster(typeid, replace=reload)(type_)
+ _cext.register_type_caster(typeid, replace=True)(type_)
for attr in cls.attributes:
typeid = ir.DynamicAttr.lookup_typeid(attr.attr_name)
- _cext.register_type_caster(typeid, replace=reload)(attr)
+ _cext.register_type_caster(typeid, replace=True)(attr)
for op in cls.operations:
- _cext.register_operation(cls, replace=reload)(op)
- _cext.register_op_adaptor(op, replace=reload)(op.Adaptor)
+ _cext.register_operation(cls, replace=True)(op)
+ _cext.register_op_adaptor(op, replace=True)(op.Adaptor)
class Pure:
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index 2c9ab4fe321a2..6791cbe7fde65 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -105,6 +105,72 @@ class AddOp(Operation, dialect=MyInt, name="add"):
print(adaptor1.rhs)
+# CHECK: TEST: testDialectLoadInMultipleContexts
+ at run
+def testDialectLoadInMultipleContexts():
+ class ContextLoadDialect(Dialect, name="ext_context_load"):
+ pass
+
+ class ContextLoadType(ContextLoadDialect.Type, name="type"):
+ value: IntegerAttr
+
+ class ContextLoadAttr(ContextLoadDialect.Attribute, name="attr"):
+ value: StringAttr
+
+ class ContextLoadOp(ContextLoadDialect.Operation, name="op"):
+ attr: ContextLoadAttr
+ result: Result[ContextLoadType]
+
+ # CHECK: same context: Dialect 'ext_context_load' has already been loaded in the current context.
+
+ def check_dialect(context_name, type_value):
+ i32 = IntegerType.get_signless(32)
+ result_type = ContextLoadType.get(IntegerAttr.get(i32, type_value))
+ attr = ContextLoadAttr.get(StringAttr.get(context_name))
+
+ module = Module.create()
+ with InsertionPoint(module.body):
+ ContextLoadOp(attr, result_type)
+
+ assert module.operation.verify()
+ module = Module.parse(str(module))
+ op = module.body.operations[0]
+ assert isinstance(op, ContextLoadOp)
+ assert isinstance(op.attr, ContextLoadAttr)
+ assert isinstance(op.result.type, ContextLoadType)
+
+ # CHECK: first context: ContextLoadOp, ContextLoadAttr, ContextLoadType
+ # CHECK: "first context"
+ # CHECK: 1 : i32
+ # CHECK: second context: ContextLoadOp, ContextLoadAttr, ContextLoadType
+ # CHECK: "second context"
+ # CHECK: 2 : i32
+ print(
+ f"{context_name}: {type(op).__name__}, "
+ f"{type(op.attr).__name__}, {type(op.result.type).__name__}"
+ )
+ print(op.attr.value)
+ print(op.result.type.value)
+
+ first_context = Context()
+ second_context = Context()
+
+ with first_context, Location.unknown():
+ ContextLoadDialect.load()
+ try:
+ ContextLoadDialect.load()
+ except RuntimeError as e:
+ print("same context:", e)
+ else:
+ raise AssertionError("loading a dialect twice in one context must fail")
+
+ check_dialect("first context", 1)
+
+ with second_context, Location.unknown():
+ ContextLoadDialect.load()
+ check_dialect("second context", 2)
+
+
# CHECK: TEST: testExtDialect
@run
def testExtDialect():
diff --git a/mlir/test/python/dialects/transform_op_interface.py b/mlir/test/python/dialects/transform_op_interface.py
index b0b416530eccc..811dd0a9149fb 100644
--- a/mlir/test/python/dialects/transform_op_interface.py
+++ b/mlir/test/python/dialects/transform_op_interface.py
@@ -25,7 +25,7 @@ def run(emit_schedule):
with ir.Context() as ctx, ir.Location.unknown():
payload = emit_payload()
- MyTransform.load(reload=True)
+ MyTransform.load()
GetNamedAttributeOp.attach_interface_impls(ctx)
PrintParamOp.attach_interface_impls(ctx)
diff --git a/mlir/test/python/dialects/transform_pattern_descriptor_op_interface.py b/mlir/test/python/dialects/transform_pattern_descriptor_op_interface.py
index 9cd73331cfdea..e194547fdc33a 100644
--- a/mlir/test/python/dialects/transform_pattern_descriptor_op_interface.py
+++ b/mlir/test/python/dialects/transform_pattern_descriptor_op_interface.py
@@ -16,7 +16,7 @@ def run(emit_schedule):
with ir.Context(), ir.Location.unknown():
payload = emit_payload()
- MyPatternDescriptors.load(reload=True)
+ MyPatternDescriptors.load()
# NB: Pattern descriptor ops have their interfaces attached
# in their respective test functions.
>From b96aa343d82a203ecec263e7b291aea0a792c24e Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sun, 19 Jul 2026 14:21:12 +0800
Subject: [PATCH 2/3] address comment
---
mlir/test/python/dialects/ext.py | 24 ++++++++++++++----------
1 file changed, 14 insertions(+), 10 deletions(-)
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index 6791cbe7fde65..707f86aaac521 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -121,8 +121,6 @@ class ContextLoadOp(ContextLoadDialect.Operation, name="op"):
attr: ContextLoadAttr
result: Result[ContextLoadType]
- # CHECK: same context: Dialect 'ext_context_load' has already been loaded in the current context.
-
def check_dialect(context_name, type_value):
i32 = IntegerType.get_signless(32)
result_type = ContextLoadType.get(IntegerAttr.get(i32, type_value))
@@ -139,12 +137,6 @@ def check_dialect(context_name, type_value):
assert isinstance(op.attr, ContextLoadAttr)
assert isinstance(op.result.type, ContextLoadType)
- # CHECK: first context: ContextLoadOp, ContextLoadAttr, ContextLoadType
- # CHECK: "first context"
- # CHECK: 1 : i32
- # CHECK: second context: ContextLoadOp, ContextLoadAttr, ContextLoadType
- # CHECK: "second context"
- # CHECK: 2 : i32
print(
f"{context_name}: {type(op).__name__}, "
f"{type(op.attr).__name__}, {type(op.result.type).__name__}"
@@ -160,16 +152,28 @@ def check_dialect(context_name, type_value):
try:
ContextLoadDialect.load()
except RuntimeError as e:
+ # CHECK: same context: Dialect 'ext_context_load' has already been loaded in the current context.
print("same context:", e)
- else:
- raise AssertionError("loading a dialect twice in one context must fail")
+ # CHECK: first context: ContextLoadOp, ContextLoadAttr, ContextLoadType
+ # CHECK: "first context"
+ # CHECK: 1 : i32
check_dialect("first context", 1)
with second_context, Location.unknown():
ContextLoadDialect.load()
+ # CHECK: second context: ContextLoadOp, ContextLoadAttr, ContextLoadType
+ # CHECK: "second context"
+ # CHECK: 2 : i32
check_dialect("second context", 2)
+ with first_context, Location.unknown():
+ try:
+ ContextLoadDialect.load()
+ except RuntimeError as e:
+ # CHECK: same context again: Dialect 'ext_context_load' has already been loaded in the current context.
+ print("same context again:", e)
+
# CHECK: TEST: testExtDialect
@run
>From 8a15f38b86b5feadb2b2c924eabaf3bd9ea21090 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Thu, 23 Jul 2026 23:55:45 +0800
Subject: [PATCH 3/3] remove useless = in lambda capture list
---
mlir/lib/Bindings/Python/IRCore.cpp | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 9ec6d048dbadf..cceef55f9b15d 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3397,7 +3397,7 @@ void populateIRCore(nb::module_ &m) {
"Alias for `dialects`.")
.def(
"get_dialect_descriptor",
- [=](PyMlirContext &self, std::string &name) {
+ [](PyMlirContext &self, std::string &name) {
MlirDialect dialect = mlirContextGetOrLoadDialect(
self.get(), {name.data(), name.size()});
if (mlirDialectIsNull(dialect)) {
@@ -3410,7 +3410,7 @@ void populateIRCore(nb::module_ &m) {
"Gets or loads a dialect by name, returning its descriptor object.")
.def(
"is_dialect_loaded",
- [=](PyMlirContext &self, std::string &name) {
+ [](PyMlirContext &self, std::string &name) {
MlirDialect dialect = mlirContextGetLoadedDialect(
self.get(), {name.data(), name.size()});
return !mlirDialectIsNull(dialect);
@@ -3545,7 +3545,7 @@ void populateIRCore(nb::module_ &m) {
nb::class_<PyDialects>(m, "Dialects")
.def(
"__getitem__",
- [=](PyDialects &self, std::string keyName) {
+ [](PyDialects &self, std::string keyName) {
MlirDialect dialect =
self.getDialectForKey(keyName, /*attrError=*/false);
nb::object descriptor =
@@ -3555,7 +3555,7 @@ void populateIRCore(nb::module_ &m) {
"Gets a dialect by name using subscript notation.")
.def(
"__getattr__",
- [=](PyDialects &self, std::string attrName) {
+ [](PyDialects &self, std::string attrName) {
MlirDialect dialect =
self.getDialectForKey(attrName, /*attrError=*/true);
nb::object descriptor =
More information about the Mlir-commits
mailing list