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  pdl.pattern @match_addi : benefit(1) {
21    %args = operands
22    %results = types
23    %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
24    rewrite %op with "transform.dialect"
25  }
26
27  sequence %arg0 {
28  ^bb1(%arg1: !pdl.operation):
29    %0 = pdl_match @match_addi in %arg1
30    // CHECK: = transform.loop.get_parent_for
31    %1 = transform.loop.get_parent_for %0
32    %2 = transform.loop.get_parent_for %0 { num_loops = 2 }
33    %3 = transform.loop.get_parent_for %0 { num_loops = 3 }
34    transform.test_print_remark_at_operand %1, "third loop"
35    transform.test_print_remark_at_operand %2, "second loop"
36    transform.test_print_remark_at_operand %3, "first loop"
37  }
38}
39
40// -----
41
42func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
43  // expected-note @below {{target op}}
44  arith.addi %arg0, %arg1 : index
45  return
46}
47
48transform.with_pdl_patterns {
49^bb0(%arg0: !pdl.operation):
50  pdl.pattern @match_addi : benefit(1) {
51    %args = operands
52    %results = types
53    %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
54    rewrite %op with "transform.dialect"
55  }
56
57  sequence %arg0 {
58  ^bb1(%arg1: !pdl.operation):
59    %0 = pdl_match @match_addi in %arg1
60    // expected-error @below {{could not find an 'scf.for' parent}}
61    %1 = transform.loop.get_parent_for %0
62  }
63}
64
65// -----
66
67// Outlined functions:
68//
69// CHECK: func @foo(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}})
70// CHECK:   scf.for
71// CHECK:     arith.addi
72//
73// CHECK: func @foo[[SUFFIX:.+]](%{{.+}}, %{{.+}}, %{{.+}})
74// CHECK:   scf.for
75// CHECK:     arith.addi
76//
77// CHECK-LABEL @loop_outline_op
78func.func @loop_outline_op(%arg0: index, %arg1: index, %arg2: index) {
79  // CHECK: scf.for
80  // CHECK-NOT: scf.for
81  // CHECK:   scf.execute_region
82  // CHECK:     func.call @foo
83  scf.for %i = %arg0 to %arg1 step %arg2 {
84    scf.for %j = %arg0 to %arg1 step %arg2 {
85      arith.addi %i, %j : index
86    }
87  }
88  // CHECK: scf.execute_region
89  // CHECK-NOT: scf.for
90  // CHECK:   func.call @foo[[SUFFIX]]
91  scf.for %j = %arg0 to %arg1 step %arg2 {
92    arith.addi %j, %j : index
93  }
94  return
95}
96
97transform.with_pdl_patterns {
98^bb0(%arg0: !pdl.operation):
99  pdl.pattern @match_addi : benefit(1) {
100    %args = operands
101    %results = types
102    %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
103    rewrite %op with "transform.dialect"
104  }
105
106  sequence %arg0 {
107  ^bb1(%arg1: !pdl.operation):
108    %0 = pdl_match @match_addi in %arg1
109    %1 = transform.loop.get_parent_for %0
110    // CHECK: = transform.loop.outline %{{.*}}
111    transform.loop.outline %1 {func_name = "foo"}
112  }
113}
114
115// -----
116
117func.func private @cond() -> i1
118func.func private @body()
119
120func.func @loop_outline_op_multi_region() {
121  // expected-note @below {{target op}}
122  scf.while : () -> () {
123    %0 = func.call @cond() : () -> i1
124    scf.condition(%0)
125  } do {
126  ^bb0:
127    func.call @body() : () -> ()
128    scf.yield
129  }
130  return
131}
132
133transform.with_pdl_patterns {
134^bb0(%arg0: !pdl.operation):
135  pdl.pattern @match_while : benefit(1) {
136    %args = operands
137    %results = types
138    %op = operation "scf.while"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
139    rewrite %op with "transform.dialect"
140  }
141
142  sequence %arg0 {
143  ^bb1(%arg1: !pdl.operation):
144    %0 = pdl_match @match_while in %arg1
145    // expected-error @below {{failed to outline}}
146    transform.loop.outline %0 {func_name = "foo"}
147  }
148}
149
150// -----
151
152// CHECK-LABEL: @loop_peel_op
153func.func @loop_peel_op() {
154  // CHECK: %[[C0:.+]] = arith.constant 0
155  // CHECK: %[[C42:.+]] = arith.constant 42
156  // CHECK: %[[C5:.+]] = arith.constant 5
157  // CHECK: %[[C40:.+]] = arith.constant 40
158  // CHECK: scf.for %{{.+}} = %[[C0]] to %[[C40]] step %[[C5]]
159  // CHECK:   arith.addi
160  // CHECK: scf.for %{{.+}} = %[[C40]] to %[[C42]] step %[[C5]]
161  // CHECK:   arith.addi
162  %0 = arith.constant 0 : index
163  %1 = arith.constant 42 : index
164  %2 = arith.constant 5 : index
165  scf.for %i = %0 to %1 step %2 {
166    arith.addi %i, %i : index
167  }
168  return
169}
170
171transform.with_pdl_patterns {
172^bb0(%arg0: !pdl.operation):
173  pdl.pattern @match_addi : benefit(1) {
174    %args = operands
175    %results = types
176    %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
177    rewrite %op with "transform.dialect"
178  }
179
180  sequence %arg0 {
181  ^bb1(%arg1: !pdl.operation):
182    %0 = pdl_match @match_addi in %arg1
183    %1 = transform.loop.get_parent_for %0
184    transform.loop.peel %1
185  }
186}
187
188// -----
189
190func.func @loop_pipeline_op(%A: memref<?xf32>, %result: memref<?xf32>) {
191  %c0 = arith.constant 0 : index
192  %c1 = arith.constant 1 : index
193  %c4 = arith.constant 4 : index
194  %cf = arith.constant 1.0 : f32
195  // CHECK: memref.load %[[MEMREF:.+]][%{{.+}}]
196  // CHECK: memref.load %[[MEMREF]]
197  // CHECK: arith.addf
198  // CHECK: scf.for
199  // CHECK:   memref.load
200  // CHECK:   arith.addf
201  // CHECK:   memref.store
202  // CHECK: arith.addf
203  // CHECK: memref.store
204  // CHECK: memref.store
205  // expected-remark @below {{transformed}}
206  scf.for %i0 = %c0 to %c4 step %c1 {
207    %A_elem = memref.load %A[%i0] : memref<?xf32>
208    %A1_elem = arith.addf %A_elem, %cf : f32
209    memref.store %A1_elem, %result[%i0] : memref<?xf32>
210  }
211  return
212}
213
214transform.with_pdl_patterns {
215^bb0(%arg0: !pdl.operation):
216  pdl.pattern @match_addf : benefit(1) {
217    %args = operands
218    %results = types
219    %op = operation "arith.addf"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
220    rewrite %op with "transform.dialect"
221  }
222
223  sequence %arg0 {
224  ^bb1(%arg1: !pdl.operation):
225    %0 = pdl_match @match_addf in %arg1
226    %1 = transform.loop.get_parent_for %0
227    %2 = transform.loop.pipeline %1
228    // Verify that the returned handle is usable.
229    transform.test_print_remark_at_operand %2, "transformed"
230  }
231}
232
233// -----
234
235// CHECK-LABEL: @loop_unroll_op
236func.func @loop_unroll_op() {
237  %c0 = arith.constant 0 : index
238  %c42 = arith.constant 42 : index
239  %c5 = arith.constant 5 : index
240  // CHECK: scf.for %[[I:.+]] =
241  scf.for %i = %c0 to %c42 step %c5 {
242    // CHECK-COUNT-4: arith.addi %[[I]]
243    arith.addi %i, %i : index
244  }
245  return
246}
247
248transform.with_pdl_patterns {
249^bb0(%arg0: !pdl.operation):
250  pdl.pattern @match_addi : benefit(1) {
251    %args = operands
252    %results = types
253    %op = operation "arith.addi"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
254    rewrite %op with "transform.dialect"
255  }
256
257  sequence %arg0 {
258  ^bb1(%arg1: !pdl.operation):
259    %0 = pdl_match @match_addi in %arg1
260    %1 = transform.loop.get_parent_for %0
261    transform.loop.unroll %1 { factor = 4 }
262  }
263}
264
265