ops.mlir
9.09 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
// RUN: mlir-opt %s | FileCheck %s
// Verify the printed output can be parsed.
// RUN: mlir-opt %s | mlir-opt | FileCheck %s
// Verify the generic form can be parsed.
// RUN: mlir-opt -mlir-print-op-generic %s | mlir-opt | FileCheck %s
func @std_for(%arg0 : index, %arg1 : index, %arg2 : index) {
scf.for %i0 = %arg0 to %arg1 step %arg2 {
scf.for %i1 = %arg0 to %arg1 step %arg2 {
%min_cmp = cmpi "slt", %i0, %i1 : index
%min = select %min_cmp, %i0, %i1 : index
%max_cmp = cmpi "sge", %i0, %i1 : index
%max = select %max_cmp, %i0, %i1 : index
scf.for %i2 = %min to %max step %i1 {
}
}
}
return
}
// CHECK-LABEL: func @std_for(
// CHECK-NEXT: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
// CHECK-NEXT: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
// CHECK-NEXT: %{{.*}} = cmpi "slt", %{{.*}}, %{{.*}} : index
// CHECK-NEXT: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : index
// CHECK-NEXT: %{{.*}} = cmpi "sge", %{{.*}}, %{{.*}} : index
// CHECK-NEXT: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : index
// CHECK-NEXT: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
func @std_if(%arg0: i1, %arg1: f32) {
scf.if %arg0 {
%0 = addf %arg1, %arg1 : f32
}
return
}
// CHECK-LABEL: func @std_if(
// CHECK-NEXT: scf.if %{{.*}} {
// CHECK-NEXT: %{{.*}} = addf %{{.*}}, %{{.*}} : f32
func @std_if_else(%arg0: i1, %arg1: f32) {
scf.if %arg0 {
%0 = addf %arg1, %arg1 : f32
} else {
%1 = addf %arg1, %arg1 : f32
}
return
}
// CHECK-LABEL: func @std_if_else(
// CHECK-NEXT: scf.if %{{.*}} {
// CHECK-NEXT: %{{.*}} = addf %{{.*}}, %{{.*}} : f32
// CHECK-NEXT: } else {
// CHECK-NEXT: %{{.*}} = addf %{{.*}}, %{{.*}} : f32
func @std_parallel_loop(%arg0 : index, %arg1 : index, %arg2 : index,
%arg3 : index, %arg4 : index) {
%step = constant 1 : index
scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
step (%arg4, %step) {
%min_cmp = cmpi "slt", %i0, %i1 : index
%min = select %min_cmp, %i0, %i1 : index
%max_cmp = cmpi "sge", %i0, %i1 : index
%max = select %max_cmp, %i0, %i1 : index
%zero = constant 0.0 : f32
%int_zero = constant 0 : i32
%red:2 = scf.parallel (%i2) = (%min) to (%max) step (%i1)
init (%zero, %int_zero) -> (f32, i32) {
%one = constant 1.0 : f32
scf.reduce(%one) : f32 {
^bb0(%lhs : f32, %rhs: f32):
%res = addf %lhs, %rhs : f32
scf.reduce.return %res : f32
}
%int_one = constant 1 : i32
scf.reduce(%int_one) : i32 {
^bb0(%lhs : i32, %rhs: i32):
%res = muli %lhs, %rhs : i32
scf.reduce.return %res : i32
}
}
}
return
}
// CHECK-LABEL: func @std_parallel_loop(
// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG3:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG4:[A-Za-z0-9]+]]:
// CHECK: %[[STEP:.*]] = constant 1 : index
// CHECK-NEXT: scf.parallel (%[[I0:.*]], %[[I1:.*]]) = (%[[ARG0]], %[[ARG1]]) to
// CHECK: (%[[ARG2]], %[[ARG3]]) step (%[[ARG4]], %[[STEP]]) {
// CHECK-NEXT: %[[MIN_CMP:.*]] = cmpi "slt", %[[I0]], %[[I1]] : index
// CHECK-NEXT: %[[MIN:.*]] = select %[[MIN_CMP]], %[[I0]], %[[I1]] : index
// CHECK-NEXT: %[[MAX_CMP:.*]] = cmpi "sge", %[[I0]], %[[I1]] : index
// CHECK-NEXT: %[[MAX:.*]] = select %[[MAX_CMP]], %[[I0]], %[[I1]] : index
// CHECK-NEXT: %[[ZERO:.*]] = constant 0.000000e+00 : f32
// CHECK-NEXT: %[[INT_ZERO:.*]] = constant 0 : i32
// CHECK-NEXT: scf.parallel (%{{.*}}) = (%[[MIN]]) to (%[[MAX]])
// CHECK-SAME: step (%[[I1]])
// CHECK-SAME: init (%[[ZERO]], %[[INT_ZERO]]) -> (f32, i32) {
// CHECK-NEXT: %[[ONE:.*]] = constant 1.000000e+00 : f32
// CHECK-NEXT: scf.reduce(%[[ONE]]) : f32 {
// CHECK-NEXT: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):
// CHECK-NEXT: %[[RES:.*]] = addf %[[LHS]], %[[RHS]] : f32
// CHECK-NEXT: scf.reduce.return %[[RES]] : f32
// CHECK-NEXT: }
// CHECK-NEXT: %[[INT_ONE:.*]] = constant 1 : i32
// CHECK-NEXT: scf.reduce(%[[INT_ONE]]) : i32 {
// CHECK-NEXT: ^bb0(%[[LHS:.*]]: i32, %[[RHS:.*]]: i32):
// CHECK-NEXT: %[[RES:.*]] = muli %[[LHS]], %[[RHS]] : i32
// CHECK-NEXT: scf.reduce.return %[[RES]] : i32
// CHECK-NEXT: }
// CHECK-NEXT: scf.yield
// CHECK-NEXT: }
// CHECK-NEXT: scf.yield
func @parallel_explicit_yield(
%arg0: index, %arg1: index, %arg2: index) {
scf.parallel (%i0) = (%arg0) to (%arg1) step (%arg2) {
scf.yield
}
return
}
// CHECK-LABEL: func @parallel_explicit_yield(
// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:
// CHECK-NEXT: scf.parallel (%{{.*}}) = (%[[ARG0]]) to (%[[ARG1]]) step (%[[ARG2]])
// CHECK-NEXT: scf.yield
// CHECK-NEXT: }
// CHECK-NEXT: return
// CHECK-NEXT: }
func @std_if_yield(%arg0: i1, %arg1: f32)
{
%x, %y = scf.if %arg0 -> (f32, f32) {
%0 = addf %arg1, %arg1 : f32
%1 = subf %arg1, %arg1 : f32
scf.yield %0, %1 : f32, f32
} else {
%0 = subf %arg1, %arg1 : f32
%1 = addf %arg1, %arg1 : f32
scf.yield %0, %1 : f32, f32
}
return
}
// CHECK-LABEL: func @std_if_yield(
// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:
// CHECK-NEXT: %{{.*}}:2 = scf.if %[[ARG0]] -> (f32, f32) {
// CHECK-NEXT: %[[T1:.*]] = addf %[[ARG1]], %[[ARG1]]
// CHECK-NEXT: %[[T2:.*]] = subf %[[ARG1]], %[[ARG1]]
// CHECK-NEXT: scf.yield %[[T1]], %[[T2]] : f32, f32
// CHECK-NEXT: } else {
// CHECK-NEXT: %[[T3:.*]] = subf %[[ARG1]], %[[ARG1]]
// CHECK-NEXT: %[[T4:.*]] = addf %[[ARG1]], %[[ARG1]]
// CHECK-NEXT: scf.yield %[[T3]], %[[T4]] : f32, f32
// CHECK-NEXT: }
func @std_for_yield(%arg0 : index, %arg1 : index, %arg2 : index) {
%s0 = constant 0.0 : f32
%result = scf.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0) -> (f32) {
%sn = addf %si, %si : f32
scf.yield %sn : f32
}
return
}
// CHECK-LABEL: func @std_for_yield(
// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:
// CHECK-NEXT: %[[INIT:.*]] = constant
// CHECK-NEXT: %{{.*}} = scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
// CHECK-SAME: iter_args(%[[ITER:.*]] = %[[INIT]]) -> (f32) {
// CHECK-NEXT: %[[NEXT:.*]] = addf %[[ITER]], %[[ITER]] : f32
// CHECK-NEXT: scf.yield %[[NEXT]] : f32
// CHECK-NEXT: }
func @std_for_yield_multi(%arg0 : index, %arg1 : index, %arg2 : index) {
%s0 = constant 0.0 : f32
%t0 = constant 1 : i32
%u0 = constant 1.0 : f32
%result1:3 = scf.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0, %ti = %t0, %ui = %u0) -> (f32, i32, f32) {
%sn = addf %si, %si : f32
%tn = addi %ti, %ti : i32
%un = subf %ui, %ui : f32
scf.yield %sn, %tn, %un : f32, i32, f32
}
return
}
// CHECK-LABEL: func @std_for_yield_multi(
// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:
// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:
// CHECK-NEXT: %[[INIT1:.*]] = constant
// CHECK-NEXT: %[[INIT2:.*]] = constant
// CHECK-NEXT: %[[INIT3:.*]] = constant
// CHECK-NEXT: %{{.*}}:3 = scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
// CHECK-SAME: iter_args(%[[ITER1:.*]] = %[[INIT1]], %[[ITER2:.*]] = %[[INIT2]], %[[ITER3:.*]] = %[[INIT3]]) -> (f32, i32, f32) {
// CHECK-NEXT: %[[NEXT1:.*]] = addf %[[ITER1]], %[[ITER1]] : f32
// CHECK-NEXT: %[[NEXT2:.*]] = addi %[[ITER2]], %[[ITER2]] : i32
// CHECK-NEXT: %[[NEXT3:.*]] = subf %[[ITER3]], %[[ITER3]] : f32
// CHECK-NEXT: scf.yield %[[NEXT1]], %[[NEXT2]], %[[NEXT3]] : f32, i32, f32
func @conditional_reduce(%buffer: memref<1024xf32>, %lb: index, %ub: index, %step: index) -> (f32) {
%sum_0 = constant 0.0 : f32
%c0 = constant 0.0 : f32
%sum = scf.for %iv = %lb to %ub step %step iter_args(%sum_iter = %sum_0) -> (f32) {
%t = load %buffer[%iv] : memref<1024xf32>
%cond = cmpf "ugt", %t, %c0 : f32
%sum_next = scf.if %cond -> (f32) {
%new_sum = addf %sum_iter, %t : f32
scf.yield %new_sum : f32
} else {
scf.yield %sum_iter : f32
}
scf.yield %sum_next : f32
}
return %sum : f32
}
// CHECK-LABEL: func @conditional_reduce(
// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]
// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]
// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]
// CHECK-SAME: %[[ARG3:[A-Za-z0-9]+]]
// CHECK-NEXT: %[[INIT:.*]] = constant
// CHECK-NEXT: %[[ZERO:.*]] = constant
// CHECK-NEXT: %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[ARG1]] to %[[ARG2]] step %[[ARG3]]
// CHECK-SAME: iter_args(%[[ITER:.*]] = %[[INIT]]) -> (f32) {
// CHECK-NEXT: %[[T:.*]] = load %[[ARG0]][%[[IV]]]
// CHECK-NEXT: %[[COND:.*]] = cmpf "ugt", %[[T]], %[[ZERO]]
// CHECK-NEXT: %[[IFRES:.*]] = scf.if %[[COND]] -> (f32) {
// CHECK-NEXT: %[[THENRES:.*]] = addf %[[ITER]], %[[T]]
// CHECK-NEXT: scf.yield %[[THENRES]] : f32
// CHECK-NEXT: } else {
// CHECK-NEXT: scf.yield %[[ITER]] : f32
// CHECK-NEXT: }
// CHECK-NEXT: scf.yield %[[IFRES]] : f32
// CHECK-NEXT: }
// CHECK-NEXT: return %[[RESULT]]