[Mlir-commits] [mlir] [MLIR][CAPI][Python] Add support for querying memory effect instances (PR #213459)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 1 09:42:16 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Twice (PragmaTwice)
<details>
<summary>Changes</summary>
The current memory-effect bindings are write-only. Python interface implementations can append effects to a temporary
`MemoryEffectInstancesList`, but callers cannot inspect an effect instance or ask an operation for its effects.
This patch adds the missing query support to the C API and exposes it in Python. Effect instances can be inspected and copied, and the Python APIs use regular `MemoryEffectInstance` objects and Python-native `list`s. The Transform helpers follow the same model.
`MemoryEffectInstancesList` is **removed** from both the C and Python APIs. Python implementations and helpers now return effects instead of mutating a supplied list, while the C API passes effects through callbacks (`MlirMemoryEffectInstancesCallback`).
This is a **breaking change**. The migration is limited to replacing output parameters with return values or callbacks.
Assited by: Codex / GPT 5.6 Sol
---
Patch is 52.62 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/213459.diff
15 Files Affected:
- (modified) mlir/include/mlir-c/Dialect/Transform.h (+10-5)
- (modified) mlir/include/mlir-c/Interfaces.h (+61-10)
- (modified) mlir/include/mlir/Bindings/Python/IRInterfaces.h (+13-13)
- (modified) mlir/include/mlir/CAPI/Interfaces.h (-3)
- (modified) mlir/lib/Bindings/Python/DialectTransform.cpp (+41-21)
- (modified) mlir/lib/Bindings/Python/IRInterfaces.cpp (+147-28)
- (modified) mlir/lib/CAPI/Dialect/Transform.cpp (+41-13)
- (modified) mlir/lib/CAPI/Interfaces/Interfaces.cpp (+74-5)
- (modified) mlir/python/mlir/dialects/ext.py (+2-2)
- (modified) mlir/test/CAPI/CMakeLists.txt (+1)
- (modified) mlir/test/CAPI/ir.c (+87-4)
- (modified) mlir/test/CAPI/transform.c (+46)
- (modified) mlir/test/python/dialects/ext.py (+2-2)
- (modified) mlir/test/python/dialects/memory_effects_op_interface.py (+97-34)
- (modified) mlir/test/python/dialects/transform_op_interface.py (+12-8)
``````````diff
diff --git a/mlir/include/mlir-c/Dialect/Transform.h b/mlir/include/mlir-c/Dialect/Transform.h
index cbda09cdbc37b..83f1ed7dc37e9 100644
--- a/mlir/include/mlir-c/Dialect/Transform.h
+++ b/mlir/include/mlir-c/Dialect/Transform.h
@@ -246,25 +246,30 @@ MLIR_CAPI_EXPORTED void mlirPatternDescriptorOpInterfaceAttachFallbackModel(
/// Helper to mark operands as only reading handles.
MLIR_CAPI_EXPORTED void
mlirTransformOnlyReadsHandle(MlirOpOperand *operands, intptr_t numOperands,
- MlirMemoryEffectInstancesList effects);
+ MlirMemoryEffectInstancesCallback callback,
+ void *userData);
/// Helper to mark operands as consuming handles.
MLIR_CAPI_EXPORTED void
mlirTransformConsumesHandle(MlirOpOperand *operands, intptr_t numOperands,
- MlirMemoryEffectInstancesList effects);
+ MlirMemoryEffectInstancesCallback callback,
+ void *userData);
/// Helper to mark results as producing handles.
MLIR_CAPI_EXPORTED void
mlirTransformProducesHandle(MlirValue *results, intptr_t numResults,
- MlirMemoryEffectInstancesList effects);
+ MlirMemoryEffectInstancesCallback callback,
+ void *userData);
/// Helper to mark potential modifications to the payload IR.
MLIR_CAPI_EXPORTED void
-mlirTransformModifiesPayload(MlirMemoryEffectInstancesList effects);
+mlirTransformModifiesPayload(MlirMemoryEffectInstancesCallback callback,
+ void *userData);
/// Helper to mark potential reads from the payload IR.
MLIR_CAPI_EXPORTED void
-mlirTransformOnlyReadsPayload(MlirMemoryEffectInstancesList effects);
+mlirTransformOnlyReadsPayload(MlirMemoryEffectInstancesCallback callback,
+ void *userData);
#ifdef __cplusplus
}
diff --git a/mlir/include/mlir-c/Interfaces.h b/mlir/include/mlir-c/Interfaces.h
index a416a6ab76f87..aa03b7f9e7b66 100644
--- a/mlir/include/mlir-c/Interfaces.h
+++ b/mlir/include/mlir-c/Interfaces.h
@@ -30,7 +30,6 @@ extern "C" {
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
@@ -159,6 +158,10 @@ MLIR_CAPI_EXPORTED MlirMemoryEffect mlirMemoryEffectsReadGet(void);
/// Returns the borrowed singleton instance of the write memory effect.
MLIR_CAPI_EXPORTED MlirMemoryEffect mlirMemoryEffectsWriteGet(void);
+/// Returns the TypeID identifying the concrete type of the given memory effect.
+MLIR_CAPI_EXPORTED MlirTypeID
+mlirMemoryEffectGetEffectID(MlirMemoryEffect effect);
+
/// Returns the borrowed singleton instance of the default side effect
/// resource.
MLIR_CAPI_EXPORTED MlirSideEffectResource
@@ -212,15 +215,53 @@ mlirMemoryEffectInstanceCreateForSymbol(MlirMemoryEffect effect,
bool effectOnFullRegion,
MlirSideEffectResource resource);
-/// Destroys a memory effect instance created by one of the functions above.
+/// Destroys an owned memory effect instance created or cloned by this API.
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);
+/// Creates an owned copy of a memory effect instance. The caller must destroy
+/// the returned instance with `mlirMemoryEffectInstanceDestroy`.
+MLIR_CAPI_EXPORTED MlirMemoryEffectInstance
+mlirMemoryEffectInstanceClone(MlirMemoryEffectInstance instance);
+
+/// Returns the memory effect of the given instance.
+MLIR_CAPI_EXPORTED MlirMemoryEffect
+mlirMemoryEffectInstanceGetEffect(MlirMemoryEffectInstance instance);
+
+/// Returns the side effect resource of the given instance.
+MLIR_CAPI_EXPORTED MlirSideEffectResource
+mlirMemoryEffectInstanceGetResource(MlirMemoryEffectInstance instance);
+
+/// Returns the stage of the given instance.
+MLIR_CAPI_EXPORTED int
+mlirMemoryEffectInstanceGetStage(MlirMemoryEffectInstance instance);
+
+/// Returns true if the given instance has effect on every single value of
+/// the resource.
+MLIR_CAPI_EXPORTED bool mlirMemoryEffectInstanceGetEffectOnFullRegion(
+ MlirMemoryEffectInstance instance);
+
+/// Returns the parameters of the given instance, or a null attribute if there
+/// are no parameters.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirMemoryEffectInstanceGetParameters(MlirMemoryEffectInstance instance);
+
+/// Returns the value (OpOperand, OpResult, or BlockArgument) of the given
+/// instance, or a null value if there is no associated value.
+MLIR_CAPI_EXPORTED MlirValue
+mlirMemoryEffectInstanceGetValue(MlirMemoryEffectInstance instance);
+
+/// Returns the symbol reference of the given instance, or a null attribute if
+/// there is no associated symbol.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirMemoryEffectInstanceGetSymbolRef(MlirMemoryEffectInstance instance);
+
+/// Callback used to return a list of memory effect instances. `effects` points
+/// to `numEffects` consecutive borrowed instances that are only valid for the
+/// duration of the callback. The caller-provided `userData` is forwarded to
+/// the callback.
+typedef void (*MlirMemoryEffectInstancesCallback)(
+ intptr_t numEffects, MlirMemoryEffectInstance *effects, void *userData);
/// Returns the interface TypeID of the MemoryEffectsOpInterface.
MLIR_CAPI_EXPORTED MlirTypeID mlirMemoryEffectsOpInterfaceTypeID(void);
@@ -231,9 +272,12 @@ typedef struct {
void (*construct)(void *userData);
/// Optional destructor for user data. Set to nullptr to disable it.
void (*destruct)(void *userData);
- /// Get memory effects callback.
- void (*getEffects)(MlirOperation op, MlirMemoryEffectInstancesList effects,
- void *userData);
+ /// Get memory effects callback. Implementations return effects by invoking
+ /// `callback` with an array of memory effect instances. The callback copies
+ /// the instances synchronously, so implementations retain ownership of them.
+ void (*getEffects)(MlirOperation op,
+ MlirMemoryEffectInstancesCallback callback,
+ void *callbackUserData, void *userData);
void *userData;
} MlirMemoryEffectsOpInterfaceCallbacks;
@@ -243,6 +287,13 @@ MLIR_CAPI_EXPORTED void mlirMemoryEffectsOpInterfaceAttachFallbackModel(
MlirContext ctx, MlirStringRef opName,
MlirMemoryEffectsOpInterfaceCallbacks callbacks);
+/// Gets the memory effects of the given operation. The operation must
+/// implement the MemoryEffectsOpInterface. Invokes `callback` once with all
+/// effects; the instances are borrowed and only valid during the callback.
+MLIR_CAPI_EXPORTED void mlirMemoryEffectsOpInterfaceGetEffects(
+ MlirOperation operation, MlirMemoryEffectInstancesCallback callback,
+ void *userData);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir/Bindings/Python/IRInterfaces.h b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
index c1047c7e20ab7..ab8c70c4a4c82 100644
--- a/mlir/include/mlir/Bindings/Python/IRInterfaces.h
+++ b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
@@ -163,6 +163,12 @@ class PyMemoryEffectInstance {
public:
explicit PyMemoryEffectInstance(MlirMemoryEffectInstance instance)
: instance(instance) {}
+ PyMemoryEffectInstance(const PyMemoryEffect &effect,
+ const nanobind::object &target,
+ const nanobind::object ¶meters, int stage,
+ bool effectOnFullRegion,
+ const PySideEffectResource &resource);
+ PyMemoryEffectInstance(const PyMemoryEffectInstance &) = delete;
PyMemoryEffectInstance(PyMemoryEffectInstance &&other) noexcept
: instance(other.instance) {
other.instance.ptr = nullptr;
@@ -173,24 +179,18 @@ class PyMemoryEffectInstance {
}
MlirMemoryEffectInstance get() const { return instance; }
+ PyMemoryEffect getEffect() const;
+ PySideEffectResource getResource() const;
+ int getStage() const;
+ bool getEffectOnFullRegion() const;
+ nanobind::object getParameters() const;
+ nanobind::object getValue() const;
+ nanobind::object getSymbolRef() const;
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; }
- operator MlirMemoryEffectInstancesList() const { return effects; }
-
-private:
- MlirMemoryEffectInstancesList effects;
-};
-
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
} // namespace python
} // namespace mlir
diff --git a/mlir/include/mlir/CAPI/Interfaces.h b/mlir/include/mlir/CAPI/Interfaces.h
index 55d850ccd01eb..1d18f6f7d64a8 100644
--- a/mlir/include/mlir/CAPI/Interfaces.h
+++ b/mlir/include/mlir/CAPI/Interfaces.h
@@ -19,9 +19,6 @@
#include "mlir/CAPI/Wrap.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
-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)
diff --git a/mlir/lib/Bindings/Python/DialectTransform.cpp b/mlir/lib/Bindings/Python/DialectTransform.cpp
index a4a91b28522cf..dd8c2d711edb2 100644
--- a/mlir/lib/Bindings/Python/DialectTransform.cpp
+++ b/mlir/lib/Bindings/Python/DialectTransform.cpp
@@ -493,36 +493,55 @@ struct ParamType : PyConcreteType<ParamType> {
//===----------------------------------------------------------------------===//
namespace {
-void onlyReadsHandle(nb::iterable &operands,
- const PyMemoryEffectsInstanceList &effects) {
+void collectMemoryEffectInstances(intptr_t numEffects,
+ MlirMemoryEffectInstance *effects,
+ void *userData) {
+ auto *result = static_cast<std::vector<PyMemoryEffectInstance> *>(userData);
+ result->reserve(result->size() + numEffects);
+ for (intptr_t i = 0; i < numEffects; ++i)
+ result->emplace_back(mlirMemoryEffectInstanceClone(effects[i]));
+}
+
+std::vector<PyMemoryEffectInstance> 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);
+ std::vector<PyMemoryEffectInstance> effects;
+ mlirTransformOnlyReadsHandle(operandsVec.data(), operandsVec.size(),
+ collectMemoryEffectInstances, &effects);
+ return effects;
};
-void consumesHandle(nb::iterable &operands,
- const PyMemoryEffectsInstanceList &effects) {
+std::vector<PyMemoryEffectInstance> 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);
+ std::vector<PyMemoryEffectInstance> effects;
+ mlirTransformConsumesHandle(operandsVec.data(), operandsVec.size(),
+ collectMemoryEffectInstances, &effects);
+ return effects;
};
-void producesHandle(nb::iterable &results,
- const PyMemoryEffectsInstanceList &effects) {
+std::vector<PyMemoryEffectInstance> 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);
+ std::vector<PyMemoryEffectInstance> effects;
+ mlirTransformProducesHandle(resultsVec.data(), resultsVec.size(),
+ collectMemoryEffectInstances, &effects);
+ return effects;
};
-void modifiesPayload(const PyMemoryEffectsInstanceList &effects) {
- mlirTransformModifiesPayload(effects);
+std::vector<PyMemoryEffectInstance> modifiesPayload() {
+ std::vector<PyMemoryEffectInstance> effects;
+ mlirTransformModifiesPayload(collectMemoryEffectInstances, &effects);
+ return effects;
}
-void onlyReadsPayload(const PyMemoryEffectsInstanceList &effects) {
- mlirTransformOnlyReadsPayload(effects);
+std::vector<PyMemoryEffectInstance> onlyReadsPayload() {
+ std::vector<PyMemoryEffectInstance> effects;
+ mlirTransformOnlyReadsPayload(collectMemoryEffectInstances, &effects);
+ return effects;
}
} // namespace
@@ -546,21 +565,22 @@ static void populateDialectTransformSubmodule(nb::module_ &m) {
PyPatternDescriptorOpInterface::bind(m);
m.def("only_reads_handle", onlyReadsHandle,
- "Mark operands as only reading handles.", nb::arg("operands"),
- nb::arg("effects"));
+ "Returns effects marking operands as only reading handles.",
+ nb::arg("operands"));
m.def("consumes_handle", consumesHandle,
- "Mark operands as consuming handles.", nb::arg("operands"),
- nb::arg("effects"));
+ "Returns effects marking operands as consuming handles.",
+ nb::arg("operands"));
- m.def("produces_handle", producesHandle, "Mark results as producing handles.",
- nb::arg("results"), nb::arg("effects"));
+ m.def("produces_handle", producesHandle,
+ "Returns effects marking results as producing handles.",
+ nb::arg("results"));
m.def("modifies_payload", modifiesPayload,
- "Mark the transform as modifying the payload.", nb::arg("effects"));
+ "Returns effects marking potential payload modifications.");
m.def("only_reads_payload", onlyReadsPayload,
- "Mark the transform as only reading the payload.", nb::arg("effects"));
+ "Returns effects marking payload reads.");
}
} // namespace transform
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index d16015acaf3dc..762dffc827043 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -44,13 +44,10 @@ MlirAttribute unwrapOptionalAttribute(const nb::object &attribute) {
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();
+MlirMemoryEffectInstance createMemoryEffectInstance(
+ const PyMemoryEffect &effect, const nb::object &target,
+ const nb::object ¶meters, int stage, bool effectOnFullRegion,
+ const PySideEffectResource &resource) {
MlirAttribute unwrappedParameters = unwrapOptionalAttribute(parameters);
MlirMemoryEffectInstance rawInstance{nullptr};
@@ -93,9 +90,7 @@ void appendMemoryEffectInstance(PyMemoryEffectsInstanceList &effects,
"SymbolRefAttr, or None");
}
}
-
- PyMemoryEffectInstance instance(rawInstance);
- mlirMemoryEffectInstancesListAppend(list, instance.get());
+ return rawInstance;
}
/// Takes in an optional ist of operands and converts them into a std::vector
@@ -176,6 +171,61 @@ wrapRegions(std::optional<std::vector<PyRegion>> regions) {
} // namespace
+PyMemoryEffectInstance::PyMemoryEffectInstance(
+ const PyMemoryEffect &effect, const nb::object &target,
+ const nb::object ¶meters, int stage, bool effectOnFullRegion,
+ const PySideEffectResource &resource)
+ : PyMemoryEffectInstance(createMemoryEffectInstance(
+ effect, target, parameters, stage, effectOnFullRegion, resource)) {}
+
+PyMemoryEffect PyMemoryEffectInstance::getEffect() const {
+ return PyMemoryEffect(mlirMemoryEffectInstanceGetEffect(instance));
+}
+
+PySideEffectResource PyMemoryEffectInstance::getResource() const {
+ return PySideEffectResource(mlirMemoryEffectInstanceGetResource(instance));
+}
+
+int PyMemoryEffectInstance::getStage() const {
+ return mlirMemoryEffectInstanceGetStage(instance);
+}
+
+bool PyMemoryEffectInstance::getEffectOnFullRegion() const {
+ return mlirMemoryEffectInstanceGetEffectOnFullRegion(instance);
+}
+
+nb::object PyMemoryEffectInstance::getParameters() const {
+ MlirAttribute parameters = mlirMemoryEffectInstanceGetParameters(instance);
+ if (mlirAttributeIsNull(parameters))
+ return nb::none();
+ PyMlirContextRef context =
+ PyMlirContext::forContext(mlirAttributeGetContext(parameters));
+ return PyAttribute(context, parameters).maybeDownCast();
+}
+
+nb::object PyMemoryEffectInstance::getValue() const {
+ MlirValue value = mlirMemoryEffectInstanceGetValue(instance);
+ if (mlirValueIsNull(value))
+ return nb::none();
+ MlirOperation owner =
+ mlirValueIsAOpResult(value)
+ ? mlirOpResultGetOwner(value)
+ : mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(value));
+ PyMlirContextRef context =
+ PyMlirContext::forContext(mlirOperationGetContext(owner));
+ return PyValue(PyOperation::forOperation(context, owner), value)
+ .maybeDownCast();
+}
+
+nb::object PyMemoryEffectInstance::getSymbolRef() const {
+ MlirAttribute symbol = mlirMemoryEffectInstanceGetSymbolRef(instance);
+ if (mlirAttributeIsNull(symbol))
+ return nb::none();
+ PyMlirContextRef context =
+ PyMlirContext::forContext(mlirAttributeGetContext(symbol));
+ return PyAttribute(context, symbol).maybeDownCast();
+}
+
/// Python wrapper for InferTypeOpInterface. This interface has only static
/// methods.
class PyInferTypeOpInterface
@@ -494,22 +544,36 @@ class PyMemoryEffectsOpInterface
nb::handle(static_cast<PyObject *>(userData)).dec_ref();
};
callbacks.getEffects = [](MlirOperation op,
- MlirMemoryEffectInstancesList effects,
- void *userData) {
+ MlirMemoryEffectInstancesCallback callback,
+ void *callbackUserData, void *userData) {
nb::handle pyClass(static_cast<PyObject *>(userData));
// Get the 'get_effects' method from the Python class.
auto pyGetEffects =
nb::cast<nb::callable>(nb::getattr(pyClass, "get_effects"));
- PyMemoryEffectsInstanceList effectsWrapper{effects};
-
PyMlirContextRef context =
PyMlirContext::forContext(mlirOperationGetContext(op));
auto opview = PyOperation::forOperation(context, op)->createOpView();
- // Invoke `pyClass.get_effects(op, effects)`.
- pyGetEffects(opview, effectsWrapper);
+ // Invoke `pyClass.get_effects(op)` and pass the resulting instances back
+ // to the C++ interface as a borrowed array.
+ nb::object result = pyGetEffects(opview);
+ nb::iterable iterable;
+ if (!nb::try_cast<nb::iterable>(result, iterable))
+ throw nb::type_error("get_effects must return an iterable");
+
+ std::vector<nb::object> effectObjects;
+ std::vector<MlirMemoryEffectInstance> effects;
+ for (nb::handle object : iterable) {
+ PyMemoryEffectInstance *effect = nullptr;
+ if (!nb::try_cast<PyMemoryEffectInstance *>(object, effect) || !effect)
+ throw nb::type_error(
+ "get_effects must return MemoryEffectInstance objects");
+ effectObjects.push_back(nb::borrow<nb::object>(object));
+ effects.push_back(effect->get());
+ }
+ callback(effects.size(), effects.data(), callbackUserData);
};
mlirMemoryEffectsOpInterfaceAttachFallbackModel(
@@ -517,7 +581,33...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/213459
More information about the Mlir-commits
mailing list