[Mlir-commits] [mlir] [MLIR][CAPI][Python] Add support for constructing memory effect instances (PR #210586)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Jul 19 07:57:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Twice (PragmaTwice)
<details>
<summary>Changes</summary>
Python implementations of `MemoryEffectsOpInterface` receive a `MemoryEffectInstancesList` (added in #<!-- -->176920), but currently have no generic way to populate it.
This adds the C API for constructing and appending memory effect instances and exposes it in Python through `MemoryEffect`, `SideEffectResource`, and `MemoryEffectInstancesList.append`. The tests use CSE and DCE to verify that the declared effects are observed by MLIR passes.
Assisted-by: GPT 5.6 Sol
---
Patch is 30.95 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/210586.diff
8 Files Affected:
- (modified) mlir/include/mlir-c/Interfaces.h (+78)
- (modified) mlir/include/mlir/Bindings/Python/IRInterfaces.h (+52-1)
- (modified) mlir/include/mlir/CAPI/Interfaces.h (+4)
- (modified) mlir/lib/Bindings/Python/DialectTransform.cpp (+10-10)
- (modified) mlir/lib/Bindings/Python/IRInterfaces.cpp (+104-3)
- (modified) mlir/lib/CAPI/Interfaces/Interfaces.cpp (+72)
- (modified) mlir/test/CAPI/ir.c (+43)
- (added) mlir/test/python/dialects/memory_effects_op_interface.py (+289)
``````````diff
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...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/210586
More information about the Mlir-commits
mailing list