[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:
1 parent
040c2b2101
commit
7c2c06ccfc
7 files changed
+453
-1
No files matched your search
@@ -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
|
||||
//===---------------------------------------------------------------------===//
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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
|
||||
//===---------------------------------------------------------------------===//
|
||||
|
||||
@@ -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)
|
||||
@@ -43,6 +43,7 @@ endif()
|
||||
_add_capi_test_executable(mlir-capi-ir-test
|
||||
ir.c
|
||||
LINK_LIBS PRIVATE
|
||||
MLIRCAPIInterfaces
|
||||
MLIRCAPIIR
|
||||
MLIRCAPIFunc
|
||||
MLIRCAPIRegisterEverything
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"):
|
||||
|
||||
Reference in new issue
Block a user