maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

Add an eager API version for BR and COND_BR 

When building unstructured control-flow there is a need to construct mlir::Block* before being able to fill them. This invites goto-style programming.

This CL introduces an alternative eager API for BR and COND_BR in which blocks are created eagerly and captured on the fly.

This allows reducing the number of calls to `BlockBuilder` from 4 to 2 in the `builder_blocks_eager` test and from 3 to 2 in the `builder_cond_branch_eager` test.

PiperOrigin-RevId: 237046114
This commit is contained in:
Nicolas Vasilache 2019-03-06 07:46:36 -08:00 committed by jpienaar
parent 38f1d2d77e
commit ee4a80bbd6
3 changed files with 174 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
mlir/include/mlir/EDSC/Intrinsics.h
View File

@ -43,6 +43,22 @@ namespace intrinsics {
/// All Handles have already captured previously constructed IR objects.
ValueHandle BR(BlockHandle bh, ArrayRef<ValueHandle> operands);
/// Creates a new mlir::Block* and branches to it from the current block.
/// Argument types are specified by `operands`.
/// Captures the new block in `bh` and the actual `operands` in `captures`. To
/// insert the new mlir::Block*, a local ScopedContext is constructed and
/// released to the current block. The branch instruction is then added to the
/// new block.
///
/// Prerequisites:
/// `b` has not yet captured an mlir::Block*.
/// No `captures` have captured any mlir::Value*.
/// All `operands` have already captured an mlir::Value*
/// captures.size() == operands.size()
/// captures and operands are pairwise of the same type.
ValueHandle BR(BlockHandle *bh, ArrayRef<ValueHandle *> captures,
ArrayRef<ValueHandle> operands);
/// Branches into the mlir::Block* captured by BlockHandle `trueBranch` with
/// `trueOperands` if `cond` evaluates to `true` (resp. `falseBranch` and
/// `falseOperand` if `cond` evaluates to `false`).
@ -53,6 +69,29 @@ ValueHandle COND_BR(ValueHandle cond, BlockHandle trueBranch,
ArrayRef<ValueHandle> trueOperands, BlockHandle falseBranch,
ArrayRef<ValueHandle> falseOperands);
/// Eagerly creates new mlir::Block* with argument types specified by
/// `trueOperands`/`falseOperands`.
/// Captures the new blocks in `trueBranch`/`falseBranch` and the arguments in
/// `trueCaptures/falseCaptures`.
/// To insert the new mlir::Block*, a local ScopedContext is constructed and
/// released. The branch instruction is then added in the original location and
/// targeting the eagerly constructed blocks.
///
/// Prerequisites:
/// `trueBranch`/`falseBranch` has not yet captured an mlir::Block*.
/// No `trueCaptures`/`falseCaptures` have captured any mlir::Value*.
/// All `trueOperands`/`trueOperands` have already captured an mlir::Value*
/// `trueCaptures`.size() == `trueOperands`.size()
/// `falseCaptures`.size() == `falseOperands`.size()
/// `trueCaptures` and `trueOperands` are pairwise of the same type
/// `falseCaptures` and `falseOperands` are pairwise of the same type.
ValueHandle COND_BR(ValueHandle cond, BlockHandle *trueBranch,
ArrayRef<ValueHandle *> trueCaptures,
ArrayRef<ValueHandle> trueOperands,
BlockHandle *falseBranch,
ArrayRef<ValueHandle *> falseCaptures,
ArrayRef<ValueHandle> falseOperands);
////////////////////////////////////////////////////////////////////////////////
// TODO(ntv): Intrinsics below this line should be TableGen'd.
////////////////////////////////////////////////////////////////////////////////

54
mlir/lib/EDSC/Intrinsics.cpp
View File

@ -31,6 +31,36 @@ ValueHandle mlir::edsc::intrinsics::BR(BlockHandle bh,
SmallVector<Value *, 4> ops(operands.begin(), operands.end());
return ValueHandle::create<BranchOp>(bh.getBlock(), ops);
}
static void enforceEmptyCapturesMatchOperands(ArrayRef<ValueHandle *> captures,
ArrayRef<ValueHandle> operands) {
assert(captures.size() == operands.size() &&
"Expected same number of captures as operands");
for (auto it : llvm::zip(captures, operands)) {
(void)it;
assert(!std::get<0>(it)->hasValue() &&
"Unexpected already captured ValueHandle");
assert(std::get<1>(it) && "Expected already captured ValueHandle");
assert(std::get<0>(it)->getType() == std::get<1>(it).getType() &&
"Expected the same type for capture and operand");
}
}
ValueHandle mlir::edsc::intrinsics::BR(BlockHandle *bh,
ArrayRef<ValueHandle *> captures,
ArrayRef<ValueHandle> operands) {
assert(!*bh && "Unexpected already captured BlockHandle");
enforceEmptyCapturesMatchOperands(captures, operands);
{ // Clone the scope explicitly to avoid modifying the insertion point in the
// current scope which result in surprising usage.
auto *currentB = ScopedContext::getBuilder();
FuncBuilder b(currentB->getInsertionBlock(), currentB->getInsertionPoint());
Location loc = ScopedContext::getLocation();
ScopedContext scope(b, loc);
BlockBuilder(bh, captures)({/* no body */});
} // Release before adding the branch to the eagerly created block.
SmallVector<Value *, 4> ops(operands.begin(), operands.end());
return ValueHandle::create<BranchOp>(bh->getBlock(), ops);
}
ValueHandle
mlir::edsc::intrinsics::COND_BR(ValueHandle cond, BlockHandle trueBranch,
@ -43,9 +73,33 @@ mlir::edsc::intrinsics::COND_BR(ValueHandle cond, BlockHandle trueBranch,
falseBranch.getBlock(), falseOps);
}
ValueHandle mlir::edsc::intrinsics::COND_BR(
ValueHandle cond, BlockHandle *trueBranch,
ArrayRef<ValueHandle *> trueCaptures, ArrayRef<ValueHandle> trueOperands,
BlockHandle *falseBranch, ArrayRef<ValueHandle *> falseCaptures,
ArrayRef<ValueHandle> falseOperands) {
assert(!*trueBranch && "Unexpected already captured BlockHandle");
assert(!*falseBranch && "Unexpected already captured BlockHandle");
enforceEmptyCapturesMatchOperands(trueCaptures, trueOperands);
enforceEmptyCapturesMatchOperands(falseCaptures, falseOperands);
{ // Clone the scope explicitly.
auto *currentB = ScopedContext::getBuilder();
FuncBuilder b(currentB->getInsertionBlock(), currentB->getInsertionPoint());
Location loc = ScopedContext::getLocation();
ScopedContext scope(b, loc);
BlockBuilder(trueBranch, trueCaptures)({/* no body */});
BlockBuilder(falseBranch, falseCaptures)({/* no body */});
} // Release before adding the branch to the eagerly created block.
SmallVector<Value *, 4> trueOps(trueOperands.begin(), trueOperands.end());
SmallVector<Value *, 4> falseOps(falseOperands.begin(), falseOperands.end());
return ValueHandle::create<CondBranchOp>(
cond, trueBranch->getBlock(), trueOps, falseBranch->getBlock(), falseOps);
}
////////////////////////////////////////////////////////////////////////////////
// TODO(ntv): Intrinsics below this line should be TableGen'd.
////////////////////////////////////////////////////////////////////////////////
ValueHandle mlir::edsc::intrinsics::RETURN(ArrayRef<ValueHandle> operands) {
SmallVector<Value *, 4> ops(operands.begin(), operands.end());
return ValueHandle::create<ReturnOp>(ops);

82
mlir/test/EDSC/builder-api-test.cpp
View File

@ -197,6 +197,50 @@ TEST_FUNC(builder_blocks) {
f->print(llvm::outs());
}
TEST_FUNC(builder_blocks_eager) {
using namespace edsc;
using namespace edsc::intrinsics;
using namespace edsc::op;
auto f = makeFunction("builder_blocks_eager");
ScopedContext scope(f.get());
ValueHandle c1(ValueHandle::create<ConstantIntOp>(42, 32)),
c2(ValueHandle::create<ConstantIntOp>(1234, 32));
ValueHandle arg1(c1.getType()), arg2(c1.getType()), arg3(c1.getType()),
arg4(c1.getType()), r(c1.getType());
// clang-format off
BlockHandle b1, b2;
{ // Toplevel function scope.
BR(&b1, {&arg1, &arg2}, {c1, c2}); // eagerly builds a new block for b1
// We cannot construct b2 eagerly with a `BR(&b2, ...)` call from within b1
// because it would result in b2 being nested under b1 which is not what we
// want in this test.
BlockBuilder(&b2, {&arg3, &arg4})({
// Instead, construct explicitly
BR(b1, {arg3, arg4}),
});
/// And come back to append into b1 once b2 exists.
BlockBuilder(b1, Append())({
r = arg1 + arg2,
BR(b2, {arg1, r}),
});
}
// CHECK-LABEL: @builder_blocks_eager
// 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: %2 = addi %0, %1 : i32
// CHECK-NEXT: br ^bb2(%0, %2 : i32, i32)
// CHECK-NEXT: ^bb2(%3: i32, %4: i32): // pred: ^bb1
// CHECK-NEXT: br ^bb1(%3, %4 : i32, i32)
// CHECK-NEXT: }
// clang-format on
f->print(llvm::outs());
}
TEST_FUNC(builder_cond_branch) {
using namespace edsc;
using namespace edsc::intrinsics;
@ -211,7 +255,6 @@ TEST_FUNC(builder_cond_branch) {
ValueHandle arg1(c32.getType()), arg2(c64.getType()), arg3(c32.getType());
BlockHandle b1, b2, functionBlock(&f->front());
;
BlockBuilder(&b1, {&arg1})({
RETURN({}),
});
@ -237,6 +280,43 @@ TEST_FUNC(builder_cond_branch) {
f->print(llvm::outs());
}
TEST_FUNC(builder_cond_branch_eager) {
using namespace edsc;
using namespace edsc::intrinsics;
using namespace edsc::op;
auto f = makeFunction("builder_cond_branch_eager", {},
{IntegerType::get(1, &globalContext())});
ScopedContext scope(f.get());
ValueHandle funcArg(f->getArgument(0));
ValueHandle c32(ValueHandle::create<ConstantIntOp>(32, 32)),
c64(ValueHandle::create<ConstantIntOp>(64, 64)),
c42(ValueHandle::create<ConstantIntOp>(42, 32));
ValueHandle arg1(c32.getType()), arg2(c64.getType()), arg3(c32.getType());
// clang-format off
BlockHandle b1, b2;
COND_BR(funcArg, &b1, {&arg1}, {c32}, &b2, {&arg2, &arg3}, {c64, c42});
BlockBuilder(b1, Append())({
RETURN({}),
});
BlockBuilder(b2, Append())({
RETURN({}),
});
// CHECK-LABEL: @builder_cond_branch_eager
// CHECK: %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());
}
int main() {
RUN_TESTS();
return 0;