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