[Mlir-commits] [mlir] [MLIR][Python] Add effect and speculatability specifiers for Python-defined ops (PR #216773)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 22 08:47:50 PDT 2026
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/216773
>From ee3f1af865c8ceefada8040041443e04e4f9b5f0 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Mon, 17 Aug 2026 23:33:33 +0800
Subject: [PATCH 1/5] [MLIR][Python] Add effect and speculatability specifiers
for Python-defined ops
---
mlir/include/mlir-c/ExtensibleDialect.h | 10 ++
mlir/include/mlir/Bindings/Python/IRCore.h | 6 +
mlir/lib/Bindings/Python/IRCore.cpp | 23 ++++
mlir/lib/CAPI/IR/ExtensibleDialect.cpp | 15 +++
mlir/python/mlir/dialects/ext.py | 53 +++++---
mlir/test/python/dialects/ext.py | 74 ++++++++++++
.../memory_effects_op_interface.py | 114 ++++++++++++++++++
7 files changed, 279 insertions(+), 16 deletions(-)
rename mlir/test/python/{dialects => ir}/memory_effects_op_interface.py (75%)
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..7bc310b8db2e6 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1997,6 +1997,12 @@ 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..93db76e12c0e0 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -1035,6 +1035,69 @@ 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 +1142,17 @@ 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 {
>From 89fe9cfd0bb36e3f6f2f337dbdf594e155322a7b Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Tue, 18 Aug 2026 01:04:30 +0800
Subject: [PATCH 2/5] format
---
mlir/include/mlir/Bindings/Python/IRCore.h | 3 ++-
mlir/test/python/dialects/ext.py | 7 ++-----
2 files changed, 4 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 7bc310b8db2e6..52f69c45cea46 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1997,7 +1997,8 @@ class MLIR_PYTHON_API_EXPORTED IsIsolatedFromAbove : public PyDynamicOpTrait {
static void bind(nanobind::module_ &m);
};
-class MLIR_PYTHON_API_EXPORTED RecursiveMemoryEffects : public PyDynamicOpTrait {
+class MLIR_PYTHON_API_EXPORTED RecursiveMemoryEffects
+ : public PyDynamicOpTrait {
public:
static bool attach(const nanobind::object &opName, PyMlirContext &context);
static void bind(nanobind::module_ &m);
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index 93db76e12c0e0..d46060c654a7e 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -1075,9 +1075,7 @@ class RecursivelySpeculatableOp(
assert module.operation.verify()
no_effect_iface = ir.MemoryEffectsOpInterface(no_effect)
- always_speculatable_iface = ir.ConditionallySpeculatable(
- always_speculatable
- )
+ always_speculatable_iface = ir.ConditionallySpeculatable(always_speculatable)
recursively_speculatable_iface = ir.ConditionallySpeculatable(
recursively_speculatable
)
@@ -1149,8 +1147,7 @@ class NoPureOp(TestPure.Operation, name="no_pure"):
# CHECK: pure speculatable: True
print(
"pure speculatable:",
- pure_spec_iface.getSpeculatability()
- == ir.Speculatability.Speculatable,
+ pure_spec_iface.getSpeculatability() == ir.Speculatability.Speculatable,
)
patterns = RewritePatternSet()
>From e71593f44de50a31d4b04a9a510d00c970102e8c Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 22 Aug 2026 15:02:11 +0800
Subject: [PATCH 3/5] address comments
---
mlir/lib/Bindings/Python/IRCore.cpp | 40 ++++++++++++++++++++---------
1 file changed, 28 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 348250207df73..62ac6b974aeda 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2614,8 +2614,10 @@ bool PyDynamicOpTrait::attach(const nb::object &opName,
void PyDynamicOpTrait::bind(nb::module_ &m) {
nb::class_<PyDynamicOpTrait> cls(m, "DynamicOpTrait");
cls.attr("attach") = classmethod(
- [](const nb::object &cls, const nb::object &opName, nb::object target,
- DefaultingPyMlirContext context) {
+ [](const nb::object &cls,
+ const nb::typed<nb::object, std::variant<nb::type_object, nb::str>>
+ &opName,
+ nb::object target, DefaultingPyMlirContext context) {
if (target.is_none())
target = cls;
return PyDynamicOpTrait::attach(opName, target, *context.get());
@@ -2634,9 +2636,13 @@ bool PyDynamicOpTraits::IsTerminator::attach(const nb::object &opName,
void PyDynamicOpTraits::IsTerminator::bind(nb::module_ &m) {
nb::class_<PyDynamicOpTraits::IsTerminator, PyDynamicOpTrait> cls(
m, "IsTerminatorTrait");
- cls.attr(typeIDAttr) = PyTypeID(mlirDynamicOpTraitIsTerminatorGetTypeID());
+ cls.def_prop_ro_static(typeIDAttr, [](nanobind::object & /*class*/) {
+ return PyTypeID(mlirDynamicOpTraitIsTerminatorGetTypeID());
+ });
cls.attr("attach") = classmethod(
- [](const nb::object &cls, const nb::object &opName,
+ [](const nb::object &cls,
+ const nb::typed<nb::object, std::variant<nb::type_object, nb::str>>
+ &opName,
DefaultingPyMlirContext context) {
return PyDynamicOpTraits::IsTerminator::attach(opName, *context.get());
},
@@ -2653,9 +2659,13 @@ bool PyDynamicOpTraits::NoTerminator::attach(const nb::object &opName,
void PyDynamicOpTraits::NoTerminator::bind(nb::module_ &m) {
nb::class_<PyDynamicOpTraits::NoTerminator, PyDynamicOpTrait> cls(
m, "NoTerminatorTrait");
- cls.attr(typeIDAttr) = PyTypeID(mlirDynamicOpTraitNoTerminatorGetTypeID());
+ cls.def_prop_ro_static(typeIDAttr, [](nanobind::object & /*class*/) {
+ return PyTypeID(mlirDynamicOpTraitNoTerminatorGetTypeID());
+ });
cls.attr("attach") = classmethod(
- [](const nb::object &cls, const nb::object &opName,
+ [](const nb::object &cls,
+ const nb::typed<nb::object, std::variant<nb::type_object, nb::str>>
+ &opName,
DefaultingPyMlirContext context) {
return PyDynamicOpTraits::NoTerminator::attach(opName, *context.get());
},
@@ -2672,10 +2682,13 @@ bool PyDynamicOpTraits::IsIsolatedFromAbove::attach(const nb::object &opName,
void PyDynamicOpTraits::IsIsolatedFromAbove::bind(nb::module_ &m) {
nb::class_<PyDynamicOpTraits::IsIsolatedFromAbove, PyDynamicOpTrait> cls(
m, "IsIsolatedFromAboveTrait");
- cls.attr(typeIDAttr) =
- PyTypeID(mlirDynamicOpTraitIsIsolatedFromAboveGetTypeID());
+ cls.def_prop_ro_static(typeIDAttr, [](nanobind::object & /*class*/) {
+ return PyTypeID(mlirDynamicOpTraitIsIsolatedFromAboveGetTypeID());
+ });
cls.attr("attach") = classmethod(
- [](const nb::object &cls, const nb::object &opName,
+ [](const nb::object &cls,
+ const nb::typed<nb::object, std::variant<nb::type_object, nb::str>>
+ &opName,
DefaultingPyMlirContext context) {
return PyDynamicOpTraits::IsIsolatedFromAbove::attach(opName,
*context.get());
@@ -2694,10 +2707,13 @@ bool PyDynamicOpTraits::RecursiveMemoryEffects::attach(const nb::object &opName,
void PyDynamicOpTraits::RecursiveMemoryEffects::bind(nb::module_ &m) {
nb::class_<PyDynamicOpTraits::RecursiveMemoryEffects, PyDynamicOpTrait> cls(
m, "RecursiveMemoryEffectsTrait");
- cls.attr(typeIDAttr) =
- PyTypeID(mlirDynamicOpTraitRecursiveMemoryEffectsGetTypeID());
+ cls.def_prop_ro_static(typeIDAttr, [](nanobind::object & /*class*/) {
+ return PyTypeID(mlirDynamicOpTraitRecursiveMemoryEffectsGetTypeID());
+ });
cls.attr("attach") = classmethod(
- [](const nb::object &cls, const nb::object &opName,
+ [](const nb::object &cls,
+ const nb::typed<nb::object, std::variant<nb::type_object, nb::str>>
+ &opName,
DefaultingPyMlirContext context) {
return PyDynamicOpTraits::RecursiveMemoryEffects::attach(
opName, *context.get());
>From 2de8447ebbaee583a355134a7f11a168de9958a9 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 22 Aug 2026 15:10:10 +0800
Subject: [PATCH 4/5] add more tests
---
mlir/test/python/ir/memory_effects_op_interface.py | 6 +++++-
1 file changed, 5 insertions(+), 1 deletion(-)
diff --git a/mlir/test/python/ir/memory_effects_op_interface.py b/mlir/test/python/ir/memory_effects_op_interface.py
index 5545ad626ca89..2211649899e99 100644
--- a/mlir/test/python/ir/memory_effects_op_interface.py
+++ b/mlir/test/python/ir/memory_effects_op_interface.py
@@ -187,11 +187,15 @@ def run_pass(source, pipeline):
with ir.Context(), ir.Location.unknown():
MemoryEffectsTest.load()
- # CHECK: recursive memory effects traits: False True
+ from mlir.dialects import scf, arith
+
+ # CHECK: recursive memory effects traits: False True True False
print(
"recursive memory effects traits:",
RegionOp.has_trait(ir.RecursiveMemoryEffectsTrait),
RecursiveRegionOp.has_trait(ir.RecursiveMemoryEffectsTrait),
+ scf.IfOp.has_trait(ir.RecursiveMemoryEffectsTrait),
+ arith.AddIOp.has_trait(ir.RecursiveMemoryEffectsTrait),
)
# CHECK: memory effect properties: True True True True
>From 5f3ad3ea25393215ce6fa69fe37b6d888f58a1ce Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 22 Aug 2026 23:47:24 +0800
Subject: [PATCH 5/5] [MLIR][Python] Refactor memory effects op interface tests
---
.../python/ir/memory_effects_op_interface.py | 329 +++++++++---------
1 file changed, 174 insertions(+), 155 deletions(-)
diff --git a/mlir/test/python/ir/memory_effects_op_interface.py b/mlir/test/python/ir/memory_effects_op_interface.py
index 2211649899e99..4d979c61d34bc 100644
--- a/mlir/test/python/ir/memory_effects_op_interface.py
+++ b/mlir/test/python/ir/memory_effects_op_interface.py
@@ -1,9 +1,10 @@
# RUN: env PYTHONUNBUFFERED=1 %PYTHON %s 2>&1 | FileCheck %s
+from contextlib import contextmanager
from typing import Any
from mlir import ir
-from mlir.dialects import ext, func
+from mlir.dialects import arith, ext, func, scf
from mlir.passmanager import PassManager
@@ -178,17 +179,33 @@ class RecursiveRegionOp(
body: ext.Region
-def run_pass(source, pipeline):
- module = ir.Module.parse(source)
+def run(test):
+ print("\nTEST:", test.__name__)
+ with ir.Context(), ir.Location.unknown():
+ MemoryEffectsTest.load()
+ test()
+
+
+ at contextmanager
+def function_module(inputs=(), results=()):
+ module = ir.Module.create()
+ with ir.InsertionPoint(module.body):
+ function = func.FuncOp("test", (inputs, results))
+ with ir.InsertionPoint(function.add_entry_block()):
+ yield module, function
+
+
+def run_pass(module, pipeline):
PassManager.parse(pipeline).run(module.operation)
- return str(module)
-with ir.Context(), ir.Location.unknown():
- MemoryEffectsTest.load()
+def count_ops(module, op_type):
+ return len(ir.get_ops_of_type(module, op_type))
- from mlir.dialects import scf, arith
+# CHECK-LABEL: TEST: testRecursiveMemoryEffectsTraits
+ at run
+def testRecursiveMemoryEffectsTraits():
# CHECK: recursive memory effects traits: False True True False
print(
"recursive memory effects traits:",
@@ -198,6 +215,10 @@ def run_pass(source, pipeline):
arith.AddIOp.has_trait(ir.RecursiveMemoryEffectsTrait),
)
+
+# CHECK-LABEL: TEST: testMemoryEffectProperties
+ at run
+def testMemoryEffectProperties():
# CHECK: memory effect properties: True True True True
print(
"memory effect properties:",
@@ -222,17 +243,15 @@ def run_pass(source, pipeline):
isinstance(ir.SideEffectResource.Default, ir.SideEffectResource),
)
- query_module = ir.Module.parse(
- """
- module {
- func.func @test(%arg0: i32) -> i32 {
- %0 = "memory_effects_test.read"(%arg0) : (i32) -> i32
- return %0 : i32
- }
- }
- """
- )
- read_op = query_module.body.operations[0].regions[0].blocks[0].operations[0]
+
+# CHECK-LABEL: TEST: testQueryMemoryEffects
+ at run
+def testQueryMemoryEffects():
+ i32 = ir.IntegerType.get_signless(32)
+ with function_module(inputs=[i32], results=[i32]) as (_, function):
+ read_op = ReadOp(function.arguments[0], i32)
+ func.ReturnOp([read_op.result])
+
read_effects = ir.MemoryEffectsOpInterface(read_op).get_effects()
read_effect = read_effects[0]
# CHECK: queried effects: True 1 True True 1 True True True
@@ -254,6 +273,10 @@ def run_pass(source, pipeline):
read_effect.symbol_ref is None,
)
+
+# CHECK-LABEL: TEST: testSymbolEffectProperties
+ at run
+def testSymbolEffectProperties():
symbol_effect = ir.MemoryEffectInstance(
ir.MemoryEffect.Read, ir.FlatSymbolRefAttr.get("global")
)
@@ -265,109 +288,105 @@ def run_pass(source, pipeline):
symbol_effect.parameters is None,
)
- read_cse = run_pass(
- """
- module {
- func.func @test(%arg0: i32) -> (i32, i32) {
- %0 = "memory_effects_test.read"(%arg0) : (i32) -> i32
- %1 = "memory_effects_test.read"(%arg0) : (i32) -> i32
- return %0, %1 : i32, i32
- }
- }
- """,
- "builtin.module(func.func(cse))",
- )
+
+# CHECK-LABEL: TEST: testMemoryEffectsCSE
+ at run
+def testMemoryEffectsCSE():
+ i32 = ir.IntegerType.get_signless(32)
+
+ with function_module(inputs=[i32], results=[i32, i32]) as (
+ read_cse,
+ function,
+ ):
+ read0 = ReadOp(function.arguments[0], i32)
+ read1 = ReadOp(function.arguments[0], i32)
+ func.ReturnOp([read0.result, read1.result])
+ run_pass(read_cse, "builtin.module(func.func(cse))")
+
# A single Read effect remains CSE-eligible.
# CHECK: CSE read count: 1
- print("CSE read count:", read_cse.count('"memory_effects_test.read"'))
-
- write_cse = run_pass(
- """
- module {
- func.func @test(%arg0: i32) -> (i32, i32) {
- %0 = "memory_effects_test.write"(%arg0) : (i32) -> i32
- %1 = "memory_effects_test.write"(%arg0) : (i32) -> i32
- return %0, %1 : i32, i32
- }
- }
- """,
- "builtin.module(func.func(cse))",
- )
+ print("CSE read count:", count_ops(read_cse, ReadOp))
+
+ with function_module(inputs=[i32], results=[i32, i32]) as (
+ write_cse,
+ function,
+ ):
+ write0 = WriteOp(function.arguments[0], i32)
+ write1 = WriteOp(function.arguments[0], i32)
+ func.ReturnOp([write0.result, write1.result])
+ run_pass(write_cse, "builtin.module(func.func(cse))")
+
# Writes cannot be CSE'd.
# CHECK: CSE write count: 2
- print("CSE write count:", write_cse.count('"memory_effects_test.write"'))
-
- read_across_write = run_pass(
- """
- module {
- func.func @test(%arg0: i32) -> (i32, i32) {
- %0 = "memory_effects_test.read"(%arg0) : (i32) -> i32
- "memory_effects_test.write_barrier"(%arg0) : (i32) -> ()
- %1 = "memory_effects_test.read"(%arg0) : (i32) -> i32
- return %0, %1 : i32, i32
- }
- }
- """,
- "builtin.module(func.func(cse))",
- )
+ print("CSE write count:", count_ops(write_cse, WriteOp))
+
+ with function_module(inputs=[i32], results=[i32, i32]) as (
+ read_across_write,
+ function,
+ ):
+ read0 = ReadOp(function.arguments[0], i32)
+ WriteBarrierOp(function.arguments[0])
+ read1 = ReadOp(function.arguments[0], i32)
+ func.ReturnOp([read0.result, read1.result])
+ run_pass(read_across_write, "builtin.module(func.func(cse))")
+
# A potentially-aliasing Write on the default resource blocks Read CSE.
# CHECK: CSE read across write count: 2
print(
"CSE read across write count:",
- read_across_write.count('"memory_effects_test.read"'),
+ count_ops(read_across_write, ReadOp),
)
- 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))",
- )
+
+# CHECK-LABEL: TEST: testRecursiveMemoryEffectsCSE
+ at run
+def testRecursiveMemoryEffectsCSE():
+ i32 = ir.IntegerType.get_signless(32)
+ with function_module(results=[i32, i32, i32, i32]) as (
+ recursive_cse,
+ _,
+ ):
+ region_ops = [
+ RegionOp(i32),
+ RegionOp(i32),
+ RecursiveRegionOp(i32),
+ RecursiveRegionOp(i32),
+ ]
+ for region_op in region_ops:
+ region_op.body.blocks.append()
+ with ir.InsertionPoint(region_op.body.blocks[0]):
+ NoEffectOp()
+ func.ReturnOp([region_op.result for region_op in region_ops])
+ run_pass(recursive_cse, "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"'),
+ count_ops(recursive_cse, RegionOp),
)
print(
"CSE recursive region count:",
- recursive_cse.count('"memory_effects_test.recursive_region"'),
+ count_ops(recursive_cse, RecursiveRegionOp),
)
- dead_code = run_pass(
- """
- module {
- func.func @test() {
- "memory_effects_test.no_effect"() : () -> ()
- "memory_effects_test.read_dead"() : () -> ()
- "memory_effects_test.write_dead"() : () -> ()
- "memory_effects_test.free_dead"() : () -> ()
- "memory_effects_test.allocate_dead"() : () -> ()
- %0 = "memory_effects_test.allocate_result"() : () -> i32
- return
- }
- }
- """,
- "builtin.module(func.func(trivial-dce))",
- )
+
+# CHECK-LABEL: TEST: testMemoryEffectsDCE
+ at run
+def testMemoryEffectsDCE():
+ i32 = ir.IntegerType.get_signless(32)
+ with function_module() as (dead_code, _):
+ NoEffectOp()
+ ReadDeadOp()
+ WriteDeadOp()
+ FreeDeadOp()
+ AllocateDeadOp()
+ AllocateResultOp(i32)
+ func.ReturnOp([])
+ run_pass(dead_code, "builtin.module(func.func(trivial-dce))")
+
# Empty and Read-only effect lists are dead. Write, Free and untargeted
# Allocate effects are observable. An Allocate targeting its own unused
# result is dead.
@@ -377,87 +396,87 @@ def run_pass(source, pipeline):
# CHECK: DCE free count: 1
# CHECK: DCE untargeted allocate count: 1
# CHECK: DCE result allocate count: 0
- print("DCE no effect count:", dead_code.count('"memory_effects_test.no_effect"'))
- print("DCE read count:", dead_code.count('"memory_effects_test.read_dead"'))
- print("DCE write count:", dead_code.count('"memory_effects_test.write_dead"'))
- print("DCE free count:", dead_code.count('"memory_effects_test.free_dead"'))
+ print("DCE no effect count:", count_ops(dead_code, NoEffectOp))
+ print("DCE read count:", count_ops(dead_code, ReadDeadOp))
+ print("DCE write count:", count_ops(dead_code, WriteDeadOp))
+ print("DCE free count:", count_ops(dead_code, FreeDeadOp))
print(
"DCE untargeted allocate count:",
- dead_code.count('"memory_effects_test.allocate_dead"'),
+ count_ops(dead_code, AllocateDeadOp),
)
print(
"DCE result allocate count:",
- dead_code.count('"memory_effects_test.allocate_result"'),
+ count_ops(dead_code, AllocateResultOp),
)
- 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))",
- )
+
+# CHECK-LABEL: TEST: testRecursiveReadDCE
+ at run
+def testRecursiveReadDCE():
+ i32 = ir.IntegerType.get_signless(32)
+ with function_module() as (recursive_read_dce, _):
+ region_op = RegionOp(i32)
+ region_op.body.blocks.append()
+ with ir.InsertionPoint(region_op.body.blocks[0]):
+ ReadDeadOp()
+
+ recursive_region_op = RecursiveRegionOp(i32)
+ recursive_region_op.body.blocks.append()
+ with ir.InsertionPoint(recursive_region_op.body.blocks[0]):
+ ReadDeadOp()
+ func.ReturnOp([])
+ run_pass(recursive_read_dce, "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"'),
+ count_ops(recursive_read_dce, RegionOp),
)
print(
"DCE recursive read region count:",
- recursive_read_dce.count('"memory_effects_test.recursive_region"'),
+ count_ops(recursive_read_dce, RecursiveRegionOp),
)
- 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))",
- )
+
+# CHECK-LABEL: TEST: testRecursiveWriteDCE
+ at run
+def testRecursiveWriteDCE():
+ i32 = ir.IntegerType.get_signless(32)
+ with function_module() as (recursive_write_dce, _):
+ recursive_region_op = RecursiveRegionOp(i32)
+ recursive_region_op.body.blocks.append()
+ with ir.InsertionPoint(recursive_region_op.body.blocks[0]):
+ WriteDeadOp()
+ func.ReturnOp([])
+ run_pass(recursive_write_dce, "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"'),
+ count_ops(recursive_write_dce, RecursiveRegionOp),
)
print(
"DCE nested write count:",
- recursive_write_dce.count('"memory_effects_test.write_dead"'),
+ count_ops(recursive_write_dce, WriteDeadOp),
)
- target_variants = run_pass(
- """
- module {
- func.func @test() {
- "memory_effects_test.block_argument_target"() ({
- ^bb0(%arg0: i32):
- }) : () -> ()
- "memory_effects_test.symbol_target"() : () -> ()
- return
- }
- }
- """,
- "builtin.module(func.func(trivial-dce))",
- )
+
+# CHECK-LABEL: TEST: testMemoryEffectTargets
+ at run
+def testMemoryEffectTargets():
+ i32 = ir.IntegerType.get_signless(32)
+ with function_module() as (target_variants, _):
+ block_argument_target = BlockArgumentTargetOp()
+ block_argument_target.body.blocks.append(i32)
+ SymbolTargetOp()
+ func.ReturnOp([])
+ run_pass(target_variants, "builtin.module(func.func(trivial-dce))")
+
# These Read effects exercise BlockArgument and SymbolRefAttr targets and
# remain removable by trivial-dce.
# CHECK: invalid symbol target: target Attribute must be a SymbolRefAttr
@@ -467,9 +486,9 @@ def run_pass(source, pipeline):
# CHECK: DCE symbol target count: 0
print(
"DCE block argument target count:",
- target_variants.count('"memory_effects_test.block_argument_target"'),
+ count_ops(target_variants, BlockArgumentTargetOp),
)
print(
"DCE symbol target count:",
- target_variants.count('"memory_effects_test.symbol_target"'),
+ count_ops(target_variants, SymbolTargetOp),
)
More information about the Mlir-commits
mailing list