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

[mlir] scf::ForOp: provide builders with callbacks for loop body 

Thanks to a recent change that made `::build` functions take an instance of
`OpBuilder`, it is now possible to build operations within a region attached to
the operation about to be created. Exercise this on `scf::ForOp` by taking a
callback that populates the loop body while the loop is being created.

Additionally, provide helper functions to build perfect nests of `ForOp`s,
with support for iteration arguments. These functions provide the same
functionality as EDSC LoopNestBuilder with simpler implementation, without
relying on edsc::ScopedContext, and using `OpBuilder` in an unambiguous way.
Compatibility functions for EDSC are provided, but may be removed in the
future.

Differential Revision: https://reviews.llvm.org/D79688
This commit is contained in:
Alex Zinenko 2020-05-19 16:22:00 +02:00
parent 1aadd6ce61
commit d1560f3956
8 changed files with 219 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
View File

@ -71,6 +71,17 @@ private:
SmallVector<LoopBuilder, 4> loops;
};
/// Adapters for building loop nests using the builder and the location stored
/// in ScopedContext. Actual builders are in scf::buildLoopNest.
scf::ValueVector loopNestBuilder(ValueRange lbs, ValueRange ubs,
ValueRange steps,
function_ref<void(ValueRange)> fun = nullptr);
scf::ValueVector loopNestBuilder(Value lb, Value ub, Value step,
function_ref<void(Value)> fun = nullptr);
scf::ValueVector loopNestBuilder(
Value lb, Value ub, Value step, ValueRange iterArgInitValues,
function_ref<scf::ValueVector(Value, ValueRange)> fun = nullptr);
} // namespace edsc
} // namespace mlir

36
mlir/include/mlir/Dialect/SCF/SCF.h
View File

@ -40,9 +40,43 @@ void ensureLoopTerminator(Region &region, Builder &builder, Location loc);
ForOp getForInductionVarOwner(Value val);
/// Returns the parallel loop parent of an induction variable. If the provided
// value is not an induction variable, then return nullptr.
/// value is not an induction variable, then return nullptr.
ParallelOp getParallelForInductionVarOwner(Value val);
/// An owning vector of values, handy to return from functions.
using ValueVector = std::vector<Value>;
/// Creates a perfect nest of "for" loops, i.e. all loops but the innermost
/// contain only another loop and a terminator. The lower, upper bounds and
/// steps are provided as `lbs`, `ubs` and `steps`, which are expected to be of
/// the same size. `iterArgs` points to the initial values of the loop iteration
/// arguments, which will be forwarded through the nest to the innermost loop.
/// The body of the loop is populated using `bodyBuilder`, which accepts an
/// ordered list of induction variables of all loops, followed by a list of
/// iteration arguments of the innermost loop, in the same order as provided to
/// `iterArgs`. This function is expected to return as many values as
/// `iterArgs`, of the same type and in the same order, that will be treated as
/// yielded from the loop body and forwarded back through the loop nest. If the
/// function is not provided, the loop nest is not expected to have iteration
/// arguments, the body of the innermost loop will be left empty, containing
/// only the zero-operand terminator. Returns the values yielded by the
/// outermost loop. If bound arrays are empty, the body builder will be called
/// once to construct the IR outside of the loop with an empty list of induction
/// variables.
ValueVector buildLoopNest(
OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
ValueRange steps, ValueRange iterArgs,
function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
bodyBuilder = nullptr);
/// A convenience version for building loop nests without iteration arguments
/// (like for reductions). Does not take the initial value of reductions or
/// expect the body building functions to return their current value.
ValueVector buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
ValueRange ubs, ValueRange steps,
function_ref<void(OpBuilder &, Location, ValueRange)>
bodyBuilder = nullptr);
} // end namespace scf
} // end namespace mlir
#endif // MLIR_DIALECT_SCF_H_

7
mlir/include/mlir/Dialect/SCF/SCFOps.td
View File

@ -137,10 +137,15 @@ def ForOp : SCF_Op<"for",
let builders = [
OpBuilder<"OpBuilder &builder, OperationState &result, "
"Value lowerBound, Value upperBound, Value step, "
"ValueRange iterArgs = llvm::None">
"ValueRange iterArgs = llvm::None, "
"function_ref<void(OpBuilder &, Location, Value, ValueRange)>"
" = nullptr">
];
let extraClassDeclaration = [{
using BodyBuilderFn =
function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
Value getInductionVar() { return getBody()->getArgument(0); }
Block::BlockArgListType getRegionIterArgs() {
return getBody()->getArguments().drop_front();

8
mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
View File

@ -512,8 +512,8 @@ LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
Value tmp = std_alloc(tmpMemRefType(transfer));
StdIndexedValue local(tmp);
Value vec = vector_type_cast(tmp);
SmallVector<Value, 8> ivs(lbs.size());
LoopNestBuilder(ivs, lbs, ubs, steps)([&] {
loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
auto ivs = llvm::to_vector<8>(loopIvs);
// Swap the ivs which will reorder memory accesses.
if (coalescedIdx >= 0)
std::swap(ivs.back(), ivs[coalescedIdx]);
@ -586,8 +586,8 @@ LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
StdIndexedValue local(tmp);
Value vec = vector_type_cast(tmp);
std_store(vectorValue, vec);
SmallVector<Value, 8> ivs(lbs.size());
LoopNestBuilder(ivs, lbs, ubs, steps)([&] {
loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
auto ivs = llvm::to_vector<8>(loopIvs);
// Swap the ivs which will reorder memory accesses.
if (coalescedIdx >= 0)
std::swap(ivs.back(), ivs[coalescedIdx]);

4
mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
View File

@ -79,11 +79,13 @@ static void insertCopyLoops(OpBuilder &builder, Location loc,
// Produce the loop nest with copies.
SmallVector<Value, 8> ivs(lbs.size());
LoopNestBuilder(ivs, lbs, ubs, steps)([&]() {
loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
ivs.assign(loopIvs.begin(), loopIvs.end());
auto activeIvs = llvm::makeArrayRef(ivs).take_back(rank);
StdIndexedValue fromHandle(from), toHandle(to);
toHandle(activeIvs) = fromHandle(activeIvs);
});
ivs[0].getParentBlock()->dump();
// Map the innermost loops to threads in reverse order.
for (auto en :

48
mlir/lib/Dialect/SCF/EDSC/Builders.cpp
View File

@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/SCF/EDSC/Builders.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
@ -110,3 +111,50 @@ mlir::edsc::makeLoopBuilder(Value *iv, Value lb, Value ub, Value step,
result.enter(body);
return result;
}
mlir::scf::ValueVector
mlir::edsc::loopNestBuilder(ValueRange lbs, ValueRange ubs, ValueRange steps,
function_ref<void(ValueRange)> fun) {
// Delegates actual construction to scf::buildLoopNest by wrapping `fun` into
// the expected function interface.
assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
return mlir::scf::buildLoopNest(
ScopedContext::getBuilderRef(), ScopedContext::getLocation(), lbs, ubs,
steps, [&](OpBuilder &builder, Location loc, ValueRange ivs) {
ScopedContext context(builder, loc);
if (fun)
fun(ivs);
});
}
mlir::scf::ValueVector
mlir::edsc::loopNestBuilder(Value lb, Value ub, Value step,
function_ref<void(Value)> fun) {
// Delegates to the ValueRange-based version by wrapping the lambda.
auto wrapper = [&](ValueRange ivs) {
assert(ivs.size() == 1);
if (fun)
fun(ivs[0]);
};
return loopNestBuilder(ValueRange(lb), ValueRange(ub), ValueRange(step),
wrapper);
}
mlir::scf::ValueVector mlir::edsc::loopNestBuilder(
Value lb, Value ub, Value step, ValueRange iterArgInitValues,
function_ref<scf::ValueVector(Value, ValueRange)> fun) {
// Delegates actual construction to scf::buildLoopNest by wrapping `fun` into
// the expected function interface.
assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
return mlir::scf::buildLoopNest(
ScopedContext::getBuilderRef(), ScopedContext::getLocation(), lb, ub,
step, iterArgInitValues,
[&](OpBuilder &builder, Location loc, ValueRange ivs, ValueRange args) {
assert(ivs.size() == 1 && "expected one induction variable");
ScopedContext context(builder, loc);
if (fun)
return fun(ivs[0], args);
return scf::ValueVector(iterArgInitValues.begin(),
iterArgInitValues.end());
});
}

110
mlir/lib/Dialect/SCF/SCF.cpp
View File

@ -40,18 +40,30 @@ SCFDialect::SCFDialect(MLIRContext *context)
//===----------------------------------------------------------------------===//
void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
Value ub, Value step, ValueRange iterArgs) {
Value ub, Value step, ValueRange iterArgs,
BodyBuilderFn bodyBuilder) {
result.addOperands({lb, ub, step});
result.addOperands(iterArgs);
for (Value v : iterArgs)
result.addTypes(v.getType());
Region *bodyRegion = result.addRegion();
bodyRegion->push_back(new Block());
if (iterArgs.empty())
ForOp::ensureTerminator(*bodyRegion, builder, result.location);
bodyRegion->front().addArgument(builder.getIndexType());
bodyRegion->push_back(new Block);
Block &bodyBlock = bodyRegion->front();
bodyBlock.addArgument(builder.getIndexType());
for (Value v : iterArgs)
bodyRegion->front().addArgument(v.getType());
bodyBlock.addArgument(v.getType());
// Create the default terminator if the builder is not provided and if the
// iteration arguments are not provided. Otherwise, leave this to the caller
// because we don't know which values to return from the loop.
if (iterArgs.empty() && !bodyBuilder) {
ForOp::ensureTerminator(*bodyRegion, builder, result.location);
} else if (bodyBuilder) {
OpBuilder::InsertionGuard guard(builder);
builder.setInsertionPointToStart(&bodyBlock);
bodyBuilder(builder, result.location, bodyBlock.getArgument(0),
bodyBlock.getArguments().drop_front());
}
}
static LogicalResult verify(ForOp op) {
@ -229,6 +241,92 @@ void ForOp::getSuccessorRegions(Optional<unsigned> index,
regions.push_back(RegionSuccessor(getResults()));
}
ValueVector mlir::scf::buildLoopNest(
OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
ValueRange steps, ValueRange iterArgs,
function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
bodyBuilder) {
assert(lbs.size() == ubs.size() &&
"expected the same number of lower and upper bounds");
assert(lbs.size() == steps.size() &&
"expected the same number of lower bounds and steps");
// If there are no bounds, call the body-building function and return early.
if (lbs.empty()) {
ValueVector results =
bodyBuilder ? bodyBuilder(builder, loc, ValueRange(), iterArgs)
: ValueVector();
assert(results.size() == iterArgs.size() &&
"loop nest body must return as many values as loop has iteration "
"arguments");
return results;
}
// First, create the loop structure iteratively using the body-builder
// callback of `ForOp::build`. Do not create `YieldOp`s yet.
OpBuilder::InsertionGuard guard(builder);
SmallVector<scf::ForOp, 4> loops;
SmallVector<Value, 4> ivs;
loops.reserve(lbs.size());
ivs.reserve(lbs.size());
ValueRange currentIterArgs = iterArgs;
Location currentLoc = loc;
for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
auto loop = builder.create<scf::ForOp>(
currentLoc, lbs[i], ubs[i], steps[i], currentIterArgs,
[&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
ValueRange args) {
ivs.push_back(iv);
// It is safe to store ValueRange args because it points to block
// arguments of a loop operation that we also own.
currentIterArgs = args;
currentLoc = nestedLoc;
});
// Set the builder to point to the body of the newly created loop. We don't
// do this in the callback beacause the builder is reset when the callback
// returns.
builder.setInsertionPointToStart(loop.getBody());
loops.push_back(loop);
}
// For all loops but the innermost, yield the results of the nested loop.
for (unsigned i = 0, e = loops.size() - 1; i < e; ++i) {
builder.setInsertionPointToEnd(loops[i].getBody());
builder.create<scf::YieldOp>(loc, loops[i + 1].getResults());
}
// In the body of the innermost loop, call the body building function if any
// and yield its results.
builder.setInsertionPointToStart(loops.back().getBody());
ValueVector results = bodyBuilder
? bodyBuilder(builder, currentLoc, ivs,
loops.back().getRegionIterArgs())
: ValueVector();
assert(results.size() == iterArgs.size() &&
"loop nest body must return as many values as loop has iteration "
"arguments");
builder.setInsertionPointToEnd(loops.back().getBody());
builder.create<scf::YieldOp>(loc, results);
// Return the results of the outermost loop.
return ValueVector(loops.front().result_begin(), loops.front().result_end());
}
ValueVector mlir::scf::buildLoopNest(
OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
ValueRange steps,
function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
// Delegate to the main function by wrapping the body builder.
return buildLoopNest(builder, loc, lbs, ubs, steps, llvm::None,
[&bodyBuilder](OpBuilder &nestedBuilder,
Location nestedLoc, ValueRange ivs,
ValueRange) -> ValueVector {
if (bodyBuilder)
bodyBuilder(nestedBuilder, nestedLoc, ivs);
return {};
});
}
//===----------------------------------------------------------------------===//
// IfOp
//===----------------------------------------------------------------------===//

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

@ -139,10 +139,10 @@ TEST_FUNC(builder_loop_for) {
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
Value i, a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
d(f.getArgument(3));
using namespace edsc::op;
LoopNestBuilder(&i, a - b, c + d, a)();
loopNestBuilder(a - b, c + d, a);
// clang-format off
// CHECK-LABEL: func @builder_loop_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
@ -1076,16 +1076,14 @@ TEST_FUNC(builder_loop_for_yield) {
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 i, a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
d(f.getArgument(3));
Value args01[2];
Value &arg0 = args01[0], &arg1 = args01[1];
using namespace edsc::op;
auto results =
LoopNestBuilder(&i, a - b, c + d, a, args01, {init0, init1})([&] {
auto sum = arg0 + arg1;
loop_yield(ArrayRef<Value>{arg1, sum});
});
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};
});
results[0] + results[1];
// clang-format off