[Mlir-commits] [mlir] 98195e0 - [MLIR][Python] Make Python-defined dialect loading context-aware (#210501)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 24 04:34:30 PDT 2026
Author: Twice
Date: 2026-07-24T19:34:24+08:00
New Revision: 98195e0458dd5ec037785a6f861c6402bec778b3
URL: https://github.com/llvm/llvm-project/commit/98195e0458dd5ec037785a6f861c6402bec778b3
DIFF: https://github.com/llvm/llvm-project/commit/98195e0458dd5ec037785a6f861c6402bec778b3.diff
LOG: [MLIR][Python] Make Python-defined dialect loading context-aware (#210501)
Python-defined dialect loading currently relies on
`Dialect._mlir_module` to infer whether a dialect has already been
loaded. This state belongs to the Python dialect class rather than an
MLIR context.
Consequently, loading the same dialect after switching contexts requires
`reload=True`, while reloading it in a context where it is already
present can hit the operation registration assertion reported in
#210053.
This patch adds `mlirContextGetLoadedDialect` (following
https://github.com/llvm/lighthouse/pull/228#discussion_r3589792891) to
the C API and exposes it as `Context.is_dialect_loaded`.
`Dialect.load()` now queries the active context:
- loading a dialect more than once in the same context raises a
`RuntimeError`;
- loading the same Python-defined dialect in another context succeeds
without a `reload` flag.
The `reload` parameter is semantically wrong (mentioned in
https://github.com/llvm/lighthouse/pull/228#discussion_r3586672116) thus
removed from the API as a small breaking change. Usually, users who are
using `reload` can just remove it and everything will work fine.
Fixes #210053.
Assisted by GPT 5.6 Sol (for writing test cases).
Added:
Modified:
mlir/include/mlir-c/IR.h
mlir/lib/Bindings/Python/IRCore.cpp
mlir/lib/CAPI/IR/IR.cpp
mlir/python/mlir/dialects/ext.py
mlir/test/python/dialects/ext.py
mlir/test/python/dialects/transform_op_interface.py
mlir/test/python/dialects/transform_pattern_descriptor_op_interface.py
Removed:
################################################################################
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..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)) {
@@ -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 {
@@ -3537,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 =
@@ -3547,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 =
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..4e8d30d82d5f8 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -29,6 +29,7 @@
__all__ = [
"Dialect",
+ "DialectAlreadyLoadedError",
"Operation",
"Operand",
"Result",
@@ -47,6 +48,10 @@
Region = ir.Region
+class DialectAlreadyLoadedError(RuntimeError):
+ """Raised when a dialect is loaded more than once in the current context."""
+
+
def construct_instance(origin, args):
if not issubclass(origin, ir.Type | ir.Attribute):
raise TypeError(f"unsupported type in constraints: {origin}")
@@ -957,18 +962,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
diff erent 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 DialectAlreadyLoadedError(
+ f"Dialect '{cls.DIALECT_NAMESPACE}' has already been loaded in the current context."
+ )
cls._mlir_module = cls._emit_module()
pm = PassManager()
@@ -980,19 +978,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..dde980df53dba 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -105,6 +105,81 @@ 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]
+
+ 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)
+
+ 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 DialectAlreadyLoadedError as e:
+ assert isinstance(e, RuntimeError)
+ # CHECK: same context: Dialect 'ext_context_load' has already been loaded in the current context.
+ print("same context:", e)
+ else:
+ raise AssertionError("expected DialectAlreadyLoadedError")
+
+ # 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 DialectAlreadyLoadedError as e:
+ # CHECK: same context again: Dialect 'ext_context_load' has already been loaded in the current context.
+ print("same context again:", e)
+ else:
+ raise AssertionError("expected DialectAlreadyLoadedError")
+
+
# 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