EDSC: move FileCheck tests into the source file
EDSC provide APIs for constructing and modifying the IR. These APIs are
currently tested by a "test" module pass that reads the dummy IR (empty
functions), recognizes certain function names and injects the IR into those
functions based on their name. This situation is unsatisfactory because the
expected outcome of the test lives in a different file than the input to the
test, i.e. the API calls.
Create a new binary for tests that constructs the IR from scratch using EDSC
APIs and prints it. Put FileCheck comments next to the printing. This removes
the need to have a file with dummy inputs and assert on its contents in the
test driver. The test source includes a simplistic test harness that runs all
functions marked as TEST_FUNC but intentionally does not include any
value-testing functionality.
PiperOrigin-RevId: 235886629
diff --git a/lib/EDSC/LowerEDSCTestPass.cpp b/lib/EDSC/LowerEDSCTestPass.cpp
index 5c51745..cdbe6e5 100644
--- a/lib/EDSC/LowerEDSCTestPass.cpp
+++ b/lib/EDSC/LowerEDSCTestPass.cpp
@@ -44,267 +44,6 @@
#include "mlir/EDSC/reference-impl.inc"
PassResult LowerEDSCTestPass::runOnFunction(Function *f) {
- // Inject a EDSC-constructed infinite loop implemented by mutual branching
- // between two blocks, following the pattern:
- //
- // br ^bb1
- // ^bb1:
- // br ^bb2
- // ^bb2:
- // br ^bb1
- //
- // Use blocks with arguments.
- if (f->getName().strref() == "blocks") {
- using namespace edsc::op;
-
- FuncBuilder builder(f);
- edsc::ScopedEDSCContext context;
- // Declare two blocks. Note that we must declare the blocks before creating
- // branches to them.
- auto type = builder.getIntegerType(32);
- edsc::Expr arg1(type), arg2(type), arg3(type), arg4(type), r(type);
- edsc::StmtBlock b1 = edsc::block({arg1, arg2}, {}),
- b2 = edsc::block({arg3, arg4}, {});
- auto c1 = edsc::constantInteger(type, 42);
- auto c2 = edsc::constantInteger(type, 1234);
-
- // Make an infinite loops by branching between the blocks. Note that copy-
- // assigning a block won't work well with branches, update the body instead.
- b1.set({r = arg1 + arg2, edsc::Branch(b2, {arg1, r})});
- b2.set({edsc::Branch(b1, {arg3, arg4})});
- auto instr = edsc::Branch(b2, {c1, c2});
-
- // Remove the existing 'return' from the function, reset the builder after
- // the instruction iterator invalidation and emit a branch to b2. This
- // should also emit blocks b2 and b1 that appear as successors to the
- // current block after the branch instruction is insterted.
- f->begin()->clear();
- builder.setInsertionPoint(&*f->begin(), f->begin()->begin());
- edsc::MLIREmitter(&builder, f->getLoc()).emitStmt(instr);
- }
-
- // Inject two EDSC-constructed blocks with arguments and a conditional branch
- // instruction that transfers control to these blocks.
- if (f->getName().strref() == "cond_branch") {
- FuncBuilder builder(f);
- edsc::ScopedEDSCContext context;
- auto i1 = builder.getIntegerType(1);
- auto i32 = builder.getIntegerType(32);
- auto i64 = builder.getIntegerType(64);
- edsc::Expr arg1(i32), arg2(i64), arg3(i32);
- // Declare two blocks with different numbers of arguments.
- edsc::StmtBlock b1 = edsc::block({arg1}, {edsc::Return()}),
- b2 = edsc::block({arg2, arg3}, {edsc::Return()});
- edsc::Expr funcArg(i1);
-
- // Inject the conditional branch.
- auto condBranch = edsc::CondBranch(
- funcArg, b1, {edsc::constantInteger(i32, 32)}, b2,
- {edsc::constantInteger(i64, 64), edsc::constantInteger(i32, 42)});
-
- assert(f->getNumArguments() == 1 && "cond_branch must have 1 argument");
- assert(f->getArgument(0)->getType() == i1 &&
- "the argument of cond_branch must have i1 type");
-
- // Remove the existing `return` instruction from the entry block of the
- // function. It will be replaced by the conditional branch.
- f->begin()->clear();
- builder.setInsertionPoint(&*f->begin(), f->begin()->begin());
- edsc::MLIREmitter(&builder, f->getLoc())
- .bind(edsc::Bindable(funcArg), f->getArgument(0))
- .emitStmt(condBranch);
- return success();
- }
-
- // Inject a EDSC-constructed `for` loop with bounds coming from function
- // arguments.
- if (f->getName().strref() == "dynamic_for_func_args") {
- assert(!f->getBlocks().empty() && "dynamic_for should not be empty");
- FuncBuilder builder(&f->getBlocks().front(),
- f->getBlocks().front().begin());
- assert(f->getNumArguments() == 2 && "dynamic_for expected 4 arguments");
- for (const auto *arg : f->getArguments()) {
- (void)arg;
- assert(arg->getType().isIndex() &&
- "dynamic_for expected index arguments");
- }
-
- using namespace edsc::op;
- Type index = IndexType::get(f->getContext());
- edsc::ScopedEDSCContext context;
- edsc::Expr lb(index), ub(index), step(index);
- step = edsc::constantInteger(index, 3);
- auto loop = edsc::For(lb, ub, step, {lb * step + ub, step + lb});
- edsc::MLIREmitter(&builder, f->getLoc())
- .bind(edsc::Bindable(lb), f->getArgument(0))
- .bind(edsc::Bindable(ub), f->getArgument(1))
- .emitStmt(loop);
- return success();
- }
-
- // Inject a EDSC-constructed `for` loop with non-constant bounds that are
- // obtained from AffineApplyOp (also constructed using EDSC operator
- // overloads).
- if (f->getName().strref() == "dynamic_for") {
- assert(!f->getBlocks().empty() && "dynamic_for should not be empty");
- FuncBuilder builder(&f->getBlocks().front(),
- f->getBlocks().front().begin());
- assert(f->getNumArguments() == 4 && "dynamic_for expected 4 arguments");
- for (const auto *arg : f->getArguments()) {
- (void)arg;
- assert(arg->getType().isIndex() &&
- "dynamic_for expected index arguments");
- }
-
- Type index = IndexType::get(f->getContext());
- edsc::ScopedEDSCContext context;
- edsc::Expr lb1(index), lb2(index), ub1(index), ub2(index), step(index);
- using namespace edsc::op;
- auto lb = lb1 - lb2;
- auto ub = ub1 + ub2;
- auto loop = edsc::For(lb, ub, step, {});
- edsc::MLIREmitter(&builder, f->getLoc())
- .bind(edsc::Bindable(lb1), f->getArgument(0))
- .bind(edsc::Bindable(lb2), f->getArgument(1))
- .bind(edsc::Bindable(ub1), f->getArgument(2))
- .bind(edsc::Bindable(ub2), f->getArgument(3))
- .bindConstant<ConstantIndexOp>(edsc::Bindable(step), 2)
- .emitStmt(loop);
-
- return success();
- }
- if (f->getName().strref() == "max_min_for") {
- assert(!f->getBlocks().empty() && "max_min_for should not be empty");
- FuncBuilder builder(&f->getBlocks().front(),
- f->getBlocks().front().begin());
- assert(f->getNumArguments() == 4 && "max_min_for expected 4 arguments");
- assert(std::all_of(f->args_begin(), f->args_end(),
- [](const Value *s) { return s->getType().isIndex(); }) &&
- "max_min_for expected index arguments");
-
- edsc::ScopedEDSCContext context;
- edsc::Expr lb1(f->getArgument(0)->getType());
- edsc::Expr lb2(f->getArgument(1)->getType());
- edsc::Expr ub1(f->getArgument(2)->getType());
- edsc::Expr ub2(f->getArgument(3)->getType());
- edsc::Expr iv(builder.getIndexType());
- edsc::Expr step = edsc::constantInteger(builder.getIndexType(), 1);
- auto loop =
- edsc::MaxMinFor(edsc::Bindable(iv), {lb1, lb2}, {ub1, ub2}, step, {});
- edsc::MLIREmitter(&builder, f->getLoc())
- .bind(edsc::Bindable(lb1), f->getArgument(0))
- .bind(edsc::Bindable(lb2), f->getArgument(1))
- .bind(edsc::Bindable(ub1), f->getArgument(2))
- .bind(edsc::Bindable(ub2), f->getArgument(3))
- .emitStmt(loop);
-
- return success();
- }
- if (f->getName().strref() == "call_indirect") {
- assert(!f->getBlocks().empty() && "call_indirect should not be empty");
- FuncBuilder builder(&f->getBlocks().front(),
- f->getBlocks().front().begin());
- Function *callee = f->getModule()->getNamedFunction("callee");
- Function *calleeArgs = f->getModule()->getNamedFunction("callee_args");
- Function *secondOrderCallee =
- f->getModule()->getNamedFunction("second_order_callee");
- assert(callee && calleeArgs && secondOrderCallee &&
- "could not find required declarations");
-
- auto funcRetIndexType = builder.getFunctionType({}, builder.getIndexType());
-
- edsc::ScopedEDSCContext context;
- edsc::Expr func(callee->getType()), funcArgs(calleeArgs->getType()),
- secondOrderFunc(secondOrderCallee->getType());
- auto stmt = edsc::call(func, {});
- auto chainedCallResult =
- edsc::call(edsc::call(secondOrderFunc, funcRetIndexType, {func}),
- builder.getIndexType(), {});
- auto argsCall =
- edsc::call(funcArgs, {chainedCallResult, chainedCallResult});
- edsc::MLIREmitter(&builder, f->getLoc())
- .bindConstant<ConstantOp>(edsc::Bindable(func),
- builder.getFunctionAttr(callee))
- .bindConstant<ConstantOp>(edsc::Bindable(funcArgs),
- builder.getFunctionAttr(calleeArgs))
- .bindConstant<ConstantOp>(edsc::Bindable(secondOrderFunc),
- builder.getFunctionAttr(secondOrderCallee))
- .emitStmt(stmt)
- .emitStmt(chainedCallResult)
- .emitStmt(argsCall);
-
- return success();
- }
-
- // Inject an EDSC-constructed computation that assigns Stmt and uses the LHS.
- if (f->getName().strref().contains("assignments")) {
- FuncBuilder builder(f);
- edsc::ScopedEDSCContext context;
- edsc::MLIREmitter emitter(&builder, f->getLoc());
-
- edsc::Expr zero = emitter.zero();
- edsc::Expr one = emitter.one();
- auto args = emitter.makeBoundFunctionArguments(f);
- auto views = emitter.makeBoundMemRefViews(args.begin(), args.end());
-
- Type indexType = builder.getIndexType();
- edsc::Expr i(indexType);
- edsc::Expr A = args[0], B = args[1], C = args[2];
- edsc::Expr M = views[0].dim(0);
- // clang-format off
- using namespace edsc::op;
- edsc::Stmt scalarA, scalarB, tmp;
- auto block = edsc::block({
- For(i, zero, M, one, {
- scalarA = load(A, {i}),
- scalarB = load(B, {i}),
- tmp = scalarA * scalarB,
- store(tmp, C, {i})
- }),
- });
- // clang-format on
- emitter.emitStmts(block.getBody());
- }
-
- // Inject an EDSC-constructed computation to exercise imperfectly nested 2-d
- // tiling.
- if (f->getName().strref().contains("tile_2d")) {
- FuncBuilder builder(f);
- edsc::ScopedEDSCContext context;
- edsc::MLIREmitter emitter(&builder, f->getLoc());
-
- edsc::Expr zero = emitter.zero();
- edsc::Expr one = emitter.one();
- auto args = emitter.makeBoundFunctionArguments(f);
- auto views = emitter.makeBoundMemRefViews(args.begin(), args.end());
-
- Type indexType = builder.getIndexType();
- edsc::Expr i(indexType), j(indexType), k1(indexType), k2(indexType);
- edsc::Indexed A(args[0]), B(args[1]), C(args[2]);
- edsc::Expr M = views[0].dim(0), N = views[0].dim(1), O = views[0].dim(2);
- // clang-format off
- using namespace edsc::op;
- edsc::Stmt scalarA, scalarB, tmp;
- auto block = edsc::block({
- For(ArrayRef<edsc::Expr>{i, j}, {zero, zero}, {M, N}, {one, one}, {
- For(k1, zero, O, one, {
- C({i, j, k1}) = A({i, j, k1}) + B({i, j, k1})
- }),
- For(k2, zero, O, one, {
- C({i, j, k2}) = A({i, j, k2}) + B({i, j, k2})
- }),
- }),
- });
- // clang-format on
- emitter.emitStmts(block.getBody());
-
- auto li = emitter.getAffineForOp(i), lj = emitter.getAffineForOp(j),
- lk1 = emitter.getAffineForOp(k1), lk2 = emitter.getAffineForOp(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);
- }
-
f->walk([](Instruction *op) {
if (op->getName().getStringRef() == "print") {
auto opName = op->getAttrOfType<StringAttr>("op");
diff --git a/test/EDSC/Test.h b/test/EDSC/Test.h
new file mode 100644
index 0000000..6b02108
--- /dev/null
+++ b/test/EDSC/Test.h
@@ -0,0 +1,72 @@
+//===- Test.h - Simple macros for API unit tests ----------------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file define simple macros for declaring test functions and running them.
+// The actual checking must be performed on the outputs with FileCheck.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TEST_TEST_H_
+#define MLIR_TEST_TEST_H_
+
+#include <functional>
+#include <vector>
+
+namespace test_detail {
+// Returns a mutable list of known test functions. Used internally by test
+// macros to add and run tests. This function is static to ensure it creates a
+// new list in each test file.
+static std::vector<std::function<void()>> &tests() {
+ static std::vector<std::function<void()>> list;
+ return list;
+}
+
+// Test registration class. Used internally by test macros to register tests
+// during static allocation.
+struct TestRegistration {
+ explicit TestRegistration(std::function<void()> func) {
+ test_detail::tests().push_back(func);
+ }
+};
+} // end namespace test_detail
+
+/// Declares a test function with the given name and adds it to the list of
+/// known tets. The body of the function must follow immediately. Example:
+///
+/// TEST_FUNC(mytest) {
+/// // CHECK: expected-output-here
+/// emitSomethingToStdOut();
+/// }
+///
+#define TEST_FUNC(name) \
+ void name(); \
+ static test_detail::TestRegistration name##Registration(name); \
+ void name()
+
+/// Runs all registered tests. Example:
+///
+/// int main() {
+/// RUN_TESTS();
+/// return 0;
+/// }
+#define RUN_TESTS \
+ []() { \
+ for (auto f : test_detail::tests()) \
+ f(); \
+ }
+
+#endif // MLIR_TEST_TEST_H_
diff --git a/test/EDSC/api-test.cpp b/test/EDSC/api-test.cpp
new file mode 100644
index 0000000..09d3751
--- /dev/null
+++ b/test/EDSC/api-test.cpp
@@ -0,0 +1,423 @@
+//===- APITest.cpp - Test for EDSC APIs -----------------------------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+// RUN: %p/api-test | FileCheck %s
+
+#include "mlir/AffineOps/AffineOps.h"
+#include "mlir/EDSC/MLIREmitter.h"
+#include "mlir/EDSC/Types.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/StandardOps/StandardOps.h"
+#include "mlir/Transforms/LoopUtils.h"
+
+#include "Test.h"
+
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+static MLIRContext &globalContext() {
+ static thread_local MLIRContext context;
+ return context;
+}
+
+static std::unique_ptr<Function> makeFunction(StringRef name,
+ ArrayRef<Type> results = {},
+ ArrayRef<Type> args = {}) {
+ auto &ctx = globalContext();
+ auto function = llvm::make_unique<Function>(
+ UnknownLoc::get(&ctx), name, FunctionType::get(args, results, &ctx));
+ function->addEntryBlock();
+ return function;
+}
+
+// Inject a EDSC-constructed infinite loop implemented by mutual branching
+// between two blocks, following the pattern:
+//
+// br ^bb1
+// ^bb1:
+// br ^bb2
+// ^bb2:
+// br ^bb1
+//
+// Use blocks with arguments.
+TEST_FUNC(blocks) {
+ using namespace edsc::op;
+
+ auto f = makeFunction("blocks");
+ FuncBuilder builder(f.get());
+ edsc::ScopedEDSCContext context;
+ // Declare two blocks. Note that we must declare the blocks before creating
+ // branches to them.
+ auto type = builder.getIntegerType(32);
+ edsc::Expr arg1(type), arg2(type), arg3(type), arg4(type), r(type);
+ edsc::StmtBlock b1 = edsc::block({arg1, arg2}, {}),
+ b2 = edsc::block({arg3, arg4}, {});
+ auto c1 = edsc::constantInteger(type, 42);
+ auto c2 = edsc::constantInteger(type, 1234);
+
+ // Make an infinite loops by branching between the blocks. Note that copy-
+ // assigning a block won't work well with branches, update the body instead.
+ b1.set({r = arg1 + arg2, edsc::Branch(b2, {arg1, r})});
+ b2.set({edsc::Branch(b1, {arg3, arg4})});
+ auto instr = edsc::Branch(b2, {c1, c2});
+
+ // Emit a branch to b2. This should also emit blocks b2 and b1 that appear as
+ // successors to the current block after the branch instruction is insterted.
+ edsc::MLIREmitter(&builder, f->getLoc()).emitStmt(instr);
+
+ // clang-format off
+ // CHECK-LABEL: @blocks
+ // CHECK: %c42_i32 = constant 42 : i32
+ // CHECK-NEXT: %c1234_i32 = constant 1234 : i32
+ // CHECK-NEXT: br ^bb1(%c42_i32, %c1234_i32 : i32, i32)
+ // CHECK-NEXT: ^bb1(%0: i32, %1: i32): // 2 preds: ^bb0, ^bb2
+ // CHECK-NEXT: br ^bb2(%0, %1 : i32, i32)
+ // CHECK-NEXT: ^bb2(%2: i32, %3: i32): // pred: ^bb1
+ // CHECK-NEXT: %4 = addi %2, %3 : i32
+ // CHECK-NEXT: br ^bb1(%2, %4 : i32, i32)
+ // CHECK-NEXT: }
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+// Inject two EDSC-constructed blocks with arguments and a conditional branch
+// instruction that transfers control to these blocks.
+TEST_FUNC(cond_branch) {
+ auto f =
+ makeFunction("cond_branch", {}, {IntegerType::get(1, &globalContext())});
+
+ FuncBuilder builder(f.get());
+ edsc::ScopedEDSCContext context;
+ auto i1 = builder.getIntegerType(1);
+ auto i32 = builder.getIntegerType(32);
+ auto i64 = builder.getIntegerType(64);
+ edsc::Expr arg1(i32), arg2(i64), arg3(i32);
+ // Declare two blocks with different numbers of arguments.
+ edsc::StmtBlock b1 = edsc::block({arg1}, {edsc::Return()}),
+ b2 = edsc::block({arg2, arg3}, {edsc::Return()});
+ edsc::Expr funcArg(i1);
+
+ // Inject the conditional branch.
+ auto condBranch = edsc::CondBranch(
+ funcArg, b1, {edsc::constantInteger(i32, 32)}, b2,
+ {edsc::constantInteger(i64, 64), edsc::constantInteger(i32, 42)});
+
+ builder.setInsertionPoint(&*f->begin(), f->begin()->begin());
+ edsc::MLIREmitter(&builder, f->getLoc())
+ .bind(edsc::Bindable(funcArg), f->getArgument(0))
+ .emitStmt(condBranch);
+
+ // clang-format off
+ // CHECK-LABEL: @cond_branch
+ // CHECK: %c0 = constant 0 : index
+ // CHECK-NEXT: %c1 = constant 1 : index
+ // CHECK-NEXT: %c32_i32 = constant 32 : i32
+ // CHECK-NEXT: %c64_i64 = constant 64 : i64
+ // CHECK-NEXT: %c42_i32 = constant 42 : i32
+ // CHECK-NEXT: cond_br %arg0, ^bb1(%c32_i32 : i32), ^bb2(%c64_i64, %c42_i32 : i64, i32)
+ // CHECK-NEXT: ^bb1(%0: i32): // pred: ^bb0
+ // CHECK-NEXT: return
+ // CHECK-NEXT: ^bb2(%1: i64, %2: i32): // pred: ^bb0
+ // CHECK-NEXT: return
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+// Inject a EDSC-constructed `for` loop with bounds coming from function
+// arguments.
+TEST_FUNC(dynamic_for_func_args) {
+ auto indexType = IndexType::get(&globalContext());
+ auto f = makeFunction("dynamic_for_func_args", {}, {indexType, indexType});
+ FuncBuilder builder(f.get());
+
+ using namespace edsc::op;
+ Type index = IndexType::get(f->getContext());
+ edsc::ScopedEDSCContext context;
+ edsc::Expr lb(index), ub(index), step(index);
+ step = edsc::constantInteger(index, 3);
+ auto loop = edsc::For(lb, ub, step, {lb * step + ub, step + lb});
+ edsc::MLIREmitter(&builder, f->getLoc())
+ .bind(edsc::Bindable(lb), f->getArgument(0))
+ .bind(edsc::Bindable(ub), f->getArgument(1))
+ .emitStmt(loop)
+ .emitStmt(edsc::Return());
+
+ // clang-format off
+ // CHECK-LABEL: func @dynamic_for_func_args(%arg0: index, %arg1: index) {
+ // CHECK: for %i0 = (d0) -> (d0)(%arg0) to (d0) -> (d0)(%arg1) step 3 {
+ // CHECK: {{.*}} = affine.apply (d0) -> (d0 * 3)(%arg0)
+ // CHECK: {{.*}} = affine.apply (d0, d1) -> (d0 * 3 + d1)(%arg0, %arg1)
+ // CHECK: {{.*}} = affine.apply (d0) -> (d0 + 3)(%arg0)
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+// Inject a EDSC-constructed `for` loop with non-constant bounds that are
+// obtained from AffineApplyOp (also constructed using EDSC operator
+// overloads).
+TEST_FUNC(dynamic_for) {
+ auto indexType = IndexType::get(&globalContext());
+ auto f = makeFunction("dynamic_for", {},
+ {indexType, indexType, indexType, indexType});
+ FuncBuilder builder(f.get());
+
+ edsc::ScopedEDSCContext context;
+ edsc::Expr lb1(indexType), lb2(indexType), ub1(indexType), ub2(indexType),
+ step(indexType);
+ using namespace edsc::op;
+ auto lb = lb1 - lb2;
+ auto ub = ub1 + ub2;
+ auto loop = edsc::For(lb, ub, step, {});
+ edsc::MLIREmitter(&builder, f->getLoc())
+ .bind(edsc::Bindable(lb1), f->getArgument(0))
+ .bind(edsc::Bindable(lb2), f->getArgument(1))
+ .bind(edsc::Bindable(ub1), f->getArgument(2))
+ .bind(edsc::Bindable(ub2), f->getArgument(3))
+ .bindConstant<ConstantIndexOp>(edsc::Bindable(step), 2)
+ .emitStmt(loop);
+
+ // clang-format off
+ // CHECK-LABEL: func @dynamic_for(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
+ // CHECK: %0 = affine.apply (d0, d1) -> (d0 - d1)(%arg0, %arg1)
+ // CHECK-NEXT: %1 = affine.apply (d0, d1) -> (d0 + d1)(%arg2, %arg3)
+ // CHECK-NEXT: for %i0 = (d0) -> (d0)(%0) to (d0) -> (d0)(%1) step 2 {
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+// Inject a EDSC-constructed empty `for` loop with max/min bounds that
+// corresponds to
+//
+// for max(%arg0, %arg1) to (%arg2, %arg3) step 1
+//
+TEST_FUNC(max_min_for) {
+ auto indexType = IndexType::get(&globalContext());
+ auto f = makeFunction("max_min_for", {},
+ {indexType, indexType, indexType, indexType});
+ FuncBuilder builder(f.get());
+
+ edsc::ScopedEDSCContext context;
+ edsc::Expr lb1(f->getArgument(0)->getType());
+ edsc::Expr lb2(f->getArgument(1)->getType());
+ edsc::Expr ub1(f->getArgument(2)->getType());
+ edsc::Expr ub2(f->getArgument(3)->getType());
+ edsc::Expr iv(builder.getIndexType());
+ edsc::Expr step = edsc::constantInteger(builder.getIndexType(), 1);
+ auto loop =
+ edsc::MaxMinFor(edsc::Bindable(iv), {lb1, lb2}, {ub1, ub2}, step, {});
+ edsc::MLIREmitter(&builder, f->getLoc())
+ .bind(edsc::Bindable(lb1), f->getArgument(0))
+ .bind(edsc::Bindable(lb2), f->getArgument(1))
+ .bind(edsc::Bindable(ub1), f->getArgument(2))
+ .bind(edsc::Bindable(ub2), f->getArgument(3))
+ .emitStmt(loop);
+
+ // clang-format off
+ // CHECK-LABEL: func @max_min_for(%arg0: index, %arg1: index, %arg2: index,
+ // %arg3: index) { CHECK: for %i0 = max (d0, d1) -> (d0, d1)(%arg0, %arg1) to
+ // min (d0, d1) -> (d0, d1)(%arg2, %arg3) {
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+// Inject EDSC-constructed chain of indirect calls that corresponds to
+//
+// @callee()
+// var x = @second_order_callee(@callee)
+// @callee_args(x, x)
+//
+TEST_FUNC(call_indirect) {
+ auto indexType = IndexType::get(&globalContext());
+ auto callee = makeFunction("callee");
+ auto calleeArgs = makeFunction("callee_args", {}, {indexType, indexType});
+ auto secondOrderCallee =
+ makeFunction("second_order_callee",
+ {FunctionType::get({}, {indexType}, &globalContext())},
+ {FunctionType::get({}, {}, &globalContext())});
+ auto f = makeFunction("call_indirect");
+ FuncBuilder builder(f.get());
+
+ auto funcRetIndexType = builder.getFunctionType({}, builder.getIndexType());
+
+ edsc::ScopedEDSCContext context;
+ edsc::Expr func(callee->getType()), funcArgs(calleeArgs->getType()),
+ secondOrderFunc(secondOrderCallee->getType());
+ auto stmt = edsc::call(func, {});
+ auto chainedCallResult =
+ edsc::call(edsc::call(secondOrderFunc, funcRetIndexType, {func}),
+ builder.getIndexType(), {});
+ auto argsCall = edsc::call(funcArgs, {chainedCallResult, chainedCallResult});
+ edsc::MLIREmitter(&builder, f->getLoc())
+ .bindConstant<ConstantOp>(edsc::Bindable(func),
+ builder.getFunctionAttr(callee.get()))
+ .bindConstant<ConstantOp>(edsc::Bindable(funcArgs),
+ builder.getFunctionAttr(calleeArgs.get()))
+ .bindConstant<ConstantOp>(
+ edsc::Bindable(secondOrderFunc),
+ builder.getFunctionAttr(secondOrderCallee.get()))
+ .emitStmt(stmt)
+ .emitStmt(chainedCallResult)
+ .emitStmt(argsCall);
+
+ // clang-format off
+ // CHECK-LABEL: @call_indirect
+ // CHECK: %f = constant @callee : () -> ()
+ // CHECK: %f_0 = constant @callee_args : (index, index) -> ()
+ // CHECK: %f_1 = constant @second_order_callee : (() -> ()) -> (() -> index)
+ // CHECK: call_indirect %f() : () -> ()
+ // CHECK: %0 = call_indirect %f_1(%f) : (() -> ()) -> (() -> index)
+ // CHECK: %1 = call_indirect %0() : () -> index
+ // CHECK: call_indirect %f_0(%1, %1) : (index, index) -> ()
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+// Inject EDSC-constructed 1-D pointwise-add loop with assignments to scalars,
+// `dim` indicates the shape of the memref storing the values.
+static std::unique_ptr<Function> makeAssignmentsFunction(int dim) {
+ auto memrefType =
+ MemRefType::get({dim}, FloatType::getF32(&globalContext()), {}, 0);
+ auto f =
+ makeFunction("assignments", {}, {memrefType, memrefType, memrefType});
+ FuncBuilder builder(f.get());
+
+ edsc::ScopedEDSCContext context;
+ edsc::MLIREmitter emitter(&builder, f->getLoc());
+
+ edsc::Expr zero = emitter.zero();
+ edsc::Expr one = emitter.one();
+ auto args = emitter.makeBoundFunctionArguments(f.get());
+ auto views = emitter.makeBoundMemRefViews(args.begin(), args.end());
+
+ Type indexType = builder.getIndexType();
+ edsc::Expr i(indexType);
+ edsc::Expr A = args[0], B = args[1], C = args[2];
+ edsc::Expr M = views[0].dim(0);
+ // clang-format off
+ using namespace edsc::op;
+ edsc::Stmt scalarA, scalarB, tmp;
+ auto block = edsc::block({
+ For(i, zero, M, one, {
+ scalarA = load(A, {i}),
+ scalarB = load(B, {i}),
+ tmp = scalarA * scalarB,
+ store(tmp, C, {i})
+ }),
+ });
+ // clang-format on
+ emitter.emitStmts(block.getBody());
+
+ return f;
+}
+
+TEST_FUNC(assignments_1) {
+ auto f = makeAssignmentsFunction(4);
+
+ // clang-format off
+ // CHECK-LABEL: func @assignments(%arg0: memref<4xf32>, %arg1: memref<4xf32>, %arg2: memref<4xf32>) {
+ // CHECK: for %[[iv:.*]] = 0 to 4 {
+ // CHECK: %[[a:.*]] = load %arg0[%[[iv]]] : memref<4xf32>
+ // CHECK: %[[b:.*]] = load %arg1[%[[iv]]] : memref<4xf32>
+ // CHECK: %[[tmp:.*]] = mulf %[[a]], %[[b]] : f32
+ // CHECK: store %[[tmp]], %arg2[%[[iv]]] : memref<4xf32>
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+TEST_FUNC(assignments_2) {
+ auto f = makeAssignmentsFunction(-1);
+
+ // clang-format off
+ // CHECK-LABEL: func @assignments(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
+ // CHECK: for %[[iv:.*]] = {{.*}} to {{.*}} {
+ // CHECK: %[[a:.*]] = load %arg0[%[[iv]]] : memref<?xf32>
+ // CHECK: %[[b:.*]] = load %arg1[%[[iv]]] : memref<?xf32>
+ // CHECK: %[[tmp:.*]] = mulf %[[a]], %[[b]] : f32
+ // CHECK: store %[[tmp]], %arg2[%[[iv]]] : memref<?xf32>
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+// Inject an EDSC-constructed computation to exercise imperfectly nested 2-d
+// tiling.
+TEST_FUNC(tile_2d) {
+ auto memrefType =
+ MemRefType::get({-1, -1, -1}, FloatType::getF32(&globalContext()), {}, 0);
+ auto f = makeFunction("tile_2d", {}, {memrefType, memrefType, memrefType});
+ FuncBuilder builder(f.get());
+
+ edsc::ScopedEDSCContext context;
+ edsc::MLIREmitter emitter(&builder, f->getLoc());
+
+ edsc::Expr zero = emitter.zero();
+ edsc::Expr one = emitter.one();
+ auto args = emitter.makeBoundFunctionArguments(f.get());
+ auto views = emitter.makeBoundMemRefViews(args.begin(), args.end());
+
+ Type indexType = builder.getIndexType();
+ edsc::Expr i(indexType), j(indexType), k1(indexType), k2(indexType);
+ edsc::Indexed A(args[0]), B(args[1]), C(args[2]);
+ edsc::Expr M = views[0].dim(0), N = views[0].dim(1), O = views[0].dim(2);
+ // clang-format off
+ using namespace edsc::op;
+ edsc::Stmt scalarA, scalarB, tmp;
+ auto block = edsc::block({
+ For(ArrayRef<edsc::Expr>{i, j}, {zero, zero}, {M, N}, {one, one}, {
+ For(k1, zero, O, one, {
+ C({i, j, k1}) = A({i, j, k1}) + B({i, j, k1})
+ }),
+ For(k2, zero, O, one, {
+ C({i, j, k2}) = A({i, j, k2}) + B({i, j, k2})
+ }),
+ }),
+ });
+ // clang-format on
+ emitter.emitStmts(block.getBody());
+
+ auto li = emitter.getAffineForOp(i), lj = emitter.getAffineForOp(j),
+ lk1 = emitter.getAffineForOp(k1), lk2 = emitter.getAffineForOp(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: for %i0 = (d0) -> (d0)({{.*}}) to (d0) -> (d0)({{.*}}) step 512 {
+ // CHECK: for %i1 = (d0) -> (d0)({{.*}}) to (d0) -> (d0)({{.*}}) step 1024 {
+ // CHECK: for %i2 = (d0) -> (d0)({{.*}}) to (d0) -> (d0)({{.*}}) {
+ // CHECK: for %i3 = max {{.*}}, %i0) to min {{.*}}, %i0) step 16 {
+ // CHECK: for %i4 = max {{.*}}, %i1) to min {{.*}}, %i1) step 32 {
+ // CHECK: for %i5 = max {{.*}}, %i1, %i4) to min {{.*}}, %i1, %i4) {
+ // CHECK: for %i6 = max {{.*}}, %i0, %i3) to min {{.*}}, %i0, %i3) {
+ // CHECK: for %i7 = (d0) -> (d0)({{.*}}) to (d0) -> (d0)({{.*}}) {
+ // CHECK: for %i8 = max {{.*}}, %i0) to min {{.*}}, %i0) {
+ // CHECK: for %i9 = max {{.*}}, %i1) to min {{.*}}, %i1) {
+ // clang-format on
+ f->print(llvm::outs());
+}
+
+int main() {
+ RUN_TESTS();
+ return 0;
+}
diff --git a/test/EDSC/for-loops.mlir b/test/EDSC/for-loops.mlir
deleted file mode 100644
index 3e6711b..0000000
--- a/test/EDSC/for-loops.mlir
+++ /dev/null
@@ -1,148 +0,0 @@
-// RUN: mlir-opt -lower-edsc-test %s | FileCheck %s
-
-// Maps used in dynamic_for below.
-// CHECK-DAG: #[[idmap:.*]] = (d0) -> (d0)
-// CHECK-DAG: #[[diffmap:.*]] = (d0, d1) -> (d0 - d1)
-// CHECK-DAG: #[[addmap:.*]] = (d0, d1) -> (d0 + d1)
-// CHECK-DAG: #[[prodconstmap:.*]] = (d0) -> (d0 * 3)
-// CHECK-DAG: #[[addconstmap:.*]] = (d0) -> (d0 + 3)
-// CHECK-DAG: #[[composedmap:.*]] = (d0, d1) -> (d0 * 3 + d1)
-// CHECK-DAG: #[[id2dmap:.*]] = (d0, d1) -> (d0, d1)
-
-// This function will be detected by the test pass that will insert
-// EDSC-constructed blocks with arguments forming an infinite loop.
-// CHECK-LABEL: @blocks
-func @blocks() {
- return
-//CHECK: %c42_i32 = constant 42 : i32
-//CHECK-NEXT: %c1234_i32 = constant 1234 : i32
-//CHECK-NEXT: br ^bb1(%c42_i32, %c1234_i32 : i32, i32)
-//CHECK-NEXT: ^bb1(%0: i32, %1: i32): // 2 preds: ^bb0, ^bb2
-//CHECK-NEXT: br ^bb2(%0, %1 : i32, i32)
-//CHECK-NEXT: ^bb2(%2: i32, %3: i32): // pred: ^bb1
-//CHECK-NEXT: %4 = addi %2, %3 : i32
-//CHECK-NEXT: br ^bb1(%2, %4 : i32, i32)
-//CHECK-NEXT: }
-}
-
-// This function will be detected by the test pass that will insert two
-// EDSC-constructed blocks with arguments and a conditional branch that goes to
-// both of them.
-func @cond_branch(%arg0: i1) {
- return
-// CHECK-LABEL: @cond_branch
-// CHECK-NEXT: %c0 = constant 0 : index
-// CHECK-NEXT: %c1 = constant 1 : index
-// CHECK-NEXT: %c32_i32 = constant 32 : i32
-// CHECK-NEXT: %c64_i64 = constant 64 : i64
-// CHECK-NEXT: %c42_i32 = constant 42 : i32
-// CHECK-NEXT: cond_br %arg0, ^bb1(%c32_i32 : i32), ^bb2(%c64_i64, %c42_i32 : i64, i32)
-// CHECK-NEXT: ^bb1(%0: i32): // pred: ^bb0
-// CHECK-NEXT: return
-// CHECK-NEXT: ^bb2(%1: i64, %2: i32): // pred: ^bb0
-// CHECK-NEXT: return
-}
-
-// This function will be detected by the test pass that will insert an
-// EDSC-constructed empty `for` loop that corresponds to
-// for %arg0 to %arg1 step 2
-// before the `return` instruction.
-// CHECK-LABEL: func @dynamic_for_func_args(%arg0: index, %arg1: index) {
-// CHECK: for %i0 = #[[idmap]](%arg0) to #[[idmap]](%arg1) step 3 {
-// CHECK: {{.*}} = affine.apply #[[prodconstmap]](%arg0)
-// CHECK: {{.*}} = affine.apply #[[composedmap]](%arg0, %arg1)
-// CHECK: {{.*}} = affine.apply #[[addconstmap]](%arg0)
-func @dynamic_for_func_args(%arg0 : index, %arg1 : index) {
- return
-}
-
-// This function will be detected by the test pass that will insert an
-// EDSC-constructed empty `for` loop that corresponds to
-// for (%arg0 - %arg1) to (%arg2 + %arg3) step 2
-// before the `return` instruction.
-//
-// CHECK-LABEL: func @dynamic_for(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
-// CHECK: %[[step:.*]] = constant 2 : index
-// CHECK: %[[lb:.*]] = affine.apply #[[diffmap]](%arg0, %arg1)
-// CHECK: %[[ub:.*]] = affine.apply #[[addmap]](%arg2, %arg3)
-// CHECK: for %i0 = #[[idmap]](%[[lb]]) to #[[idmap]](%[[ub]]) step 2 {
-func @dynamic_for(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) {
- return
-}
-
-// These functions will be detected by the test pass that will insert an
-// EDSC-constructed 1-D pointwise-add loop with assignments to scalars before
-// the `return` instruction.
-//
-// CHECK-LABEL: func @assignments_1(%arg0: memref<4xf32>, %arg1: memref<4xf32>, %arg2: memref<4xf32>) {
-// CHECK: for %[[iv:.*]] = 0 to 4 {
-// CHECK: %[[a:.*]] = load %arg0[%[[iv]]] : memref<4xf32>
-// CHECK: %[[b:.*]] = load %arg1[%[[iv]]] : memref<4xf32>
-// CHECK: %[[tmp:.*]] = mulf %[[a]], %[[b]] : f32
-// CHECK: store %[[tmp]], %arg2[%[[iv]]] : memref<4xf32>
-func @assignments_1(%arg0: memref<4xf32>, %arg1: memref<4xf32>, %arg2: memref<4xf32>) {
- return
-}
-
-// CHECK-LABEL: func @assignments_2(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
-// CHECK: for %[[iv:.*]] = {{.*}} to {{.*}} {
-// CHECK: %[[a:.*]] = load %arg0[%[[iv]]] : memref<?xf32>
-// CHECK: %[[b:.*]] = load %arg1[%[[iv]]] : memref<?xf32>
-// CHECK: %[[tmp:.*]] = mulf %[[a]], %[[b]] : f32
-// CHECK: store %[[tmp]], %arg2[%[[iv]]] : memref<?xf32>
-func @assignments_2(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
- return
-}
-
-// This function will be detected by the test pass that will insert an
-// EDSC-constructed empty `for` loop with max/min bounds that correspond to
-// for max(%arg0, %arg1) to (%arg2, %arg3) step 1
-// before the `return` instruction.
-//
-// CHECK-LABEL: func @max_min_for(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
-// CHECK: for %i0 = max #[[id2dmap]](%arg0, %arg1) to min #[[id2dmap]](%arg2, %arg3) {
-func @max_min_for(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) {
- return
-}
-
-func @callee()
-func @callee_args(index, index)
-func @second_order_callee(() -> ()) -> (() -> (index))
-
-// This function will be detected by the test pass that will insert an
-// EDSC-constructed chain of indirect calls that corresponds to
-// @callee()
-// var x = @second_order_callee(@callee)
-// @callee_args(x, x)
-// before the `return` instruction.
-//
-// CHECK-LABEL: @call_indirect
-// CHECK: %f = constant @callee : () -> ()
-// CHECK: %f_0 = constant @callee_args : (index, index) -> ()
-// CHECK: %f_1 = constant @second_order_callee : (() -> ()) -> (() -> index)
-// CHECK: call_indirect %f() : () -> ()
-// CHECK: %0 = call_indirect %f_1(%f) : (() -> ()) -> (() -> index)
-// CHECK: %1 = call_indirect %0() : () -> index
-// CHECK: call_indirect %f_0(%1, %1) : (index, index) -> ()
-func @call_indirect() {
- return
-}
-
-// This function will be detected by the test pass that will insert an
-// EDSC-constructed chain of indirect calls that corresponds to an imperfectly
-// nested loop nest with 2 common outer loops and 2 inner 1-d loops.
-// CHECK-LABEL: func @tile_2d
-// CHECK: for %i0 = #[[idmap]]({{.*}}) to #[[idmap]]({{.*}}) step 512 {
-// CHECK: for %i1 = #[[idmap]]({{.*}}) to #[[idmap]]({{.*}}) step 1024 {
-// CHECK: for %i2 = #[[idmap]]({{.*}}) to #[[idmap]]({{.*}}) {
-// CHECK: for %i3 = max #{{.*}}, %i0) to min #{{.*}}, %i0) step 16 {
-// CHECK: for %i4 = max #{{.*}}, %i1) to min #{{.*}}, %i1) step 32 {
-// CHECK: for %i5 = max #{{.*}}, %i1, %i4) to min #{{.*}}, %i1, %i4) {
-// CHECK: for %i6 = max #{{.*}}, %i0, %i3) to min #{{.*}}, %i0, %i3) {
-// CHECK: for %i7 = #[[idmap]]({{.*}}) to #[[idmap]]({{.*}}) {
-// CHECK: for %i8 = max #{{.*}}, %i0) to min #{{.*}}, %i0) {
-// CHECK: for %i9 = max #{{.*}}, %i1) to min #{{.*}}, %i1) {
-func @tile_2d(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?x?xf32>) {
- return
-}
-
