From 8647a9bc511bd47ca80ae48761662436066a09bb Mon Sep 17 00:00:00 2001 From: Alex Zinenko Date: Thu, 18 Jun 2020 17:53:50 +0200 Subject: [PATCH] [mlir] Refactor affine loop nest builders Existing implementation of affine loop nest builders relies on EDSC ScopedContext, which is not used pervasively. Provide a common OpBuilder-based helper function to construct a perfect nest of affine loops with the body of the innermost loop populated by a callback. Use this function to implement the EDSC version. Affine "for" loops differ from SCF "for" loops by (1) not allowing to yield values and (2) supporting short-hand form for constant bounds, which justifies a separate implementation of the loop nest builder for the same of simplicity. Differential Revision: https://reviews.llvm.org/D81955 --- mlir/include/mlir/Dialect/Affine/IR/AffineOps.h | 15 +++++ mlir/lib/Dialect/Affine/EDSC/Builders.cpp | 55 ++++------------- mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 80 +++++++++++++++++++++++++ 3 files changed, 108 insertions(+), 42 deletions(-) diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h index 808fe89..93bac5a 100644 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h @@ -381,6 +381,21 @@ AffineForOp getForInductionVarOwner(Value val); void extractForInductionVars(ArrayRef forInsts, SmallVectorImpl *ivs); +/// Builds a perfect nest of affine "for" loops, i.e. each loop except the +/// innermost only contains another loop and a terminator. The loops iterate +/// from "lbs" to "ubs" with "steps". The body of the innermost loop is +/// populated by calling "bodyBuilderFn" and providing it with an OpBuilder, a +/// Location and a list of loop induction variables. +void buildAffineLoopNest(OpBuilder &builder, Location loc, + ArrayRef lbs, ArrayRef ubs, + ArrayRef steps, + function_ref + bodyBuilderFn = nullptr); +void buildAffineLoopNest(OpBuilder &builder, Location loc, ValueRange lbs, + ValueRange ubs, ArrayRef steps, + function_ref + bodyBuilderFn = nullptr); + /// AffineBound represents a lower or upper bound in the for operation. /// This class does not own the underlying operands. Instead, it refers /// to the operands stored in the AffineForOp. Its life span should not exceed diff --git a/mlir/lib/Dialect/Affine/EDSC/Builders.cpp b/mlir/lib/Dialect/Affine/EDSC/Builders.cpp index 5c18119..87a0b9d 100644 --- a/mlir/lib/Dialect/Affine/EDSC/Builders.cpp +++ b/mlir/lib/Dialect/Affine/EDSC/Builders.cpp @@ -68,51 +68,22 @@ 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. + // Wrap the body builder function into an interface compatible with the main + // builder. + auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc, + ValueRange ivs) { + ScopedContext context(nestedBuilder, nestedLoc); + bodyBuilderFn(ivs); + }; + function_ref wrapper; + if (bodyBuilderFn) + wrapper = wrappedBuilderFn; + + // Extract the builder, location and construct the loop nest. 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()); - } + buildAffineLoopNest(builder, loc, lbs, ubs, steps, wrapper); } void mlir::edsc::affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step, diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index dd01b37..ed3623a 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -1688,6 +1688,86 @@ void mlir::extractForInductionVars(ArrayRef forInsts, ivs->push_back(forInst.getInductionVar()); } +/// Builds an affine loop nest, using "loopCreatorFn" to create individual loop +/// operations. +template +void buildAffineLoopNestImpl( + OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs, + ArrayRef steps, + function_ref bodyBuilderFn, + LoopCreatorTy &&loopCreatorFn) { + 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. + OpBuilder::InsertionGuard guard(builder); + if (lbs.empty()) { + if (bodyBuilderFn) + bodyBuilderFn(builder, loc, ValueRange()); + return; + } + + // 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) { + OpBuilder::InsertionGuard nestedGuard(nestedBuilder); + bodyBuilderFn(nestedBuilder, nestedLoc, ivs); + } + nestedBuilder.create(nestedLoc); + }; + + // Delegate actual loop creation to the callback in order to dispatch + // between constant- and variable-bound loops. + auto loop = loopCreatorFn(builder, loc, lbs[i], ubs[i], steps[i], loopBody); + builder.setInsertionPointToStart(loop.getBody()); + } +} + +/// Creates an affine loop from the bounds known to be constants. +static AffineForOp buildAffineLoopFromConstants( + OpBuilder &builder, Location loc, int64_t lb, int64_t ub, int64_t step, + function_ref bodyBuilderFn) { + return builder.create(loc, lb, ub, step, bodyBuilderFn); +} + +/// Creates an affine loop from the bounds that may or may not be constants. +static AffineForOp buildAffineLoopFromValues( + OpBuilder &builder, Location loc, Value lb, Value ub, int64_t step, + function_ref bodyBuilderFn) { + auto lbConst = lb.getDefiningOp(); + auto ubConst = ub.getDefiningOp(); + if (lbConst && ubConst) + return buildAffineLoopFromConstants(builder, loc, lbConst.getValue(), + ubConst.getValue(), step, + bodyBuilderFn); + return builder.create(loc, lb, builder.getDimIdentityMap(), ub, + builder.getDimIdentityMap(), step, + bodyBuilderFn); +} + +void mlir::buildAffineLoopNest( + OpBuilder &builder, Location loc, ArrayRef lbs, + ArrayRef ubs, ArrayRef steps, + function_ref bodyBuilderFn) { + buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn, + buildAffineLoopFromConstants); +} + +void mlir::buildAffineLoopNest( + OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs, + ArrayRef steps, + function_ref bodyBuilderFn) { + buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn, + buildAffineLoopFromValues); +} + //===----------------------------------------------------------------------===// // AffineIfOp //===----------------------------------------------------------------------===// -- 2.7.4