[mlir] Provide OpBuilder-based replacements for edsc::BlockBuilder

The ScopedBuilder class in EDSC is being gradually phased out in favor of core
OpBuilder-based helpers with callbacks. Provide helper functions that are
compatible with `edsc::ScopedContext` and can be used to create and populate
blocks using callbacks that take block arguments as callback arguments. This
removes the need for `edsc::BlockHandle`, forward-declaration of `Value`s used
for block arguments and the tag `edsc::Append` class, leading to noticable
reduction in the verbosity of the code using helper functions.

Remove "eager mode" construction tests that are only relevant to the
`BlockBuilder`-based approach.

`edsc::BlockHandle` and `edsc::BlockBuilder` are now deprecated and will be
removed soon.

Differential Revision: https://reviews.llvm.org/D82008
This commit is contained in:
Alex Zinenko 2020-06-17 16:46:45 +02:00
parent e31e8f1ed5
commit a75e09372e
7 changed files with 139 additions and 139 deletions

View File

@ -55,7 +55,7 @@ private:
std::unique_ptr<BuilderType> builder;
};
inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
inline void defaultRegionBuilder(ValueRange args) {}
/// Build a `linalg.generic` op with the specified `inputs`, `outputs` and
/// `region`.
@ -76,8 +76,7 @@ inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
Operation *makeGenericLinalgOp(
ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
ArrayRef<StructuredIndexed> outputs,
function_ref<void(ArrayRef<BlockArgument>)> regionBuilder =
defaultRegionBuilder,
function_ref<void(ValueRange)> regionBuilder = defaultRegionBuilder,
ArrayRef<Value> otherValues = {}, ArrayRef<Attribute> otherAttributes = {});
namespace ops {
@ -89,11 +88,11 @@ using edsc::StructuredIndexed;
/// Build the body of a region to compute a scalar multiply, under the current
/// ScopedContext, at the current insert point.
void mulRegionBuilder(ArrayRef<BlockArgument> args);
void mulRegionBuilder(ValueRange args);
/// Build the body of a region to compute a scalar multiply-accumulate, under
/// the current ScopedContext, at the current insert point.
void macRegionBuilder(ArrayRef<BlockArgument> args);
void macRegionBuilder(ValueRange args);
/// TODO(ntv): In the future we should tie these implementations to something in
/// Tablegen that generates the proper interfaces and the proper sugared named
@ -149,7 +148,7 @@ Operation *linalg_generic_pointwise_max(StructuredIndexed I1,
// TODO(ntv): Implement more useful pointwise operations on a per-need basis.
using MatmulRegionBuilder = function_ref<void(ArrayRef<BlockArgument> args)>;
using MatmulRegionBuilder = function_ref<void(ValueRange args)>;
/// Build a linalg.generic, under the current ScopedContext, at the current
/// insert point, that computes:

View File

@ -52,6 +52,7 @@ using std_sign_extendi = ValueBuilder<SignExtendIOp>;
/// Prerequisites:
/// All Handles have already captured previously constructed IR objects.
BranchOp std_br(BlockHandle bh, ValueRange operands);
BranchOp std_br(Block *block, ValueRange operands);
/// Creates a new mlir::Block* and branches to it from the current block.
/// Argument types are specified by `operands`.
@ -78,6 +79,8 @@ BranchOp std_br(BlockHandle *bh, ArrayRef<Type> types,
CondBranchOp std_cond_br(Value cond, BlockHandle trueBranch,
ValueRange trueOperands, BlockHandle falseBranch,
ValueRange falseOperands);
CondBranchOp std_cond_br(Value cond, Block *trueBranch, ValueRange trueOperands,
Block *falseBranch, ValueRange falseOperands);
/// Eagerly creates new mlir::Block* with argument types specified by
/// `trueOperands`/`falseOperands`.

View File

@ -187,6 +187,8 @@ private:
// trying to enter a Block that has already been constructed.
class Append {};
/// Deprecated. Use buildInNewBlock or appendToBlock instead.
///
/// A BlockBuilder is a NestedBuilder for mlir::Block*.
/// This exists by opposition to LoopBuilder which is not related to an
/// mlir::Block* but to a mlir::Value.
@ -231,6 +233,8 @@ private:
BlockBuilder &operator=(BlockBuilder &other) = delete;
};
/// Deprecated. Use Block * instead.
///
/// A BlockHandle represents a (potentially "delayed") Block abstraction.
/// This extra abstraction is necessary because an mlir::Block is not an
/// mlir::Value.
@ -269,6 +273,35 @@ private:
mlir::Block *block;
};
/// Creates a block in the region that contains the insertion block of the
/// OpBuilder currently at the top of ScopedContext stack (appends the block to
/// the region). Be aware that this will NOT update the insertion point of the
/// builder to insert into the newly constructed block.
Block *createBlock(TypeRange argTypes = llvm::None);
/// Creates a block in the specified region using OpBuilder at the top of
/// ScopedContext stack (appends the block to the region). Be aware that this
/// will NOT update the insertion point of the builder to insert into the newly
/// constructed block.
Block *createBlockInRegion(Region &region, TypeRange argTypes = llvm::None);
/// Calls "builderFn" with ScopedContext reconfigured to insert into "block" and
/// passes in the block arguments. If the block has a terminator, the operations
/// are inserted before the terminator, otherwise appended to the block.
void appendToBlock(Block *block, function_ref<void(ValueRange)> builderFn);
/// Creates a block in the region that contains the insertion block of the
/// OpBuilder currently at the top of ScopedContext stack, and calls "builderFn"
/// to populate the body of the block while passing it the block arguments.
Block *buildInNewBlock(TypeRange argTypes,
function_ref<void(ValueRange)> builderFn);
/// Creates a block in the specified region using OpBuilder at the top of
/// ScopedContext stack, and calls "builderFn" to populate the body of the block
/// while passing it the block arguments.
Block *buildInNewBlock(Region &region, TypeRange argTypes,
function_ref<void(ValueRange)> builderFn);
/// A StructuredIndexed represents an indexable quantity that is either:
/// 1. a captured value, which is suitable for buffer and tensor operands, or;
/// 2. a captured type, which is suitable for tensor return values.

View File

@ -73,8 +73,8 @@ GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
Operation *mlir::edsc::makeGenericLinalgOp(
ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
ArrayRef<StructuredIndexed> outputs,
function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
function_ref<void(ValueRange)> regionBuilder, ArrayRef<Value> otherValues,
ArrayRef<Attribute> otherAttributes) {
for (unsigned i = 0, e = outputs.size(); i + 1 < e; ++i)
assert(!(outputs[i].getType().isa<RankedTensorType>() &&
outputs[i + 1].getType().isa<MemRefType>()) &&
@ -136,15 +136,12 @@ Operation *mlir::edsc::makeGenericLinalgOp(
assert(op->getRegion(0).empty());
OpBuilder opBuilder(op);
ScopedContext scope(opBuilder, op->getLoc());
BlockHandle b;
SmallVector<Value, 8> handles(blockTypes.size());
BlockBuilder(&b, op->getRegion(0), blockTypes,
handles)([&] { regionBuilder(b.getBlock()->getArguments()); });
buildInNewBlock(op->getRegion(0), blockTypes, regionBuilder);
assert(llvm::hasSingleElement(op->getRegion(0)));
return op;
}
void mlir::edsc::ops::mulRegionBuilder(ArrayRef<BlockArgument> args) {
void mlir::edsc::ops::mulRegionBuilder(ValueRange args) {
using edsc::op::operator+;
using edsc::op::operator*;
assert(args.size() == 2 && "expected 2 block arguments");
@ -152,7 +149,7 @@ void mlir::edsc::ops::mulRegionBuilder(ArrayRef<BlockArgument> args) {
linalg_yield(a * b);
}
void mlir::edsc::ops::macRegionBuilder(ArrayRef<BlockArgument> args) {
void mlir::edsc::ops::macRegionBuilder(ValueRange args) {
using edsc::op::operator+;
using edsc::op::operator*;
assert(args.size() == 3 && "expected 3 block arguments");
@ -165,14 +162,14 @@ Operation *mlir::edsc::ops::linalg_generic_pointwise(
SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
IteratorType::Parallel);
if (O.getType().isa<RankedTensorType>()) {
auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
auto fun = [&unaryOp](ValueRange args) {
assert(args.size() == 1 && "expected 1 block arguments");
Value a(args[0]);
linalg_yield(unaryOp(a));
};
return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
}
auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
auto fun = [&unaryOp](ValueRange args) {
assert(args.size() == 2 && "expected 2 block arguments");
Value a(args[0]);
linalg_yield(unaryOp(a));
@ -193,14 +190,14 @@ Operation *mlir::edsc::ops::linalg_generic_pointwise(
SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
IteratorType::Parallel);
if (O.getType().isa<RankedTensorType>()) {
auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
auto fun = [&binaryOp](ValueRange args) {
assert(args.size() == 2 && "expected 2 block arguments");
Value a(args[0]), b(args[1]);
linalg_yield(binaryOp(a, b));
};
return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, fun);
}
auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
auto fun = [&binaryOp](ValueRange args) {
assert(args.size() == 3 && "expected 3 block arguments");
Value a(args[0]), b(args[1]);
linalg_yield(binaryOp(a, b));

View File

@ -18,6 +18,10 @@ BranchOp mlir::edsc::intrinsics::std_br(BlockHandle bh, ValueRange operands) {
return OperationBuilder<BranchOp>(bh.getBlock(), ops);
}
BranchOp mlir::edsc::intrinsics::std_br(Block *block, ValueRange operands) {
return OperationBuilder<BranchOp>(block, operands);
}
BranchOp mlir::edsc::intrinsics::std_br(BlockHandle *bh, ArrayRef<Type> types,
MutableArrayRef<Value> captures,
ValueRange operands) {
@ -27,6 +31,14 @@ BranchOp mlir::edsc::intrinsics::std_br(BlockHandle *bh, ArrayRef<Type> types,
return OperationBuilder<BranchOp>(bh->getBlock(), ops);
}
CondBranchOp mlir::edsc::intrinsics::std_cond_br(Value cond, Block *trueBranch,
ValueRange trueOperands,
Block *falseBranch,
ValueRange falseOperands) {
return OperationBuilder<CondBranchOp>(cond, trueBranch, trueOperands,
falseBranch, falseOperands);
}
CondBranchOp mlir::edsc::intrinsics::std_cond_br(Value cond,
BlockHandle trueBranch,
ValueRange trueOperands,

View File

@ -87,6 +87,59 @@ BlockHandle mlir::edsc::BlockHandle::createInRegion(Region &region,
return res;
}
Block *mlir::edsc::createBlock(TypeRange argTypes) {
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
OpBuilder &builder = ScopedContext::getBuilderRef();
Block *block = builder.getInsertionBlock();
assert(block != nullptr &&
"insertion point not set up in the builder within ScopedContext");
return createBlockInRegion(*block->getParent(), argTypes);
}
Block *mlir::edsc::createBlockInRegion(Region &region, TypeRange argTypes) {
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
OpBuilder &builder = ScopedContext::getBuilderRef();
OpBuilder::InsertionGuard guard(builder);
return builder.createBlock(&region, {}, argTypes);
}
void mlir::edsc::appendToBlock(Block *block,
function_ref<void(ValueRange)> builderFn) {
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
OpBuilder &builder = ScopedContext::getBuilderRef();
OpBuilder::InsertionGuard guard(builder);
if (block->empty() || block->back().isKnownNonTerminator())
builder.setInsertionPointToEnd(block);
else
builder.setInsertionPoint(&block->back());
builderFn(block->getArguments());
}
Block *mlir::edsc::buildInNewBlock(TypeRange argTypes,
function_ref<void(ValueRange)> builderFn) {
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
OpBuilder &builder = ScopedContext::getBuilderRef();
Block *block = builder.getInsertionBlock();
assert(block != nullptr &&
"insertion point not set up in the builder within ScopedContext");
return buildInNewBlock(*block->getParent(), argTypes, builderFn);
}
Block *mlir::edsc::buildInNewBlock(Region &region, TypeRange argTypes,
function_ref<void(ValueRange)> builderFn) {
assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
OpBuilder &builder = ScopedContext::getBuilderRef();
Block *block = createBlockInRegion(region, argTypes);
OpBuilder::InsertionGuard guard(builder);
builder.setInsertionPointToStart(block);
builderFn(block->getArguments());
return block;
}
void mlir::edsc::LoopBuilder::operator()(function_ref<void(void)> fun) {
// Call to `exit` must be explicit and asymmetric (cannot happen in the
// destructor) because of ordering wrt comma operator.

View File

@ -182,12 +182,12 @@ TEST_FUNC(builder_block_append) {
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
BlockHandle b1, functionBlock(&f.front());
BlockBuilder(&b1, {}, {})([&] { std_constant_index(0); });
BlockBuilder(b1, Append())([&] { std_constant_index(1); });
BlockBuilder(b1, Append())([&] { std_ret(); });
// Get back to entry block and add a branch into b1
BlockBuilder(functionBlock, Append())([&] { std_br(b1, {}); });
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
@ -211,28 +211,18 @@ TEST_FUNC(builder_blocks) {
Value c1(std_constant_int(42, 32)), c2(std_constant_int(1234, 32));
ReturnOp ret = std_ret();
Value r;
Value args12[2];
Value &arg1 = args12[0], &arg2 = args12[1];
Value args34[2];
Value &arg3 = args34[0], &arg4 = args34[1];
BlockHandle b1, b2, functionBlock(&f.front());
BlockBuilder(&b1, {c1.getType(), c1.getType()}, args12)(
// b2 has not yet been constructed, need to come back later.
// This is a byproduct of non-structured control-flow.
);
BlockBuilder(&b2, {c1.getType(), c1.getType()}, args34)([&] {
std_br(b1, {arg3, arg4});
});
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..
BlockBuilder(b1, Append())([&] {
r = arg1 + arg2;
std_br(b2, {arg1, r});
// 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
BlockBuilder(functionBlock, Append())([&] { std_br(b1, {c1, c2}); });
// 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
@ -251,68 +241,22 @@ TEST_FUNC(builder_blocks) {
f.erase();
}
TEST_FUNC(builder_blocks_eager) {
using namespace edsc::op;
auto f = makeFunction("builder_blocks_eager");
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
Value c1(std_constant_int(42, 32)), c2(std_constant_int(1234, 32));
Value res;
Value args1And2[2], args3And4[2];
Value &arg1 = args1And2[0], &arg2 = args1And2[1], &arg3 = args3And4[0],
&arg4 = args3And4[1];
// clang-format off
BlockHandle b1, b2;
{ // Toplevel function scope.
// Build a new block for b1 eagerly.
std_br(&b1, {c1.getType(), c1.getType()}, args1And2, {c1, c2});
// Construct a new block b2 explicitly with a branch into b1.
BlockBuilder(&b2, {c1.getType(), c1.getType()}, args3And4)([&]{
std_br(b1, {arg3, arg4});
});
/// And come back to append into b1 once b2 exists.
BlockBuilder(b1, Append())([&]{
res = arg1 + arg2;
std_br(b2, {arg1, res});
});
}
// CHECK-LABEL: @builder_blocks_eager
// 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());
f.erase();
}
TEST_FUNC(builder_cond_branch) {
auto f = makeFunction("builder_cond_branch", {},
{IntegerType::get(1, &globalContext())});
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
Value funcArg(f.getArgument(0));
Value c32(std_constant_int(32, 32)), c64(std_constant_int(64, 64)),
c42(std_constant_int(42, 32));
ReturnOp ret = std_ret();
Value arg1;
Value args23[2];
BlockHandle b1, b2, functionBlock(&f.front());
BlockBuilder(&b1, c32.getType(), arg1)([&] { std_ret(); });
BlockBuilder(&b2, {c64.getType(), c32.getType()}, args23)([&] { std_ret(); });
// Get back to entry block and add a conditional branch
BlockBuilder(functionBlock, Append())([&] {
std_cond_br(funcArg, b1, {c32}, b2, {c64, c42});
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();
@ -331,44 +275,6 @@ TEST_FUNC(builder_cond_branch) {
f.erase();
}
TEST_FUNC(builder_cond_branch_eager) {
using namespace edsc::op;
auto f = makeFunction("builder_cond_branch_eager", {},
{IntegerType::get(1, &globalContext())});
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
Value arg0(f.getArgument(0));
Value c32(std_constant_int(32, 32)), c64(std_constant_int(64, 64)),
c42(std_constant_int(42, 32));
// clang-format off
BlockHandle b1, b2;
Value arg1[1], args2And3[2];
std_cond_br(arg0,
&b1, c32.getType(), arg1, c32,
&b2, {c64.getType(), c32.getType()}, args2And3, {c64, c42});
BlockBuilder(b1, Append())([]{
std_ret();
});
BlockBuilder(b2, Append())([]{
std_ret();
});
// CHECK-LABEL: @builder_cond_branch_eager
// 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());
f.erase();
}
TEST_FUNC(builder_helpers) {
using namespace edsc::op;
auto f32Type = FloatType::getF32(&globalContext());
@ -433,13 +339,10 @@ TEST_FUNC(insertion_in_block) {
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
BlockHandle b1;
// clang-format off
std_constant_int(0, 32);
(BlockBuilder(&b1))([]{
std_constant_int(1, 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
@ -1057,7 +960,7 @@ TEST_FUNC(memref_vector_matmul_test) {
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
Value A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
auto contractionBuilder = [](ArrayRef<BlockArgument> args) {
auto contractionBuilder = [](ValueRange args) {
assert(args.size() == 3 && "expected 3 block arguments");
(linalg_yield(vector_contraction_matmul(args[0], args[1], args[2])));
};