1*5f0d4f20SAlex Zinenko# RUN: %PYTHON %s | FileCheck %s 2*5f0d4f20SAlex Zinenko 3*5f0d4f20SAlex Zinenkofrom mlir.ir import * 4*5f0d4f20SAlex Zinenkofrom mlir.dialects import transform 5*5f0d4f20SAlex Zinenkofrom mlir.dialects import pdl 6*5f0d4f20SAlex Zinenkofrom mlir.dialects.transform import loop 7*5f0d4f20SAlex Zinenko 8*5f0d4f20SAlex Zinenko 9*5f0d4f20SAlex Zinenkodef run(f): 10*5f0d4f20SAlex Zinenko with Context(), Location.unknown(): 11*5f0d4f20SAlex Zinenko module = Module.create() 12*5f0d4f20SAlex Zinenko with InsertionPoint(module.body): 13*5f0d4f20SAlex Zinenko print("\nTEST:", f.__name__) 14*5f0d4f20SAlex Zinenko f() 15*5f0d4f20SAlex Zinenko print(module) 16*5f0d4f20SAlex Zinenko return f 17*5f0d4f20SAlex Zinenko 18*5f0d4f20SAlex Zinenko 19*5f0d4f20SAlex Zinenko@run 20*5f0d4f20SAlex Zinenkodef getParentLoop(): 21*5f0d4f20SAlex Zinenko sequence = transform.SequenceOp() 22*5f0d4f20SAlex Zinenko with InsertionPoint(sequence.body): 23*5f0d4f20SAlex Zinenko loop.GetParentForOp(sequence.bodyTarget, num_loops=2) 24*5f0d4f20SAlex Zinenko transform.YieldOp() 25*5f0d4f20SAlex Zinenko # CHECK-LABEL: TEST: getParentLoop 26*5f0d4f20SAlex Zinenko # CHECK: = transform.loop.get_parent_for % 27*5f0d4f20SAlex Zinenko # CHECK: num_loops = 2 28*5f0d4f20SAlex Zinenko 29*5f0d4f20SAlex Zinenko 30*5f0d4f20SAlex Zinenko@run 31*5f0d4f20SAlex Zinenkodef loopOutline(): 32*5f0d4f20SAlex Zinenko sequence = transform.SequenceOp() 33*5f0d4f20SAlex Zinenko with InsertionPoint(sequence.body): 34*5f0d4f20SAlex Zinenko loop.LoopOutlineOp(sequence.bodyTarget, func_name="foo") 35*5f0d4f20SAlex Zinenko transform.YieldOp() 36*5f0d4f20SAlex Zinenko # CHECK-LABEL: TEST: loopOutline 37*5f0d4f20SAlex Zinenko # CHECK: = transform.loop.outline % 38*5f0d4f20SAlex Zinenko # CHECK: func_name = "foo" 39*5f0d4f20SAlex Zinenko 40*5f0d4f20SAlex Zinenko 41*5f0d4f20SAlex Zinenko@run 42*5f0d4f20SAlex Zinenkodef loopPeel(): 43*5f0d4f20SAlex Zinenko sequence = transform.SequenceOp() 44*5f0d4f20SAlex Zinenko with InsertionPoint(sequence.body): 45*5f0d4f20SAlex Zinenko loop.LoopPeelOp(sequence.bodyTarget) 46*5f0d4f20SAlex Zinenko transform.YieldOp() 47*5f0d4f20SAlex Zinenko # CHECK-LABEL: TEST: loopPeel 48*5f0d4f20SAlex Zinenko # CHECK: = transform.loop.peel % 49*5f0d4f20SAlex Zinenko 50*5f0d4f20SAlex Zinenko 51*5f0d4f20SAlex Zinenko@run 52*5f0d4f20SAlex Zinenkodef loopPipeline(): 53*5f0d4f20SAlex Zinenko sequence = transform.SequenceOp() 54*5f0d4f20SAlex Zinenko with InsertionPoint(sequence.body): 55*5f0d4f20SAlex Zinenko loop.LoopPipelineOp(sequence.bodyTarget, iteration_interval=3) 56*5f0d4f20SAlex Zinenko transform.YieldOp() 57*5f0d4f20SAlex Zinenko # CHECK-LABEL: TEST: loopPipeline 58*5f0d4f20SAlex Zinenko # CHECK: = transform.loop.pipeline % 59*5f0d4f20SAlex Zinenko # CHECK-DAG: iteration_interval = 3 60*5f0d4f20SAlex Zinenko # CHECK-DAG: read_latency = 10 61*5f0d4f20SAlex Zinenko 62*5f0d4f20SAlex Zinenko 63*5f0d4f20SAlex Zinenko@run 64*5f0d4f20SAlex Zinenkodef loopUnroll(): 65*5f0d4f20SAlex Zinenko sequence = transform.SequenceOp() 66*5f0d4f20SAlex Zinenko with InsertionPoint(sequence.body): 67*5f0d4f20SAlex Zinenko loop.LoopUnrollOp(sequence.bodyTarget, factor=42) 68*5f0d4f20SAlex Zinenko transform.YieldOp() 69*5f0d4f20SAlex Zinenko # CHECK-LABEL: TEST: loopUnroll 70*5f0d4f20SAlex Zinenko # CHECK: transform.loop.unroll % 71*5f0d4f20SAlex Zinenko # CHECK: factor = 42 72