[Mlir-commits] [mlir] [MLIR][Python] Add effect and speculatability specifiers for Python-defined ops (PR #216773)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 17 10:08:24 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Twice (PragmaTwice)
<details>
<summary>Changes</summary>
This PR adds standalone effect and speculatability specifiers for Python-defined operations.
`NoMemoryEffect` and `AlwaysSpeculatable`, previously nested under `Pure`, are now public, and `RecursivelySpeculatable` is added. `Pure` remains a shorthand for attaching `NoMemoryEffect` and `AlwaysSpeculatable`.
This also exposes `OpTrait::HasRecursiveMemoryEffects` through the C API and Python bindings as `ir.RecursiveMemoryEffectsTrait`, allowing region-bearing Python-defined operations to derive their memory effects from nested operations:
```python
class LeafOp(
TestDialect.Operation,
name="leaf",
traits=[NoMemoryEffect, AlwaysSpeculatable],
):
pass
class RegionOp(
TestDialect.Operation,
name="region",
traits=[
ir.NoTerminatorTrait,
ir.RecursiveMemoryEffectsTrait,
RecursivelySpeculatable,
],
):
body: Region
```
Tests cover direct interface queries and validate recursive memory effects with CSE and trivial DCE.
Related to #<!-- -->177735 and #<!-- -->195505.
Assisted-by: Codex / GPT-5.6 Sol
---
Full diff: https://github.com/llvm/llvm-project/pull/216773.diff
7 Files Affected:
- (modified) mlir/include/mlir-c/ExtensibleDialect.h (+10)
- (modified) mlir/include/mlir/Bindings/Python/IRCore.h (+7)
- (modified) mlir/lib/Bindings/Python/IRCore.cpp (+23)
- (modified) mlir/lib/CAPI/IR/ExtensibleDialect.cpp (+15)
- (modified) mlir/python/mlir/dialects/ext.py (+37-16)
- (modified) mlir/test/python/dialects/ext.py (+71)
- (renamed) mlir/test/python/ir/memory_effects_op_interface.py (+114)
``````````diff
diff --git a/mlir/include/mlir-c/ExtensibleDialect.h b/mlir/include/mlir-c/ExtensibleDialect.h
index f8dcb3f6800d3..f787ff8dd20c1 100644
--- a/mlir/include/mlir-c/ExtensibleDialect.h
+++ b/mlir/include/mlir-c/ExtensibleDialect.h
@@ -71,6 +71,16 @@ mlirDynamicOpTraitNoTerminatorCreate(void);
/// terminator.
MLIR_CAPI_EXPORTED MlirTypeID mlirDynamicOpTraitNoTerminatorGetTypeID(void);
+/// Get the dynamic op trait that indicates memory effects of an operation
+/// includes the effects of operations nested within its regions.
+MLIR_CAPI_EXPORTED MlirDynamicOpTrait
+mlirDynamicOpTraitRecursiveMemoryEffectsCreate(void);
+
+/// Get the type ID of the dynamic op trait that indicates memory effects of an
+/// operation includes the effects of operations nested within its regions.
+MLIR_CAPI_EXPORTED MlirTypeID
+mlirDynamicOpTraitRecursiveMemoryEffectsGetTypeID(void);
+
/// Destroy the dynamic op trait.
MLIR_CAPI_EXPORTED void
mlirDynamicOpTraitDestroy(MlirDynamicOpTrait dynamicOpTrait);
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 3314e0b2a8fcf..52f69c45cea46 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1997,6 +1997,13 @@ class MLIR_PYTHON_API_EXPORTED IsIsolatedFromAbove : public PyDynamicOpTrait {
static void bind(nanobind::module_ &m);
};
+class MLIR_PYTHON_API_EXPORTED RecursiveMemoryEffects
+ : public PyDynamicOpTrait {
+public:
+ static bool attach(const nanobind::object &opName, PyMlirContext &context);
+ static void bind(nanobind::module_ &m);
+};
+
} // namespace PyDynamicOpTraits
MLIR_PYTHON_API_EXPORTED MlirValue getUniqueResult(MlirOperation operation);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 75cfd2a0a1c0b..348250207df73 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2685,6 +2685,28 @@ void PyDynamicOpTraits::IsIsolatedFromAbove::bind(nb::module_ &m) {
nb::arg("context").none() = nb::none());
}
+bool PyDynamicOpTraits::RecursiveMemoryEffects::attach(const nb::object &opName,
+ PyMlirContext &context) {
+ MlirDynamicOpTrait trait = mlirDynamicOpTraitRecursiveMemoryEffectsCreate();
+ return attachOpTrait(opName, trait, context);
+}
+
+void PyDynamicOpTraits::RecursiveMemoryEffects::bind(nb::module_ &m) {
+ nb::class_<PyDynamicOpTraits::RecursiveMemoryEffects, PyDynamicOpTrait> cls(
+ m, "RecursiveMemoryEffectsTrait");
+ cls.attr(typeIDAttr) =
+ PyTypeID(mlirDynamicOpTraitRecursiveMemoryEffectsGetTypeID());
+ cls.attr("attach") = classmethod(
+ [](const nb::object &cls, const nb::object &opName,
+ DefaultingPyMlirContext context) {
+ return PyDynamicOpTraits::RecursiveMemoryEffects::attach(
+ opName, *context.get());
+ },
+ "Attach RecursiveMemoryEffects trait to the given operation name.",
+ nb::arg("cls"), nb::arg("op_name"),
+ nb::arg("context").none() = nb::none());
+}
+
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
} // namespace python
} // namespace mlir
@@ -5320,6 +5342,7 @@ void populateIRCore(nb::module_ &m) {
PyDynamicOpTraits::IsTerminator::bind(m);
PyDynamicOpTraits::NoTerminator::bind(m);
PyDynamicOpTraits::IsIsolatedFromAbove::bind(m);
+ PyDynamicOpTraits::RecursiveMemoryEffects::bind(m);
// MLIRError exception.
MLIRError::bind(m);
diff --git a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
index 2b4d5af0c9f13..2b4187da9eabf 100644
--- a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
+++ b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
@@ -54,6 +54,21 @@ MlirTypeID mlirDynamicOpTraitNoTerminatorGetTypeID() {
return wrap(DynamicOpTraits::NoTerminator::getStaticTypeID());
}
+namespace mlir::DynamicOpTraits {
+
+class RecursiveMemoryEffects
+ : public DynamicOpTraitImpl<OpTrait::HasRecursiveMemoryEffects> {};
+
+} // namespace mlir::DynamicOpTraits
+
+MlirDynamicOpTrait mlirDynamicOpTraitRecursiveMemoryEffectsCreate(void) {
+ return wrap(new DynamicOpTraits::RecursiveMemoryEffects());
+}
+
+MlirTypeID mlirDynamicOpTraitRecursiveMemoryEffectsGetTypeID(void) {
+ return wrap(DynamicOpTraits::RecursiveMemoryEffects::getStaticTypeID());
+}
+
MlirDynamicOpTrait mlirDynamicOpTraitIsIsolatedFromAboveCreate() {
return wrap(new DynamicOpTraits::IsIsolatedFromAbove());
}
diff --git a/mlir/python/mlir/dialects/ext.py b/mlir/python/mlir/dialects/ext.py
index bd59772d8a9e1..ba321ffc32c32 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -30,17 +30,24 @@
__all__ = [
"Dialect",
"DialectAlreadyLoadedError",
+ # components of dialects
"Operation",
+ "Type",
+ "Attribute",
+ # types for operation fields
"Operand",
"Result",
"Region",
- "Type",
- "Attribute",
- "Pure",
+ # specifiers for operation fields
"result",
- "infer_result",
"operand",
"attribute",
+ "infer_result",
+ # interfaces and traits
+ "Pure",
+ "NoMemoryEffect",
+ "AlwaysSpeculatable",
+ "RecursivelySpeculatable",
]
Operand = ir.Value
@@ -995,20 +1002,34 @@ def load(cls) -> None:
_cext.register_op_adaptor(op, replace=True)(op.Adaptor)
-class Pure:
- """Always speculatable operation that does not touch memory."""
+class NoMemoryEffect(ir.MemoryEffectsOpInterface):
+ """Operation that has no effect on memory."""
+
+ @staticmethod
+ def get_effects(op):
+ return []
+
+
+class AlwaysSpeculatable(ir.ConditionallySpeculatable):
+ """Operation that is always speculatable."""
+
+ @staticmethod
+ def get_speculatability(op):
+ return ir.Speculatability.Speculatable
+
- class NoMemoryEffect(ir.MemoryEffectsOpInterface):
- @staticmethod
- def get_effects(op):
- return []
+class RecursivelySpeculatable(ir.ConditionallySpeculatable):
+ """Operation that is speculatable if all operations in all its regions are speculatable."""
- class AlwaysSpeculatable(ir.ConditionallySpeculatable):
- @staticmethod
- def get_speculatability(op):
- return ir.Speculatability.Speculatable
+ @staticmethod
+ def get_speculatability(op):
+ return ir.Speculatability.RecursivelySpeculatable
+
+
+class Pure:
+ """Always speculatable operation that does not touch memory."""
@staticmethod
def attach(op_name):
- Pure.NoMemoryEffect.attach(op_name)
- Pure.AlwaysSpeculatable.attach(op_name)
+ NoMemoryEffect.attach(op_name)
+ AlwaysSpeculatable.attach(op_name)
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index b98b36ed07e0e..d46060c654a7e 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -1035,6 +1035,67 @@ class PureOp(TestIface.Operation, name="pure"):
print("static spec query error:", e)
+# CHECK: TEST: testExtDialectWithPublicInterfaces
+ at run
+def testExtDialectWithPublicInterfaces():
+ class TestPublicIface(Dialect, name="ext_public_iface"):
+ pass
+
+ class NoEffectOp(
+ TestPublicIface.Operation, name="no_effect", traits=[NoMemoryEffect]
+ ):
+ pass
+
+ class AlwaysSpeculatableOp(
+ TestPublicIface.Operation,
+ name="always_speculatable",
+ traits=[AlwaysSpeculatable],
+ ):
+ pass
+
+ class RecursivelySpeculatableOp(
+ TestPublicIface.Operation,
+ name="recursively_speculatable",
+ traits=[NoTerminatorTrait, RecursivelySpeculatable],
+ ):
+ body: Region
+
+ with Context(), Location.unknown():
+ TestPublicIface.load()
+
+ module = Module.create()
+ with InsertionPoint(module.body):
+ no_effect = NoEffectOp()
+ always_speculatable = AlwaysSpeculatableOp()
+ recursively_speculatable = RecursivelySpeculatableOp()
+ recursively_speculatable.body.blocks.append()
+ with InsertionPoint(recursively_speculatable.body.blocks[0]):
+ AlwaysSpeculatableOp()
+
+ assert module.operation.verify()
+
+ no_effect_iface = ir.MemoryEffectsOpInterface(no_effect)
+ always_speculatable_iface = ir.ConditionallySpeculatable(always_speculatable)
+ recursively_speculatable_iface = ir.ConditionallySpeculatable(
+ recursively_speculatable
+ )
+
+ # CHECK: public no memory effects: 0
+ print("public no memory effects:", len(no_effect_iface.get_effects()))
+ # CHECK: public always speculatable: True
+ print(
+ "public always speculatable:",
+ always_speculatable_iface.getSpeculatability()
+ == ir.Speculatability.Speculatable,
+ )
+ # CHECK: public recursively speculatable: True
+ print(
+ "public recursively speculatable:",
+ recursively_speculatable_iface.getSpeculatability()
+ == ir.Speculatability.RecursivelySpeculatable,
+ )
+
+
# CHECK: TEST: testExtDialectWithPure
@run
def testExtDialectWithPure():
@@ -1079,6 +1140,16 @@ class NoPureOp(TestPure.Operation, name="no_pure"):
# CHECK: }
print(module)
+ pure_memory_iface = ir.MemoryEffectsOpInterface(p1)
+ pure_spec_iface = ir.ConditionallySpeculatable(p1)
+ # CHECK: pure memory effects: 0
+ print("pure memory effects:", len(pure_memory_iface.get_effects()))
+ # CHECK: pure speculatable: True
+ print(
+ "pure speculatable:",
+ pure_spec_iface.getSpeculatability() == ir.Speculatability.Speculatable,
+ )
+
patterns = RewritePatternSet()
apply_patterns_and_fold_greedily(module, patterns.freeze())
# CHECK: module {
diff --git a/mlir/test/python/dialects/memory_effects_op_interface.py b/mlir/test/python/ir/memory_effects_op_interface.py
similarity index 75%
rename from mlir/test/python/dialects/memory_effects_op_interface.py
rename to mlir/test/python/ir/memory_effects_op_interface.py
index d5370ea5358cb..5545ad626ca89 100644
--- a/mlir/test/python/dialects/memory_effects_op_interface.py
+++ b/mlir/test/python/ir/memory_effects_op_interface.py
@@ -160,6 +160,24 @@ class SymbolTargetOp(
pass
+class RegionOp(
+ MemoryEffectsTest.Operation,
+ name="region",
+ traits=[ir.NoTerminatorTrait],
+):
+ result: ext.Result[Any]
+ body: ext.Region
+
+
+class RecursiveRegionOp(
+ MemoryEffectsTest.Operation,
+ name="recursive_region",
+ traits=[ir.NoTerminatorTrait, ir.RecursiveMemoryEffectsTrait],
+):
+ result: ext.Result[Any]
+ body: ext.Region
+
+
def run_pass(source, pipeline):
module = ir.Module.parse(source)
PassManager.parse(pipeline).run(module.operation)
@@ -169,6 +187,13 @@ def run_pass(source, pipeline):
with ir.Context(), ir.Location.unknown():
MemoryEffectsTest.load()
+ # CHECK: recursive memory effects traits: False True
+ print(
+ "recursive memory effects traits:",
+ RegionOp.has_trait(ir.RecursiveMemoryEffectsTrait),
+ RecursiveRegionOp.has_trait(ir.RecursiveMemoryEffectsTrait),
+ )
+
# CHECK: memory effect properties: True True True True
print(
"memory effect properties:",
@@ -288,6 +313,41 @@ def run_pass(source, pipeline):
read_across_write.count('"memory_effects_test.read"'),
)
+ recursive_cse = run_pass(
+ """
+ module {
+ func.func @test() -> (i32, i32, i32, i32) {
+ %0 = "memory_effects_test.region"() ({
+ "memory_effects_test.no_effect"() : () -> ()
+ }) : () -> i32
+ %1 = "memory_effects_test.region"() ({
+ "memory_effects_test.no_effect"() : () -> ()
+ }) : () -> i32
+ %2 = "memory_effects_test.recursive_region"() ({
+ "memory_effects_test.no_effect"() : () -> ()
+ }) : () -> i32
+ %3 = "memory_effects_test.recursive_region"() ({
+ "memory_effects_test.no_effect"() : () -> ()
+ }) : () -> i32
+ return %0, %1, %2, %3 : i32, i32, i32, i32
+ }
+ }
+ """,
+ "builtin.module(func.func(cse))",
+ )
+ # An op without RecursiveMemoryEffects has unknown effects and cannot be
+ # CSE'd. The trait makes the other op's empty nested effects visible.
+ # CHECK: CSE non-recursive region count: 2
+ # CHECK: CSE recursive region count: 1
+ print(
+ "CSE non-recursive region count:",
+ recursive_cse.count('"memory_effects_test.region"'),
+ )
+ print(
+ "CSE recursive region count:",
+ recursive_cse.count('"memory_effects_test.recursive_region"'),
+ )
+
dead_code = run_pass(
"""
module {
@@ -326,6 +386,60 @@ def run_pass(source, pipeline):
dead_code.count('"memory_effects_test.allocate_result"'),
)
+ recursive_read_dce = run_pass(
+ """
+ module {
+ func.func @test() {
+ %0 = "memory_effects_test.region"() ({
+ "memory_effects_test.read_dead"() : () -> ()
+ }) : () -> i32
+ %1 = "memory_effects_test.recursive_region"() ({
+ "memory_effects_test.read_dead"() : () -> ()
+ }) : () -> i32
+ return
+ }
+ }
+ """,
+ "builtin.module(func.func(trivial-dce))",
+ )
+ # The non-recursive op has unknown effects and remains. The recursive op is
+ # removable because all nested effects are reads.
+ # CHECK: DCE non-recursive read region count: 1
+ # CHECK: DCE recursive read region count: 0
+ print(
+ "DCE non-recursive read region count:",
+ recursive_read_dce.count('"memory_effects_test.region"'),
+ )
+ print(
+ "DCE recursive read region count:",
+ recursive_read_dce.count('"memory_effects_test.recursive_region"'),
+ )
+
+ recursive_write_dce = run_pass(
+ """
+ module {
+ func.func @test() {
+ %0 = "memory_effects_test.recursive_region"() ({
+ "memory_effects_test.write_dead"() : () -> ()
+ }) : () -> i32
+ return
+ }
+ }
+ """,
+ "builtin.module(func.func(trivial-dce))",
+ )
+ # A nested Write remains observable through RecursiveMemoryEffects.
+ # CHECK: DCE recursive write region count: 1
+ # CHECK: DCE nested write count: 1
+ print(
+ "DCE recursive write region count:",
+ recursive_write_dce.count('"memory_effects_test.recursive_region"'),
+ )
+ print(
+ "DCE nested write count:",
+ recursive_write_dce.count('"memory_effects_test.write_dead"'),
+ )
+
target_variants = run_pass(
"""
module {
``````````
</details>
https://github.com/llvm/llvm-project/pull/216773
More information about the Mlir-commits
mailing list