[Mlir-commits] [mlir] 8647a9b - [mlir] Refactor affine loop nest builders
Alex Zinenko
llvmlistbot at llvm.org
Thu Jun 18 14:03:30 PDT 2020
Author: Alex Zinenko
Date: 2020-06-18T23:03:13+02:00
New Revision: 8647a9bc511bd47ca80ae48761662436066a09bb
URL: https://github.com/llvm/llvm-project/commit/8647a9bc511bd47ca80ae48761662436066a09bb
DIFF: https://github.com/llvm/llvm-project/commit/8647a9bc511bd47ca80ae48761662436066a09bb.diff
LOG: [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
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/lib/Dialect/Affine/EDSC/Builders.cpp
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index 808fe892b4c1..93bac5a4a29f 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<AffineForOp> forInsts,
SmallVectorImpl<Value> *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<int64_t> lbs, ArrayRef<int64_t> ubs,
+ ArrayRef<int64_t> steps,
+ function_ref<void(OpBuilder &, Location, ValueRange)>
+ bodyBuilderFn = nullptr);
+void buildAffineLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
+ ValueRange ubs, ArrayRef<int64_t> steps,
+ function_ref<void(OpBuilder &, Location, ValueRange)>
+ 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 5c18119dcd74..87a0b9defdc3 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<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.
+ // 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<void(OpBuilder &, Location, ValueRange)> 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<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());
- }
+ 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 dd01b3783d1f..ed3623ac0a43 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<AffineForOp> forInsts,
ivs->push_back(forInst.getInductionVar());
}
+/// Builds an affine loop nest, using "loopCreatorFn" to create individual loop
+/// operations.
+template <typename BoundListTy, typename LoopCreatorTy>
+void buildAffineLoopNestImpl(
+ OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs,
+ ArrayRef<int64_t> steps,
+ function_ref<void(OpBuilder &, Location, ValueRange)> 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<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) {
+ OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
+ bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
+ }
+ nestedBuilder.create<AffineTerminatorOp>(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<void(OpBuilder &, Location, Value)> bodyBuilderFn) {
+ return builder.create<AffineForOp>(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<void(OpBuilder &, Location, Value)> bodyBuilderFn) {
+ auto lbConst = lb.getDefiningOp<ConstantIndexOp>();
+ auto ubConst = ub.getDefiningOp<ConstantIndexOp>();
+ if (lbConst && ubConst)
+ return buildAffineLoopFromConstants(builder, loc, lbConst.getValue(),
+ ubConst.getValue(), step,
+ bodyBuilderFn);
+ return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub,
+ builder.getDimIdentityMap(), step,
+ bodyBuilderFn);
+}
+
+void mlir::buildAffineLoopNest(
+ OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs,
+ ArrayRef<int64_t> ubs, ArrayRef<int64_t> steps,
+ function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
+ buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
+ buildAffineLoopFromConstants);
+}
+
+void mlir::buildAffineLoopNest(
+ OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
+ ArrayRef<int64_t> steps,
+ function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
+ buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
+ buildAffineLoopFromValues);
+}
+
//===----------------------------------------------------------------------===//
// AffineIfOp
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list