[Mlir-commits] [mlir] [MLIR][Python] Make Python-defined dialect loading context-aware (PR #210501)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jul 18 03:59:57 PDT 2026
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/210501
>From 17d00d79d53674da4c0d54188ad2afe70dc60bf0 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 18 Jul 2026 18:37:27 +0800
Subject: [PATCH] [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..3ea214da40531 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.
More information about the Mlir-commits
mailing list