1 //===- builder-api-test.cpp - Tests for Declarative Builder APIs ----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 // RUN: mlir-edsc-builder-api-test | FileCheck %s
10
11 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
12 #include "mlir/Dialect/Linalg/EDSC/Builders.h"
13 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
14 #include "mlir/Dialect/SCF/EDSC/Intrinsics.h"
15 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
16 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
17 #include "mlir/EDSC/Builders.h"
18 #include "mlir/IR/AffineExpr.h"
19 #include "mlir/IR/Builders.h"
20 #include "mlir/IR/BuiltinOps.h"
21 #include "mlir/IR/BuiltinTypes.h"
22 #include "mlir/IR/IntegerSet.h"
23 #include "mlir/IR/MLIRContext.h"
24 #include "mlir/IR/Types.h"
25 #include "mlir/Pass/Pass.h"
26 #include "mlir/Pass/PassManager.h"
27 #include "mlir/Transforms/LoopUtils.h"
28 #include "mlir/Transforms/Passes.h"
29
30 #include "APITest.h"
31
32 #include "llvm/Support/raw_ostream.h"
33
34 using namespace mlir;
35 using namespace mlir::edsc;
36 using namespace mlir::edsc::intrinsics;
37
globalContext()38 static MLIRContext &globalContext() {
39 static thread_local MLIRContext context;
40 static thread_local bool initOnce = [&]() {
41 // clang-format off
42 context.loadDialect<AffineDialect,
43 scf::SCFDialect,
44 linalg::LinalgDialect,
45 StandardOpsDialect,
46 vector::VectorDialect>();
47 // clang-format on
48 return true;
49 }();
50 (void)initOnce;
51 context.allowUnregisteredDialects();
52 return context;
53 }
54
makeFunction(StringRef name,ArrayRef<Type> results={},ArrayRef<Type> args={})55 static FuncOp makeFunction(StringRef name, ArrayRef<Type> results = {},
56 ArrayRef<Type> args = {}) {
57 auto &ctx = globalContext();
58 auto function = FuncOp::create(UnknownLoc::get(&ctx), name,
59 FunctionType::get(args, results, &ctx));
60 function.addEntryBlock();
61 return function;
62 }
63
TEST_FUNC(builder_dynamic_for_func_args)64 TEST_FUNC(builder_dynamic_for_func_args) {
65 auto indexType = IndexType::get(&globalContext());
66 auto f32Type = FloatType::getF32(&globalContext());
67 auto f =
68 makeFunction("builder_dynamic_for_func_args", {}, {indexType, indexType});
69
70 OpBuilder builder(f.getBody());
71 ScopedContext scope(builder, f.getLoc());
72 Value lb(f.getArgument(0)), ub(f.getArgument(1));
73 Value f7(std_constant_float(llvm::APFloat(7.0f), f32Type));
74 Value f13(std_constant_float(llvm::APFloat(13.0f), f32Type));
75 Value i7(std_constant_int(7, 32));
76 Value i13(std_constant_int(13, 32));
77 affineLoopBuilder(lb, ub, 3, [&](Value i) {
78 using namespace edsc::op;
79 lb *std_constant_index(3) + ub;
80 lb + std_constant_index(3);
81 affineLoopBuilder(lb, ub, 2, [&](Value j) {
82 ceilDiv(std_constant_index(31) * floorDiv(i + j * std_constant_index(3),
83 std_constant_index(32)),
84 std_constant_index(32));
85 ((f7 + f13) / f7) % f13 - f7 *f13;
86 ((i7 + i13) / i7) % i13 - i7 *i13;
87 });
88 });
89
90 // clang-format off
91 // CHECK-LABEL: func @builder_dynamic_for_func_args(%{{.*}}: index, %{{.*}}: index) {
92 // CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) step 3 {
93 // CHECK: {{.*}} = affine.apply affine_map<()[s0] -> (s0 * 3)>()[%{{.*}}]
94 // CHECK: {{.*}} = affine.apply affine_map<()[s0, s1] -> (s1 + s0 * 3)>()[%{{.*}}, %{{.*}}]
95 // CHECK: {{.*}} = affine.apply affine_map<()[s0] -> (s0 + 3)>()[%{{.*}}]
96 // CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) step 2 {
97 // CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> ((d0 + d1 * 3) floordiv 32)>(%{{.*}}, %{{.*}})
98 // CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> (((d0 + d1 * 3) floordiv 32) * 31)>(%{{.*}}, %{{.*}})
99 // CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> ((((d0 + d1 * 3) floordiv 32) * 31) ceildiv 32)>(%{{.*}}, %{{.*}})
100 // CHECK-DAG: [[rf1:%[0-9]+]] = addf {{.*}}, {{.*}} : f32
101 // CHECK-DAG: [[rf2:%[0-9]+]] = divf [[rf1]], {{.*}} : f32
102 // CHECK-DAG: [[rf3:%[0-9]+]] = remf [[rf2]], {{.*}} : f32
103 // CHECK-DAG: [[rf4:%[0-9]+]] = mulf {{.*}}, {{.*}} : f32
104 // CHECK: {{.*}} = subf [[rf3]], [[rf4]] : f32
105 // CHECK-DAG: [[ri1:%[0-9]+]] = addi {{.*}}, {{.*}} : i32
106 // CHECK-DAG: [[ri2:%[0-9]+]] = divi_signed [[ri1]], {{.*}} : i32
107 // CHECK-DAG: [[ri3:%[0-9]+]] = remi_signed [[ri2]], {{.*}} : i32
108 // CHECK-DAG: [[ri4:%[0-9]+]] = muli {{.*}}, {{.*}} : i32
109 // CHECK: {{.*}} = subi [[ri3]], [[ri4]] : i32
110 // clang-format on
111 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
112 f.erase();
113 }
114
TEST_FUNC(builder_dynamic_for)115 TEST_FUNC(builder_dynamic_for) {
116 auto indexType = IndexType::get(&globalContext());
117 auto f = makeFunction("builder_dynamic_for", {},
118 {indexType, indexType, indexType, indexType});
119
120 OpBuilder builder(f.getBody());
121 ScopedContext scope(builder, f.getLoc());
122 Value i, a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
123 d(f.getArgument(3));
124 using namespace edsc::op;
125 affineLoopBuilder(a - b, c + d, 2);
126
127 // clang-format off
128 // CHECK-LABEL: func @builder_dynamic_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
129 // CHECK-DAG: [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
130 // CHECK-DAG: [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
131 // CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>([[r0]]) to affine_map<(d0) -> (d0)>([[r1]]) step 2 {
132 // clang-format on
133 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
134 f.erase();
135 }
136
TEST_FUNC(builder_loop_for)137 TEST_FUNC(builder_loop_for) {
138 auto indexType = IndexType::get(&globalContext());
139 auto f = makeFunction("builder_loop_for", {},
140 {indexType, indexType, indexType, indexType});
141
142 OpBuilder builder(f.getBody());
143 ScopedContext scope(builder, f.getLoc());
144 Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
145 d(f.getArgument(3));
146 using namespace edsc::op;
147 loopNestBuilder(a - b, c + d, a);
148
149 // clang-format off
150 // CHECK-LABEL: func @builder_loop_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
151 // CHECK-DAG: [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
152 // CHECK-DAG: [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
153 // CHECK-NEXT: scf.for %{{.*}} = [[r0]] to [[r1]] step {{.*}} {
154 // clang-format on
155 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
156 f.erase();
157 }
158
TEST_FUNC(builder_max_min_for)159 TEST_FUNC(builder_max_min_for) {
160 auto indexType = IndexType::get(&globalContext());
161 auto f = makeFunction("builder_max_min_for", {},
162 {indexType, indexType, indexType, indexType});
163
164 OpBuilder builder(f.getBody());
165 ScopedContext scope(builder, f.getLoc());
166 Value lb1(f.getArgument(0)), lb2(f.getArgument(1)), ub1(f.getArgument(2)),
167 ub2(f.getArgument(3));
168 affineLoopBuilder({lb1, lb2}, {ub1, ub2}, 1);
169 std_ret();
170
171 // clang-format off
172 // CHECK-LABEL: func @builder_max_min_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
173 // CHECK: affine.for %{{.*}} = max affine_map<(d0, d1) -> (d0, d1)>(%{{.*}}, %{{.*}}) to min affine_map<(d0, d1) -> (d0, d1)>(%{{.*}}, %{{.*}}) {
174 // CHECK: return
175 // clang-format on
176 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
177 f.erase();
178 }
179
TEST_FUNC(builder_affine_for_iter_args)180 TEST_FUNC(builder_affine_for_iter_args) {
181 auto indexType = IndexType::get(&globalContext());
182 auto f = makeFunction("builder_affine_for_iter_args", {},
183 {indexType, indexType, indexType});
184
185 OpBuilder builder(f.getBody());
186 ScopedContext scope(builder, f.getLoc());
187 Value i, lb_1(f.getArgument(0)), ub_1(f.getArgument(1)),
188 ub_2(f.getArgument(2));
189 Value c32(std_constant_int(32, 32));
190 Value c42(std_constant_int(42, 32));
191 using namespace edsc::op;
192 affineLoopBuilder(
193 lb_1, {ub_1, ub_2}, 2, {c32, c42}, [&](Value iv, ValueRange args) {
194 Value sum(args[0] + args[1]);
195 builder.create<AffineYieldOp>(f.getLoc(), ValueRange({args[1], sum}));
196 });
197
198 // clang-format off
199 // CHECK-LABEL: func @builder_affine_for_iter_args
200 // CHECK: (%[[lb_1:.*]]: index, %[[ub_1:.*]]: index, %[[ub_2:.*]]: index) {
201 // CHECK-NEXT: %[[c32:.*]] = constant 32 : i32
202 // CHECK-NEXT: %[[c42:.*]] = constant 42 : i32
203 // CHECK-NEXT: %{{.*}} = affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to min affine_map<(d0, d1) -> (d0, d1)>(%[[ub_1]], %[[ub_2]]) step 2 iter_args(%[[iarg_1:.*]] = %[[c32]], %[[iarg_2:.*]] = %[[c42]]) -> (i32, i32) {
204 // CHECK-NEXT: %[[sum:.*]] = addi %[[iarg_1]], %[[iarg_2]] : i32
205 // CHECK-NEXT: affine.yield %[[iarg_2]], %[[sum]] : i32, i32
206 // CHECK-NEXT: }
207 // clang-format on
208 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
209 f.erase();
210 }
211
TEST_FUNC(builder_block_append)212 TEST_FUNC(builder_block_append) {
213 using namespace edsc::op;
214 auto f = makeFunction("builder_blocks");
215
216 OpBuilder builder(f.getBody());
217 ScopedContext scope(builder, f.getLoc());
218
219 Block *b =
220 buildInNewBlock(TypeRange(), [&](ValueRange) { std_constant_index(0); });
221 appendToBlock(b, [&](ValueRange) { std_constant_index(1); });
222 appendToBlock(b, [&](ValueRange) { std_ret(); });
223 // Get back to entry block and add a branch into "b".
224 appendToBlock(&f.front(), [&](ValueRange) { std_br(b, {}); });
225
226 // clang-format off
227 // CHECK-LABEL: @builder_blocks
228 // CHECK-NEXT: br ^bb1
229 // CHECK-NEXT: ^bb1: // pred: ^bb0
230 // CHECK-NEXT: constant 0 : index
231 // CHECK-NEXT: constant 1 : index
232 // CHECK-NEXT: return
233 // CHECK-NEXT: }
234 // clang-format on
235 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
236 f.erase();
237 }
238
TEST_FUNC(builder_blocks)239 TEST_FUNC(builder_blocks) {
240 using namespace edsc::op;
241 auto f = makeFunction("builder_blocks");
242
243 OpBuilder builder(f.getBody());
244 ScopedContext scope(builder, f.getLoc());
245 Value c1(std_constant_int(42, 32)), c2(std_constant_int(1234, 32));
246 ReturnOp ret = std_ret();
247
248 Block *b1 = createBlock({c1.getType(), c1.getType()});
249 Block *b2 = buildInNewBlock({c1.getType(), c1.getType()},
250 [&](ValueRange args) { std_br(b1, args); });
251 // The insertion point within the toplevel function is now past b2, we will
252 // need to get back the entry block.
253 // This is what happens with unstructured control-flow.
254 appendToBlock(b1, [&](ValueRange args) {
255 Value r = args[0] + args[1];
256 std_br(b2, {args[0], r});
257 });
258 // Get back to entry block and add a branch into b1.
259 appendToBlock(&f.front(), [&](ValueRange) { std_br(b1, {c1, c2}); });
260 ret.erase();
261
262 // clang-format off
263 // CHECK-LABEL: @builder_blocks
264 // CHECK: %{{.*}} = constant 42 : i32
265 // CHECK-NEXT: %{{.*}} = constant 1234 : i32
266 // CHECK-NEXT: br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
267 // CHECK-NEXT: ^bb1(%{{.*}}: i32, %{{.*}}: i32): // 2 preds: ^bb0, ^bb2
268 // CHECK-NEXT: %{{.*}} = addi %{{.*}}, %{{.*}} : i32
269 // CHECK-NEXT: br ^bb2(%{{.*}}, %{{.*}} : i32, i32)
270 // CHECK-NEXT: ^bb2(%{{.*}}: i32, %{{.*}}: i32): // pred: ^bb1
271 // CHECK-NEXT: br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
272 // CHECK-NEXT: }
273 // clang-format on
274 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
275 f.erase();
276 }
277
TEST_FUNC(builder_cond_branch)278 TEST_FUNC(builder_cond_branch) {
279 auto f = makeFunction("builder_cond_branch", {},
280 {IntegerType::get(1, &globalContext())});
281
282 OpBuilder builder(f.getBody());
283 ScopedContext scope(builder, f.getLoc());
284 Value c32(std_constant_int(32, 32)), c64(std_constant_int(64, 64)),
285 c42(std_constant_int(42, 32));
286 ReturnOp ret = std_ret();
287
288 Block *b1 = buildInNewBlock(c32.getType(), [&](ValueRange) { std_ret(); });
289 Block *b2 = buildInNewBlock({c64.getType(), c32.getType()},
290 [&](ValueRange) { std_ret(); });
291 // Get back to entry block and add a conditional branch.
292 appendToBlock(&f.front(), [&](ValueRange args) {
293 std_cond_br(args[0], b1, {c32}, b2, {c64, c42});
294 });
295 ret.erase();
296
297 // clang-format off
298 // CHECK-LABEL: @builder_cond_branch
299 // CHECK: %{{.*}} = constant 32 : i32
300 // CHECK-NEXT: %{{.*}} = constant 64 : i64
301 // CHECK-NEXT: %{{.*}} = constant 42 : i32
302 // CHECK-NEXT: cond_br %{{.*}}, ^bb1(%{{.*}} : i32), ^bb2(%{{.*}}, %{{.*}} : i64, i32)
303 // CHECK-NEXT: ^bb1(%{{.*}}: i32): // pred: ^bb0
304 // CHECK-NEXT: return
305 // CHECK-NEXT: ^bb2(%{{.*}}: i64, %{{.*}}: i32): // pred: ^bb0
306 // CHECK-NEXT: return
307 // clang-format on
308 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
309 f.erase();
310 }
311
TEST_FUNC(builder_helpers)312 TEST_FUNC(builder_helpers) {
313 using namespace edsc::op;
314 auto f32Type = FloatType::getF32(&globalContext());
315 auto memrefType =
316 MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
317 ShapedType::kDynamicSize},
318 f32Type, {}, 0);
319 auto f =
320 makeFunction("builder_helpers", {}, {memrefType, memrefType, memrefType});
321
322 OpBuilder builder(f.getBody());
323 ScopedContext scope(builder, f.getLoc());
324 // clang-format off
325 Value f7 = std_constant_float(llvm::APFloat(7.0f), f32Type);
326 MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1)),
327 vC(f.getArgument(2));
328 AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
329 Value lb0, lb1, lb2, ub0, ub1, ub2;
330 int64_t step0, step1, step2;
331 std::tie(lb0, ub0, step0) = vA.range(0);
332 std::tie(lb1, ub1, step1) = vA.range(1);
333 lb2 = vA.lb(2);
334 ub2 = vA.ub(2);
335 step2 = vA.step(2);
336 affineLoopNestBuilder({lb0, lb1}, {ub0, ub1}, {step0, step1}, [&](ValueRange ivs) {
337 Value i = ivs[0];
338 Value j = ivs[1];
339 affineLoopBuilder(lb2, ub2, step2, [&](Value k1){
340 C(i, j, k1) = f7 + A(i, j, k1) + B(i, j, k1);
341 });
342 affineLoopBuilder(lb2, ub2, step2, [&](Value k2){
343 C(i, j, k2) += A(i, j, k2) + B(i, j, k2);
344 });
345 });
346
347 // clang-format off
348 // CHECK-LABEL: @builder_helpers
349 // CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
350 // CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
351 // CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
352 // CHECK-DAG: [[a:%.*]] = affine.load %arg0[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
353 // CHECK-DAG: [[b:%.*]] = addf {{.*}}, [[a]] : f32
354 // CHECK-DAG: [[c:%.*]] = affine.load %arg1[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
355 // CHECK-DAG: [[d:%.*]] = addf [[b]], [[c]] : f32
356 // CHECK-NEXT: affine.store [[d]], %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
357 // CHECK-NEXT: }
358 // CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) {
359 // CHECK-DAG: [[a:%.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
360 // CHECK-DAG: [[b:%.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
361 // CHECK-DAG: [[c:%.*]] = addf [[b]], [[a]] : f32
362 // CHECK-DAG: [[d:%.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
363 // CHECK-DAG: [[e:%.*]] = addf [[d]], [[c]] : f32
364 // CHECK-NEXT: affine.store [[e]], %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
365 // clang-format on
366 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
367 f.erase();
368 }
369
TEST_FUNC(insertion_in_block)370 TEST_FUNC(insertion_in_block) {
371 using namespace edsc::op;
372 auto indexType = IndexType::get(&globalContext());
373 auto f = makeFunction("insertion_in_block", {}, {indexType, indexType});
374
375 OpBuilder builder(f.getBody());
376 ScopedContext scope(builder, f.getLoc());
377 std_constant_int(0, 32);
378 buildInNewBlock({}, [&](ValueRange) { std_constant_int(1, 32); });
379 std_constant_int(2, 32);
380 // clang-format off
381 // CHECK-LABEL: @insertion_in_block
382 // CHECK: {{.*}} = constant 0 : i32
383 // CHECK: {{.*}} = constant 2 : i32
384 // CHECK: ^bb1: // no predecessors
385 // CHECK: {{.*}} = constant 1 : i32
386 // clang-format on
387 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
388 f.erase();
389 }
390
TEST_FUNC(zero_and_std_sign_extendi_op_i1_to_i8)391 TEST_FUNC(zero_and_std_sign_extendi_op_i1_to_i8) {
392 using namespace edsc::op;
393 auto i1Type = IntegerType::get(1, &globalContext());
394 auto i8Type = IntegerType::get(8, &globalContext());
395 auto memrefType = MemRefType::get({}, i1Type, {}, 0);
396 auto f = makeFunction("zero_and_std_sign_extendi_op", {},
397 {memrefType, memrefType});
398
399 OpBuilder builder(f.getBody());
400 ScopedContext scope(builder, f.getLoc());
401 AffineIndexedValue A(f.getArgument(0));
402 AffineIndexedValue B(f.getArgument(1));
403 // clang-format off
404 edsc::intrinsics::std_zero_extendi(A, i8Type);
405 edsc::intrinsics::std_sign_extendi(B, i8Type);
406 // CHECK-LABEL: @zero_and_std_sign_extendi_op
407 // CHECK: %[[SRC1:.*]] = affine.load
408 // CHECK: zexti %[[SRC1]] : i1 to i8
409 // CHECK: %[[SRC2:.*]] = affine.load
410 // CHECK: sexti %[[SRC2]] : i1 to i8
411 // clang-format on
412 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
413 f.erase();
414 }
415
TEST_FUNC(operator_or)416 TEST_FUNC(operator_or) {
417 auto i1Type = IntegerType::get(/*width=*/1, &globalContext());
418 auto f = makeFunction("operator_or", {}, {i1Type, i1Type});
419
420 OpBuilder builder(f.getBody());
421 ScopedContext scope(builder, f.getLoc());
422
423 using op::operator||;
424 Value lhs(f.getArgument(0));
425 Value rhs(f.getArgument(1));
426 lhs || rhs;
427
428 // clang-format off
429 // CHECK-LABEL: @operator_or
430 // CHECK: [[ARG0:%.*]]: i1, [[ARG1:%.*]]: i1
431 // CHECK: or [[ARG0]], [[ARG1]]
432 // clang-format on
433 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
434 f.erase();
435 }
436
TEST_FUNC(operator_and)437 TEST_FUNC(operator_and) {
438 auto i1Type = IntegerType::get(/*width=*/1, &globalContext());
439 auto f = makeFunction("operator_and", {}, {i1Type, i1Type});
440
441 OpBuilder builder(f.getBody());
442 ScopedContext scope(builder, f.getLoc());
443
444 using op::operator&&;
445 using op::negate;
446 Value lhs(f.getArgument(0));
447 Value rhs(f.getArgument(1));
448 negate(lhs && rhs);
449
450 // clang-format off
451 // CHECK-LABEL: @operator_and
452 // CHECK: [[ARG0:%.*]]: i1, [[ARG1:%.*]]: i1
453 // CHECK: [[AND:%.*]] = and [[ARG0]], [[ARG1]]
454 // CHECK: [[TRUE:%.*]] = constant true
455 // CHECK: subi [[TRUE]], [[AND]] : i1
456 // clang-format on
457 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
458 f.erase();
459 }
460
TEST_FUNC(divis_op_i32)461 TEST_FUNC(divis_op_i32) {
462 using namespace edsc::op;
463 auto f = makeFunction("divis_op", {}, {});
464
465 OpBuilder builder(f.getBody());
466 ScopedContext scope(builder, f.getLoc());
467 auto i32Type = builder.getI32Type();
468 std_divis(std_constant_int(10, i32Type), std_constant_int(2, i32Type));
469
470 // clang-format off
471 // CHECK-LABEL: @divis_op
472 // CHECK-DAG: {{.*}} = constant 10
473 // CHECK-DAG: {{.*}} = constant 2
474 // CHECK-NEXT: {{.*}} = divi_signed
475 // clang-format on
476 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
477 f.erase();
478 }
479
TEST_FUNC(diviu_op_i32)480 TEST_FUNC(diviu_op_i32) {
481 using namespace edsc::op;
482 auto f = makeFunction("diviu_op", {}, {});
483
484 OpBuilder builder(f.getBody());
485 ScopedContext scope(builder, f.getLoc());
486 auto i32Type = builder.getI32Type();
487 std_diviu(std_constant_int(10, i32Type), std_constant_int(2, i32Type));
488
489 // clang-format off
490 // CHECK-LABEL: @diviu_op
491 // CHECK-DAG: {{.*}} = constant 10
492 // CHECK-DAG: {{.*}} = constant 2
493 // CHECK-NEXT: {{.*}} = divi_unsigned
494 // clang-format on
495 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
496 f.erase();
497 }
498
TEST_FUNC(fpext_f32_f64)499 TEST_FUNC(fpext_f32_f64) {
500 using namespace edsc::op;
501 auto f = makeFunction("fpext", {}, {});
502
503 OpBuilder builder(f.getBody());
504 ScopedContext scope(builder, f.getLoc());
505 auto f32Type = builder.getF32Type();
506 auto f64Type = builder.getF64Type();
507 std_fpext(std_constant_float(llvm::APFloat(10.0f), f32Type), f64Type);
508
509 // clang-format off
510 // CHECK-LABEL: @fpext
511 // CHECK: {{.*}} = constant 1.0
512 // CHECK-NEXT: {{.*}} = fpext
513 // clang-format on
514 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
515 f.erase();
516 }
517
TEST_FUNC(fptrunc_f32_bf16)518 TEST_FUNC(fptrunc_f32_bf16) {
519 using namespace edsc::op;
520 auto f = makeFunction("fptrunc", {}, {});
521
522 OpBuilder builder(f.getBody());
523 ScopedContext scope(builder, f.getLoc());
524 auto f32Type = builder.getF32Type();
525 auto bf16Type = builder.getBF16Type();
526 std_fptrunc(std_constant_float(llvm::APFloat(10.0f), f32Type), bf16Type);
527
528 // clang-format off
529 // CHECK-LABEL: @fptrunc
530 // CHECK: {{.*}} = constant 1.0
531 // CHECK-NEXT: {{.*}} = fptrunc
532 // clang-format on
533 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
534 f.erase();
535 }
536
TEST_FUNC(select_op_i32)537 TEST_FUNC(select_op_i32) {
538 using namespace edsc::op;
539 auto i32Type = IntegerType::get(32, &globalContext());
540 auto memrefType = MemRefType::get(
541 {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, i32Type, {}, 0);
542 auto f = makeFunction("select_op", {}, {memrefType});
543
544 OpBuilder builder(f.getBody());
545 ScopedContext scope(builder, f.getLoc());
546 Value zero = std_constant_index(0), one = std_constant_index(1);
547 MemRefBoundsCapture vA(f.getArgument(0));
548 AffineIndexedValue A(f.getArgument(0));
549 affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
550 using namespace edsc::op;
551 Value i = ivs[0], j = ivs[1];
552 std_select(eq(i, zero), A(zero, zero), A(i, j));
553 std_select(ne(i, zero), A(zero, zero), A(i, j));
554 std_select(slt(i, zero), A(zero, zero), A(i, j));
555 std_select(sle(i, zero), A(zero, zero), A(i, j));
556 std_select(sgt(i, zero), A(zero, zero), A(i, j));
557 std_select(sge(i, zero), A(zero, zero), A(i, j));
558 std_select(ult(i, zero), A(zero, zero), A(i, j));
559 std_select(ule(i, zero), A(zero, zero), A(i, j));
560 std_select(ugt(i, zero), A(zero, zero), A(i, j));
561 std_select(uge(i, zero), A(zero, zero), A(i, j));
562 });
563
564 // clang-format off
565 // CHECK-LABEL: @select_op
566 // CHECK: affine.for %{{.*}} = 0 to 1 {
567 // CHECK-NEXT: affine.for %{{.*}} = 0 to 1 {
568 // CHECK-DAG: {{.*}} = cmpi "eq"
569 // CHECK-DAG: {{.*}} = affine.load
570 // CHECK-DAG: {{.*}} = affine.load
571 // CHECK-NEXT: {{.*}} = select
572 // CHECK-DAG: {{.*}} = cmpi "ne"
573 // CHECK-DAG: {{.*}} = affine.load
574 // CHECK-DAG: {{.*}} = affine.load
575 // CHECK-NEXT: {{.*}} = select
576 // CHECK-DAG: {{.*}} = cmpi "slt"
577 // CHECK-DAG: {{.*}} = affine.load
578 // CHECK-DAG: {{.*}} = affine.load
579 // CHECK-NEXT: {{.*}} = select
580 // CHECK-DAG: {{.*}} = cmpi "sle"
581 // CHECK-DAG: {{.*}} = affine.load
582 // CHECK-DAG: {{.*}} = affine.load
583 // CHECK-NEXT: {{.*}} = select
584 // CHECK-DAG: {{.*}} = cmpi "sgt"
585 // CHECK-DAG: {{.*}} = affine.load
586 // CHECK-DAG: {{.*}} = affine.load
587 // CHECK-NEXT: {{.*}} = select
588 // CHECK-DAG: {{.*}} = cmpi "sge"
589 // CHECK-DAG: {{.*}} = affine.load
590 // CHECK-DAG: {{.*}} = affine.load
591 // CHECK-NEXT: {{.*}} = select
592 // CHECK-DAG: {{.*}} = cmpi "ult"
593 // CHECK-DAG: {{.*}} = affine.load
594 // CHECK-DAG: {{.*}} = affine.load
595 // CHECK-NEXT: {{.*}} = select
596 // CHECK-DAG: {{.*}} = cmpi "ule"
597 // CHECK-DAG: {{.*}} = affine.load
598 // CHECK-DAG: {{.*}} = affine.load
599 // CHECK-NEXT: {{.*}} = select
600 // CHECK-DAG: {{.*}} = cmpi "ugt"
601 // CHECK-DAG: {{.*}} = affine.load
602 // CHECK-DAG: {{.*}} = affine.load
603 // CHECK-NEXT: {{.*}} = select
604 // CHECK-DAG: {{.*}} = cmpi "uge"
605 // CHECK-DAG: {{.*}} = affine.load
606 // CHECK-DAG: {{.*}} = affine.load
607 // CHECK-NEXT: {{.*}} = select
608 // clang-format on
609 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
610 f.erase();
611 }
612
TEST_FUNC(select_op_f32)613 TEST_FUNC(select_op_f32) {
614 auto f32Type = FloatType::getF32(&globalContext());
615 auto memrefType = MemRefType::get(
616 {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
617 auto f = makeFunction("select_op", {}, {memrefType, memrefType});
618
619 OpBuilder builder(f.getBody());
620 ScopedContext scope(builder, f.getLoc());
621 // clang-format off
622 Value zero = std_constant_index(0), one = std_constant_index(1);
623 MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1));
624 AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1));
625 affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
626 using namespace edsc::op;
627 Value i = ivs[0], j = ivs[1];
628 std_select(eq(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
629 std_select(ne(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
630 std_select(sge(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
631 std_select(sle(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
632 std_select(slt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
633 std_select(sgt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
634 std_select(uge(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
635 std_select(ule(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
636 std_select(ult(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
637 std_select(ugt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
638 });
639
640 // clang-format off
641 // CHECK-LABEL: @select_op
642 // CHECK: affine.for %{{.*}} = 0 to 1 {
643 // CHECK-NEXT: affine.for %{{.*}} = 0 to 1 {
644 // CHECK-DAG: cmpf "oeq"
645 // CHECK-DAG: affine.load
646 // CHECK-DAG: affine.load
647 // CHECK-DAG: affine.load
648 // CHECK-DAG: affine.load
649 // CHECK-DAG: affine.apply
650 // CHECK-NEXT: select
651 // CHECK-DAG: cmpf "one"
652 // CHECK-DAG: affine.load
653 // CHECK-DAG: affine.load
654 // CHECK-DAG: affine.load
655 // CHECK-DAG: affine.load
656 // CHECK-DAG: affine.apply
657 // CHECK-NEXT: select
658 // CHECK-DAG: cmpf "oge"
659 // CHECK-DAG: affine.load
660 // CHECK-DAG: affine.load
661 // CHECK-DAG: affine.load
662 // CHECK-DAG: affine.load
663 // CHECK-DAG: affine.apply
664 // CHECK-NEXT: select
665 // CHECK-DAG: cmpf "ole"
666 // CHECK-DAG: affine.load
667 // CHECK-DAG: affine.load
668 // CHECK-DAG: affine.load
669 // CHECK-DAG: affine.load
670 // CHECK-DAG: affine.apply
671 // CHECK-NEXT: select
672 // CHECK-DAG: cmpf "olt"
673 // CHECK-DAG: affine.load
674 // CHECK-DAG: affine.load
675 // CHECK-DAG: affine.load
676 // CHECK-DAG: affine.load
677 // CHECK-DAG: affine.apply
678 // CHECK-NEXT: select
679 // CHECK-DAG: cmpf "ogt"
680 // CHECK-DAG: affine.load
681 // CHECK-DAG: affine.load
682 // CHECK-DAG: affine.load
683 // CHECK-DAG: affine.load
684 // CHECK-DAG: affine.apply
685 // CHECK-NEXT: select
686 // CHECK-DAG: cmpf "oge"
687 // CHECK-DAG: affine.load
688 // CHECK-DAG: affine.load
689 // CHECK-DAG: affine.load
690 // CHECK-DAG: affine.load
691 // CHECK-DAG: affine.apply
692 // CHECK-NEXT: select
693 // CHECK-DAG: cmpf "ole"
694 // CHECK-DAG: affine.load
695 // CHECK-DAG: affine.load
696 // CHECK-DAG: affine.load
697 // CHECK-DAG: affine.load
698 // CHECK-DAG: affine.apply
699 // CHECK-NEXT: select
700 // CHECK-DAG: cmpf "olt"
701 // CHECK-DAG: affine.load
702 // CHECK-DAG: affine.load
703 // CHECK-DAG: affine.load
704 // CHECK-DAG: affine.load
705 // CHECK-DAG: affine.apply
706 // CHECK-NEXT: select
707 // CHECK-DAG: cmpf "ogt"
708 // CHECK-DAG: affine.load
709 // CHECK-DAG: affine.load
710 // CHECK-DAG: affine.load
711 // CHECK-DAG: affine.load
712 // CHECK-DAG: affine.apply
713 // CHECK-NEXT: select
714 // clang-format on
715 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
716 f.erase();
717 }
718
719 // Inject an EDSC-constructed computation to exercise imperfectly nested 2-d
720 // tiling.
TEST_FUNC(tile_2d)721 TEST_FUNC(tile_2d) {
722 auto memrefType =
723 MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
724 ShapedType::kDynamicSize},
725 FloatType::getF32(&globalContext()), {}, 0);
726 auto f = makeFunction("tile_2d", {}, {memrefType, memrefType, memrefType});
727
728 OpBuilder builder(f.getBody());
729 ScopedContext scope(builder, f.getLoc());
730 Value zero = std_constant_index(0);
731 MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1)),
732 vC(f.getArgument(2));
733 AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1)),
734 C(f.getArgument(2));
735 Value i, j, k1, k2;
736 Value M(vC.ub(0)), N(vC.ub(1)), O(vC.ub(2));
737
738 // clang-format off
739 using namespace edsc::op;
740 affineLoopNestBuilder({zero, zero}, {M, N}, {1, 1}, [&](ValueRange ivs) {
741 i = ivs[0];
742 j = ivs[1];
743 affineLoopBuilder(zero, O, 1, [&](Value k) {
744 k1 = k;
745 C(i, j, k1) = A(i, j, k1) + B(i, j, k1);
746 });
747 affineLoopBuilder(zero, O, 1, [&](Value k) {
748 k2 = k;
749 C(i, j, k2) = A(i, j, k2) + B(i, j, k2);
750 });
751 });
752 // clang-format on
753
754 auto li = getForInductionVarOwner(i), lj = getForInductionVarOwner(j),
755 lk1 = getForInductionVarOwner(k1), lk2 = getForInductionVarOwner(k2);
756 auto indicesL1 = mlir::tile({li, lj}, {512, 1024}, {lk1, lk2});
757 auto lii1 = indicesL1[0][0], ljj1 = indicesL1[1][0];
758 mlir::tile({ljj1, lii1}, {32, 16}, ljj1);
759
760 // clang-format off
761 // CHECK-LABEL: func @tile_2d
762 // CHECK: %[[ZERO:.*]] = constant 0 : index
763 // CHECK: %[[M:[0-9]+]] = dim %arg2, %c0{{[_0-9]*}} : memref<?x?x?xf32>
764 // CHECK: %[[N:[0-9]+]] = dim %arg2, %c1{{[_0-9]*}} : memref<?x?x?xf32>
765 // CHECK: %[[P:[0-9]+]] = dim %arg2, %c2{{[_0-9]*}} : memref<?x?x?xf32>
766 // CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[M]]) step 512 {
767 // CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[N]]) step 1024 {
768 // CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[P]]) {
769 // CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 512)>(%{{.*}})[%[[M]]] step 16 {
770 // CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 1024)>(%{{.*}})[%[[N]]] step 32 {
771 // CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0, d1) -> (0, d0, d1)>(%{{.*}}, %{{.*}}) to min affine_map<(d0, d1)[s0] -> (s0, d0 + 1024, d1 + 32)>(%{{.*}}, %{{.*}})[%[[N]]] {
772 // CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0, d1) -> (0, d0, d1)>(%{{.*}}, %{{.*}}) to min affine_map<(d0, d1)[s0] -> (s0, d0 + 512, d1 + 16)>(%{{.*}}, %{{.*}})[%[[M]]] {
773 // CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
774 // CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
775 // CHECK-NEXT: {{.*}} = addf {{.*}}, {{.*}} : f32
776 // CHECK-NEXT: affine.store {{.*}}, {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
777 // CHECK: }
778 // CHECK-NEXT: }
779 // CHECK-NEXT: }
780 // CHECK-NEXT: }
781 // CHECK-NEXT: }
782 // CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[P]]) {
783 // CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 512)>(%{{.*}})[%[[M]]] {
784 // CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 1024)>(%{{.*}})[%[[N]]] {
785 // CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
786 // CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
787 // CHECK-NEXT: {{.*}}= addf {{.*}}, {{.*}} : f32
788 // CHECK-NEXT: affine.store {{.*}}, {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
789 // clang-format on
790 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
791 f.erase();
792 }
793
794 // Exercise StdIndexedValue for loads and stores.
TEST_FUNC(indirect_access)795 TEST_FUNC(indirect_access) {
796 using namespace edsc::op;
797 auto memrefType = MemRefType::get({ShapedType::kDynamicSize},
798 FloatType::getF32(&globalContext()), {}, 0);
799 auto f = makeFunction("indirect_access", {},
800 {memrefType, memrefType, memrefType, memrefType});
801
802 OpBuilder builder(f.getBody());
803 ScopedContext scope(builder, f.getLoc());
804 Value zero = std_constant_index(0);
805 MemRefBoundsCapture vC(f.getArgument(2));
806 AffineIndexedValue B(f.getArgument(1)), D(f.getArgument(3));
807 StdIndexedValue A(f.getArgument(0)), C(f.getArgument(2));
808 Value N(vC.ub(0));
809
810 // clang-format off
811 affineLoopBuilder(zero, N, 1, [&](Value i) {
812 C((Value)D(i)) = A((Value)B(i));
813 });
814 // clang-format on
815
816 // clang-format off
817 // CHECK-LABEL: func @indirect_access
818 // CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>, %[[ARG2:.*]]: memref<?xf32>, %[[ARG3:.*]]: memref<?xf32>)
819 // CHECK-DAG: [[B:%.*]] = affine.load %[[ARG1]]
820 // CHECK-DAG: [[D:%.*]] = affine.load %[[ARG3]]
821 // CHECK: load %{{.*}}{{\[}}[[B]]{{\]}}
822 // CHECK: store %{{.*}}, %{{.*}}{{\[}}[[D]]{{\]}}
823 // clang-format on
824 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
825 f.erase();
826 }
827
828 // Exercise affine loads and stores build with empty maps.
TEST_FUNC(empty_map_load_store)829 TEST_FUNC(empty_map_load_store) {
830 using namespace edsc::op;
831 auto memrefType =
832 MemRefType::get({}, FloatType::getF32(&globalContext()), {}, 0);
833 auto f = makeFunction("empty_map_load_store", {},
834 {memrefType, memrefType, memrefType, memrefType});
835
836 OpBuilder builder(f.getBody());
837 ScopedContext scope(builder, f.getLoc());
838 Value zero = std_constant_index(0);
839 Value one = std_constant_index(1);
840 AffineIndexedValue input(f.getArgument(0)), res(f.getArgument(1));
841
842 // clang-format off
843 affineLoopBuilder(zero, one, 1, [&](Value) {
844 res() = input();
845 });
846 // clang-format on
847
848 // clang-format off
849 // CHECK-LABEL: func @empty_map_load_store(
850 // CHECK: [[A:%.*]] = affine.load %{{.*}}[]
851 // CHECK: affine.store [[A]], %{{.*}}[]
852 // clang-format on
853 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
854 f.erase();
855 }
856
857 // clang-format off
858 // CHECK-LABEL: func @affine_if_op
859 // CHECK: affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}[[s0:.*]], [[s1:.*]]{{\]}}
860 // CHECK-NOT: else
861 // CHECK: affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}[[s0:.*]], [[s1:.*]]{{\]}}
862 // CHECK-NEXT: } else {
863 // clang-format on
TEST_FUNC(affine_if_op)864 TEST_FUNC(affine_if_op) {
865 using namespace edsc::op;
866 auto f32Type = FloatType::getF32(&globalContext());
867 auto memrefType = MemRefType::get(
868 {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
869 auto f = makeFunction("affine_if_op", {}, {memrefType});
870
871 OpBuilder builder(f.getBody());
872 ScopedContext scope(builder, f.getLoc());
873
874 Value zero = std_constant_index(0), ten = std_constant_index(10);
875
876 SmallVector<bool, 4> isEq = {false, false, false, false};
877 SmallVector<AffineExpr, 4> affineExprs = {
878 builder.getAffineDimExpr(0), // d0 >= 0
879 builder.getAffineDimExpr(1), // d1 >= 0
880 builder.getAffineSymbolExpr(0), // s0 >= 0
881 builder.getAffineSymbolExpr(1) // s1 >= 0
882 };
883 auto intSet = IntegerSet::get(2, 2, affineExprs, isEq);
884
885 SmallVector<Value, 4> affineIfArgs = {zero, zero, ten, ten};
886 intrinsics::affine_if(intSet, affineIfArgs, /*withElseRegion=*/false);
887 intrinsics::affine_if(intSet, affineIfArgs, /*withElseRegion=*/true);
888
889 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
890 f.erase();
891 }
892
893 // clang-format off
894 // CHECK-LABEL: func @linalg_generic_pointwise
895 // CHECK: linalg.generic {
896 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
897 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
898 // CHECK-SAME: ins({{.*}}memref<?x?xf32>, memref<?x?xf32>)
899 // CHECK-SAME: outs({{.*}}memref<?x?xf32>)
900 // CHECK: addf
901 // CHECK: linalg.generic {
902 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
903 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
904 // CHECK-SAME: ins({{.*}}memref<?x?xf32>, memref<?x?xf32>)
905 // CHECK-SAME: outs({{.*}}memref<?x?xf32>)
906 // CHECK: cmpf "ogt"
907 // CHECK: select
908 // CHECK: linalg.generic {
909 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
910 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
911 // CHECK-SAME: ins(%{{[a-z0-9]*}} : memref<?x?xf32>)
912 // CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?xf32>)
913 // CHECK: tanh
914 // clang-format on
TEST_FUNC(linalg_generic_pointwise_test)915 TEST_FUNC(linalg_generic_pointwise_test) {
916 using namespace edsc;
917 using namespace edsc::ops;
918
919 auto f32Type = FloatType::getF32(&globalContext());
920 auto memrefType = MemRefType::get(
921 {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
922 auto f = makeFunction("linalg_generic_pointwise", {},
923 {memrefType, memrefType, memrefType});
924
925 OpBuilder builder(f.getBody());
926 ScopedContext scope(builder, f.getLoc());
927 Value A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
928 AffineExpr i, j;
929 bindDims(&globalContext(), i, j);
930 StructuredIndexed SA(A), SB(B), SC(C);
931 linalg_generic_pointwise_add(SA({i, j}), SB({i, j}), SC({i, j}));
932 linalg_generic_pointwise_max(SA({i, j}), SB({i, j}), SC({i, j}));
933 linalg_generic_pointwise_tanh(SA({i, j}), SC({i, j}));
934
935 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
936 f.erase();
937 }
938
939 // clang-format off
940 // CHECK-LABEL: func @linalg_generic_matmul
941 // CHECK: linalg.generic {
942 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
943 // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
944 // CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : memref<?x?xf32>, memref<?x?xf32>)
945 // CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?xf32>)
946 /// CHECK: ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
947 // CHECK: %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
948 // CHECK: %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
949 // CHECK: linalg.yield %[[a4]] : f32
950 // CHECK: }
951 // clang-format on
TEST_FUNC(linalg_generic_matmul_test)952 TEST_FUNC(linalg_generic_matmul_test) {
953 using namespace edsc;
954 using namespace edsc::ops;
955
956 auto f32Type = FloatType::getF32(&globalContext());
957 auto memrefType = MemRefType::get(
958 {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
959 auto f = makeFunction("linalg_generic_matmul", {},
960 {memrefType, memrefType, memrefType});
961
962 OpBuilder builder(f.getBody());
963 ScopedContext scope(builder, f.getLoc());
964 linalg_generic_matmul(f.getArguments());
965
966 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
967 f.erase();
968 }
969
970 // clang-format off
971 // CHECK-LABEL: func @linalg_generic_conv_nhwc
972 // CHECK: linalg.generic {
973 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2 * 3 + d4 * 5, d3 * 4 + d5 * 6, d6)>,
974 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d1)>,
975 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2, d3, d1)>],
976 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]}
977 // CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
978 // CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?x?x?xf32>)
979 /// CHECK: ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
980 // CHECK: %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
981 // CHECK: %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
982 // CHECK: linalg.yield %[[a4]] : f32
983 // CHECK: }
984 // clang-format on
TEST_FUNC(linalg_generic_conv_nhwc)985 TEST_FUNC(linalg_generic_conv_nhwc) {
986 using namespace edsc;
987 using namespace edsc::ops;
988
989 auto f32Type = FloatType::getF32(&globalContext());
990 auto memrefType =
991 MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
992 ShapedType::kDynamicSize, ShapedType::kDynamicSize},
993 f32Type, {}, 0);
994 auto f = makeFunction("linalg_generic_conv_nhwc", {},
995 {memrefType, memrefType, memrefType});
996
997 OpBuilder builder(f.getBody());
998 ScopedContext scope(builder, f.getLoc());
999 linalg_generic_conv_nhwc(f.getArguments(),
1000 /*strides=*/{3, 4}, /*dilations=*/{5, 6});
1001
1002 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
1003 f.erase();
1004 }
1005
1006 // clang-format off
1007 // CHECK-LABEL: func @linalg_generic_dilated_conv_nhwc
1008 // CHECK: linalg.generic {
1009 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d3 * 3 + d5 * 5, d4 * 4 + d6 * 6, d2)>,
1010 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d2, d1)>,
1011 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d3, d4, d1 + d2 * 7)>],
1012 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
1013 // CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
1014 // CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?x?x?xf32>)
1015 // CHECK: ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
1016 // CHECK: %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
1017 // CHECK: %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
1018 // CHECK: linalg.yield %[[a4]] : f32
1019 // CHECK: }
1020 // clang-format on
TEST_FUNC(linalg_generic_dilated_conv_nhwc)1021 TEST_FUNC(linalg_generic_dilated_conv_nhwc) {
1022 using namespace edsc;
1023 using namespace edsc::ops;
1024
1025 auto f32Type = FloatType::getF32(&globalContext());
1026 auto memrefType =
1027 MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
1028 ShapedType::kDynamicSize, ShapedType::kDynamicSize},
1029 f32Type, {}, 0);
1030 auto f = makeFunction("linalg_generic_dilated_conv_nhwc", {},
1031 {memrefType, memrefType, memrefType});
1032
1033 OpBuilder builder(f.getBody());
1034 ScopedContext scope(builder, f.getLoc());
1035 linalg_generic_dilated_conv_nhwc(f.getArguments(),
1036 /*depth_multiplier=*/7,
1037 /*strides=*/{3, 4}, /*dilations=*/{5, 6});
1038
1039 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
1040 f.erase();
1041 }
1042
1043 // clang-format off
1044 // CHECK-LABEL: func @linalg_metadata_ops
1045 // CHECK: linalg.reshape {{.*}} [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d2)>] : memref<4x8x16xf32> into memref<32x16xf32>
1046 // CHECK: linalg.reshape {{.*}} [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d2)>] : memref<32x16xf32> into memref<4x8x16xf32>
1047 // clang-format on
TEST_FUNC(linalg_metadata_ops)1048 TEST_FUNC(linalg_metadata_ops) {
1049 using linalg::ReassociationExprs;
1050
1051 auto f32Type = FloatType::getF32(&globalContext());
1052 auto memrefType = MemRefType::get({4, 8, 16}, f32Type, {}, 0);
1053 auto f = makeFunction("linalg_metadata_ops", {}, {memrefType});
1054
1055 OpBuilder builder(f.getBody());
1056 ScopedContext scope(builder, f.getLoc());
1057 AffineExpr i, j, k;
1058 bindDims(&globalContext(), i, j, k);
1059 Value v(f.getArgument(0));
1060 SmallVector<ReassociationExprs, 2> maps = {ReassociationExprs({i, j}),
1061 ReassociationExprs({k})};
1062 auto reshaped = linalg_reshape(v, maps);
1063 linalg_reshape(memrefType, reshaped, maps);
1064
1065 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
1066 f.erase();
1067 }
1068
1069 // clang-format off
1070 // CHECK-LABEL: func @linalg_tensors
1071 // CHECK: linalg.generic {
1072 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
1073 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
1074 // CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, memref<?x?xf32>)
1075 // CHECK: addf
1076 // CHECK: } -> tensor<?x?xf32>
1077 // CHECK: linalg.generic {
1078 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
1079 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
1080 // CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, tensor<?x?xf32>)
1081 // CHECK: cmpf "ogt"
1082 // CHECK: select
1083 // CHECK: } -> tensor<?x?xf32>
1084 // CHECK: linalg.generic {
1085 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
1086 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
1087 // CHECK-SAME: ins(%{{[a-z0-9]*}} : tensor<?x?xf32>)
1088 // CHECK: tanh
1089 // CHECK: } -> tensor<?x?xf32>
1090 // CHECK: linalg.generic {
1091 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
1092 // CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>,
1093 // CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>],
1094 // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
1095 // CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, memref<?x?xf32>)
1096 // CHECK: mulf
1097 // CHECK: } -> tensor<?x?xf32>
1098 // CHECK: linalg.generic {
1099 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
1100 // CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>,
1101 // CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>],
1102 // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
1103 // CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, memref<?x?xf32>)
1104 // CHECK-SAME: init(%{{[a-z0-9]*}} : tensor<?x?xf32>)
1105 // CHECK: mulf
1106 // CHECK: addf
1107 // CHECK: } -> tensor<?x?xf32>
1108 // clang-format on
TEST_FUNC(linalg_tensors_test)1109 TEST_FUNC(linalg_tensors_test) {
1110 using namespace edsc;
1111 using namespace edsc::ops;
1112
1113 auto f32Type = FloatType::getF32(&globalContext());
1114 auto memrefType = MemRefType::get(
1115 {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
1116 auto tensorType = RankedTensorType::get(
1117 {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type);
1118 auto f = makeFunction("linalg_tensors", {}, {tensorType, memrefType});
1119
1120 OpBuilder builder(f.getBody());
1121 ScopedContext scope(builder, f.getLoc());
1122 Value A(f.getArgument(0)), B(f.getArgument(1));
1123 AffineExpr i, j;
1124 bindDims(&globalContext(), i, j);
1125 StructuredIndexed SA(A), SB(B), SC(tensorType);
1126 Value added = linalg_generic_pointwise_add(SA({i, j}), SB({i, j}), SC({i, j}))
1127 ->getResult(0);
1128 Value maxed = linalg_generic_pointwise_max(
1129 SA({i, j}), StructuredIndexed(added)({i, j}), SC({i, j}))
1130 ->getResult(0);
1131 Value tanhed = linalg_generic_pointwise_tanh(StructuredIndexed(maxed)({i, j}),
1132 SC({i, j}))
1133 ->getResult(0);
1134 Value o1 = linalg_generic_matmul(A, B, tanhed, tensorType)->getResult(0);
1135 linalg_generic_matmul(A, B, o1, tensorType);
1136
1137 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
1138 f.erase();
1139 }
1140
TEST_FUNC(vector_extractelement_op_i32)1141 TEST_FUNC(vector_extractelement_op_i32) {
1142 using namespace edsc::op;
1143 auto f = makeFunction("vector_extractelement_op", {}, {});
1144
1145 OpBuilder builder(f.getBody());
1146 ScopedContext scope(builder, f.getLoc());
1147 auto i32Type = builder.getI32Type();
1148 auto vectorType = VectorType::get(/*shape=*/{8}, i32Type);
1149 vector_extract_element(
1150 i32Type, std_constant(vectorType, builder.getI32VectorAttr({10})),
1151 std_constant_int(0, i32Type));
1152
1153 // clang-format off
1154 // CHECK-LABEL: @vector_extractelement_op
1155 // CHECK-DAG: {{.*}} = constant dense<10>
1156 // CHECK-DAG: {{.*}} = constant 0
1157 // CHECK-NEXT: {{.*}} = vector.extractelement
1158 // clang-format on
1159 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
1160 f.erase();
1161 }
1162
1163 // clang-format off
1164 // CHECK-LABEL: func @memref_vector_matmul_test(
1165 // CHECK-SAME: %[[A:.*]]: memref<?x?xvector<4x16xf32>>,
1166 // CHECK-SAME: %[[B:.*]]: memref<?x?xvector<16x8xf32>>,
1167 // CHECK-SAME: %[[C:.*]]: memref<?x?xvector<4x8xf32>>)
1168 // CHECK: linalg.generic {{{.*}}}
1169 // CHECK-SAME: ins(%[[A]], %[[B]] : memref<?x?xvector<4x16xf32>>, memref<?x?xvector<16x8xf32>>)
1170 // CHECK-SAME: outs(%[[C]] : memref<?x?xvector<4x8xf32>>)
1171 // CHECK: vector.contract{{.*}}[affine_map<(d0, d1, d2) -> (d0, d2)>,
1172 // CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>,
1173 // CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>],
1174 // CHECK-SAME: {{.*}}["parallel", "parallel", "reduction"]
1175 // clang-format on
TEST_FUNC(memref_vector_matmul_test)1176 TEST_FUNC(memref_vector_matmul_test) {
1177 using namespace edsc;
1178 using namespace edsc::ops;
1179
1180 int64_t M = 4, N = 8, K = 16;
1181 auto f32Type = FloatType::getF32(&globalContext());
1182 auto mkVectorType = VectorType::get({M, K}, f32Type);
1183 auto knVectorType = VectorType::get({K, N}, f32Type);
1184 auto mnVectorType = VectorType::get({M, N}, f32Type);
1185 auto typeA =
1186 MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
1187 mkVectorType, {}, 0);
1188 auto typeB =
1189 MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
1190 knVectorType, {}, 0);
1191 auto typeC =
1192 MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
1193 mnVectorType, {}, 0);
1194 auto f = makeFunction("memref_vector_matmul_test", {}, {typeA, typeB, typeC});
1195
1196 OpBuilder builder(f.getBody());
1197 ScopedContext scope(builder, f.getLoc());
1198 Value A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
1199 auto contractionBuilder = [](ValueRange args) {
1200 assert(args.size() == 3 && "expected 3 block arguments");
1201 (linalg_yield(vector_contraction_matmul(args[0], args[1], args[2])));
1202 };
1203 linalg_generic_matmul(A, B, C, contractionBuilder);
1204
1205 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
1206 f.erase();
1207 }
1208
TEST_FUNC(builder_loop_for_yield)1209 TEST_FUNC(builder_loop_for_yield) {
1210 auto indexType = IndexType::get(&globalContext());
1211 auto f32Type = FloatType::getF32(&globalContext());
1212 auto f = makeFunction("builder_loop_for_yield", {},
1213 {indexType, indexType, indexType, indexType});
1214
1215 OpBuilder builder(f.getBody());
1216 ScopedContext scope(builder, f.getLoc());
1217 Value init0 = std_constant_float(llvm::APFloat(1.0f), f32Type);
1218 Value init1 = std_constant_float(llvm::APFloat(2.0f), f32Type);
1219 Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
1220 d(f.getArgument(3));
1221 using namespace edsc::op;
1222 auto results = loopNestBuilder(a - b, c + d, a, {init0, init1},
1223 [&](Value iv, ValueRange args) {
1224 Value sum = args[0] + args[1];
1225 return scf::ValueVector{args[1], sum};
1226 }).getResults();
1227 results[0] + results[1];
1228
1229 // clang-format off
1230 // CHECK-LABEL: func @builder_loop_for_yield(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
1231 // CHECK: [[init0:%.*]] = constant
1232 // CHECK: [[init1:%.*]] = constant
1233 // CHECK-DAG: [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
1234 // CHECK-DAG: [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
1235 // CHECK-NEXT: [[res:%[0-9]+]]:2 = scf.for %{{.*}} = [[r0]] to [[r1]] step {{.*}} iter_args([[arg0:%.*]] = [[init0]], [[arg1:%.*]] = [[init1]]) -> (f32, f32) {
1236 // CHECK: [[sum:%[0-9]+]] = addf [[arg0]], [[arg1]] : f32
1237 // CHECK: scf.yield [[arg1]], [[sum]] : f32, f32
1238 // CHECK: addf [[res]]#0, [[res]]#1 : f32
1239 // clang-format on
1240
1241 f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
1242 f.erase();
1243 }
1244
main()1245 int main() {
1246 RUN_TESTS();
1247 return 0;
1248 }
1249