102 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import transform5from mlir.dialects import pdl6from mlir.dialects.transform import loop7 8 9def run(f):10 with Context(), Location.unknown():11 module = Module.create()12 with InsertionPoint(module.body):13 print("\nTEST:", f.__name__)14 f()15 print(module)16 return f17 18 19@run20def loopOutline():21 sequence = transform.SequenceOp(22 transform.FailurePropagationMode.Propagate,23 [],24 transform.OperationType.get("scf.for"),25 )26 with InsertionPoint(sequence.body):27 loop.LoopOutlineOp(28 transform.AnyOpType.get(),29 transform.AnyOpType.get(),30 sequence.bodyTarget,31 func_name="foo",32 )33 transform.YieldOp()34 # CHECK-LABEL: TEST: loopOutline35 # CHECK: = transform.loop.outline %36 # CHECK: func_name = "foo"37 38 39@run40def loopPeel():41 sequence = transform.SequenceOp(42 transform.FailurePropagationMode.Propagate,43 [],44 transform.OperationType.get("scf.for"),45 )46 with InsertionPoint(sequence.body):47 loop.LoopPeelOp(transform.AnyOpType.get(), transform.AnyOpType.get(), sequence.bodyTarget)48 transform.YieldOp()49 # CHECK-LABEL: TEST: loopPeel50 # CHECK: = transform.loop.peel %51 52@run53def loopPeel_peel_front():54 sequence = transform.SequenceOp(55 transform.FailurePropagationMode.Propagate,56 [],57 transform.OperationType.get("scf.for"),58 )59 with InsertionPoint(sequence.body):60 loop.LoopPeelOp(61 transform.AnyOpType.get(),62 transform.AnyOpType.get(),63 sequence.bodyTarget,64 peel_front=True,65 )66 transform.YieldOp()67 # CHECK-LABEL: TEST: loopPeel_peel_front68 # CHECK: = transform.loop.peel %[[ARG0:.*]] {peel_front = true}69 70 71@run72def loopPipeline():73 sequence = transform.SequenceOp(74 transform.FailurePropagationMode.Propagate,75 [],76 transform.OperationType.get("scf.for"),77 )78 with InsertionPoint(sequence.body):79 loop.LoopPipelineOp(80 pdl.OperationType.get(), sequence.bodyTarget, iteration_interval=381 )82 transform.YieldOp()83 # CHECK-LABEL: TEST: loopPipeline84 # CHECK: = transform.loop.pipeline %85 # CHECK-DAG: iteration_interval = 386 # (read_latency has default value and is not printed)87 88 89@run90def loopUnroll():91 sequence = transform.SequenceOp(92 transform.FailurePropagationMode.Propagate,93 [],94 transform.OperationType.get("scf.for"),95 )96 with InsertionPoint(sequence.body):97 loop.LoopUnrollOp(sequence.bodyTarget, factor=42)98 transform.YieldOp()99 # CHECK-LABEL: TEST: loopUnroll100 # CHECK: transform.loop.unroll %101 # CHECK: factor = 42102