[Mlir-commits] [mlir] 3f5bd53 - [mlir] Introduce callback-based builders for AffineForOp
Alex Zinenko
llvmlistbot at llvm.org
Mon Jun 15 08:55:59 PDT 2020
Author: Alex Zinenko
Date: 2020-06-15T17:55:49+02:00
New Revision: 3f5bd53eb1044b166fd9d2771ca7135d1b29ee6f
URL: https://github.com/llvm/llvm-project/commit/3f5bd53eb1044b166fd9d2771ca7135d1b29ee6f
DIFF: https://github.com/llvm/llvm-project/commit/3f5bd53eb1044b166fd9d2771ca7135d1b29ee6f.diff
LOG: [mlir] Introduce callback-based builders for AffineForOp
Similarly to `scf::ForOp`, introduce additional `function_ref` arguments to
`AffineForOp::build` that can be used to populate the body of the loop during
its construction. Provide compatibility functions for constructing affine loop
nests using `edsc::ScopedContext`.
`edsc::AffineLoopNestBuilder` and reletad functionality is now deprecated and
will be removed soon, users are expected to switch to `affineLoopNestBuilder`
that provides similar functionality with a simpler OpBuilder-based
implementation.
Differential Revision: https://reviews.llvm.org/D81754
Added:
Modified:
mlir/include/mlir/Dialect/Affine/EDSC/Builders.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
mlir/lib/Dialect/Affine/EDSC/Builders.cpp
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/test/EDSC/builder-api-test.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/EDSC/Builders.h b/mlir/include/mlir/Dialect/Affine/EDSC/Builders.h
index 0d9da1fe0dce..3cb326d2fd9e 100644
--- a/mlir/include/mlir/Dialect/Affine/EDSC/Builders.h
+++ b/mlir/include/mlir/Dialect/Affine/EDSC/Builders.h
@@ -28,6 +28,8 @@ namespace edsc {
LoopBuilder makeAffineLoopBuilder(Value *iv, ArrayRef<Value> lbs,
ArrayRef<Value> ubs, int64_t step);
+/// Deprecated. Use affineLoopNestBuilder instead.
+///
/// Explicit nested LoopBuilder. Offers a compressed multi-loop builder to avoid
/// explicitly writing all the loops in a nest. This simple functionality is
/// also useful to write rank-agnostic custom ops.
@@ -67,6 +69,31 @@ class AffineLoopNestBuilder {
SmallVector<LoopBuilder, 4> loops;
};
+/// Creates a perfect nest of affine "for" loops, given the list of lower
+/// bounds, upper bounds and steps. The three lists are expected to contain the
+/// same number of elements. Uses the OpBuilder and Location stored in
+/// ScopedContext and assumes they are non-null. The optional "bodyBuilderFn"
+/// callback is called to construct the body of the innermost loop and is passed
+/// the list of loop induction variables, in order from outermost to innermost.
+/// The function is expected to use the builder and location stored in
+/// ScopedContext at the moment of the call. The function should not create
+/// the affine terminator op, which will be added regardless of the
+/// "bodyBuilderFn" being present.
+void affineLoopNestBuilder(
+ ValueRange lbs, ValueRange ubs, ArrayRef<int64_t> steps,
+ function_ref<void(ValueRange)> bodyBuilderFn = nullptr);
+
+/// Creates a single affine "for" loop, iterating from max(lbs) to min(ubs) with
+/// the given step. Uses the OpBuilder and Location stored in ScopedContext and
+/// assumes they are non-null. The optional "bodyBuilderFn" callback is called
+/// to construct the body of the loop and is passed the induction variable. The
+/// function is expected to use the builder and location stored in ScopedContext
+/// at the moment of the call. The function should not create the affine
+/// terminator op, which will be added regardless of the "bodyBuilderFn" being
+/// present.
+void affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step,
+ function_ref<void(Value)> bodyBuilderFn = nullptr);
+
namespace op {
Value operator+(Value lhs, Value rhs);
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index eaf1fada9362..78a5e773a050 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -179,11 +179,15 @@ def AffineForOp : Affine_Op<"for",
let skipDefaultBuilders = 1;
let builders = [
OpBuilder<"OpBuilder &builder, OperationState &result, "
- "int64_t lowerBound, int64_t upperBound, int64_t step = 1">,
+ "int64_t lowerBound, int64_t upperBound, int64_t step = 1, "
+ "function_ref<void(OpBuilder &, Location, Value)> bodyBuilder "
+ " = nullptr">,
OpBuilder<"OpBuilder &builder, OperationState &result, "
"ValueRange lbOperands, AffineMap lbMap, "
"ValueRange ubOperands, AffineMap ubMap, "
- "int64_t step = 1">
+ "int64_t step = 1, "
+ "function_ref<void(OpBuilder &, Location, Value)> bodyBuilder "
+ " = nullptr">
];
let extraClassDeclaration = [{
diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
index f7bdd53233ae..99ded0686a54 100644
--- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
+++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
@@ -164,13 +164,13 @@ void NDTransferOpHelper<ConcreteOp>::emitLoops(Lambda loopBodyBuilder) {
auto majorLbs = vectorBoundsCapture.getLbs();
auto majorUbs = vectorBoundsCapture.getUbs();
auto majorSteps = vectorBoundsCapture.getSteps();
- SmallVector<Value, 8> majorIvs(vectorBoundsCapture.rank());
- AffineLoopNestBuilder(majorIvs, majorLbs, majorUbs, majorSteps)([&] {
- ValueRange indices(xferOp.indices());
- loopBodyBuilder(majorIvs, indices.take_front(leadingRank),
- indices.drop_front(leadingRank).take_front(majorRank),
- indices.take_back(minorRank), memrefBoundsCapture);
- });
+ affineLoopNestBuilder(
+ majorLbs, majorUbs, majorSteps, [&](ValueRange majorIvs) {
+ ValueRange indices(xferOp.indices());
+ loopBodyBuilder(majorIvs, indices.take_front(leadingRank),
+ indices.drop_front(leadingRank).take_front(majorRank),
+ indices.take_back(minorRank), memrefBoundsCapture);
+ });
}
}
diff --git a/mlir/lib/Dialect/Affine/EDSC/Builders.cpp b/mlir/lib/Dialect/Affine/EDSC/Builders.cpp
index 98e6be955cba..5c18119dcd74 100644
--- a/mlir/lib/Dialect/Affine/EDSC/Builders.cpp
+++ b/mlir/lib/Dialect/Affine/EDSC/Builders.cpp
@@ -64,6 +64,79 @@ mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(Value *iv,
loops.emplace_back(makeAffineLoopBuilder(iv, lbs, ubs, step));
}
+void mlir::edsc::affineLoopNestBuilder(
+ ValueRange lbs, ValueRange ubs, ArrayRef<int64_t> steps,
+ function_ref<void(ValueRange)> bodyBuilderFn) {
+ assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+ assert(lbs.size() == ubs.size() && "Mismatch in number of arguments");
+ assert(lbs.size() == steps.size() && "Mismatch in number of arguments");
+
+ // If there are no loops to be constructed, construct the body anyway.
+ if (lbs.empty()) {
+ if (bodyBuilderFn)
+ bodyBuilderFn(ValueRange());
+ return;
+ }
+
+ // Fetch the builder and location.
+ OpBuilder &builder = ScopedContext::getBuilderRef();
+ OpBuilder::InsertionGuard guard(builder);
+ Location loc = ScopedContext::getLocation();
+ AffineMap identity = builder.getDimIdentityMap();
+
+ // Create the loops iteratively and store the induction variables.
+ SmallVector<Value, 4> ivs;
+ ivs.reserve(lbs.size());
+ for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
+ // Callback for creating the loop body, always creates the terminator.
+ auto loopBody = [&](OpBuilder &nestedBuilder, Location nestedLoc,
+ Value iv) {
+ ivs.push_back(iv);
+ // In the innermost loop, call the body builder.
+ if (i == e - 1 && bodyBuilderFn) {
+ ScopedContext nestedContext(nestedBuilder, loc);
+ OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
+ bodyBuilderFn(ivs);
+ }
+ nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+ };
+
+ // Create the loop. If the bounds are known to be constants, use the
+ // constant form of the loop.
+ auto lbConst = lbs[i].getDefiningOp<ConstantIndexOp>();
+ auto ubConst = ubs[i].getDefiningOp<ConstantIndexOp>();
+ auto loop = lbConst && ubConst
+ ? builder.create<AffineForOp>(loc, lbConst.getValue(),
+ ubConst.getValue(), steps[i],
+ loopBody)
+ : builder.create<AffineForOp>(loc, lbs[i], identity, ubs[i],
+ identity, steps[i], loopBody);
+ builder.setInsertionPointToStart(loop.getBody());
+ }
+}
+
+void mlir::edsc::affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step,
+ function_ref<void(Value)> bodyBuilderFn) {
+ // Fetch the builder and location.
+ assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+ OpBuilder &builder = ScopedContext::getBuilderRef();
+ Location loc = ScopedContext::getLocation();
+
+ // Create the actual loop and call the body builder, if provided, after
+ // updating the scoped context.
+ builder.create<AffineForOp>(
+ loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,
+ builder.getMultiDimIdentityMap(ubs.size()), step,
+ [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv) {
+ if (bodyBuilderFn) {
+ ScopedContext nestedContext(nestedBuilder, nestedLoc);
+ OpBuilder::InsertionGuard guard(nestedBuilder);
+ bodyBuilderFn(iv);
+ }
+ nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+ });
+}
+
mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
MutableArrayRef<Value> ivs, ArrayRef<Value> lbs, ArrayRef<Value> ubs,
ArrayRef<int64_t> steps) {
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 1aae5cf45ae2..dd01b3783d1f 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1176,9 +1176,10 @@ LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
// AffineForOp
//===----------------------------------------------------------------------===//
-void AffineForOp::build(OpBuilder &builder, OperationState &result,
- ValueRange lbOperands, AffineMap lbMap,
- ValueRange ubOperands, AffineMap ubMap, int64_t step) {
+void AffineForOp::build(
+ OpBuilder &builder, OperationState &result, ValueRange lbOperands,
+ AffineMap lbMap, ValueRange ubOperands, AffineMap ubMap, int64_t step,
+ function_ref<void(OpBuilder &, Location, Value)> bodyBuilder) {
assert(((!lbMap && lbOperands.empty()) ||
lbOperands.size() == lbMap.getNumInputs()) &&
"lower bound operand count does not match the affine map");
@@ -1202,17 +1203,25 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result,
// Create a region and a block for the body. The argument of the region is
// the loop induction variable.
Region *bodyRegion = result.addRegion();
- Block *body = new Block();
- body->addArgument(IndexType::get(builder.getContext()));
+ Block *body = new Block;
+ Value inductionVar = body->addArgument(IndexType::get(builder.getContext()));
bodyRegion->push_back(body);
- ensureTerminator(*bodyRegion, builder, result.location);
+ if (bodyBuilder) {
+ OpBuilder::InsertionGuard guard(builder);
+ builder.setInsertionPointToStart(body);
+ bodyBuilder(builder, result.location, inductionVar);
+ } else {
+ ensureTerminator(*bodyRegion, builder, result.location);
+ }
}
-void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb,
- int64_t ub, int64_t step) {
+void AffineForOp::build(
+ OpBuilder &builder, OperationState &result, int64_t lb, int64_t ub,
+ int64_t step,
+ function_ref<void(OpBuilder &, Location, Value)> bodyBuilder) {
auto lbMap = AffineMap::getConstantMap(lb, builder.getContext());
auto ubMap = AffineMap::getConstantMap(ub, builder.getContext());
- return build(builder, result, {}, lbMap, {}, ubMap, step);
+ return build(builder, result, {}, lbMap, {}, ubMap, step, bodyBuilder);
}
static LogicalResult verify(AffineForOp op) {
diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp
index 3435926f867e..6bb7316c0aa4 100644
--- a/mlir/test/EDSC/builder-api-test.cpp
+++ b/mlir/test/EDSC/builder-api-test.cpp
@@ -67,16 +67,16 @@ TEST_FUNC(builder_dynamic_for_func_args) {
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
- Value i, j, lb(f.getArgument(0)), ub(f.getArgument(1));
+ 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));
- AffineLoopNestBuilder(&i, lb, ub, 3)([&] {
+ affineLoopBuilder(lb, ub, 3, [&](Value i) {
using namespace edsc::op;
lb *std_constant_index(3) + ub;
lb + std_constant_index(3);
- AffineLoopNestBuilder(&j, lb, ub, 2)([&] {
+ 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));
@@ -120,7 +120,7 @@ TEST_FUNC(builder_dynamic_for) {
Value i, a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
d(f.getArgument(3));
using namespace edsc::op;
- AffineLoopNestBuilder(&i, a - b, c + d, 2)();
+ affineLoopBuilder(a - b, c + d, 2);
// clang-format off
// CHECK-LABEL: func @builder_dynamic_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
@@ -161,9 +161,9 @@ TEST_FUNC(builder_max_min_for) {
OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
- Value i, lb1(f.getArgument(0)), lb2(f.getArgument(1)), ub1(f.getArgument(2)),
+ Value lb1(f.getArgument(0)), lb2(f.getArgument(1)), ub1(f.getArgument(2)),
ub2(f.getArgument(3));
- AffineLoopNestBuilder(&i, {lb1, lb2}, {ub1, ub2}, 1)();
+ affineLoopBuilder({lb1, lb2}, {ub1, ub2}, 1);
std_ret();
// clang-format off
@@ -386,20 +386,20 @@ TEST_FUNC(builder_helpers) {
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 ivs[2];
- Value &i = ivs[0], &j = ivs[1];
- Value k1, k2, lb0, lb1, lb2, ub0, ub1, ub2;
+ 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(ivs, {lb0, lb1}, {ub0, ub1}, {step0, step1})([&]{
- AffineLoopNestBuilder(&k1, lb2, ub2, step2)([&]{
+ 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);
});
- AffineLoopNestBuilder(&k2, lb2, ub2, step2)([&]{
+ affineLoopBuilder(lb2, ub2, step2, [&](Value k2){
C(i, j, k2) += A(i, j, k2) + B(i, j, k2);
});
});
@@ -525,16 +525,14 @@ TEST_FUNC(select_op_i32) {
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));
AffineIndexedValue A(f.getArgument(0));
- Value ivs[2];
- Value &i = ivs[0], &j = ivs[1];
- AffineLoopNestBuilder(ivs, {zero, zero}, {one, one}, {1, 1})([&]{
- std_select(eq(i, zero), A(zero, zero), A(i, j));
+ affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
+ std_select(eq(ivs[0], zero), A(zero, zero), A(ivs[0], ivs[1]));
});
+ // clang-format off
// CHECK-LABEL: @select_op
// CHECK: affine.for %{{.*}} = 0 to 1 {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 1 {
@@ -559,10 +557,9 @@ TEST_FUNC(select_op_f32) {
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));
- Value ivs[2];
- Value &i = ivs[0], &j = ivs[1];
- AffineLoopNestBuilder(ivs, {zero, zero}, {one, one}, {1, 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(B(i, j) >= B(i + one, j), A(zero, zero), A(i, j));
@@ -637,18 +634,20 @@ TEST_FUNC(tile_2d) {
vC(f.getArgument(2));
AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1)),
C(f.getArgument(2));
- Value ivs[2];
- Value &i = ivs[0], &j = ivs[1];
- Value k1, k2;
+ 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(ivs, {zero, zero}, {M, N}, {1, 1})([&]{
- AffineLoopNestBuilder(&k1, zero, O, 1)([&]{
+ 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);
});
- AffineLoopNestBuilder(&k2, zero, O, 1)([&]{
+ affineLoopBuilder(zero, O, 1, [&](Value k) {
+ k2 = k;
C(i, j, k2) = A(i, j, k2) + B(i, j, k2);
});
});
@@ -708,10 +707,10 @@ TEST_FUNC(indirect_access) {
MemRefBoundsCapture vC(f.getArgument(2));
AffineIndexedValue B(f.getArgument(1)), D(f.getArgument(3));
StdIndexedValue A(f.getArgument(0)), C(f.getArgument(2));
- Value i, N(vC.ub(0));
+ Value N(vC.ub(0));
// clang-format off
- AffineLoopNestBuilder(&i, zero, N, 1)([&]{
+ affineLoopBuilder(zero, N, 1, [&](Value i) {
C((Value)D(i)) = A((Value)B(i));
});
// clang-format on
@@ -743,8 +742,7 @@ TEST_FUNC(empty_map_load_store) {
AffineIndexedValue input(f.getArgument(0)), res(f.getArgument(1));
// clang-format off
- Value iv;
- AffineLoopNestBuilder(&iv, zero, one, 1)([&]{
+ affineLoopBuilder(zero, one, 1, [&](Value) {
res() = input();
});
// clang-format on
More information about the Mlir-commits
mailing list