[Mlir-commits] [mlir] a4402bd - [MLIR][Python] Add effect and speculatability specifiers for Python-defined ops (#216773)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 22 09:27:42 PDT 2026
Author: Twice
Date: 2026-08-23T00:27:37+08:00
New Revision: a4402bd0cc5fc5f187e264217a349f853a40ec50
URL: https://github.com/llvm/llvm-project/commit/a4402bd0cc5fc5f187e264217a349f853a40ec50
DIFF: https://github.com/llvm/llvm-project/commit/a4402bd0cc5fc5f187e264217a349f853a40ec50.diff
LOG: [MLIR][Python] Add effect and speculatability specifiers for Python-defined ops (#216773)
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
Added:
mlir/test/python/ir/memory_effects_op_interface.py
Modified:
mlir/include/mlir-c/ExtensibleDialect.h
mlir/include/mlir/Bindings/Python/IRCore.h
mlir/lib/Bindings/Python/IRCore.cpp
mlir/lib/CAPI/IR/ExtensibleDialect.cpp
mlir/python/mlir/dialects/ext.py
mlir/test/python/dialects/ext.py
Removed:
mlir/test/python/dialects/memory_effects_op_interface.py
################################################################################
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 05c7d63ccf766..b28e6ec606d52 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -2007,6 +2007,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 a19a2c8b1c7f6..0adad4a349478 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());
@@ -2685,6 +2698,31 @@ 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.def_prop_ro_static(typeIDAttr, [](nanobind::object & /*class*/) {
+ return PyTypeID(mlirDynamicOpTraitRecursiveMemoryEffectsGetTypeID());
+ });
+ cls.attr("attach") = classmethod(
+ [](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());
+ },
+ "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
@@ -5382,6 +5420,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 50%
rename from mlir/test/python/dialects/memory_effects_op_interface.py
rename to mlir/test/python/ir/memory_effects_op_interface.py
index d5370ea5358cb..4d979c61d34bc 100644
--- a/mlir/test/python/dialects/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
@@ -160,15 +161,64 @@ class SymbolTargetOp(
pass
-def run_pass(source, pipeline):
- module = ir.Module.parse(source)
+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(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))
+
+
+# CHECK-LABEL: TEST: testRecursiveMemoryEffectsTraits
+ at run
+def testRecursiveMemoryEffectsTraits():
+ # 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-LABEL: TEST: testMemoryEffectProperties
+ at run
+def testMemoryEffectProperties():
# CHECK: memory effect properties: True True True True
print(
"memory effect properties:",
@@ -193,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
@@ -225,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")
)
@@ -236,74 +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),
)
- 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: 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:",
+ count_ops(recursive_cse, RegionOp),
+ )
+ print(
+ "CSE recursive region count:",
+ count_ops(recursive_cse, RecursiveRegionOp),
)
+
+
+# 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.
@@ -313,33 +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),
)
- 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: 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:",
+ count_ops(recursive_read_dce, RegionOp),
+ )
+ print(
+ "DCE recursive read region count:",
+ count_ops(recursive_read_dce, RecursiveRegionOp),
)
+
+
+# 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:",
+ count_ops(recursive_write_dce, RecursiveRegionOp),
+ )
+ print(
+ "DCE nested write count:",
+ count_ops(recursive_write_dce, WriteDeadOp),
+ )
+
+
+# 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
@@ -349,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