[MLIR][Python] Add ConditionallySpeculatable interface and Pure specifier (#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:
```python
class PureOp(
    TestPure.Operation,
    name="pure",
    traits=[Pure]  # just like in the ODS!
):
    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/GPT5.4
This commit is contained in:
Twice authored and GitHub committed 2026-05-05 11:22:33 +08:00
1 parent 040c2b2101
commit 7c2c06ccfc
7 files changed
+453 -1

No files matched your search

+43
View File
@@ -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
//===---------------------------------------------------------------------===//
+77 -1
View File
@@ -338,6 +338,74 @@ public:
}
};
/// 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 @@ public:
};
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);
+97
View File
@@ -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
//===---------------------------------------------------------------------===//
+20
View File
@@ -35,6 +35,7 @@ __all__ = [
"Region",
"Type",
"Attribute",
"Pure",
"result",
"infer_result",
"operand",
@@ -992,3 +993,22 @@ class Dialect(ir.Dialect):
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)
+1
View File
@@ -43,6 +43,7 @@ endif()
_add_capi_test_executable(mlir-capi-ir-test
ir.c
LINK_LIBS PRIVATE
MLIRCAPIInterfaces
MLIRCAPIIR
MLIRCAPIFunc
MLIRCAPIRegisterEverything
+99
View File
@@ -18,6 +18,7 @@
#include "mlir-c/Diagnostics.h"
#include "mlir-c/Dialect/Func.h"
#include "mlir-c/IntegerSet.h"
#include "mlir-c/Interfaces.h"
#include "mlir-c/RegisterEverything.h"
#include "mlir-c/Support.h"
@@ -2526,6 +2527,102 @@ int testBlockPredecessorsSuccessors(MlirContext ctx) {
return 0;
}
typedef struct {
intptr_t callbackCount;
} SpeculatabilityCallbackData;
static MlirSpeculatability conditionallySpeculatableCallback(MlirOperation op,
void *userData) {
SpeculatabilityCallbackData *data = (SpeculatabilityCallbackData *)userData;
MlirStringRef opName = mlirIdentifierStr(mlirOperationGetName(op));
if (!mlirStringRefEqual(opName,
mlirStringRefCreateFromCString("memref.store")))
return MlirSpeculatabilityNotSpeculatable;
++data->callbackCount;
return MlirSpeculatabilityRecursivelySpeculatable;
}
int testInterfaces(MlirContext ctx) {
// CHECK-LABEL: @testInterfaces
fprintf(stderr, "@testInterfaces\n");
MlirTypeID condSpecTypeID = mlirConditionallySpeculatableOpInterfaceTypeID();
if (mlirTypeIDIsNull(condSpecTypeID)) {
fprintf(stderr, "ERROR: Expected ConditionallySpeculatable type id\n");
return 1;
}
MlirStringRef constantName = mlirStringRefCreateFromCString("arith.constant");
MlirStringRef storeName = mlirStringRefCreateFromCString("memref.store");
if (!mlirOperationImplementsInterfaceStatic(constantName, ctx,
condSpecTypeID)) {
fprintf(stderr, "ERROR: Expected arith.constant to implement "
"ConditionallySpeculatable\n");
return 2;
}
if (mlirOperationImplementsInterfaceStatic(storeName, ctx, condSpecTypeID)) {
fprintf(stderr, "ERROR: Expected memref.store to not implement "
"ConditionallySpeculatable before attach\n");
return 3;
}
MlirLocation loc = mlirLocationUnknownGet(ctx);
MlirType i32 = mlirIntegerTypeGet(ctx, 32);
MlirAttribute zero =
mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : i32"));
MlirNamedAttribute valueAttr = mlirNamedAttributeGet(
mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("value")), zero);
MlirOperationState constantState = mlirOperationStateGet(constantName, loc);
mlirOperationStateAddResults(&constantState, 1, &i32);
mlirOperationStateAddAttributes(&constantState, 1, &valueAttr);
MlirOperation constantOp = mlirOperationCreate(&constantState);
if (!mlirOperationImplementsInterface(constantOp, condSpecTypeID)) {
fprintf(stderr, "ERROR: Expected arith.constant instance to implement "
"ConditionallySpeculatable\n");
return 4;
}
fprintf(
stderr, "arith.constant speculatability: %d\n",
mlirConditionallySpeculatableOpInterfaceGetSpeculatability(constantOp));
// CHECK: arith.constant speculatability: 1
MlirOperationState storeState = mlirOperationStateGet(storeName, loc);
MlirOperation storeOp = mlirOperationCreate(&storeState);
if (mlirOperationImplementsInterface(storeOp, condSpecTypeID)) {
fprintf(stderr, "ERROR: Expected memref.store instance to not implement "
"ConditionallySpeculatable before attach\n");
return 5;
}
SpeculatabilityCallbackData callbackData = {0};
MlirConditionallySpeculatableOpInterfaceCallbacks callbacks = {
.construct = NULL,
.destruct = NULL,
.getSpeculatability = conditionallySpeculatableCallback,
.userData = &callbackData,
};
mlirConditionallySpeculatableOpInterfaceAttachFallbackModel(ctx, storeName,
callbacks);
fprintf(
stderr, "memref.store static after attach: %d\n",
mlirOperationImplementsInterfaceStatic(storeName, ctx, condSpecTypeID));
fprintf(stderr, "memref.store instance after attach: %d\n",
mlirOperationImplementsInterface(storeOp, condSpecTypeID));
fprintf(stderr, "memref.store speculatability: %d\n",
mlirConditionallySpeculatableOpInterfaceGetSpeculatability(storeOp));
fprintf(stderr, "callback count: %" PRIdPTR "\n", callbackData.callbackCount);
// CHECK: memref.store static after attach: 1
// CHECK: memref.store instance after attach: 1
// CHECK: memref.store speculatability: 2
// CHECK: callback count: 1
mlirOperationDestroy(storeOp);
mlirOperationDestroy(constantOp);
return 0;
}
int main(void) {
MlirContext ctx = mlirContextCreate();
registerAllUpstreamDialects(ctx);
@@ -2575,6 +2672,8 @@ int main(void) {
if (testBlockPredecessorsSuccessors(ctx))
return 17;
if (testInterfaces(ctx))
return 18;
// CHECK: DESTROY MAIN CONTEXT
// CHECK: reportResourceDelete: resource_i64_blob
+116
View File
@@ -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
@@ -840,6 +841,121 @@ def testExtDialectWithAttrInOp():
print(module)
# CHECK: TEST: testExtDialectWithInterfaces
@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
@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"):