2022-05-30 15:14:02 +02:00
|
|
|
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
|
|
|
|
|
|
from mlir.ir import *
|
|
|
|
|
from mlir.dialects import transform
|
|
|
|
|
from mlir.dialects import pdl
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def run(f):
|
|
|
|
|
with Context(), Location.unknown():
|
|
|
|
|
module = Module.create()
|
|
|
|
|
with InsertionPoint(module.body):
|
|
|
|
|
print("\nTEST:", f.__name__)
|
|
|
|
|
f()
|
|
|
|
|
print(module)
|
|
|
|
|
return f
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@run
|
|
|
|
|
def testSequenceOp():
|
2022-08-12 13:53:46 +00:00
|
|
|
sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
|
|
|
|
|
[pdl.OperationType.get()])
|
2022-05-30 15:14:02 +02:00
|
|
|
with InsertionPoint(sequence.body):
|
|
|
|
|
transform.YieldOp([sequence.bodyTarget])
|
|
|
|
|
# CHECK-LABEL: TEST: testSequenceOp
|
2022-08-12 13:53:46 +00:00
|
|
|
# CHECK: = transform.sequence failures(propagate) {
|
2022-05-30 15:14:02 +02:00
|
|
|
# CHECK: ^{{.*}}(%[[ARG0:.+]]: !pdl.operation):
|
|
|
|
|
# CHECK: yield %[[ARG0]] : !pdl.operation
|
|
|
|
|
# CHECK: } : !pdl.operation
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@run
|
|
|
|
|
def testNestedSequenceOp():
|
2022-08-12 13:53:46 +00:00
|
|
|
sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE)
|
2022-05-30 15:14:02 +02:00
|
|
|
with InsertionPoint(sequence.body):
|
2022-08-12 13:53:46 +00:00
|
|
|
nested = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, sequence.bodyTarget)
|
2022-05-30 15:14:02 +02:00
|
|
|
with InsertionPoint(nested.body):
|
2022-08-12 13:53:46 +00:00
|
|
|
doubly_nested = transform.SequenceOp(
|
|
|
|
|
transform.FailurePropagationMode.PROPAGATE,
|
|
|
|
|
[pdl.OperationType.get()], nested.bodyTarget)
|
2022-05-30 15:14:02 +02:00
|
|
|
with InsertionPoint(doubly_nested.body):
|
|
|
|
|
transform.YieldOp([doubly_nested.bodyTarget])
|
|
|
|
|
transform.YieldOp()
|
|
|
|
|
transform.YieldOp()
|
|
|
|
|
# CHECK-LABEL: TEST: testNestedSequenceOp
|
2022-08-12 13:53:46 +00:00
|
|
|
# CHECK: transform.sequence failures(propagate) {
|
2022-05-30 15:14:02 +02:00
|
|
|
# CHECK: ^{{.*}}(%[[ARG0:.+]]: !pdl.operation):
|
2022-08-12 13:53:46 +00:00
|
|
|
# CHECK: sequence %[[ARG0]] failures(propagate) {
|
2022-05-30 15:14:02 +02:00
|
|
|
# CHECK: ^{{.*}}(%[[ARG1:.+]]: !pdl.operation):
|
2022-08-12 13:53:46 +00:00
|
|
|
# CHECK: = sequence %[[ARG1]] failures(propagate) {
|
2022-05-30 15:14:02 +02:00
|
|
|
# CHECK: ^{{.*}}(%[[ARG2:.+]]: !pdl.operation):
|
|
|
|
|
# CHECK: yield %[[ARG2]] : !pdl.operation
|
|
|
|
|
# CHECK: } : !pdl.operation
|
|
|
|
|
# CHECK: }
|
|
|
|
|
# CHECK: }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@run
|
|
|
|
|
def testTransformPDLOps():
|
|
|
|
|
withPdl = transform.WithPDLPatternsOp()
|
|
|
|
|
with InsertionPoint(withPdl.body):
|
2022-08-12 13:53:46 +00:00
|
|
|
sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
|
|
|
|
|
[pdl.OperationType.get()],
|
2022-05-30 15:14:02 +02:00
|
|
|
withPdl.bodyTarget)
|
|
|
|
|
with InsertionPoint(sequence.body):
|
|
|
|
|
match = transform.PDLMatchOp(sequence.bodyTarget, "pdl_matcher")
|
|
|
|
|
transform.YieldOp(match)
|
|
|
|
|
# CHECK-LABEL: TEST: testTransformPDLOps
|
|
|
|
|
# CHECK: transform.with_pdl_patterns {
|
|
|
|
|
# CHECK: ^{{.*}}(%[[ARG0:.+]]: !pdl.operation):
|
2022-08-12 13:53:46 +00:00
|
|
|
# CHECK: = sequence %[[ARG0]] failures(propagate) {
|
2022-05-30 15:14:02 +02:00
|
|
|
# CHECK: ^{{.*}}(%[[ARG1:.+]]: !pdl.operation):
|
|
|
|
|
# CHECK: %[[RES:.+]] = pdl_match @pdl_matcher in %[[ARG1]]
|
|
|
|
|
# CHECK: yield %[[RES]] : !pdl.operation
|
|
|
|
|
# CHECK: } : !pdl.operation
|
|
|
|
|
# CHECK: }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@run
|
|
|
|
|
def testGetClosestIsolatedParentOp():
|
2022-08-12 13:53:46 +00:00
|
|
|
sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE)
|
2022-05-30 15:14:02 +02:00
|
|
|
with InsertionPoint(sequence.body):
|
|
|
|
|
transform.GetClosestIsolatedParentOp(sequence.bodyTarget)
|
|
|
|
|
transform.YieldOp()
|
|
|
|
|
# CHECK-LABEL: TEST: testGetClosestIsolatedParentOp
|
|
|
|
|
# CHECK: transform.sequence
|
|
|
|
|
# CHECK: ^{{.*}}(%[[ARG1:.+]]: !pdl.operation):
|
|
|
|
|
# CHECK: = get_closest_isolated_parent %[[ARG1]]
|
2022-07-07 13:11:34 +02:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@run
|
|
|
|
|
def testMergeHandlesOp():
|
2022-08-12 13:53:46 +00:00
|
|
|
sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE)
|
2022-07-07 13:11:34 +02:00
|
|
|
with InsertionPoint(sequence.body):
|
|
|
|
|
transform.MergeHandlesOp([sequence.bodyTarget])
|
|
|
|
|
transform.YieldOp()
|
|
|
|
|
# CHECK-LABEL: TEST: testMergeHandlesOp
|
|
|
|
|
# CHECK: transform.sequence
|
|
|
|
|
# CHECK: ^{{.*}}(%[[ARG1:.+]]: !pdl.operation):
|
|
|
|
|
# CHECK: = merge_handles %[[ARG1]]
|
2022-07-07 15:55:23 +02:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@run
|
|
|
|
|
def testReplicateOp():
|
|
|
|
|
with_pdl = transform.WithPDLPatternsOp()
|
|
|
|
|
with InsertionPoint(with_pdl.body):
|
2022-08-12 13:53:46 +00:00
|
|
|
sequence = transform.SequenceOp(
|
|
|
|
|
transform.FailurePropagationMode.PROPAGATE, with_pdl.bodyTarget)
|
2022-07-07 15:55:23 +02:00
|
|
|
with InsertionPoint(sequence.body):
|
|
|
|
|
m1 = transform.PDLMatchOp(sequence.bodyTarget, "first")
|
|
|
|
|
m2 = transform.PDLMatchOp(sequence.bodyTarget, "second")
|
|
|
|
|
transform.ReplicateOp(m1, [m2])
|
|
|
|
|
transform.YieldOp()
|
|
|
|
|
# CHECK-LABEL: TEST: testReplicateOp
|
|
|
|
|
# CHECK: %[[FIRST:.+]] = pdl_match
|
|
|
|
|
# CHECK: %[[SECOND:.+]] = pdl_match
|
|
|
|
|
# CHECK: %{{.*}} = replicate num(%[[FIRST]]) %[[SECOND]]
|