[Mlir-commits] [mlir] [MLIR][Python] Add `ConditionallySpeculatable` interface and `Pure` specifier (PR #195505)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun May 3 00:00:32 PDT 2026
https://github.com/PragmaTwice created https://github.com/llvm/llvm-project/pull/195505
This PR brings two features: the `ConditionallySpeculatable` op interface and the `Pure` specifier for Python-defined ops.
The result is that you can mark an op as pure like: (just like in ODS!)
```python
class PureOp(TestPure.Operation, name="pure", traits=[Pure]):
a: Operand[IntegerType[32]]
b: Operand[IntegerType[32]]
res: Result[IntegerType[32]] = infer_result()
```
Then this op is both `NoMemoryEffect` and `AlwaysSpeculatable`.
Assisted-by: Copilot/GPT 5.5
>From 1df66434232244d5370ca52e0b73a5d75139a701 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sun, 3 May 2026 14:50:57 +0800
Subject: [PATCH 1/2] [MLIR][Python] Add ConditionallySpeculatable interface
and Pure specifier
---
mlir/include/mlir-c/Interfaces.h | 43 ++++++++
mlir/lib/Bindings/Python/IRInterfaces.cpp | 80 +++++++++++++-
mlir/lib/CAPI/Interfaces/Interfaces.cpp | 97 +++++++++++++++++
mlir/python/mlir/dialects/ext.py | 20 ++++
mlir/test/python/dialects/ext.py | 122 +++++++++++++++++++++-
5 files changed, 357 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir-c/Interfaces.h b/mlir/include/mlir-c/Interfaces.h
index 17a812dcd86a9..a5251dc78471f 100644
--- a/mlir/include/mlir-c/Interfaces.h
+++ b/mlir/include/mlir-c/Interfaces.h
@@ -97,6 +97,49 @@ mlirInferShapedTypeOpInterfaceInferReturnTypes(
void *properties, intptr_t nRegions, MlirRegion *regions,
MlirShapedTypeComponentsCallback callback, void *userData);
+//===---------------------------------------------------------------------===//
+// ConditionallySpeculatable
+//===---------------------------------------------------------------------===//
+
+/// Enum representing the speculatability of an operation.
+typedef enum {
+ /// The operation is not speculatable.
+ MlirSpeculatabilityNotSpeculatable,
+ /// The operation is speculatable.
+ MlirSpeculatabilitySpeculatable,
+ /// The operation is speculatable if all nested operations are speculatable.
+ MlirSpeculatabilityRecursivelySpeculatable
+} MlirSpeculatability;
+
+/// Returns the interface TypeID of the ConditionallySpeculatable interface.
+MLIR_CAPI_EXPORTED MlirTypeID
+mlirConditionallySpeculatableOpInterfaceTypeID(void);
+
+/// Callbacks for implementing ConditionallySpeculatable from external code.
+typedef struct {
+ /// Optional constructor for user data. Set to nullptr to disable it.
+ void (*construct)(void *userData);
+ /// Optional destructor for user data. Set to nullptr to disable it.
+ void (*destruct)(void *userData);
+ /// Returns the speculatability of the given operation.
+ MlirSpeculatability (*getSpeculatability)(MlirOperation op, void *userData);
+ void *userData;
+} MlirConditionallySpeculatableOpInterfaceCallbacks;
+
+/// Attach a new FallbackModel for the ConditionallySpeculatable interface to
+/// the named operation. The FallbackModel will call the provided callbacks.
+MLIR_CAPI_EXPORTED void
+mlirConditionallySpeculatableOpInterfaceAttachFallbackModel(
+ MlirContext ctx, MlirStringRef opName,
+ MlirConditionallySpeculatableOpInterfaceCallbacks callbacks);
+
+/// Returns the speculatability of the given operation.
+///
+/// The operation must implement the ConditionallySpeculatable interface.
+MLIR_CAPI_EXPORTED MlirSpeculatability
+mlirConditionallySpeculatableOpInterfaceGetSpeculatability(
+ MlirOperation operation);
+
//===---------------------------------------------------------------------===//
// MemoryEffectsOpInterface
//===---------------------------------------------------------------------===//
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index 561316d863f75..c4e87246bd68c 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -338,6 +338,74 @@ class PyInferShapedTypeOpInterface
}
};
+/// Wrapper around the ConditionallySpeculatable interface.
+class PyConditionallySpeculatableOpInterface
+ : public PyConcreteOpInterface<PyConditionallySpeculatableOpInterface> {
+public:
+ using PyConcreteOpInterface<
+ PyConditionallySpeculatableOpInterface>::PyConcreteOpInterface;
+
+ constexpr static const char *pyClassName = "ConditionallySpeculatable";
+ constexpr static GetTypeIDFunctionTy getInterfaceID =
+ &mlirConditionallySpeculatableOpInterfaceTypeID;
+
+ /// Attach a new ConditionallySpeculatable FallbackModel to the named
+ /// operation. The FallbackModel acts as a trampoline for callbacks on the
+ /// Python class.
+ static void attach(nb::object &target, const std::string &opName,
+ DefaultingPyMlirContext ctx) {
+ MlirConditionallySpeculatableOpInterfaceCallbacks callbacks;
+ callbacks.userData = target.ptr();
+ nb::handle(static_cast<PyObject *>(callbacks.userData)).inc_ref();
+ callbacks.construct = nullptr;
+ callbacks.destruct = [](void *userData) {
+ nb::handle(static_cast<PyObject *>(userData)).dec_ref();
+ };
+ callbacks.getSpeculatability = [](MlirOperation op, void *userData) {
+ nb::handle pyClass(static_cast<PyObject *>(userData));
+
+ auto pyGetSpeculatability =
+ nb::cast<nb::callable>(nb::getattr(pyClass, "get_speculatability"));
+
+ PyMlirContextRef context =
+ PyMlirContext::forContext(mlirOperationGetContext(op));
+ auto opview = PyOperation::forOperation(context, op)->createOpView();
+
+ return nb::cast<MlirSpeculatability>(pyGetSpeculatability(opview));
+ };
+
+ mlirConditionallySpeculatableOpInterfaceAttachFallbackModel(
+ ctx->get(), mlirStringRefCreate(opName.c_str(), opName.size()),
+ callbacks);
+ }
+
+ static void bindDerived(ClassTy &cls) {
+ cls.def(
+ "getSpeculatability",
+ [](PyConditionallySpeculatableOpInterface &self) {
+ if (self.isStatic())
+ throw nb::type_error(
+ "Cannot query speculatability on a static interface");
+ auto operation = self.getOperationObject();
+ auto *pyOperation = nb::cast<PyOperation *>(operation);
+ return mlirConditionallySpeculatableOpInterfaceGetSpeculatability(
+ pyOperation->get());
+ },
+ "Returns the speculatability of the given operation.");
+ cls.attr("attach") = classmethod(
+ [](const nb::object &cls, const nb::object &opName, nb::object target,
+ DefaultingPyMlirContext context) {
+ if (target.is_none())
+ target = cls;
+ return attach(target, nb::cast<std::string>(opName), context);
+ },
+ nb::arg("cls"), nb::arg("op_name"), nb::kw_only(),
+ nb::arg("target").none() = nb::none(),
+ nb::arg("context").none() = nb::none(),
+ "Attach the interface subclass to the given operation name.");
+ }
+};
+
/// Wrapper around the MemoryEffectsOpInterface.
class PyMemoryEffectsOpInterface
: public PyConcreteOpInterface<PyMemoryEffectsOpInterface> {
@@ -401,8 +469,16 @@ class PyMemoryEffectsOpInterface
};
void populateIRInterfaces(nb::module_ &m) {
- nb::class_<PyMemoryEffectsInstanceList>(m, "MemoryEffectInstancesList");
-
+ nb::enum_<MlirSpeculatability>(m, "Speculatability")
+ .value("NotSpeculatable", MlirSpeculatabilityNotSpeculatable)
+ .value("Speculatable", MlirSpeculatabilitySpeculatable)
+ .value("RecursivelySpeculatable",
+ MlirSpeculatabilityRecursivelySpeculatable);
+ auto memoryEffectsInstanceListClass =
+ nb::class_<PyMemoryEffectsInstanceList>(m, "MemoryEffectInstancesList");
+ (void)memoryEffectsInstanceListClass;
+
+ PyConditionallySpeculatableOpInterface::bind(m);
PyInferShapedTypeOpInterface::bind(m);
PyInferTypeOpInterface::bind(m);
PyMemoryEffectsOpInterface::bind(m);
diff --git a/mlir/lib/CAPI/Interfaces/Interfaces.cpp b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
index bfa9fbf3e6217..35a2bd562a8a1 100644
--- a/mlir/lib/CAPI/Interfaces/Interfaces.cpp
+++ b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
@@ -179,6 +179,103 @@ MlirLogicalResult mlirInferShapedTypeOpInterfaceInferReturnTypes(
return mlirLogicalResultSuccess();
}
+//===---------------------------------------------------------------------===//
+// ConditionallySpeculatable
+//===---------------------------------------------------------------------===//
+
+MlirTypeID mlirConditionallySpeculatableOpInterfaceTypeID() {
+ return wrap(ConditionallySpeculatable::getInterfaceID());
+}
+
+/// Fallback model for the ConditionallySpeculatable interface that uses C API
+/// callbacks.
+class ConditionallySpeculatableOpInterfaceFallbackModel
+ : public mlir::ConditionallySpeculatable::FallbackModel<
+ ConditionallySpeculatableOpInterfaceFallbackModel> {
+public:
+ /// Sets the callbacks that this FallbackModel will use.
+ /// NB: the callbacks can only be set through this method as the
+ /// RegisteredOperationName::attachInterface mechanism default-constructs
+ /// the FallbackModel without being able to provide arguments.
+ void
+ setCallbacks(MlirConditionallySpeculatableOpInterfaceCallbacks callbacks) {
+ this->callbacks = callbacks;
+ }
+
+ ~ConditionallySpeculatableOpInterfaceFallbackModel() {
+ if (callbacks.destruct)
+ callbacks.destruct(callbacks.userData);
+ }
+
+ static TypeID getInterfaceID() {
+ return ConditionallySpeculatable::getInterfaceID();
+ }
+
+ static bool classof(const mlir::ConditionallySpeculatable::Concept *op) {
+ // Enable casting back to the FallbackModel from the Interface. This is
+ // necessary as attachInterface(...) default-constructs the FallbackModel
+ // without being able to pass in the callbacks and returns just the Concept.
+ return true;
+ }
+
+ Speculation::Speculatability getSpeculatability(Operation *op) const {
+ assert(callbacks.getSpeculatability &&
+ "getSpeculatability callback not set");
+
+ switch (callbacks.getSpeculatability(wrap(op), callbacks.userData)) {
+ case MlirSpeculatabilityNotSpeculatable:
+ return Speculation::NotSpeculatable;
+ case MlirSpeculatabilitySpeculatable:
+ return Speculation::Speculatable;
+ case MlirSpeculatabilityRecursivelySpeculatable:
+ return Speculation::RecursivelySpeculatable;
+ }
+ llvm_unreachable("unknown speculatability");
+ }
+
+private:
+ MlirConditionallySpeculatableOpInterfaceCallbacks callbacks;
+};
+
+/// Attach a ConditionallySpeculatable FallbackModel to the given named op.
+/// The FallbackModel uses the provided callbacks to implement the interface.
+void mlirConditionallySpeculatableOpInterfaceAttachFallbackModel(
+ MlirContext ctx, MlirStringRef opName,
+ MlirConditionallySpeculatableOpInterfaceCallbacks callbacks) {
+ // Look up the operation definition in the context.
+ std::optional<RegisteredOperationName> opInfo =
+ RegisteredOperationName::lookup(unwrap(opName), unwrap(ctx));
+
+ assert(opInfo.has_value() && "operation not found in context");
+
+ // NB: the following default-constructs the FallbackModel _without_ being able
+ // to provide arguments.
+ opInfo->attachInterface<ConditionallySpeculatableOpInterfaceFallbackModel>();
+ // Cast to get the underlying FallbackModel and set the callbacks.
+ auto *model = cast<ConditionallySpeculatableOpInterfaceFallbackModel>(
+ opInfo
+ ->getInterface<ConditionallySpeculatableOpInterfaceFallbackModel>());
+ assert(model &&
+ "Failed to get ConditionallySpeculatableOpInterfaceFallbackModel");
+ model->setCallbacks(callbacks);
+}
+
+MlirSpeculatability mlirConditionallySpeculatableOpInterfaceGetSpeculatability(
+ MlirOperation operation) {
+ auto iface = dyn_cast<ConditionallySpeculatable>(unwrap(operation));
+ assert(iface && "operation does not implement ConditionallySpeculatable");
+
+ switch (iface.getSpeculatability()) {
+ case Speculation::NotSpeculatable:
+ return MlirSpeculatabilityNotSpeculatable;
+ case Speculation::Speculatable:
+ return MlirSpeculatabilitySpeculatable;
+ case Speculation::RecursivelySpeculatable:
+ return MlirSpeculatabilityRecursivelySpeculatable;
+ }
+ llvm_unreachable("unknown speculatability");
+}
+
//===---------------------------------------------------------------------===//
// MemoryEffectOpInterface
//===---------------------------------------------------------------------===//
diff --git a/mlir/python/mlir/dialects/ext.py b/mlir/python/mlir/dialects/ext.py
index c2efa9bb773cc..10ac002c75d96 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -35,6 +35,7 @@
"Region",
"Type",
"Attribute",
+ "Pure",
"result",
"infer_result",
"operand",
@@ -992,3 +993,22 @@ def load(
for op in cls.operations:
_cext.register_operation(cls, replace=reload)(op)
_cext.register_op_adaptor(op, replace=reload)(op.Adaptor)
+
+
+class Pure:
+ """Always speculatable operation that does not touch memory."""
+
+ class NoMemoryEffect(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ pass
+
+ class AlwaysSpeculatable(ir.ConditionallySpeculatable):
+ @staticmethod
+ def get_speculatability(op):
+ return ir.Speculatability.Speculatable
+
+ @staticmethod
+ def attach(op_name):
+ Pure.NoMemoryEffect.attach(op_name)
+ Pure.AlwaysSpeculatable.attach(op_name)
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index cfdbeb3362735..1362a6fa5b033 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -3,6 +3,7 @@
from mlir.ir import *
from mlir.dialects import arith
from mlir.dialects.ext import *
+from mlir.rewrite import *
from mlir import ir
from typing import Any, Optional, Sequence, TypeVar, Union
import sys
@@ -575,9 +576,9 @@ class MakeArrayOp(TestType.Operation, name="make_array"):
arr: Result[Array]
class MakeArray3Op(TestType.Operation, name="make_array3"):
- arr: Result[
- Array[IntegerType[32], IntegerAttr[IntegerType[32], 3]]
- ] = infer_result()
+ arr: Result[Array[IntegerType[32], IntegerAttr[IntegerType[32], 3]]] = (
+ infer_result()
+ )
with Context(), Location.unknown():
TestType.load()
@@ -840,6 +841,121 @@ class OpWithAttr(TestAttrInOp.Operation, name="op_with_attr"):
print(module)
+# CHECK: TEST: testExtDialectWithInterfaces
+ at run
+def testExtDialectWithInterfaces():
+ class TestIface(Dialect, name="ext_iface"):
+ pass
+
+ class NoMemoryEffectModel(ir.MemoryEffectsOpInterface):
+ @staticmethod
+ def get_effects(op, effects):
+ pass
+
+ class AlwaysSpeculatableModel(ir.ConditionallySpeculatable):
+ @staticmethod
+ def get_speculatability(op):
+ print("get_speculatability opview:", type(op).__name__)
+ return ir.Speculatability.Speculatable
+
+ class PureOp(TestIface.Operation, name="pure"):
+ pass
+
+ with Context(), Location.unknown():
+ TestIface.load()
+ NoMemoryEffectModel.attach(PureOp.OPERATION_NAME)
+ AlwaysSpeculatableModel.attach(PureOp.OPERATION_NAME)
+
+ memory_static = ir.MemoryEffectsOpInterface(PureOp)
+ spec_static = ir.ConditionallySpeculatable(PureOp)
+ # CHECK: static memory iface: MemoryEffectsOpInterface
+ print("static memory iface:", type(memory_static).__name__)
+ # CHECK: static spec iface: ConditionallySpeculatable
+ print("static spec iface:", type(spec_static).__name__)
+
+ module = Module.create()
+ with InsertionPoint(module.body):
+ pure = PureOp()
+
+ memory_iface = ir.MemoryEffectsOpInterface(pure)
+ spec_iface = ir.ConditionallySpeculatable(pure)
+ # CHECK: instance memory iface: MemoryEffectsOpInterface
+ print("instance memory iface:", type(memory_iface).__name__)
+ # CHECK: instance spec iface: ConditionallySpeculatable
+ print("instance spec iface:", type(spec_iface).__name__)
+ # CHECK: get_speculatability opview: PureOp
+ # CHECK: speculatability equals: True
+ print(
+ "speculatability equals:",
+ spec_iface.getSpeculatability() == ir.Speculatability.Speculatable,
+ )
+
+ try:
+ spec_static.getSpeculatability()
+ except TypeError as e:
+ # CHECK: static spec query error: Cannot query speculatability on a static interface
+ print("static spec query error:", e)
+
+
+# CHECK: TEST: testExtDialectWithPure
+ at run
+def testExtDialectWithPure():
+ class TestPure(Dialect, name="ext_pure"):
+ pass
+
+ class PureOp(TestPure.Operation, name="pure", traits=[Pure]):
+ a: Operand[IntegerType[32]]
+ b: Operand[IntegerType[32]]
+ res: Result[IntegerType[32]] = infer_result()
+
+ class NoPureOp(TestPure.Operation, name="no_pure"):
+ a: Operand[IntegerType[32]]
+ b: Operand[IntegerType[32]]
+ res: Result[IntegerType[32]] = infer_result()
+
+ with Context(), Location.unknown():
+ TestPure.load()
+
+ i32 = IntegerType.get(32)
+ module = Module.create()
+ with InsertionPoint(module.body):
+ c1 = arith.constant(i32, 1)
+ c2 = arith.constant(i32, 2)
+ c3 = arith.constant(i32, 3)
+ c4 = arith.constant(i32, 4)
+ p1 = PureOp(c1, c2)
+ p2 = PureOp(c2, c3)
+ p3 = PureOp(c1, c4)
+ np = NoPureOp(p1, p2)
+
+ assert module.operation.verify()
+ # CHECK: module {
+ # CHECK: %c1_i32 = arith.constant 1 : i32
+ # CHECK: %c2_i32 = arith.constant 2 : i32
+ # CHECK: %c3_i32 = arith.constant 3 : i32
+ # CHECK: %c4_i32 = arith.constant 4 : i32
+ # CHECK: %[[P0:.*]] = "ext_pure.pure"(%c1_i32, %c2_i32) : (i32, i32) -> i32
+ # CHECK: %[[P1:.*]] = "ext_pure.pure"(%c2_i32, %c3_i32) : (i32, i32) -> i32
+ # CHECK: %[[P2:.*]] = "ext_pure.pure"(%c1_i32, %c4_i32) : (i32, i32) -> i32
+ # CHECK: %[[NP:.*]] = "ext_pure.no_pure"(%[[P0]], %[[P1]]) : (i32, i32) -> i32
+ # CHECK: }
+ print(module)
+
+ patterns = RewritePatternSet()
+ apply_patterns_and_fold_greedily(module, patterns.freeze())
+ # CHECK: module {
+ # CHECK: %c1_i32 = arith.constant 1 : i32
+ # CHECK: %c2_i32 = arith.constant 2 : i32
+ # CHECK: %c3_i32 = arith.constant 3 : i32
+ # CHECK-NOT: %c4_i32 = arith.constant 4 : i32
+ # CHECK: %[[P0_FOLDED:.*]] = "ext_pure.pure"(%c1_i32, %c2_i32) : (i32, i32) -> i32
+ # CHECK: %[[P1_FOLDED:.*]] = "ext_pure.pure"(%c2_i32, %c3_i32) : (i32, i32) -> i32
+ # CHECK-NOT: "ext_pure.pure"(%c1_i32, %c4_i32)
+ # CHECK: %[[NP_FOLDED:.*]] = "ext_pure.no_pure"(%[[P0_FOLDED]], %[[P1_FOLDED]]) : (i32, i32) -> i32
+ # CHECK: }
+ print(module)
+
+
@run
def testExtDialectFieldSpecifiers():
class TestFieldSpecifiers(Dialect, name="ext_field_specifiers"):
>From 317e32ad071e536657187046799a09980d0cc96c Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sun, 3 May 2026 15:00:13 +0800
Subject: [PATCH 2/2] format
---
mlir/test/python/dialects/ext.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index 1362a6fa5b033..2c9ab4fe321a2 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -576,9 +576,9 @@ class MakeArrayOp(TestType.Operation, name="make_array"):
arr: Result[Array]
class MakeArray3Op(TestType.Operation, name="make_array3"):
- arr: Result[Array[IntegerType[32], IntegerAttr[IntegerType[32], 3]]] = (
- infer_result()
- )
+ arr: Result[
+ Array[IntegerType[32], IntegerAttr[IntegerType[32], 3]]
+ ] = infer_result()
with Context(), Location.unknown():
TestType.load()
More information about the Mlir-commits
mailing list