1// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
2
3// CHECK-LABEL: @get_parent_for_op
4func.func @get_parent_for_op(%arg0: index, %arg1: index, %arg2: index) {
5  // expected-remark @below {{first loop}}
6  scf.for %i = %arg0 to %arg1 step %arg2 {
7    // expected-remark @below {{second loop}}
8    scf.for %j = %arg0 to %arg1 step %arg2 {
9      // expected-remark @below {{third loop}}
10      scf.for %k = %arg0 to %arg1 step %arg2 {
11        arith.addi %i, %j : index
12      }
13    }
14  }
15  return
16}
17
18transform.with_pdl_patterns {
19^bb0(%arg0: !pdl.operation):
20  sequence %arg0 {
21  ^bb1(%arg1: !pdl.operation):
22    %0 = transform.structured.match ops{["arith.addi"]} in %arg1
23    // CHECK: = transform.loop.get_parent_for
24    %1 = transform.loop.get_parent_for %0
25    %2 = transform.loop.get_parent_for %0 { num_loops = 2 }
26    %3 = transform.loop.get_parent_for %0 { num_loops = 3 }
27    transform.test_print_remark_at_operand %1, "third loop"
28    transform.test_print_remark_at_operand %2, "second loop"
29    transform.test_print_remark_at_operand %3, "first loop"
30  }
31}
32
33// -----
34
35func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
36  // expected-note @below {{target op}}
37  arith.addi %arg0, %arg1 : index
38  return
39}
40
41transform.with_pdl_patterns {
42^bb0(%arg0: !pdl.operation):
43  sequence %arg0 {
44  ^bb1(%arg1: !pdl.operation):
45    %0 = transform.structured.match ops{["arith.addi"]} in %arg1
46    // expected-error @below {{could not find an 'scf.for' parent}}
47    %1 = transform.loop.get_parent_for %0
48  }
49}
50
51// -----
52
53// Outlined functions:
54//
55// CHECK: func @foo(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}})
56// CHECK:   scf.for
57// CHECK:     arith.addi
58//
59// CHECK: func @foo[[SUFFIX:.+]](%{{.+}}, %{{.+}}, %{{.+}})
60// CHECK:   scf.for
61// CHECK:     arith.addi
62//
63// CHECK-LABEL @loop_outline_op
64func.func @loop_outline_op(%arg0: index, %arg1: index, %arg2: index) {
65  // CHECK: scf.for
66  // CHECK-NOT: scf.for
67  // CHECK:   scf.execute_region
68  // CHECK:     func.call @foo
69  scf.for %i = %arg0 to %arg1 step %arg2 {
70    scf.for %j = %arg0 to %arg1 step %arg2 {
71      arith.addi %i, %j : index
72    }
73  }
74  // CHECK: scf.execute_region
75  // CHECK-NOT: scf.for
76  // CHECK:   func.call @foo[[SUFFIX]]
77  scf.for %j = %arg0 to %arg1 step %arg2 {
78    arith.addi %j, %j : index
79  }
80  return
81}
82
83transform.with_pdl_patterns {
84^bb0(%arg0: !pdl.operation):
85  sequence %arg0 {
86  ^bb1(%arg1: !pdl.operation):
87    %0 = transform.structured.match ops{["arith.addi"]} in %arg1
88    %1 = transform.loop.get_parent_for %0
89    // CHECK: = transform.loop.outline %{{.*}}
90    transform.loop.outline %1 {func_name = "foo"}
91  }
92}
93
94// -----
95
96func.func private @cond() -> i1
97func.func private @body()
98
99func.func @loop_outline_op_multi_region() {
100  // expected-note @below {{target op}}
101  scf.while : () -> () {
102    %0 = func.call @cond() : () -> i1
103    scf.condition(%0)
104  } do {
105  ^bb0:
106    func.call @body() : () -> ()
107    scf.yield
108  }
109  return
110}
111
112transform.with_pdl_patterns {
113^bb0(%arg0: !pdl.operation):
114  sequence %arg0 {
115  ^bb1(%arg1: !pdl.operation):
116    %0 = transform.structured.match ops{["scf.while"]} in %arg1
117    // expected-error @below {{failed to outline}}
118    transform.loop.outline %0 {func_name = "foo"}
119  }
120}
121
122// -----
123
124// CHECK-LABEL: @loop_peel_op
125func.func @loop_peel_op() {
126  // CHECK: %[[C0:.+]] = arith.constant 0
127  // CHECK: %[[C42:.+]] = arith.constant 42
128  // CHECK: %[[C5:.+]] = arith.constant 5
129  // CHECK: %[[C40:.+]] = arith.constant 40
130  // CHECK: scf.for %{{.+}} = %[[C0]] to %[[C40]] step %[[C5]]
131  // CHECK:   arith.addi
132  // CHECK: scf.for %{{.+}} = %[[C40]] to %[[C42]] step %[[C5]]
133  // CHECK:   arith.addi
134  %0 = arith.constant 0 : index
135  %1 = arith.constant 42 : index
136  %2 = arith.constant 5 : index
137  scf.for %i = %0 to %1 step %2 {
138    arith.addi %i, %i : index
139  }
140  return
141}
142
143transform.with_pdl_patterns {
144^bb0(%arg0: !pdl.operation):
145  sequence %arg0 {
146  ^bb1(%arg1: !pdl.operation):
147    %0 = transform.structured.match ops{["arith.addi"]} in %arg1
148    %1 = transform.loop.get_parent_for %0
149    transform.loop.peel %1
150  }
151}
152
153// -----
154
155func.func @loop_pipeline_op(%A: memref<?xf32>, %result: memref<?xf32>) {
156  %c0 = arith.constant 0 : index
157  %c1 = arith.constant 1 : index
158  %c4 = arith.constant 4 : index
159  %cf = arith.constant 1.0 : f32
160  // CHECK: memref.load %[[MEMREF:.+]][%{{.+}}]
161  // CHECK: memref.load %[[MEMREF]]
162  // CHECK: arith.addf
163  // CHECK: scf.for
164  // CHECK:   memref.load
165  // CHECK:   arith.addf
166  // CHECK:   memref.store
167  // CHECK: arith.addf
168  // CHECK: memref.store
169  // CHECK: memref.store
170  // expected-remark @below {{transformed}}
171  scf.for %i0 = %c0 to %c4 step %c1 {
172    %A_elem = memref.load %A[%i0] : memref<?xf32>
173    %A1_elem = arith.addf %A_elem, %cf : f32
174    memref.store %A1_elem, %result[%i0] : memref<?xf32>
175  }
176  return
177}
178
179transform.with_pdl_patterns {
180^bb0(%arg0: !pdl.operation):
181  sequence %arg0 {
182  ^bb1(%arg1: !pdl.operation):
183    %0 = transform.structured.match ops{["arith.addf"]} in %arg1
184    %1 = transform.loop.get_parent_for %0
185    %2 = transform.loop.pipeline %1
186    // Verify that the returned handle is usable.
187    transform.test_print_remark_at_operand %2, "transformed"
188  }
189}
190
191// -----
192
193// CHECK-LABEL: @loop_unroll_op
194func.func @loop_unroll_op() {
195  %c0 = arith.constant 0 : index
196  %c42 = arith.constant 42 : index
197  %c5 = arith.constant 5 : index
198  // CHECK: scf.for %[[I:.+]] =
199  scf.for %i = %c0 to %c42 step %c5 {
200    // CHECK-COUNT-4: arith.addi %[[I]]
201    arith.addi %i, %i : index
202  }
203  return
204}
205
206transform.with_pdl_patterns {
207^bb0(%arg0: !pdl.operation):
208  sequence %arg0 {
209  ^bb1(%arg1: !pdl.operation):
210    %0 = transform.structured.match ops{["arith.addi"]} in %arg1
211    %1 = transform.loop.get_parent_for %0
212    transform.loop.unroll %1 { factor = 4 }
213  }
214}
215
216