[Mlir-commits] [mlir] [MLIR][CAPI][Python] Add support for constructing memory effect instances (PR #210586)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 23 08:48:10 PDT 2026
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/210586
>From ee4b1d4aa4c3c0afaf896ace40531d702551cb77 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sun, 19 Jul 2026 18:07:37 +0800
Subject: [PATCH 1/4] [MLIR][CAPI][Python] Add support for constructing memory
effect instances
---
mlir/include/mlir-c/Interfaces.h | 78 +++++
.../mlir/Bindings/Python/IRInterfaces.h | 53 +++-
mlir/include/mlir/CAPI/Interfaces.h | 4 +
mlir/lib/Bindings/Python/DialectTransform.cpp | 20 +-
mlir/lib/Bindings/Python/IRInterfaces.cpp | 107 ++++++-
mlir/lib/CAPI/Interfaces/Interfaces.cpp | 72 +++++
mlir/test/CAPI/ir.c | 43 +++
.../dialects/memory_effects_op_interface.py | 291 ++++++++++++++++++
8 files changed, 654 insertions(+), 14 deletions(-)
create mode 100644 mlir/test/python/dialects/memory_effects_op_interface.py
diff --git a/mlir/include/mlir-c/Interfaces.h b/mlir/include/mlir-c/Interfaces.h
index a5251dc78471f..a416a6ab76f87 100644
--- a/mlir/include/mlir-c/Interfaces.h
+++ b/mlir/include/mlir-c/Interfaces.h
@@ -28,7 +28,10 @@ extern "C" {
}; \
typedef struct name name
+DEFINE_C_API_STRUCT(MlirMemoryEffect, void);
+DEFINE_C_API_STRUCT(MlirMemoryEffectInstance, void);
DEFINE_C_API_STRUCT(MlirMemoryEffectInstancesList, void);
+DEFINE_C_API_STRUCT(MlirSideEffectResource, void);
#undef DEFINE_C_API_STRUCT
@@ -144,6 +147,81 @@ mlirConditionallySpeculatableOpInterfaceGetSpeculatability(
// MemoryEffectsOpInterface
//===---------------------------------------------------------------------===//
+/// Returns the borrowed singleton instance of the allocate memory effect.
+MLIR_CAPI_EXPORTED MlirMemoryEffect mlirMemoryEffectsAllocateGet(void);
+
+/// Returns the borrowed singleton instance of the free memory effect.
+MLIR_CAPI_EXPORTED MlirMemoryEffect mlirMemoryEffectsFreeGet(void);
+
+/// Returns the borrowed singleton instance of the read memory effect.
+MLIR_CAPI_EXPORTED MlirMemoryEffect mlirMemoryEffectsReadGet(void);
+
+/// Returns the borrowed singleton instance of the write memory effect.
+MLIR_CAPI_EXPORTED MlirMemoryEffect mlirMemoryEffectsWriteGet(void);
+
+/// Returns the borrowed singleton instance of the default side effect
+/// resource.
+MLIR_CAPI_EXPORTED MlirSideEffectResource
+mlirSideEffectsDefaultResourceGet(void);
+
+/// Creates a memory effect instance without an associated IR entity.
+/// `parameters` may be a null attribute. The caller owns the returned instance
+/// and must destroy it with `mlirMemoryEffectInstanceDestroy`.
+MLIR_CAPI_EXPORTED MlirMemoryEffectInstance mlirMemoryEffectInstanceCreate(
+ MlirMemoryEffect effect, MlirAttribute parameters, int stage,
+ bool effectOnFullRegion, MlirSideEffectResource resource);
+
+/// Creates a memory effect instance associated with an operation operand.
+/// `parameters` may be a null attribute. The caller owns the returned instance
+/// and must destroy it with `mlirMemoryEffectInstanceDestroy`.
+MLIR_CAPI_EXPORTED MlirMemoryEffectInstance
+mlirMemoryEffectInstanceCreateForOpOperand(MlirMemoryEffect effect,
+ MlirOpOperand opOperand,
+ MlirAttribute parameters, int stage,
+ bool effectOnFullRegion,
+ MlirSideEffectResource resource);
+
+/// Creates a memory effect instance associated with an operation result.
+/// `result` must wrap an OpResult. `parameters` may be a null attribute. The
+/// caller owns the returned instance and must destroy it with
+/// `mlirMemoryEffectInstanceDestroy`.
+MLIR_CAPI_EXPORTED MlirMemoryEffectInstance
+mlirMemoryEffectInstanceCreateForOpResult(MlirMemoryEffect effect,
+ MlirValue result,
+ MlirAttribute parameters, int stage,
+ bool effectOnFullRegion,
+ MlirSideEffectResource resource);
+
+/// Creates a memory effect instance associated with a block argument.
+/// `blockArgument` must wrap a BlockArgument. `parameters` may be a null
+/// attribute. The caller owns the returned instance and must destroy it with
+/// `mlirMemoryEffectInstanceDestroy`.
+MLIR_CAPI_EXPORTED MlirMemoryEffectInstance
+mlirMemoryEffectInstanceCreateForBlockArgument(
+ MlirMemoryEffect effect, MlirValue blockArgument, MlirAttribute parameters,
+ int stage, bool effectOnFullRegion, MlirSideEffectResource resource);
+
+/// Creates a memory effect instance associated with a symbol. `symbol` must be
+/// a SymbolRefAttr. `parameters` may be a null attribute. The caller owns the
+/// returned instance and must destroy it with
+/// `mlirMemoryEffectInstanceDestroy`.
+MLIR_CAPI_EXPORTED MlirMemoryEffectInstance
+mlirMemoryEffectInstanceCreateForSymbol(MlirMemoryEffect effect,
+ MlirAttribute symbol,
+ MlirAttribute parameters, int stage,
+ bool effectOnFullRegion,
+ MlirSideEffectResource resource);
+
+/// Destroys a memory effect instance created by one of the functions above.
+MLIR_CAPI_EXPORTED void
+mlirMemoryEffectInstanceDestroy(MlirMemoryEffectInstance instance);
+
+/// Appends a copy of `instance` to the given list. This does not take ownership
+/// of `instance`; the caller remains responsible for destroying it.
+MLIR_CAPI_EXPORTED void
+mlirMemoryEffectInstancesListAppend(MlirMemoryEffectInstancesList list,
+ MlirMemoryEffectInstance instance);
+
/// Returns the interface TypeID of the MemoryEffectsOpInterface.
MLIR_CAPI_EXPORTED MlirTypeID mlirMemoryEffectsOpInterfaceTypeID(void);
diff --git a/mlir/include/mlir/Bindings/Python/IRInterfaces.h b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
index fb30e030b6c32..1ba995aa82c5e 100644
--- a/mlir/include/mlir/Bindings/Python/IRInterfaces.h
+++ b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
@@ -135,7 +135,58 @@ class PyConcreteOpInterface {
nanobind::object obj;
};
-struct PyMemoryEffectsInstanceList {
+/// A borrowed memory effect.
+class PyMemoryEffect {
+public:
+ explicit PyMemoryEffect(MlirMemoryEffect effect) : effect(effect) {}
+
+ MlirMemoryEffect get() const { return effect; }
+
+private:
+ MlirMemoryEffect effect;
+};
+
+/// A borrowed side effect resource.
+class PySideEffectResource {
+public:
+ explicit PySideEffectResource(MlirSideEffectResource resource)
+ : resource(resource) {}
+
+ MlirSideEffectResource get() const { return resource; }
+
+private:
+ MlirSideEffectResource resource;
+};
+
+/// An owning memory effect instance.
+class PyMemoryEffectInstance {
+public:
+ explicit PyMemoryEffectInstance(MlirMemoryEffectInstance instance)
+ : instance(instance) {}
+ PyMemoryEffectInstance(PyMemoryEffectInstance &&other) noexcept
+ : instance(other.instance) {
+ other.instance.ptr = nullptr;
+ }
+ ~PyMemoryEffectInstance() {
+ if (instance.ptr)
+ mlirMemoryEffectInstanceDestroy(instance);
+ }
+
+ MlirMemoryEffectInstance get() const { return instance; }
+
+private:
+ MlirMemoryEffectInstance instance;
+};
+
+/// A callback-scoped view of a list of memory effect instances.
+class PyMemoryEffectsInstanceList {
+public:
+ explicit PyMemoryEffectsInstanceList(MlirMemoryEffectInstancesList effects)
+ : effects(effects) {}
+
+ MlirMemoryEffectInstancesList get() const { return effects; }
+
+private:
MlirMemoryEffectInstancesList effects;
};
diff --git a/mlir/include/mlir/CAPI/Interfaces.h b/mlir/include/mlir/CAPI/Interfaces.h
index 15afc9fb0f18e..55d850ccd01eb 100644
--- a/mlir/include/mlir/CAPI/Interfaces.h
+++ b/mlir/include/mlir/CAPI/Interfaces.h
@@ -22,5 +22,9 @@
DEFINE_C_API_PTR_METHODS(
MlirMemoryEffectInstancesList,
llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance>)
+DEFINE_C_API_PTR_METHODS(MlirMemoryEffect, mlir::MemoryEffects::Effect)
+DEFINE_C_API_PTR_METHODS(MlirMemoryEffectInstance,
+ mlir::MemoryEffects::EffectInstance)
+DEFINE_C_API_PTR_METHODS(MlirSideEffectResource, mlir::SideEffects::Resource)
#endif // MLIR_CAPI_INTERFACES_H
diff --git a/mlir/lib/Bindings/Python/DialectTransform.cpp b/mlir/lib/Bindings/Python/DialectTransform.cpp
index bd72082cea7e8..f84189e3a22ff 100644
--- a/mlir/lib/Bindings/Python/DialectTransform.cpp
+++ b/mlir/lib/Bindings/Python/DialectTransform.cpp
@@ -494,38 +494,38 @@ struct ParamType : PyConcreteType<ParamType> {
namespace {
void onlyReadsHandle(nb::iterable &operands,
- PyMemoryEffectsInstanceList effects) {
+ const PyMemoryEffectsInstanceList &effects) {
std::vector<MlirOpOperand> operandsVec;
for (auto operand : operands)
operandsVec.push_back(nb::cast<PyOpOperand>(operand));
mlirTransformOnlyReadsHandle(operandsVec.data(), operandsVec.size(),
- effects.effects);
+ effects.get());
};
void consumesHandle(nb::iterable &operands,
- PyMemoryEffectsInstanceList effects) {
+ const PyMemoryEffectsInstanceList &effects) {
std::vector<MlirOpOperand> operandsVec;
for (auto operand : operands)
operandsVec.push_back(nb::cast<PyOpOperand>(operand));
mlirTransformConsumesHandle(operandsVec.data(), operandsVec.size(),
- effects.effects);
+ effects.get());
};
void producesHandle(nb::iterable &results,
- PyMemoryEffectsInstanceList effects) {
+ const PyMemoryEffectsInstanceList &effects) {
std::vector<MlirValue> resultsVec;
for (auto result : results)
resultsVec.push_back(nb::cast<PyOpResult>(result).get());
mlirTransformProducesHandle(resultsVec.data(), resultsVec.size(),
- effects.effects);
+ effects.get());
};
-void modifiesPayload(PyMemoryEffectsInstanceList effects) {
- mlirTransformModifiesPayload(effects.effects);
+void modifiesPayload(const PyMemoryEffectsInstanceList &effects) {
+ mlirTransformModifiesPayload(effects.get());
}
-void onlyReadsPayload(PyMemoryEffectsInstanceList effects) {
- mlirTransformOnlyReadsPayload(effects.effects);
+void onlyReadsPayload(const PyMemoryEffectsInstanceList &effects) {
+ mlirTransformOnlyReadsPayload(effects.get());
}
} // namespace
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index c4e87246bd68c..4495cb6e8397d 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -34,6 +34,70 @@ its return shaped type components. Raises ValueError on failure.)";
namespace {
+MlirAttribute unwrapOptionalAttribute(const nb::object &attribute) {
+ if (attribute.is_none())
+ return mlirAttributeGetNull();
+
+ PyAttribute *pyAttribute = nullptr;
+ if (!nb::try_cast<PyAttribute *>(attribute, pyAttribute) || !pyAttribute)
+ throw nb::type_error("parameters must be an Attribute or None");
+ return pyAttribute->get();
+}
+
+void appendMemoryEffectInstance(PyMemoryEffectsInstanceList &effects,
+ const PyMemoryEffect &effect,
+ const nb::object &target,
+ const nb::object ¶meters, int stage,
+ bool effectOnFullRegion,
+ const PySideEffectResource &resource) {
+ MlirMemoryEffectInstancesList list = effects.get();
+ MlirAttribute unwrappedParameters = unwrapOptionalAttribute(parameters);
+
+ MlirMemoryEffectInstance rawInstance{nullptr};
+ if (target.is_none()) {
+ rawInstance =
+ mlirMemoryEffectInstanceCreate(effect.get(), unwrappedParameters, stage,
+ effectOnFullRegion, resource.get());
+ } else {
+ PyOpOperand *opOperand = nullptr;
+ PyValue *value = nullptr;
+ PyAttribute *attribute = nullptr;
+ if (nb::try_cast<PyOpOperand *>(target, opOperand) && opOperand) {
+ rawInstance = mlirMemoryEffectInstanceCreateForOpOperand(
+ effect.get(), *opOperand, unwrappedParameters, stage,
+ effectOnFullRegion, resource.get());
+ } else if (nb::try_cast<PyValue *>(target, value) && value) {
+ MlirValue mlirValue = value->get();
+ if (mlirValueIsAOpResult(mlirValue)) {
+ rawInstance = mlirMemoryEffectInstanceCreateForOpResult(
+ effect.get(), mlirValue, unwrappedParameters, stage,
+ effectOnFullRegion, resource.get());
+ } else if (mlirValueIsABlockArgument(mlirValue)) {
+ rawInstance = mlirMemoryEffectInstanceCreateForBlockArgument(
+ effect.get(), mlirValue, unwrappedParameters, stage,
+ effectOnFullRegion, resource.get());
+ } else {
+ throw nb::type_error(
+ "target Value must be an OpResult or BlockArgument");
+ }
+ } else if (nb::try_cast<PyAttribute *>(target, attribute) && attribute) {
+ MlirAttribute symbol = attribute->get();
+ if (!mlirAttributeIsASymbolRef(symbol))
+ throw nb::type_error("target Attribute must be a SymbolRefAttr");
+ rawInstance = mlirMemoryEffectInstanceCreateForSymbol(
+ effect.get(), symbol, unwrappedParameters, stage, effectOnFullRegion,
+ resource.get());
+ } else {
+ throw nb::type_error(
+ "target must be an OpOperand, OpResult, BlockArgument, "
+ "SymbolRefAttr, or None");
+ }
+ }
+
+ PyMemoryEffectInstance instance(rawInstance);
+ mlirMemoryEffectInstancesListAppend(list, instance.get());
+}
+
/// Takes in an optional ist of operands and converts them into a std::vector
/// of MlirVlaues. Returns an empty std::vector if the list is empty.
std::vector<MlirValue> wrapOperands(std::optional<nb::sequence> operandList) {
@@ -474,9 +538,46 @@ void populateIRInterfaces(nb::module_ &m) {
.value("Speculatable", MlirSpeculatabilitySpeculatable)
.value("RecursivelySpeculatable",
MlirSpeculatabilityRecursivelySpeculatable);
- auto memoryEffectsInstanceListClass =
- nb::class_<PyMemoryEffectsInstanceList>(m, "MemoryEffectInstancesList");
- (void)memoryEffectsInstanceListClass;
+ nb::class_<PyMemoryEffect>(m, "MemoryEffect", "A memory effect.")
+ .def_prop_ro_static("allocate",
+ [](nb::object & /*class*/) {
+ return PyMemoryEffect(
+ mlirMemoryEffectsAllocateGet());
+ })
+ .def_prop_ro_static("free",
+ [](nb::object & /*class*/) {
+ return PyMemoryEffect(mlirMemoryEffectsFreeGet());
+ })
+ .def_prop_ro_static("read",
+ [](nb::object & /*class*/) {
+ return PyMemoryEffect(mlirMemoryEffectsReadGet());
+ })
+ .def_prop_ro_static("write", [](nb::object & /*class*/) {
+ return PyMemoryEffect(mlirMemoryEffectsWriteGet());
+ });
+
+ nb::class_<PySideEffectResource>(m, "SideEffectResource",
+ "A side effect resource.")
+ .def_prop_ro_static("default", [](nb::object & /*class*/) {
+ return PySideEffectResource(mlirSideEffectsDefaultResourceGet());
+ });
+
+ nb::class_<PyMemoryEffectsInstanceList>(
+ m, "MemoryEffectInstancesList",
+ "A memory effect list that is valid only during get_effects.")
+ .def("append", &appendMemoryEffectInstance, nb::arg("effect"),
+ nb::arg("target").none() = nb::none(), nb::kw_only(),
+ nb::arg("parameters").none() = nb::none(), nb::arg("stage") = 0,
+ nb::arg("effect_on_full_region") = false,
+ nb::arg("resource") =
+ PySideEffectResource(mlirSideEffectsDefaultResourceGet()),
+ nb::sig("def append(self, effect: MemoryEffect, target: OpOperand | "
+ "OpResult | BlockArgument | SymbolRefAttr | None = None, *, "
+ "parameters: Attribute | None = None, stage: int = 0, "
+ "effect_on_full_region: bool = False, resource: "
+ "SideEffectResource = ...) -> None"),
+ "Append a memory effect instance. The target may be an OpOperand, "
+ "OpResult, BlockArgument, SymbolRefAttr, or None.");
PyConditionallySpeculatableOpInterface::bind(m);
PyInferShapedTypeOpInterface::bind(m);
diff --git a/mlir/lib/CAPI/Interfaces/Interfaces.cpp b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
index 35a2bd562a8a1..1bf8dffa9c431 100644
--- a/mlir/lib/CAPI/Interfaces/Interfaces.cpp
+++ b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
@@ -280,6 +280,78 @@ MlirSpeculatability mlirConditionallySpeculatableOpInterfaceGetSpeculatability(
// MemoryEffectOpInterface
//===---------------------------------------------------------------------===//
+MlirMemoryEffect mlirMemoryEffectsAllocateGet() {
+ return wrap(
+ static_cast<MemoryEffects::Effect *>(MemoryEffects::Allocate::get()));
+}
+
+MlirMemoryEffect mlirMemoryEffectsFreeGet() {
+ return wrap(static_cast<MemoryEffects::Effect *>(MemoryEffects::Free::get()));
+}
+
+MlirMemoryEffect mlirMemoryEffectsReadGet() {
+ return wrap(static_cast<MemoryEffects::Effect *>(MemoryEffects::Read::get()));
+}
+
+MlirMemoryEffect mlirMemoryEffectsWriteGet() {
+ return wrap(
+ static_cast<MemoryEffects::Effect *>(MemoryEffects::Write::get()));
+}
+
+MlirSideEffectResource mlirSideEffectsDefaultResourceGet() {
+ return wrap(static_cast<SideEffects::Resource *>(
+ SideEffects::DefaultResource::get()));
+}
+
+MlirMemoryEffectInstance mlirMemoryEffectInstanceCreate(
+ MlirMemoryEffect effect, MlirAttribute parameters, int stage,
+ bool effectOnFullRegion, MlirSideEffectResource resource) {
+ return wrap(new MemoryEffects::EffectInstance(
+ unwrap(effect), unwrap(parameters), stage, effectOnFullRegion,
+ unwrap(resource)));
+}
+
+MlirMemoryEffectInstance mlirMemoryEffectInstanceCreateForOpOperand(
+ MlirMemoryEffect effect, MlirOpOperand opOperand, MlirAttribute parameters,
+ int stage, bool effectOnFullRegion, MlirSideEffectResource resource) {
+ return wrap(new MemoryEffects::EffectInstance(
+ unwrap(effect), unwrap(opOperand), unwrap(parameters), stage,
+ effectOnFullRegion, unwrap(resource)));
+}
+
+MlirMemoryEffectInstance mlirMemoryEffectInstanceCreateForOpResult(
+ MlirMemoryEffect effect, MlirValue result, MlirAttribute parameters,
+ int stage, bool effectOnFullRegion, MlirSideEffectResource resource) {
+ return wrap(new MemoryEffects::EffectInstance(
+ unwrap(effect), cast<OpResult>(unwrap(result)), unwrap(parameters), stage,
+ effectOnFullRegion, unwrap(resource)));
+}
+
+MlirMemoryEffectInstance mlirMemoryEffectInstanceCreateForBlockArgument(
+ MlirMemoryEffect effect, MlirValue blockArgument, MlirAttribute parameters,
+ int stage, bool effectOnFullRegion, MlirSideEffectResource resource) {
+ return wrap(new MemoryEffects::EffectInstance(
+ unwrap(effect), cast<BlockArgument>(unwrap(blockArgument)),
+ unwrap(parameters), stage, effectOnFullRegion, unwrap(resource)));
+}
+
+MlirMemoryEffectInstance mlirMemoryEffectInstanceCreateForSymbol(
+ MlirMemoryEffect effect, MlirAttribute symbol, MlirAttribute parameters,
+ int stage, bool effectOnFullRegion, MlirSideEffectResource resource) {
+ return wrap(new MemoryEffects::EffectInstance(
+ unwrap(effect), cast<SymbolRefAttr>(unwrap(symbol)), unwrap(parameters),
+ stage, effectOnFullRegion, unwrap(resource)));
+}
+
+void mlirMemoryEffectInstanceDestroy(MlirMemoryEffectInstance instance) {
+ delete unwrap(instance);
+}
+
+void mlirMemoryEffectInstancesListAppend(MlirMemoryEffectInstancesList list,
+ MlirMemoryEffectInstance instance) {
+ unwrap(list)->push_back(*unwrap(instance));
+}
+
MlirTypeID mlirMemoryEffectsOpInterfaceTypeID() {
return wrap(MemoryEffectOpInterface::getInterfaceID());
}
diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index 57ae8b9a2819b..3c72209ad95c3 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -2589,6 +2589,8 @@ int testInterfaces(MlirContext ctx) {
// CHECK: arith.constant speculatability: 1
MlirOperationState storeState = mlirOperationStateGet(storeName, loc);
+ MlirValue constantResult = mlirOperationGetResult(constantOp, 0);
+ mlirOperationStateAddOperands(&storeState, 1, &constantResult);
MlirOperation storeOp = mlirOperationCreate(&storeState);
if (mlirOperationImplementsInterface(storeOp, condSpecTypeID)) {
fprintf(stderr, "ERROR: Expected memref.store instance to not implement "
@@ -2619,6 +2621,47 @@ int testInterfaces(MlirContext ctx) {
// CHECK: memref.store speculatability: 2
// CHECK: callback count: 1
+ MlirMemoryEffect allocate = mlirMemoryEffectsAllocateGet();
+ MlirMemoryEffect free = mlirMemoryEffectsFreeGet();
+ MlirMemoryEffect read = mlirMemoryEffectsReadGet();
+ MlirMemoryEffect write = mlirMemoryEffectsWriteGet();
+ MlirSideEffectResource defaultResource = mlirSideEffectsDefaultResourceGet();
+ if (!allocate.ptr || !free.ptr || !read.ptr || !write.ptr ||
+ !defaultResource.ptr) {
+ fprintf(stderr, "ERROR: Expected memory effect components\n");
+ return 6;
+ }
+
+ MlirAttribute nullParameters = {NULL};
+ MlirOpOperand opOperand = mlirOperationGetOpOperand(storeOp, 0);
+ MlirBlock block = mlirBlockCreate(1, &i32, &loc);
+ MlirValue blockArgument = mlirBlockGetArgument(block, 0);
+ MlirAttribute symbol = mlirFlatSymbolRefAttrGet(
+ ctx, mlirStringRefCreateFromCString("effect_target"));
+
+ MlirMemoryEffectInstance instances[] = {
+ mlirMemoryEffectInstanceCreate(allocate, nullParameters, 0, false,
+ defaultResource),
+ mlirMemoryEffectInstanceCreateForOpOperand(read, opOperand, zero, 1,
+ false, defaultResource),
+ mlirMemoryEffectInstanceCreateForOpResult(
+ write, constantResult, nullParameters, 2, false, defaultResource),
+ mlirMemoryEffectInstanceCreateForBlockArgument(free, blockArgument, zero,
+ 3, false, defaultResource),
+ mlirMemoryEffectInstanceCreateForSymbol(read, symbol, zero, 4, true,
+ defaultResource),
+ };
+ for (intptr_t i = 0; i < 5; ++i) {
+ if (!instances[i].ptr) {
+ fprintf(stderr, "ERROR: Expected memory effect instance\n");
+ return 7;
+ }
+ mlirMemoryEffectInstanceDestroy(instances[i]);
+ }
+ mlirBlockDestroy(block);
+ fprintf(stderr, "memory effect instances constructed\n");
+ // CHECK: memory effect instances constructed
+
mlirOperationDestroy(storeOp);
mlirOperationDestroy(constantOp);
return 0;
diff --git a/mlir/test/python/dialects/memory_effects_op_interface.py b/mlir/test/python/dialects/memory_effects_op_interface.py
new file mode 100644
index 0000000000000..5689130c584b5
--- /dev/null
+++ b/mlir/test/python/dialects/memory_effects_op_interface.py
@@ -0,0 +1,291 @@
+# RUN: env PYTHONUNBUFFERED=1 %PYTHON %s 2>&1 | FileCheck %s
+
+from typing import Any
+
+from mlir import ir
+from mlir.dialects import ext, func
+from mlir.passmanager import PassManager
+
+
+class MemoryEffectsTest(ext.Dialect, name="memory_effects_test"):
+ pass
+
+
+class NoEffectModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ pass
+
+
+class ReadModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ effects.append(
+ ir.MemoryEffect.read,
+ op.op_operands[0],
+ parameters=ir.StringAttr.get("read parameter"),
+ stage=1,
+ effect_on_full_region=True,
+ resource=ir.SideEffectResource.default,
+ )
+
+
+class ReadDeadModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ effects.append(ir.MemoryEffect.read)
+
+
+class WriteModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ effects.append(ir.MemoryEffect.write)
+
+
+class FreeModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ effects.append(ir.MemoryEffect.free)
+
+
+class AllocateModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ effects.append(ir.MemoryEffect.allocate)
+
+
+class AllocateResultModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ effects.append(ir.MemoryEffect.allocate, op.results[0])
+
+
+class BlockArgumentTargetModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ effects.append(ir.MemoryEffect.read, op.regions[0].blocks[0].arguments[0])
+
+
+class SymbolTargetModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ try:
+ effects.append(ir.MemoryEffect.read, ir.StringAttr.get("not a symbol"))
+ except TypeError as error:
+ print("invalid symbol target:", error)
+ try:
+ effects.append(ir.MemoryEffect.read, parameters=42)
+ except TypeError as error:
+ print("invalid parameters:", error)
+ effects.append(
+ ir.MemoryEffect.read,
+ ir.FlatSymbolRefAttr.get("global"),
+ parameters=ir.StringAttr.get("symbol parameter"),
+ stage=2,
+ effect_on_full_region=True,
+ )
+
+
+class ReadOp(MemoryEffectsTest.Operation, name="read", traits=[ReadModel]):
+ operand: ext.Operand[Any]
+ result: ext.Result[Any]
+
+
+class WriteOp(MemoryEffectsTest.Operation, name="write", traits=[WriteModel]):
+ operand: ext.Operand[Any]
+ result: ext.Result[Any]
+
+
+class WriteBarrierOp(
+ MemoryEffectsTest.Operation, name="write_barrier", traits=[WriteModel]
+):
+ operand: ext.Operand[Any]
+
+
+class NoEffectOp(MemoryEffectsTest.Operation, name="no_effect", traits=[NoEffectModel]):
+ pass
+
+
+class ReadDeadOp(MemoryEffectsTest.Operation, name="read_dead", traits=[ReadDeadModel]):
+ pass
+
+
+class WriteDeadOp(MemoryEffectsTest.Operation, name="write_dead", traits=[WriteModel]):
+ pass
+
+
+class FreeDeadOp(MemoryEffectsTest.Operation, name="free_dead", traits=[FreeModel]):
+ pass
+
+
+class AllocateDeadOp(
+ MemoryEffectsTest.Operation, name="allocate_dead", traits=[AllocateModel]
+):
+ pass
+
+
+class AllocateResultOp(
+ MemoryEffectsTest.Operation,
+ name="allocate_result",
+ traits=[AllocateResultModel],
+):
+ result: ext.Result[Any]
+
+
+class BlockArgumentTargetOp(
+ MemoryEffectsTest.Operation,
+ name="block_argument_target",
+ traits=[ir.NoTerminatorTrait, BlockArgumentTargetModel],
+):
+ body: ext.Region
+
+
+class SymbolTargetOp(
+ MemoryEffectsTest.Operation, name="symbol_target", traits=[SymbolTargetModel]
+):
+ pass
+
+
+def run_pass(source, pipeline):
+ module = ir.Module.parse(source)
+ PassManager.parse(pipeline).run(module.operation)
+ return str(module)
+
+
+with ir.Context(), ir.Location.unknown():
+ MemoryEffectsTest.load()
+
+ # The static properties return wrappers around borrowed, statically-owned
+ # C++ singletons.
+ # CHECK: memory effect properties: True True True True
+ print(
+ "memory effect properties:",
+ isinstance(ir.MemoryEffect.allocate, ir.MemoryEffect),
+ isinstance(ir.MemoryEffect.free, ir.MemoryEffect),
+ isinstance(ir.MemoryEffect.read, ir.MemoryEffect),
+ isinstance(ir.MemoryEffect.write, ir.MemoryEffect),
+ )
+ # CHECK: default resource property: True
+ print(
+ "default resource property:",
+ isinstance(ir.SideEffectResource.default, ir.SideEffectResource),
+ )
+
+ 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))",
+ )
+ # 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))",
+ )
+ # 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))",
+ )
+ # 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"'),
+ )
+
+ 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))",
+ )
+ # 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.
+ # CHECK: DCE no effect count: 0
+ # CHECK: DCE read count: 0
+ # CHECK: DCE write count: 1
+ # 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 untargeted allocate count:",
+ dead_code.count('"memory_effects_test.allocate_dead"'),
+ )
+ print(
+ "DCE result allocate count:",
+ dead_code.count('"memory_effects_test.allocate_result"'),
+ )
+
+ 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))",
+ )
+ # These Read effects exercise BlockArgument and SymbolRefAttr targets and
+ # remain removable by trivial-dce.
+ # CHECK: invalid symbol target: target Attribute must be a SymbolRefAttr
+ # CHECK: invalid parameters: parameters must be an Attribute or None
+ # CHECK: DCE block argument target count: 0
+ # CHECK: DCE symbol target count: 0
+ print(
+ "DCE block argument target count:",
+ target_variants.count('"memory_effects_test.block_argument_target"'),
+ )
+ print(
+ "DCE symbol target count:",
+ target_variants.count('"memory_effects_test.symbol_target"'),
+ )
>From ae347ca90903a7a23e8cd26d39019ad5cd75afa9 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sun, 19 Jul 2026 22:56:28 +0800
Subject: [PATCH 2/4] fix
---
mlir/test/python/dialects/memory_effects_op_interface.py | 2 --
1 file changed, 2 deletions(-)
diff --git a/mlir/test/python/dialects/memory_effects_op_interface.py b/mlir/test/python/dialects/memory_effects_op_interface.py
index 5689130c584b5..cb043cb961bca 100644
--- a/mlir/test/python/dialects/memory_effects_op_interface.py
+++ b/mlir/test/python/dialects/memory_effects_op_interface.py
@@ -155,8 +155,6 @@ def run_pass(source, pipeline):
with ir.Context(), ir.Location.unknown():
MemoryEffectsTest.load()
- # The static properties return wrappers around borrowed, statically-owned
- # C++ singletons.
# CHECK: memory effect properties: True True True True
print(
"memory effect properties:",
>From 90d24dbc6e65f26f5d0b5824ec0cfe5abecb142a Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Thu, 23 Jul 2026 23:34:23 +0800
Subject: [PATCH 3/4] address comments
---
.../mlir/Bindings/Python/IRInterfaces.h | 7 ++--
mlir/lib/Bindings/Python/DialectTransform.cpp | 13 +++-----
mlir/lib/Bindings/Python/IRInterfaces.cpp | 10 +++---
.../dialects/memory_effects_op_interface.py | 32 +++++++++----------
4 files changed, 30 insertions(+), 32 deletions(-)
diff --git a/mlir/include/mlir/Bindings/Python/IRInterfaces.h b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
index 1ba995aa82c5e..c1047c7e20ab7 100644
--- a/mlir/include/mlir/Bindings/Python/IRInterfaces.h
+++ b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
@@ -135,7 +135,7 @@ class PyConcreteOpInterface {
nanobind::object obj;
};
-/// A borrowed memory effect.
+/// A memory effect.
class PyMemoryEffect {
public:
explicit PyMemoryEffect(MlirMemoryEffect effect) : effect(effect) {}
@@ -146,7 +146,7 @@ class PyMemoryEffect {
MlirMemoryEffect effect;
};
-/// A borrowed side effect resource.
+/// A side effect resource.
class PySideEffectResource {
public:
explicit PySideEffectResource(MlirSideEffectResource resource)
@@ -158,7 +158,7 @@ class PySideEffectResource {
MlirSideEffectResource resource;
};
-/// An owning memory effect instance.
+/// A memory effect instance.
class PyMemoryEffectInstance {
public:
explicit PyMemoryEffectInstance(MlirMemoryEffectInstance instance)
@@ -185,6 +185,7 @@ class PyMemoryEffectsInstanceList {
: effects(effects) {}
MlirMemoryEffectInstancesList get() const { return effects; }
+ operator MlirMemoryEffectInstancesList() const { return effects; }
private:
MlirMemoryEffectInstancesList effects;
diff --git a/mlir/lib/Bindings/Python/DialectTransform.cpp b/mlir/lib/Bindings/Python/DialectTransform.cpp
index f84189e3a22ff..a4a91b28522cf 100644
--- a/mlir/lib/Bindings/Python/DialectTransform.cpp
+++ b/mlir/lib/Bindings/Python/DialectTransform.cpp
@@ -498,8 +498,7 @@ void onlyReadsHandle(nb::iterable &operands,
std::vector<MlirOpOperand> operandsVec;
for (auto operand : operands)
operandsVec.push_back(nb::cast<PyOpOperand>(operand));
- mlirTransformOnlyReadsHandle(operandsVec.data(), operandsVec.size(),
- effects.get());
+ mlirTransformOnlyReadsHandle(operandsVec.data(), operandsVec.size(), effects);
};
void consumesHandle(nb::iterable &operands,
@@ -507,8 +506,7 @@ void consumesHandle(nb::iterable &operands,
std::vector<MlirOpOperand> operandsVec;
for (auto operand : operands)
operandsVec.push_back(nb::cast<PyOpOperand>(operand));
- mlirTransformConsumesHandle(operandsVec.data(), operandsVec.size(),
- effects.get());
+ mlirTransformConsumesHandle(operandsVec.data(), operandsVec.size(), effects);
};
void producesHandle(nb::iterable &results,
@@ -516,16 +514,15 @@ void producesHandle(nb::iterable &results,
std::vector<MlirValue> resultsVec;
for (auto result : results)
resultsVec.push_back(nb::cast<PyOpResult>(result).get());
- mlirTransformProducesHandle(resultsVec.data(), resultsVec.size(),
- effects.get());
+ mlirTransformProducesHandle(resultsVec.data(), resultsVec.size(), effects);
};
void modifiesPayload(const PyMemoryEffectsInstanceList &effects) {
- mlirTransformModifiesPayload(effects.get());
+ mlirTransformModifiesPayload(effects);
}
void onlyReadsPayload(const PyMemoryEffectsInstanceList &effects) {
- mlirTransformOnlyReadsPayload(effects.get());
+ mlirTransformOnlyReadsPayload(effects);
}
} // namespace
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index 4495cb6e8397d..d16015acaf3dc 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -539,26 +539,26 @@ void populateIRInterfaces(nb::module_ &m) {
.value("RecursivelySpeculatable",
MlirSpeculatabilityRecursivelySpeculatable);
nb::class_<PyMemoryEffect>(m, "MemoryEffect", "A memory effect.")
- .def_prop_ro_static("allocate",
+ .def_prop_ro_static("Allocate",
[](nb::object & /*class*/) {
return PyMemoryEffect(
mlirMemoryEffectsAllocateGet());
})
- .def_prop_ro_static("free",
+ .def_prop_ro_static("Free",
[](nb::object & /*class*/) {
return PyMemoryEffect(mlirMemoryEffectsFreeGet());
})
- .def_prop_ro_static("read",
+ .def_prop_ro_static("Read",
[](nb::object & /*class*/) {
return PyMemoryEffect(mlirMemoryEffectsReadGet());
})
- .def_prop_ro_static("write", [](nb::object & /*class*/) {
+ .def_prop_ro_static("Write", [](nb::object & /*class*/) {
return PyMemoryEffect(mlirMemoryEffectsWriteGet());
});
nb::class_<PySideEffectResource>(m, "SideEffectResource",
"A side effect resource.")
- .def_prop_ro_static("default", [](nb::object & /*class*/) {
+ .def_prop_ro_static("Default", [](nb::object & /*class*/) {
return PySideEffectResource(mlirSideEffectsDefaultResourceGet());
});
diff --git a/mlir/test/python/dialects/memory_effects_op_interface.py b/mlir/test/python/dialects/memory_effects_op_interface.py
index cb043cb961bca..3e6a1b58c55b0 100644
--- a/mlir/test/python/dialects/memory_effects_op_interface.py
+++ b/mlir/test/python/dialects/memory_effects_op_interface.py
@@ -21,64 +21,64 @@ class ReadModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
effects.append(
- ir.MemoryEffect.read,
+ ir.MemoryEffect.Read,
op.op_operands[0],
parameters=ir.StringAttr.get("read parameter"),
stage=1,
effect_on_full_region=True,
- resource=ir.SideEffectResource.default,
+ resource=ir.SideEffectResource.Default,
)
class ReadDeadModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
- effects.append(ir.MemoryEffect.read)
+ effects.append(ir.MemoryEffect.Read)
class WriteModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
- effects.append(ir.MemoryEffect.write)
+ effects.append(ir.MemoryEffect.Write)
class FreeModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
- effects.append(ir.MemoryEffect.free)
+ effects.append(ir.MemoryEffect.Free)
class AllocateModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
- effects.append(ir.MemoryEffect.allocate)
+ effects.append(ir.MemoryEffect.Allocate)
class AllocateResultModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
- effects.append(ir.MemoryEffect.allocate, op.results[0])
+ effects.append(ir.MemoryEffect.Allocate, op.results[0])
class BlockArgumentTargetModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
- effects.append(ir.MemoryEffect.read, op.regions[0].blocks[0].arguments[0])
+ effects.append(ir.MemoryEffect.Read, op.regions[0].blocks[0].arguments[0])
class SymbolTargetModel(ir.MemoryEffectsOpInterface):
@staticmethod
def get_effects(op, effects):
try:
- effects.append(ir.MemoryEffect.read, ir.StringAttr.get("not a symbol"))
+ effects.append(ir.MemoryEffect.Read, ir.StringAttr.get("not a symbol"))
except TypeError as error:
print("invalid symbol target:", error)
try:
- effects.append(ir.MemoryEffect.read, parameters=42)
+ effects.append(ir.MemoryEffect.Read, parameters=42)
except TypeError as error:
print("invalid parameters:", error)
effects.append(
- ir.MemoryEffect.read,
+ ir.MemoryEffect.Read,
ir.FlatSymbolRefAttr.get("global"),
parameters=ir.StringAttr.get("symbol parameter"),
stage=2,
@@ -158,15 +158,15 @@ def run_pass(source, pipeline):
# CHECK: memory effect properties: True True True True
print(
"memory effect properties:",
- isinstance(ir.MemoryEffect.allocate, ir.MemoryEffect),
- isinstance(ir.MemoryEffect.free, ir.MemoryEffect),
- isinstance(ir.MemoryEffect.read, ir.MemoryEffect),
- isinstance(ir.MemoryEffect.write, ir.MemoryEffect),
+ isinstance(ir.MemoryEffect.Allocate, ir.MemoryEffect),
+ isinstance(ir.MemoryEffect.Free, ir.MemoryEffect),
+ isinstance(ir.MemoryEffect.Read, ir.MemoryEffect),
+ isinstance(ir.MemoryEffect.Write, ir.MemoryEffect),
)
# CHECK: default resource property: True
print(
"default resource property:",
- isinstance(ir.SideEffectResource.default, ir.SideEffectResource),
+ isinstance(ir.SideEffectResource.Default, ir.SideEffectResource),
)
read_cse = run_pass(
>From a03c2f97ecc374ecb9563dcbbd62105d8b68c593 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Thu, 23 Jul 2026 23:47:47 +0800
Subject: [PATCH 4/4] address claude comments
---
mlir/test/python/dialects/memory_effects_op_interface.py | 5 +++++
1 file changed, 5 insertions(+)
diff --git a/mlir/test/python/dialects/memory_effects_op_interface.py b/mlir/test/python/dialects/memory_effects_op_interface.py
index 3e6a1b58c55b0..dc3b744aae2d7 100644
--- a/mlir/test/python/dialects/memory_effects_op_interface.py
+++ b/mlir/test/python/dialects/memory_effects_op_interface.py
@@ -77,6 +77,10 @@ def get_effects(op, effects):
effects.append(ir.MemoryEffect.Read, parameters=42)
except TypeError as error:
print("invalid parameters:", error)
+ try:
+ effects.append(ir.MemoryEffect.Read, 42)
+ except TypeError as error:
+ print("invalid target:", error)
effects.append(
ir.MemoryEffect.Read,
ir.FlatSymbolRefAttr.get("global"),
@@ -277,6 +281,7 @@ def run_pass(source, pipeline):
# remain removable by trivial-dce.
# CHECK: invalid symbol target: target Attribute must be a SymbolRefAttr
# CHECK: invalid parameters: parameters must be an Attribute or None
+ # CHECK: invalid target: target must be an OpOperand, OpResult, BlockArgument, SymbolRefAttr, or None
# CHECK: DCE block argument target count: 0
# CHECK: DCE symbol target count: 0
print(
More information about the Mlir-commits
mailing list