diff --git a/mlir/include/mlir/Dialect/Affine/EDSC/Builders.h b/mlir/include/mlir/Dialect/Affine/EDSC/Builders.h index 0d9da1fe0dce9..3cb326d2fd9ea 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 lbs, ArrayRef 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 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 steps, + function_ref 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 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 eaf1fada93622..78a5e773a050c 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 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 bodyBuilder " + " = nullptr"> ]; let extraClassDeclaration = [{ diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp index f7bdd53233ae2..99ded0686a544 100644 --- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp +++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp @@ -164,13 +164,13 @@ void NDTransferOpHelper::emitLoops(Lambda loopBodyBuilder) { auto majorLbs = vectorBoundsCapture.getLbs(); auto majorUbs = vectorBoundsCapture.getUbs(); auto majorSteps = vectorBoundsCapture.getSteps(); - SmallVector 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 98e6be955cba6..5c18119dcd743 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 steps, + function_ref 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 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(nestedLoc); + }; + + // Create the loop. If the bounds are known to be constants, use the + // constant form of the loop. + auto lbConst = lbs[i].getDefiningOp(); + auto ubConst = ubs[i].getDefiningOp(); + auto loop = lbConst && ubConst + ? builder.create(loc, lbConst.getValue(), + ubConst.getValue(), steps[i], + loopBody) + : builder.create(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 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( + 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(nestedLoc); + }); +} + mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder( MutableArrayRef ivs, ArrayRef lbs, ArrayRef ubs, ArrayRef steps) { diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index 1aae5cf45ae2e..dd01b3783d1f3 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 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 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 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 3435926f867e0..6bb7316c0aa4d 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