93a873dfc9
In prehistorical times, AffineApplyOp was allowed to produce multiple values. This allowed the creation of intricate SSA use-def chains. AffineApplyNormalizer was originally introduced as a means of reusing the AffineMap::compose method to write SSA use-def chains. Unfortunately, symbols that were produced by an AffineApplyOp needed to be promoted to dims and reordered for the mathematical composition to be valid. Since then, single result AffineApplyOp became the law of the land but the original assumptions were not revisited. This revision revisits these assumptions and retires AffineApplyNormalizer. Differential Revision: https://reviews.llvm.org/D94920
1251 lines
51 KiB
C++
1251 lines
51 KiB
C++
//===- builder-api-test.cpp - Tests for Declarative Builder APIs ----------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// RUN: mlir-edsc-builder-api-test | FileCheck %s
|
|
|
|
#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/Linalg/EDSC/Builders.h"
|
|
#include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/SCF/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
|
|
#include "mlir/EDSC/Builders.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/IntegerSet.h"
|
|
#include "mlir/IR/MLIRContext.h"
|
|
#include "mlir/IR/Types.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Pass/PassManager.h"
|
|
#include "mlir/Transforms/LoopUtils.h"
|
|
#include "mlir/Transforms/Passes.h"
|
|
|
|
#include "APITest.h"
|
|
|
|
#include "llvm/Support/raw_ostream.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::edsc;
|
|
using namespace mlir::edsc::intrinsics;
|
|
|
|
static MLIRContext &globalContext() {
|
|
static thread_local MLIRContext context;
|
|
static thread_local bool initOnce = [&]() {
|
|
// clang-format off
|
|
context.loadDialect<AffineDialect,
|
|
scf::SCFDialect,
|
|
linalg::LinalgDialect,
|
|
StandardOpsDialect,
|
|
vector::VectorDialect>();
|
|
// clang-format on
|
|
return true;
|
|
}();
|
|
(void)initOnce;
|
|
context.allowUnregisteredDialects();
|
|
return context;
|
|
}
|
|
|
|
static FuncOp makeFunction(StringRef name, ArrayRef<Type> results = {},
|
|
ArrayRef<Type> args = {}) {
|
|
auto &ctx = globalContext();
|
|
auto function = FuncOp::create(UnknownLoc::get(&ctx), name,
|
|
FunctionType::get(&ctx, args, results));
|
|
function.addEntryBlock();
|
|
return function;
|
|
}
|
|
|
|
TEST_FUNC(builder_dynamic_for_func_args) {
|
|
auto indexType = IndexType::get(&globalContext());
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto f =
|
|
makeFunction("builder_dynamic_for_func_args", {}, {indexType, indexType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value lb(f.getArgument(0)), ub(f.getArgument(1));
|
|
Value f7(std_constant_float(llvm::APFloat(7.0f), f32Type));
|
|
Value f13(std_constant_float(llvm::APFloat(13.0f), f32Type));
|
|
Value i7(std_constant_int(7, 32));
|
|
Value i13(std_constant_int(13, 32));
|
|
affineLoopBuilder(lb, ub, 3, [&](Value i) {
|
|
using namespace edsc::op;
|
|
lb *std_constant_index(3) + ub;
|
|
lb + std_constant_index(3);
|
|
affineLoopBuilder(lb, ub, 2, [&](Value j) {
|
|
ceilDiv(std_constant_index(31) * floorDiv(i + j * std_constant_index(3),
|
|
std_constant_index(32)),
|
|
std_constant_index(32));
|
|
((f7 + f13) / f7) % f13 - f7 *f13;
|
|
((i7 + i13) / i7) % i13 - i7 *i13;
|
|
});
|
|
});
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @builder_dynamic_for_func_args(%{{.*}}: index, %{{.*}}: index) {
|
|
// CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) step 3 {
|
|
// CHECK: {{.*}} = affine.apply affine_map<()[s0] -> (s0 * 3)>()[%{{.*}}]
|
|
// CHECK: {{.*}} = affine.apply affine_map<()[s0, s1] -> (s0 * 3 + s1)>()[%{{.*}}, %{{.*}}]
|
|
// CHECK: {{.*}} = affine.apply affine_map<()[s0] -> (s0 + 3)>()[%{{.*}}]
|
|
// CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) step 2 {
|
|
// CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> ((d0 + d1 * 3) floordiv 32)>(%{{.*}}, %{{.*}})
|
|
// CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> (((d0 + d1 * 3) floordiv 32) * 31)>(%{{.*}}, %{{.*}})
|
|
// CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> ((((d0 + d1 * 3) floordiv 32) * 31) ceildiv 32)>(%{{.*}}, %{{.*}})
|
|
// CHECK-DAG: [[rf1:%[0-9]+]] = addf {{.*}}, {{.*}} : f32
|
|
// CHECK-DAG: [[rf2:%[0-9]+]] = divf [[rf1]], {{.*}} : f32
|
|
// CHECK-DAG: [[rf3:%[0-9]+]] = remf [[rf2]], {{.*}} : f32
|
|
// CHECK-DAG: [[rf4:%[0-9]+]] = mulf {{.*}}, {{.*}} : f32
|
|
// CHECK: {{.*}} = subf [[rf3]], [[rf4]] : f32
|
|
// CHECK-DAG: [[ri1:%[0-9]+]] = addi {{.*}}, {{.*}} : i32
|
|
// CHECK-DAG: [[ri2:%[0-9]+]] = divi_signed [[ri1]], {{.*}} : i32
|
|
// CHECK-DAG: [[ri3:%[0-9]+]] = remi_signed [[ri2]], {{.*}} : i32
|
|
// CHECK-DAG: [[ri4:%[0-9]+]] = muli {{.*}}, {{.*}} : i32
|
|
// CHECK: {{.*}} = subi [[ri3]], [[ri4]] : i32
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_dynamic_for) {
|
|
auto indexType = IndexType::get(&globalContext());
|
|
auto f = makeFunction("builder_dynamic_for", {},
|
|
{indexType, indexType, indexType, indexType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value i, a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
|
|
d(f.getArgument(3));
|
|
using namespace edsc::op;
|
|
affineLoopBuilder(a - b, c + d, 2);
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @builder_dynamic_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
|
|
// CHECK-DAG: [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
|
|
// CHECK-DAG: [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
|
|
// CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>([[r0]]) to affine_map<(d0) -> (d0)>([[r1]]) step 2 {
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_loop_for) {
|
|
auto indexType = IndexType::get(&globalContext());
|
|
auto f = makeFunction("builder_loop_for", {},
|
|
{indexType, indexType, indexType, indexType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
|
|
d(f.getArgument(3));
|
|
using namespace edsc::op;
|
|
loopNestBuilder(a - b, c + d, a);
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @builder_loop_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
|
|
// CHECK-DAG: [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
|
|
// CHECK-DAG: [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
|
|
// CHECK-NEXT: scf.for %{{.*}} = [[r0]] to [[r1]] step {{.*}} {
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_max_min_for) {
|
|
auto indexType = IndexType::get(&globalContext());
|
|
auto f = makeFunction("builder_max_min_for", {},
|
|
{indexType, indexType, indexType, indexType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value lb1(f.getArgument(0)), lb2(f.getArgument(1)), ub1(f.getArgument(2)),
|
|
ub2(f.getArgument(3));
|
|
affineLoopBuilder({lb1, lb2}, {ub1, ub2}, 1);
|
|
std_ret();
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @builder_max_min_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
|
|
// CHECK: affine.for %{{.*}} = max affine_map<(d0, d1) -> (d0, d1)>(%{{.*}}, %{{.*}}) to min affine_map<(d0, d1) -> (d0, d1)>(%{{.*}}, %{{.*}}) {
|
|
// CHECK: return
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_affine_for_iter_args) {
|
|
auto indexType = IndexType::get(&globalContext());
|
|
auto f = makeFunction("builder_affine_for_iter_args", {},
|
|
{indexType, indexType, indexType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value i, lb_1(f.getArgument(0)), ub_1(f.getArgument(1)),
|
|
ub_2(f.getArgument(2));
|
|
Value c32(std_constant_int(32, 32));
|
|
Value c42(std_constant_int(42, 32));
|
|
using namespace edsc::op;
|
|
affineLoopBuilder(
|
|
lb_1, {ub_1, ub_2}, 2, {c32, c42}, [&](Value iv, ValueRange args) {
|
|
Value sum(args[0] + args[1]);
|
|
builder.create<AffineYieldOp>(f.getLoc(), ValueRange({args[1], sum}));
|
|
});
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @builder_affine_for_iter_args
|
|
// CHECK: (%[[lb_1:.*]]: index, %[[ub_1:.*]]: index, %[[ub_2:.*]]: index) {
|
|
// CHECK-NEXT: %[[c32:.*]] = constant 32 : i32
|
|
// CHECK-NEXT: %[[c42:.*]] = constant 42 : i32
|
|
// 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) {
|
|
// CHECK-NEXT: %[[sum:.*]] = addi %[[iarg_1]], %[[iarg_2]] : i32
|
|
// CHECK-NEXT: affine.yield %[[iarg_2]], %[[sum]] : i32, i32
|
|
// CHECK-NEXT: }
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_block_append) {
|
|
using namespace edsc::op;
|
|
auto f = makeFunction("builder_blocks");
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
|
|
Block *b =
|
|
buildInNewBlock(TypeRange(), [&](ValueRange) { std_constant_index(0); });
|
|
appendToBlock(b, [&](ValueRange) { std_constant_index(1); });
|
|
appendToBlock(b, [&](ValueRange) { std_ret(); });
|
|
// Get back to entry block and add a branch into "b".
|
|
appendToBlock(&f.front(), [&](ValueRange) { std_br(b, {}); });
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @builder_blocks
|
|
// CHECK-NEXT: br ^bb1
|
|
// CHECK-NEXT: ^bb1: // pred: ^bb0
|
|
// CHECK-NEXT: constant 0 : index
|
|
// CHECK-NEXT: constant 1 : index
|
|
// CHECK-NEXT: return
|
|
// CHECK-NEXT: }
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_blocks) {
|
|
using namespace edsc::op;
|
|
auto f = makeFunction("builder_blocks");
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value c1(std_constant_int(42, 32)), c2(std_constant_int(1234, 32));
|
|
ReturnOp ret = std_ret();
|
|
|
|
Block *b1 = createBlock({c1.getType(), c1.getType()});
|
|
Block *b2 = buildInNewBlock({c1.getType(), c1.getType()},
|
|
[&](ValueRange args) { std_br(b1, args); });
|
|
// The insertion point within the toplevel function is now past b2, we will
|
|
// need to get back the entry block.
|
|
// This is what happens with unstructured control-flow.
|
|
appendToBlock(b1, [&](ValueRange args) {
|
|
Value r = args[0] + args[1];
|
|
std_br(b2, {args[0], r});
|
|
});
|
|
// Get back to entry block and add a branch into b1.
|
|
appendToBlock(&f.front(), [&](ValueRange) { std_br(b1, {c1, c2}); });
|
|
ret.erase();
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @builder_blocks
|
|
// CHECK: %{{.*}} = constant 42 : i32
|
|
// CHECK-NEXT: %{{.*}} = constant 1234 : i32
|
|
// CHECK-NEXT: br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
|
|
// CHECK-NEXT: ^bb1(%{{.*}}: i32, %{{.*}}: i32): // 2 preds: ^bb0, ^bb2
|
|
// CHECK-NEXT: %{{.*}} = addi %{{.*}}, %{{.*}} : i32
|
|
// CHECK-NEXT: br ^bb2(%{{.*}}, %{{.*}} : i32, i32)
|
|
// CHECK-NEXT: ^bb2(%{{.*}}: i32, %{{.*}}: i32): // pred: ^bb1
|
|
// CHECK-NEXT: br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
|
|
// CHECK-NEXT: }
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_cond_branch) {
|
|
auto f = makeFunction("builder_cond_branch", {},
|
|
{IntegerType::get(&globalContext(), 1)});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value c32(std_constant_int(32, 32)), c64(std_constant_int(64, 64)),
|
|
c42(std_constant_int(42, 32));
|
|
ReturnOp ret = std_ret();
|
|
|
|
Block *b1 = buildInNewBlock(c32.getType(), [&](ValueRange) { std_ret(); });
|
|
Block *b2 = buildInNewBlock({c64.getType(), c32.getType()},
|
|
[&](ValueRange) { std_ret(); });
|
|
// Get back to entry block and add a conditional branch.
|
|
appendToBlock(&f.front(), [&](ValueRange args) {
|
|
std_cond_br(args[0], b1, {c32}, b2, {c64, c42});
|
|
});
|
|
ret.erase();
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @builder_cond_branch
|
|
// CHECK: %{{.*}} = constant 32 : i32
|
|
// CHECK-NEXT: %{{.*}} = constant 64 : i64
|
|
// CHECK-NEXT: %{{.*}} = constant 42 : i32
|
|
// CHECK-NEXT: cond_br %{{.*}}, ^bb1(%{{.*}} : i32), ^bb2(%{{.*}}, %{{.*}} : i64, i32)
|
|
// CHECK-NEXT: ^bb1(%{{.*}}: i32): // pred: ^bb0
|
|
// CHECK-NEXT: return
|
|
// CHECK-NEXT: ^bb2(%{{.*}}: i64, %{{.*}}: i32): // pred: ^bb0
|
|
// CHECK-NEXT: return
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_helpers) {
|
|
using namespace edsc::op;
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType =
|
|
MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
|
|
ShapedType::kDynamicSize},
|
|
f32Type, {}, 0);
|
|
auto f =
|
|
makeFunction("builder_helpers", {}, {memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
// clang-format off
|
|
Value f7 = std_constant_float(llvm::APFloat(7.0f), f32Type);
|
|
MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1)),
|
|
vC(f.getArgument(2));
|
|
AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
|
|
Value lb0, lb1, lb2, ub0, ub1, ub2;
|
|
int64_t step0, step1, step2;
|
|
std::tie(lb0, ub0, step0) = vA.range(0);
|
|
std::tie(lb1, ub1, step1) = vA.range(1);
|
|
lb2 = vA.lb(2);
|
|
ub2 = vA.ub(2);
|
|
step2 = vA.step(2);
|
|
affineLoopNestBuilder({lb0, lb1}, {ub0, ub1}, {step0, step1}, [&](ValueRange ivs) {
|
|
Value i = ivs[0];
|
|
Value j = ivs[1];
|
|
affineLoopBuilder(lb2, ub2, step2, [&](Value k1){
|
|
C(i, j, k1) = f7 + A(i, j, k1) + B(i, j, k1);
|
|
});
|
|
affineLoopBuilder(lb2, ub2, step2, [&](Value k2){
|
|
C(i, j, k2) += A(i, j, k2) + B(i, j, k2);
|
|
});
|
|
});
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @builder_helpers
|
|
// CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
|
|
// CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
|
|
// CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
|
|
// CHECK-DAG: [[a:%.*]] = affine.load %arg0[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-DAG: [[b:%.*]] = addf {{.*}}, [[a]] : f32
|
|
// CHECK-DAG: [[c:%.*]] = affine.load %arg1[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-DAG: [[d:%.*]] = addf [[b]], [[c]] : f32
|
|
// CHECK-NEXT: affine.store [[d]], %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%{{.*}}) to affine_map<(d0) -> (d0)>(%{{.*}}) {
|
|
// CHECK-DAG: [[a:%.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-DAG: [[b:%.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-DAG: [[c:%.*]] = addf [[b]], [[a]] : f32
|
|
// CHECK-DAG: [[d:%.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-DAG: [[e:%.*]] = addf [[d]], [[c]] : f32
|
|
// CHECK-NEXT: affine.store [[e]], %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(insertion_in_block) {
|
|
using namespace edsc::op;
|
|
auto indexType = IndexType::get(&globalContext());
|
|
auto f = makeFunction("insertion_in_block", {}, {indexType, indexType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
std_constant_int(0, 32);
|
|
buildInNewBlock({}, [&](ValueRange) { std_constant_int(1, 32); });
|
|
std_constant_int(2, 32);
|
|
// clang-format off
|
|
// CHECK-LABEL: @insertion_in_block
|
|
// CHECK: {{.*}} = constant 0 : i32
|
|
// CHECK: {{.*}} = constant 2 : i32
|
|
// CHECK: ^bb1: // no predecessors
|
|
// CHECK: {{.*}} = constant 1 : i32
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(zero_and_std_sign_extendi_op_i1_to_i8) {
|
|
using namespace edsc::op;
|
|
auto i1Type = IntegerType::get(&globalContext(), 1);
|
|
auto i8Type = IntegerType::get(&globalContext(), 8);
|
|
auto memrefType = MemRefType::get({}, i1Type, {}, 0);
|
|
auto f = makeFunction("zero_and_std_sign_extendi_op", {},
|
|
{memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
AffineIndexedValue A(f.getArgument(0));
|
|
AffineIndexedValue B(f.getArgument(1));
|
|
// clang-format off
|
|
edsc::intrinsics::std_zero_extendi(A, i8Type);
|
|
edsc::intrinsics::std_sign_extendi(B, i8Type);
|
|
// CHECK-LABEL: @zero_and_std_sign_extendi_op
|
|
// CHECK: %[[SRC1:.*]] = affine.load
|
|
// CHECK: zexti %[[SRC1]] : i1 to i8
|
|
// CHECK: %[[SRC2:.*]] = affine.load
|
|
// CHECK: sexti %[[SRC2]] : i1 to i8
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(operator_or) {
|
|
auto i1Type = IntegerType::get(&globalContext(), /*width=*/1);
|
|
auto f = makeFunction("operator_or", {}, {i1Type, i1Type});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
|
|
using op::operator||;
|
|
Value lhs(f.getArgument(0));
|
|
Value rhs(f.getArgument(1));
|
|
lhs || rhs;
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @operator_or
|
|
// CHECK: [[ARG0:%.*]]: i1, [[ARG1:%.*]]: i1
|
|
// CHECK: or [[ARG0]], [[ARG1]]
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(operator_and) {
|
|
auto i1Type = IntegerType::get(&globalContext(), /*width=*/1);
|
|
auto f = makeFunction("operator_and", {}, {i1Type, i1Type});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
|
|
using op::operator&&;
|
|
using op::negate;
|
|
Value lhs(f.getArgument(0));
|
|
Value rhs(f.getArgument(1));
|
|
negate(lhs && rhs);
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @operator_and
|
|
// CHECK: [[ARG0:%.*]]: i1, [[ARG1:%.*]]: i1
|
|
// CHECK: [[AND:%.*]] = and [[ARG0]], [[ARG1]]
|
|
// CHECK: [[TRUE:%.*]] = constant true
|
|
// CHECK: subi [[TRUE]], [[AND]] : i1
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(divis_op_i32) {
|
|
using namespace edsc::op;
|
|
auto f = makeFunction("divis_op", {}, {});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
auto i32Type = builder.getI32Type();
|
|
std_divis(std_constant_int(10, i32Type), std_constant_int(2, i32Type));
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @divis_op
|
|
// CHECK-DAG: {{.*}} = constant 10
|
|
// CHECK-DAG: {{.*}} = constant 2
|
|
// CHECK-NEXT: {{.*}} = divi_signed
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(diviu_op_i32) {
|
|
using namespace edsc::op;
|
|
auto f = makeFunction("diviu_op", {}, {});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
auto i32Type = builder.getI32Type();
|
|
std_diviu(std_constant_int(10, i32Type), std_constant_int(2, i32Type));
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @diviu_op
|
|
// CHECK-DAG: {{.*}} = constant 10
|
|
// CHECK-DAG: {{.*}} = constant 2
|
|
// CHECK-NEXT: {{.*}} = divi_unsigned
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(fpext_f32_f64) {
|
|
using namespace edsc::op;
|
|
auto f = makeFunction("fpext", {}, {});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
auto f32Type = builder.getF32Type();
|
|
auto f64Type = builder.getF64Type();
|
|
std_fpext(std_constant_float(llvm::APFloat(10.0f), f32Type), f64Type);
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @fpext
|
|
// CHECK: {{.*}} = constant 1.0
|
|
// CHECK-NEXT: {{.*}} = fpext
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(fptrunc_f32_bf16) {
|
|
using namespace edsc::op;
|
|
auto f = makeFunction("fptrunc", {}, {});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
auto f32Type = builder.getF32Type();
|
|
auto bf16Type = builder.getBF16Type();
|
|
std_fptrunc(std_constant_float(llvm::APFloat(10.0f), f32Type), bf16Type);
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @fptrunc
|
|
// CHECK: {{.*}} = constant 1.0
|
|
// CHECK-NEXT: {{.*}} = fptrunc
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(select_op_i32) {
|
|
using namespace edsc::op;
|
|
auto i32Type = IntegerType::get(&globalContext(), 32);
|
|
auto memrefType = MemRefType::get(
|
|
{ShapedType::kDynamicSize, ShapedType::kDynamicSize}, i32Type, {}, 0);
|
|
auto f = makeFunction("select_op", {}, {memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value zero = std_constant_index(0), one = std_constant_index(1);
|
|
MemRefBoundsCapture vA(f.getArgument(0));
|
|
AffineIndexedValue A(f.getArgument(0));
|
|
affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
|
|
using namespace edsc::op;
|
|
Value i = ivs[0], j = ivs[1];
|
|
std_select(eq(i, zero), A(zero, zero), A(i, j));
|
|
std_select(ne(i, zero), A(zero, zero), A(i, j));
|
|
std_select(slt(i, zero), A(zero, zero), A(i, j));
|
|
std_select(sle(i, zero), A(zero, zero), A(i, j));
|
|
std_select(sgt(i, zero), A(zero, zero), A(i, j));
|
|
std_select(sge(i, zero), A(zero, zero), A(i, j));
|
|
std_select(ult(i, zero), A(zero, zero), A(i, j));
|
|
std_select(ule(i, zero), A(zero, zero), A(i, j));
|
|
std_select(ugt(i, zero), A(zero, zero), A(i, j));
|
|
std_select(uge(i, zero), A(zero, zero), A(i, j));
|
|
});
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @select_op
|
|
// CHECK: affine.for %{{.*}} = 0 to 1 {
|
|
// CHECK-NEXT: affine.for %{{.*}} = 0 to 1 {
|
|
// CHECK-DAG: {{.*}} = cmpi eq
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi ne
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi slt
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi sle
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi sgt
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi sge
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi ult
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi ule
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi ugt
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// CHECK-DAG: {{.*}} = cmpi uge
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-DAG: {{.*}} = affine.load
|
|
// CHECK-NEXT: {{.*}} = select
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(select_op_f32) {
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType = MemRefType::get(
|
|
{ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
|
|
auto f = makeFunction("select_op", {}, {memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
// clang-format off
|
|
Value zero = std_constant_index(0), one = std_constant_index(1);
|
|
MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1));
|
|
AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1));
|
|
affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
|
|
using namespace edsc::op;
|
|
Value i = ivs[0], j = ivs[1];
|
|
std_select(eq(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(ne(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(sge(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(sle(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(slt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(sgt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(uge(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(ule(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(ult(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
std_select(ugt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
|
|
});
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @select_op
|
|
// CHECK: affine.for %{{.*}} = 0 to 1 {
|
|
// CHECK-NEXT: affine.for %{{.*}} = 0 to 1 {
|
|
// CHECK-DAG: cmpf oeq
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf one
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf oge
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf ole
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf olt
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf ogt
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf oge
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf ole
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf olt
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// CHECK-DAG: cmpf ogt
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.load
|
|
// CHECK-DAG: affine.apply
|
|
// CHECK-NEXT: select
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// Inject an EDSC-constructed computation to exercise imperfectly nested 2-d
|
|
// tiling.
|
|
TEST_FUNC(tile_2d) {
|
|
auto memrefType =
|
|
MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
|
|
ShapedType::kDynamicSize},
|
|
FloatType::getF32(&globalContext()), {}, 0);
|
|
auto f = makeFunction("tile_2d", {}, {memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value zero = std_constant_index(0);
|
|
MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1)),
|
|
vC(f.getArgument(2));
|
|
AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1)),
|
|
C(f.getArgument(2));
|
|
Value i, j, k1, k2;
|
|
Value M(vC.ub(0)), N(vC.ub(1)), O(vC.ub(2));
|
|
|
|
// clang-format off
|
|
using namespace edsc::op;
|
|
affineLoopNestBuilder({zero, zero}, {M, N}, {1, 1}, [&](ValueRange ivs) {
|
|
i = ivs[0];
|
|
j = ivs[1];
|
|
affineLoopBuilder(zero, O, 1, [&](Value k) {
|
|
k1 = k;
|
|
C(i, j, k1) = A(i, j, k1) + B(i, j, k1);
|
|
});
|
|
affineLoopBuilder(zero, O, 1, [&](Value k) {
|
|
k2 = k;
|
|
C(i, j, k2) = A(i, j, k2) + B(i, j, k2);
|
|
});
|
|
});
|
|
// clang-format on
|
|
|
|
auto li = getForInductionVarOwner(i), lj = getForInductionVarOwner(j),
|
|
lk1 = getForInductionVarOwner(k1), lk2 = getForInductionVarOwner(k2);
|
|
auto indicesL1 = mlir::tile({li, lj}, {512, 1024}, {lk1, lk2});
|
|
auto lii1 = indicesL1[0][0], ljj1 = indicesL1[1][0];
|
|
mlir::tile({ljj1, lii1}, {32, 16}, ljj1);
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @tile_2d
|
|
// CHECK: %[[ZERO:.*]] = constant 0 : index
|
|
// CHECK: %[[M:[0-9]+]] = dim %arg2, %c0{{[_0-9]*}} : memref<?x?x?xf32>
|
|
// CHECK: %[[N:[0-9]+]] = dim %arg2, %c1{{[_0-9]*}} : memref<?x?x?xf32>
|
|
// CHECK: %[[P:[0-9]+]] = dim %arg2, %c2{{[_0-9]*}} : memref<?x?x?xf32>
|
|
// CHECK: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[M]]) step 512 {
|
|
// CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[N]]) step 1024 {
|
|
// CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[P]]) {
|
|
// CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 512)>(%{{.*}})[%[[M]]] step 16 {
|
|
// CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 1024)>(%{{.*}})[%[[N]]] step 32 {
|
|
// 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]]] {
|
|
// 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]]] {
|
|
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-NEXT: {{.*}} = addf {{.*}}, {{.*}} : f32
|
|
// CHECK-NEXT: affine.store {{.*}}, {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: affine.for %{{.*}} = affine_map<(d0) -> (d0)>(%[[ZERO]]) to affine_map<(d0) -> (d0)>(%[[P]]) {
|
|
// CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 512)>(%{{.*}})[%[[M]]] {
|
|
// CHECK-NEXT: affine.for %{{.*}} = max affine_map<(d0) -> (0, d0)>(%{{.*}}) to min affine_map<(d0)[s0] -> (s0, d0 + 1024)>(%{{.*}})[%[[N]]] {
|
|
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-NEXT: {{.*}} = affine.load {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// CHECK-NEXT: {{.*}}= addf {{.*}}, {{.*}} : f32
|
|
// CHECK-NEXT: affine.store {{.*}}, {{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?xf32>
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// Exercise StdIndexedValue for loads and stores.
|
|
TEST_FUNC(indirect_access) {
|
|
using namespace edsc::op;
|
|
auto memrefType = MemRefType::get({ShapedType::kDynamicSize},
|
|
FloatType::getF32(&globalContext()), {}, 0);
|
|
auto f = makeFunction("indirect_access", {},
|
|
{memrefType, memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value zero = std_constant_index(0);
|
|
MemRefBoundsCapture vC(f.getArgument(2));
|
|
AffineIndexedValue B(f.getArgument(1)), D(f.getArgument(3));
|
|
StdIndexedValue A(f.getArgument(0)), C(f.getArgument(2));
|
|
Value N(vC.ub(0));
|
|
|
|
// clang-format off
|
|
affineLoopBuilder(zero, N, 1, [&](Value i) {
|
|
C((Value)D(i)) = A((Value)B(i));
|
|
});
|
|
// clang-format on
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @indirect_access
|
|
// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>, %[[ARG2:.*]]: memref<?xf32>, %[[ARG3:.*]]: memref<?xf32>)
|
|
// CHECK-DAG: [[B:%.*]] = affine.load %[[ARG1]]
|
|
// CHECK-DAG: [[D:%.*]] = affine.load %[[ARG3]]
|
|
// CHECK: load %{{.*}}{{\[}}[[B]]{{\]}}
|
|
// CHECK: store %{{.*}}, %{{.*}}{{\[}}[[D]]{{\]}}
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// Exercise affine loads and stores build with empty maps.
|
|
TEST_FUNC(empty_map_load_store) {
|
|
using namespace edsc::op;
|
|
auto memrefType =
|
|
MemRefType::get({}, FloatType::getF32(&globalContext()), {}, 0);
|
|
auto f = makeFunction("empty_map_load_store", {},
|
|
{memrefType, memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value zero = std_constant_index(0);
|
|
Value one = std_constant_index(1);
|
|
AffineIndexedValue input(f.getArgument(0)), res(f.getArgument(1));
|
|
|
|
// clang-format off
|
|
affineLoopBuilder(zero, one, 1, [&](Value) {
|
|
res() = input();
|
|
});
|
|
// clang-format on
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @empty_map_load_store(
|
|
// CHECK: [[A:%.*]] = affine.load %{{.*}}[]
|
|
// CHECK: affine.store [[A]], %{{.*}}[]
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @affine_if_op
|
|
// CHECK: affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}[[s0:.*]], [[s1:.*]]{{\]}}
|
|
// CHECK-NOT: else
|
|
// CHECK: affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}[[s0:.*]], [[s1:.*]]{{\]}}
|
|
// CHECK-NEXT: } else {
|
|
// clang-format on
|
|
TEST_FUNC(affine_if_op) {
|
|
using namespace edsc::op;
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType = MemRefType::get(
|
|
{ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
|
|
auto f = makeFunction("affine_if_op", {}, {memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
|
|
Value zero = std_constant_index(0), ten = std_constant_index(10);
|
|
|
|
SmallVector<bool, 4> isEq = {false, false, false, false};
|
|
SmallVector<AffineExpr, 4> affineExprs = {
|
|
builder.getAffineDimExpr(0), // d0 >= 0
|
|
builder.getAffineDimExpr(1), // d1 >= 0
|
|
builder.getAffineSymbolExpr(0), // s0 >= 0
|
|
builder.getAffineSymbolExpr(1) // s1 >= 0
|
|
};
|
|
auto intSet = IntegerSet::get(2, 2, affineExprs, isEq);
|
|
|
|
SmallVector<Value, 4> affineIfArgs = {zero, zero, ten, ten};
|
|
intrinsics::affine_if(intSet, affineIfArgs, /*withElseRegion=*/false);
|
|
intrinsics::affine_if(intSet, affineIfArgs, /*withElseRegion=*/true);
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @linalg_generic_pointwise
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
|
|
// CHECK-SAME: ins({{.*}}memref<?x?xf32>, memref<?x?xf32>)
|
|
// CHECK-SAME: outs({{.*}}memref<?x?xf32>)
|
|
// CHECK: addf
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
|
|
// CHECK-SAME: ins({{.*}}memref<?x?xf32>, memref<?x?xf32>)
|
|
// CHECK-SAME: outs({{.*}}memref<?x?xf32>)
|
|
// CHECK: cmpf ogt
|
|
// CHECK: select
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}} : memref<?x?xf32>)
|
|
// CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?xf32>)
|
|
// CHECK: tanh
|
|
// clang-format on
|
|
TEST_FUNC(linalg_generic_pointwise_test) {
|
|
using namespace edsc;
|
|
using namespace edsc::ops;
|
|
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType = MemRefType::get(
|
|
{ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
|
|
auto f = makeFunction("linalg_generic_pointwise", {},
|
|
{memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
|
|
AffineExpr i, j;
|
|
bindDims(&globalContext(), i, j);
|
|
StructuredIndexed SA(A), SB(B), SC(C);
|
|
linalg_generic_pointwise_add(SA({i, j}), SB({i, j}), SC({i, j}));
|
|
linalg_generic_pointwise_max(SA({i, j}), SB({i, j}), SC({i, j}));
|
|
linalg_generic_pointwise_tanh(SA({i, j}), SC({i, j}));
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @linalg_generic_matmul
|
|
// CHECK: linalg.generic {
|
|
// 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)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : memref<?x?xf32>, memref<?x?xf32>)
|
|
// CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?xf32>)
|
|
/// CHECK: ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
|
|
// CHECK: %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
|
|
// CHECK: %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
|
|
// CHECK: linalg.yield %[[a4]] : f32
|
|
// CHECK: }
|
|
// clang-format on
|
|
TEST_FUNC(linalg_generic_matmul_test) {
|
|
using namespace edsc;
|
|
using namespace edsc::ops;
|
|
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType = MemRefType::get(
|
|
{ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
|
|
auto f = makeFunction("linalg_generic_matmul", {},
|
|
{memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
linalg_generic_matmul(f.getArguments());
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @linalg_generic_conv_nhwc
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2 * 3 + d4 * 5, d3 * 4 + d5 * 6, d6)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d1)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2, d3, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
|
|
// CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?x?x?xf32>)
|
|
/// CHECK: ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
|
|
// CHECK: %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
|
|
// CHECK: %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
|
|
// CHECK: linalg.yield %[[a4]] : f32
|
|
// CHECK: }
|
|
// clang-format on
|
|
TEST_FUNC(linalg_generic_conv_nhwc) {
|
|
using namespace edsc;
|
|
using namespace edsc::ops;
|
|
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType =
|
|
MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
|
|
ShapedType::kDynamicSize, ShapedType::kDynamicSize},
|
|
f32Type, {}, 0);
|
|
auto f = makeFunction("linalg_generic_conv_nhwc", {},
|
|
{memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
linalg_generic_conv_nhwc(f.getArguments(),
|
|
/*strides=*/{3, 4}, /*dilations=*/{5, 6});
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @linalg_generic_dilated_conv_nhwc
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d3 * 3 + d5 * 5, d4 * 4 + d6 * 6, d2)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d2, d1)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d3, d4, d1 + d2 * 7)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
|
|
// CHECK-SAME: outs(%{{[a-z0-9]*}} : memref<?x?x?x?xf32>)
|
|
// CHECK: ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
|
|
// CHECK: %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
|
|
// CHECK: %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
|
|
// CHECK: linalg.yield %[[a4]] : f32
|
|
// CHECK: }
|
|
// clang-format on
|
|
TEST_FUNC(linalg_generic_dilated_conv_nhwc) {
|
|
using namespace edsc;
|
|
using namespace edsc::ops;
|
|
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType =
|
|
MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
|
|
ShapedType::kDynamicSize, ShapedType::kDynamicSize},
|
|
f32Type, {}, 0);
|
|
auto f = makeFunction("linalg_generic_dilated_conv_nhwc", {},
|
|
{memrefType, memrefType, memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
linalg_generic_dilated_conv_nhwc(f.getArguments(),
|
|
/*depth_multiplier=*/7,
|
|
/*strides=*/{3, 4}, /*dilations=*/{5, 6});
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @linalg_metadata_ops
|
|
// CHECK: linalg.reshape {{.*}} [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d2)>] : memref<4x8x16xf32> into memref<32x16xf32>
|
|
// CHECK: linalg.reshape {{.*}} [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d2)>] : memref<32x16xf32> into memref<4x8x16xf32>
|
|
// clang-format on
|
|
TEST_FUNC(linalg_metadata_ops) {
|
|
using linalg::ReassociationExprs;
|
|
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType = MemRefType::get({4, 8, 16}, f32Type, {}, 0);
|
|
auto f = makeFunction("linalg_metadata_ops", {}, {memrefType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
AffineExpr i, j, k;
|
|
bindDims(&globalContext(), i, j, k);
|
|
Value v(f.getArgument(0));
|
|
SmallVector<ReassociationExprs, 2> maps = {ReassociationExprs({i, j}),
|
|
ReassociationExprs({k})};
|
|
auto reshaped = linalg_reshape(v, maps);
|
|
linalg_reshape(memrefType, reshaped, maps);
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @linalg_tensors
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, memref<?x?xf32>)
|
|
// CHECK: addf
|
|
// CHECK: } -> tensor<?x?xf32>
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CHECK: cmpf ogt
|
|
// CHECK: select
|
|
// CHECK: } -> tensor<?x?xf32>
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}} : tensor<?x?xf32>)
|
|
// CHECK: tanh
|
|
// CHECK: } -> tensor<?x?xf32>
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, memref<?x?xf32>)
|
|
// CHECK: mulf
|
|
// CHECK: } -> tensor<?x?xf32>
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>],
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
|
|
// CHECK-SAME: ins(%{{[a-z0-9]*}}, %{{[a-z0-9]*}} : tensor<?x?xf32>, memref<?x?xf32>)
|
|
// CHECK-SAME: outs(%{{[a-z0-9]*}} : tensor<?x?xf32>)
|
|
// CHECK: mulf
|
|
// CHECK: addf
|
|
// CHECK: } -> tensor<?x?xf32>
|
|
// clang-format on
|
|
TEST_FUNC(linalg_tensors_test) {
|
|
using namespace edsc;
|
|
using namespace edsc::ops;
|
|
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto memrefType = MemRefType::get(
|
|
{ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
|
|
auto tensorType = RankedTensorType::get(
|
|
{ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type);
|
|
auto f =
|
|
makeFunction("linalg_tensors", {}, {tensorType, memrefType, tensorType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
|
|
AffineExpr i, j;
|
|
bindDims(&globalContext(), i, j);
|
|
StructuredIndexed SA(A), SB(B), SC(C);
|
|
Value added = linalg_generic_pointwise_add(SA({i, j}), SB({i, j}), SC({i, j}))
|
|
->getResult(0);
|
|
Value maxed = linalg_generic_pointwise_max(
|
|
SA({i, j}), StructuredIndexed(added)({i, j}), SC({i, j}))
|
|
->getResult(0);
|
|
Value tanhed = linalg_generic_pointwise_tanh(StructuredIndexed(maxed)({i, j}),
|
|
SC({i, j}))
|
|
->getResult(0);
|
|
Value o1 = linalg_generic_matmul(A, B, tanhed, tensorType)->getResult(0);
|
|
linalg_generic_matmul(A, B, o1, tensorType);
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(vector_extractelement_op_i32) {
|
|
using namespace edsc::op;
|
|
auto f = makeFunction("vector_extractelement_op", {}, {});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
auto i32Type = builder.getI32Type();
|
|
auto vectorType = VectorType::get(/*shape=*/{8}, i32Type);
|
|
vector_extract_element(
|
|
i32Type, std_constant(vectorType, builder.getI32VectorAttr({10})),
|
|
std_constant_int(0, i32Type));
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: @vector_extractelement_op
|
|
// CHECK-DAG: {{.*}} = constant dense<10>
|
|
// CHECK-DAG: {{.*}} = constant 0
|
|
// CHECK-NEXT: {{.*}} = vector.extractelement
|
|
// clang-format on
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @memref_vector_matmul_test(
|
|
// CHECK-SAME: %[[A:.*]]: memref<?x?xvector<4x16xf32>>,
|
|
// CHECK-SAME: %[[B:.*]]: memref<?x?xvector<16x8xf32>>,
|
|
// CHECK-SAME: %[[C:.*]]: memref<?x?xvector<4x8xf32>>)
|
|
// CHECK: linalg.generic {{{.*}}}
|
|
// CHECK-SAME: ins(%[[A]], %[[B]] : memref<?x?xvector<4x16xf32>>, memref<?x?xvector<16x8xf32>>)
|
|
// CHECK-SAME: outs(%[[C]] : memref<?x?xvector<4x8xf32>>)
|
|
// CHECK: vector.contract{{.*}}[affine_map<(d0, d1, d2) -> (d0, d2)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2) -> (d2, d1)>,
|
|
// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1)>],
|
|
// CHECK-SAME: {{.*}}["parallel", "parallel", "reduction"]
|
|
// clang-format on
|
|
TEST_FUNC(memref_vector_matmul_test) {
|
|
using namespace edsc;
|
|
using namespace edsc::ops;
|
|
|
|
int64_t M = 4, N = 8, K = 16;
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto mkVectorType = VectorType::get({M, K}, f32Type);
|
|
auto knVectorType = VectorType::get({K, N}, f32Type);
|
|
auto mnVectorType = VectorType::get({M, N}, f32Type);
|
|
auto typeA =
|
|
MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
|
|
mkVectorType, {}, 0);
|
|
auto typeB =
|
|
MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
|
|
knVectorType, {}, 0);
|
|
auto typeC =
|
|
MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
|
|
mnVectorType, {}, 0);
|
|
auto f = makeFunction("memref_vector_matmul_test", {}, {typeA, typeB, typeC});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
|
|
auto contractionBuilder = [](ValueRange args) {
|
|
assert(args.size() == 3 && "expected 3 block arguments");
|
|
(linalg_yield(vector_contraction_matmul(args[0], args[1], args[2])));
|
|
};
|
|
linalg_generic_matmul(A, B, C, contractionBuilder);
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
TEST_FUNC(builder_loop_for_yield) {
|
|
auto indexType = IndexType::get(&globalContext());
|
|
auto f32Type = FloatType::getF32(&globalContext());
|
|
auto f = makeFunction("builder_loop_for_yield", {},
|
|
{indexType, indexType, indexType, indexType});
|
|
|
|
OpBuilder builder(f.getBody());
|
|
ScopedContext scope(builder, f.getLoc());
|
|
Value init0 = std_constant_float(llvm::APFloat(1.0f), f32Type);
|
|
Value init1 = std_constant_float(llvm::APFloat(2.0f), f32Type);
|
|
Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
|
|
d(f.getArgument(3));
|
|
using namespace edsc::op;
|
|
auto results = loopNestBuilder(a - b, c + d, a, {init0, init1},
|
|
[&](Value iv, ValueRange args) {
|
|
Value sum = args[0] + args[1];
|
|
return scf::ValueVector{args[1], sum};
|
|
})
|
|
.getResults();
|
|
results[0] + results[1];
|
|
|
|
// clang-format off
|
|
// CHECK-LABEL: func @builder_loop_for_yield(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
|
|
// CHECK: [[init0:%.*]] = constant
|
|
// CHECK: [[init1:%.*]] = constant
|
|
// CHECK-DAG: [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
|
|
// CHECK-DAG: [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
|
|
// CHECK-NEXT: [[res:%[0-9]+]]:2 = scf.for %{{.*}} = [[r0]] to [[r1]] step {{.*}} iter_args([[arg0:%.*]] = [[init0]], [[arg1:%.*]] = [[init1]]) -> (f32, f32) {
|
|
// CHECK: [[sum:%[0-9]+]] = addf [[arg0]], [[arg1]] : f32
|
|
// CHECK: scf.yield [[arg1]], [[sum]] : f32, f32
|
|
// CHECK: addf [[res]]#0, [[res]]#1 : f32
|
|
// clang-format on
|
|
|
|
f.print(llvm::outs(), OpPrintingFlags().useLocalScope());
|
|
f.erase();
|
|
}
|
|
|
|
int main() {
|
|
RUN_TESTS();
|
|
return 0;
|
|
}
|