# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
from mlir.dialects.pdl import *
def constructAndPrintInModule(f):
print("\nTEST:", f.__name__)
with Context(), Location.unknown():
module = Module.create()
with InsertionPoint(module.body):
f()
print(module)
return f
# CHECK: module {
# CHECK: pdl.pattern @operations : benefit(1) {
# CHECK: %0 = attribute
# CHECK: %1 = type
# CHECK: %2 = operation {"attr" = %0} -> (%1 : !pdl.type)
# CHECK: %3 = result 0 of %2
# CHECK: %4 = operand
# CHECK: %5 = operation(%3, %4 : !pdl.value, !pdl.value)
# CHECK: rewrite %5 with "rewriter"
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_operations():
pattern = PatternOp(1, "operations")
with InsertionPoint(pattern.body):
attr = AttributeOp()
ty = TypeOp()
op0 = OperationOp(attributes={"attr": attr}, types=[ty])
op0_result = ResultOp(op0, 0)
input = OperandOp()
root = OperationOp(args=[op0_result, input])
RewriteOp(root, "rewriter")
# CHECK: module {
# CHECK: pdl.pattern @rewrite_with_args : benefit(1) {
# CHECK: %0 = operand
# CHECK: %1 = operation(%0 : !pdl.value)
# CHECK: rewrite %1 with "rewriter"(%0 : !pdl.value)
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_rewrite_with_args():
pattern = PatternOp(1, "rewrite_with_args")
with InsertionPoint(pattern.body):
input = OperandOp()
root = OperationOp(args=[input])
RewriteOp(root, "rewriter", args=[input])
# CHECK: module {
# CHECK: pdl.pattern @rewrite_multi_root_optimal : benefit(1) {
# CHECK: %0 = operand
# CHECK: %1 = operand
# CHECK: %2 = type
# CHECK: %3 = operation(%0 : !pdl.value) -> (%2 : !pdl.type)
# CHECK: %4 = result 0 of %3
# CHECK: %5 = operation(%4 : !pdl.value)
# CHECK: %6 = operation(%1 : !pdl.value) -> (%2 : !pdl.type)
# CHECK: %7 = result 0 of %6
# CHECK: %8 = operation(%4, %7 : !pdl.value, !pdl.value)
# CHECK: rewrite with "rewriter"(%5, %8 : !pdl.operation, !pdl.operation)
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_rewrite_multi_root_optimal():
pattern = PatternOp(1, "rewrite_multi_root_optimal")
with InsertionPoint(pattern.body):
input1 = OperandOp()
input2 = OperandOp()
ty = TypeOp()
op1 = OperationOp(args=[input1], types=[ty])
val1 = ResultOp(op1, 0)
root1 = OperationOp(args=[val1])
op2 = OperationOp(args=[input2], types=[ty])
val2 = ResultOp(op2, 0)
root2 = OperationOp(args=[val1, val2])
RewriteOp(name="rewriter", args=[root1, root2])
# CHECK: module {
# CHECK: pdl.pattern @rewrite_multi_root_forced : benefit(1) {
# CHECK: %0 = operand
# CHECK: %1 = operand
# CHECK: %2 = type
# CHECK: %3 = operation(%0 : !pdl.value) -> (%2 : !pdl.type)
# CHECK: %4 = result 0 of %3
# CHECK: %5 = operation(%4 : !pdl.value)
# CHECK: %6 = operation(%1 : !pdl.value) -> (%2 : !pdl.type)
# CHECK: %7 = result 0 of %6
# CHECK: %8 = operation(%4, %7 : !pdl.value, !pdl.value)
# CHECK: rewrite %5 with "rewriter"(%8 : !pdl.operation)
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_rewrite_multi_root_forced():
pattern = PatternOp(1, "rewrite_multi_root_forced")
with InsertionPoint(pattern.body):
input1 = OperandOp()
input2 = OperandOp()
ty = TypeOp()
op1 = OperationOp(args=[input1], types=[ty])
val1 = ResultOp(op1, 0)
root1 = OperationOp(args=[val1])
op2 = OperationOp(args=[input2], types=[ty])
val2 = ResultOp(op2, 0)
root2 = OperationOp(args=[val1, val2])
RewriteOp(root1, name="rewriter", args=[root2])
# CHECK: module {
# CHECK: pdl.pattern @rewrite_add_body : benefit(1) {
# CHECK: %0 = type : i32
# CHECK: %1 = type
# CHECK: %2 = operation -> (%0, %1 : !pdl.type, !pdl.type)
# CHECK: rewrite %2 {
# CHECK: %3 = type
# CHECK: %4 = operation "foo.op" -> (%0, %3 : !pdl.type, !pdl.type)
# CHECK: replace %2 with %4
# CHECK: }
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_rewrite_add_body():
pattern = PatternOp(1, "rewrite_add_body")
with InsertionPoint(pattern.body):
ty1 = TypeOp(IntegerType.get_signless(32))
ty2 = TypeOp()
root = OperationOp(types=[ty1, ty2])
rewrite = RewriteOp(root)
with InsertionPoint(rewrite.add_body()):
ty3 = TypeOp()
newOp = OperationOp(name="foo.op", types=[ty1, ty3])
ReplaceOp(root, with_op=newOp)
# CHECK: module {
# CHECK: pdl.pattern @rewrite_type : benefit(1) {
# CHECK: %0 = type : i32
# CHECK: %1 = type
# CHECK: %2 = operation -> (%0, %1 : !pdl.type, !pdl.type)
# CHECK: rewrite %2 {
# CHECK: %3 = operation "foo.op" -> (%0, %1 : !pdl.type, !pdl.type)
# CHECK: }
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_rewrite_type():
pattern = PatternOp(1, "rewrite_type")
with InsertionPoint(pattern.body):
ty1 = TypeOp(IntegerType.get_signless(32))
ty2 = TypeOp()
root = OperationOp(types=[ty1, ty2])
rewrite = RewriteOp(root)
with InsertionPoint(rewrite.add_body()):
newOp = OperationOp(name="foo.op", types=[ty1, ty2])
# CHECK: module {
# CHECK: pdl.pattern @rewrite_types : benefit(1) {
# CHECK: %0 = types
# CHECK: %1 = operation -> (%0 : !pdl.range<type>)
# CHECK: rewrite %1 {
# CHECK: %2 = types : [i32, i64]
# CHECK: %3 = operation "foo.op" -> (%0, %2 : !pdl.range<type>, !pdl.range<type>)
# CHECK: }
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_rewrite_types():
pattern = PatternOp(1, "rewrite_types")
with InsertionPoint(pattern.body):
types = TypesOp()
root = OperationOp(types=[types])
rewrite = RewriteOp(root)
with InsertionPoint(rewrite.add_body()):
otherTypes = TypesOp(
[IntegerType.get_signless(32), IntegerType.get_signless(64)]
)
newOp = OperationOp(name="foo.op", types=[types, otherTypes])
# CHECK: module {
# CHECK: pdl.pattern @rewrite_operands : benefit(1) {
# CHECK: %0 = types
# CHECK: %1 = operands : %0
# CHECK: %2 = operation(%1 : !pdl.range<value>)
# CHECK: rewrite %2 {
# CHECK: %3 = operation "foo.op" -> (%0 : !pdl.range<type>)
# CHECK: }
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_rewrite_operands():
pattern = PatternOp(1, "rewrite_operands")
with InsertionPoint(pattern.body):
types = TypesOp()
operands = OperandsOp(types)
root = OperationOp(args=[operands])
rewrite = RewriteOp(root)
with InsertionPoint(rewrite.add_body()):
newOp = OperationOp(name="foo.op", types=[types])
# CHECK: module {
# CHECK: pdl.pattern @native_rewrite : benefit(1) {
# CHECK: %0 = operation
# CHECK: rewrite %0 {
# CHECK: apply_native_rewrite "NativeRewrite"(%0 : !pdl.operation)
# CHECK: }
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_native_rewrite():
pattern = PatternOp(1, "native_rewrite")
with InsertionPoint(pattern.body):
root = OperationOp()
rewrite = RewriteOp(root)
with InsertionPoint(rewrite.add_body()):
ApplyNativeRewriteOp([], "NativeRewrite", args=[root])
# CHECK: module {
# CHECK: pdl.pattern @attribute_with_value : benefit(1) {
# CHECK: %0 = operation
# CHECK: rewrite %0 {
# CHECK: %1 = attribute = "value"
# CHECK: apply_native_rewrite "NativeRewrite"(%1 : !pdl.attribute)
# CHECK: }
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_attribute_with_value():
pattern = PatternOp(1, "attribute_with_value")
with InsertionPoint(pattern.body):
root = OperationOp()
rewrite = RewriteOp(root)
with InsertionPoint(rewrite.add_body()):
attr = AttributeOp(value=Attribute.parse('"value"'))
ApplyNativeRewriteOp([], "NativeRewrite", args=[attr])
# CHECK: module {
# CHECK: pdl.pattern @erase : benefit(1) {
# CHECK: %0 = operation
# CHECK: rewrite %0 {
# CHECK: erase %0
# CHECK: }
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_erase():
pattern = PatternOp(1, "erase")
with InsertionPoint(pattern.body):
root = OperationOp()
rewrite = RewriteOp(root)
with InsertionPoint(rewrite.add_body()):
EraseOp(root)
# CHECK: module {
# CHECK: pdl.pattern @operation_results : benefit(1) {
# CHECK: %0 = types
# CHECK: %1 = operation -> (%0 : !pdl.range<type>)
# CHECK: %2 = results of %1
# CHECK: %3 = operation(%2 : !pdl.range<value>)
# CHECK: rewrite %3 with "rewriter"
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_operation_results():
valueRange = RangeType.get(ValueType.get())
pattern = PatternOp(1, "operation_results")
with InsertionPoint(pattern.body):
types = TypesOp()
inputOp = OperationOp(types=[types])
results = ResultsOp(valueRange, inputOp)
root = OperationOp(args=[results])
RewriteOp(root, name="rewriter")
# CHECK: module {
# CHECK: pdl.pattern : benefit(1) {
# CHECK: %0 = type
# CHECK: apply_native_constraint "typeConstraint"(%0 : !pdl.type)
# CHECK: %1 = operation -> (%0 : !pdl.type)
# CHECK: rewrite %1 with "rewrite"
# CHECK: }
# CHECK: }
@constructAndPrintInModule
def test_apply_native_constraint():
pattern = PatternOp(1)
with InsertionPoint(pattern.body):
resultType = TypeOp()
ApplyNativeConstraintOp([], "typeConstraint", args=[resultType])
root = OperationOp(types=[resultType])
RewriteOp(root, name="rewrite")
Codebase Browser
llvm