shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity
llvm-project/mlir/lib/Dialect/SCF/Utils/Utils.cpp
Mehdi Amini b78ceef43e
[mlir][scf] Fix crash in extractFixedOuterLoops with iter_args loops (#184106) 
The stripmineSink helper splices loop body operations into a new inner
scf.for that has no iter_args. When the target loop carries iter_args,
values yielded by the spliced body are moved inside the inner loop, but
the outer loop's yield terminator still references those values,
creating an SSA invariant violation. In debug builds this triggers the
assertion
  use_empty() && "Cannot destroy a value that still has uses\!"
when the outer RewriterBase tries to erase the now-broken operations.

Fix: in extractFixedOuterLoops, skip the strip-mining transformation if
any of the collected perfectly-nested loops have iter_args.

Add a regression test to parametric-tiling.mlir.

Fixes #129044

Assisted-by: Claude Code
2026-03-11 14:21:57 +01:00

1700 lines
70 KiB
C++
Raw Blame History

//===- Utils.cpp ---- Misc utilities for loop transformation ----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements miscellaneous loop transformation routines.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/Support/DebugLog.h"
#include <cstdint>
using namespace mlir;
#define DEBUG_TYPE "scf-utils"
SmallVector<scf::ForOp> mlir::replaceLoopNestWithNewYields(
RewriterBase &rewriter, MutableArrayRef<scf::ForOp> loopNest,
ValueRange newIterOperands, const NewYieldValuesFn &newYieldValuesFn,
bool replaceIterOperandsUsesInLoop) {
if (loopNest.empty())
return {};
// This method is recursive (to make it more readable). Adding an
// assertion here to limit the recursion. (See
// https://discourse.llvm.org/t/rfc-update-to-mlir-developer-policy-on-recursion/62235)
assert(loopNest.size() <= 10 &&
"exceeded recursion limit when yielding value from loop nest");
// To yield a value from a perfectly nested loop nest, the following
// pattern needs to be created, i.e. starting with
//
// ```mlir
// scf.for .. {
// scf.for .. {
// scf.for .. {
// %value = ...
// }
// }
// }
// ```
//
// needs to be modified to
//
// ```mlir
// %0 = scf.for .. iter_args(%arg0 = %init) {
// %1 = scf.for .. iter_args(%arg1 = %arg0) {
// %2 = scf.for .. iter_args(%arg2 = %arg1) {
// %value = ...
// scf.yield %value
// }
// scf.yield %2
// }
// scf.yield %1
// }
// ```
//
// The inner most loop is handled using the `replaceWithAdditionalYields`
// that works on a single loop.
if (loopNest.size() == 1) {
auto innerMostLoop =
cast<scf::ForOp>(*loopNest.back().replaceWithAdditionalYields(
rewriter, newIterOperands, replaceIterOperandsUsesInLoop,
newYieldValuesFn));
return {innerMostLoop};
}
// The outer loops are modified by calling this method recursively
// - The return value of the inner loop is the value yielded by this loop.
// - The region iter args of this loop are the init_args for the inner loop.
SmallVector<scf::ForOp> newLoopNest;
NewYieldValuesFn fn =
[&](OpBuilder &innerBuilder, Location loc,
ArrayRef<BlockArgument> innerNewBBArgs) -> SmallVector<Value> {
newLoopNest = replaceLoopNestWithNewYields(rewriter, loopNest.drop_front(),
innerNewBBArgs, newYieldValuesFn,
replaceIterOperandsUsesInLoop);
return llvm::map_to_vector(
newLoopNest.front().getResults().take_back(innerNewBBArgs.size()),
[](OpResult r) -> Value { return r; });
};
scf::ForOp outerMostLoop =
cast<scf::ForOp>(*loopNest.front().replaceWithAdditionalYields(
rewriter, newIterOperands, replaceIterOperandsUsesInLoop, fn));
newLoopNest.insert(newLoopNest.begin(), outerMostLoop);
return newLoopNest;
}
/// Outline a region with a single block into a new FuncOp.
/// Assumes the FuncOp result types is the type of the yielded operands of the
/// single block. This constraint makes it easy to determine the result.
/// This method also clones the `arith::ConstantIndexOp` at the start of
/// `outlinedFuncBody` to alloc simple canonicalizations. If `callOp` is
/// provided, it will be set to point to the operation that calls the outlined
/// function.
// TODO: support more than single-block regions.
// TODO: more flexible constant handling.
FailureOr<func::FuncOp> mlir::outlineSingleBlockRegion(RewriterBase &rewriter,
Location loc,
Region &region,
StringRef funcName,
func::CallOp *callOp) {
assert(!funcName.empty() && "funcName cannot be empty");
if (!region.hasOneBlock())
return failure();
Block *originalBlock = &region.front();
Operation *originalTerminator = originalBlock->getTerminator();
// Outline before current function.
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(region.getParentOfType<FunctionOpInterface>());
SetVector<Value> captures;
getUsedValuesDefinedAbove(region, captures);
ValueRange outlinedValues(captures.getArrayRef());
SmallVector<Type> outlinedFuncArgTypes;
SmallVector<Location> outlinedFuncArgLocs;
// Region's arguments are exactly the first block's arguments as per
// Region::getArguments().
// Func's arguments are cat(regions's arguments, captures arguments).
for (BlockArgument arg : region.getArguments()) {
outlinedFuncArgTypes.push_back(arg.getType());
outlinedFuncArgLocs.push_back(arg.getLoc());
}
for (Value value : outlinedValues) {
outlinedFuncArgTypes.push_back(value.getType());
outlinedFuncArgLocs.push_back(value.getLoc());
}
FunctionType outlinedFuncType =
FunctionType::get(rewriter.getContext(), outlinedFuncArgTypes,
originalTerminator->getOperandTypes());
auto outlinedFunc =
func::FuncOp::create(rewriter, loc, funcName, outlinedFuncType);
Block *outlinedFuncBody = outlinedFunc.addEntryBlock();
// Merge blocks while replacing the original block operands.
// Warning: `mergeBlocks` erases the original block, reconstruct it later.
int64_t numOriginalBlockArguments = originalBlock->getNumArguments();
auto outlinedFuncBlockArgs = outlinedFuncBody->getArguments();
{
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPointToEnd(outlinedFuncBody);
rewriter.mergeBlocks(
originalBlock, outlinedFuncBody,
outlinedFuncBlockArgs.take_front(numOriginalBlockArguments));
// Explicitly set up a new ReturnOp terminator.
rewriter.setInsertionPointToEnd(outlinedFuncBody);
func::ReturnOp::create(rewriter, loc, originalTerminator->getResultTypes(),
originalTerminator->getOperands());
}
// Reconstruct the block that was deleted and add a
// terminator(call_results).
Block *newBlock = rewriter.createBlock(
&region, region.begin(),
TypeRange{outlinedFuncArgTypes}.take_front(numOriginalBlockArguments),
ArrayRef<Location>(outlinedFuncArgLocs)
.take_front(numOriginalBlockArguments));
{
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPointToEnd(newBlock);
SmallVector<Value> callValues;
llvm::append_range(callValues, newBlock->getArguments());
llvm::append_range(callValues, outlinedValues);
auto call = func::CallOp::create(rewriter, loc, outlinedFunc, callValues);
if (callOp)
*callOp = call;
// `originalTerminator` was moved to `outlinedFuncBody` and is still valid.
// Clone `originalTerminator` to take the callOp results then erase it from
// `outlinedFuncBody`.
IRMapping bvm;
bvm.map(originalTerminator->getOperands(), call->getResults());
rewriter.clone(*originalTerminator, bvm);
rewriter.eraseOp(originalTerminator);
}
// Lastly, explicit RAUW outlinedValues, only for uses within `outlinedFunc`.
// Clone the `arith::ConstantIndexOp` at the start of `outlinedFuncBody`.
for (auto it : llvm::zip(outlinedValues, outlinedFuncBlockArgs.take_back(
outlinedValues.size()))) {
Value orig = std::get<0>(it);
Value repl = std::get<1>(it);
{
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPointToStart(outlinedFuncBody);
if (Operation *cst = orig.getDefiningOp<arith::ConstantIndexOp>()) {
repl = rewriter.clone(*cst)->getResult(0);
}
}
orig.replaceUsesWithIf(repl, [&](OpOperand &opOperand) {
return outlinedFunc->isProperAncestor(opOperand.getOwner());
});
}
return outlinedFunc;
}
LogicalResult mlir::outlineIfOp(RewriterBase &b, scf::IfOp ifOp,
func::FuncOp *thenFn, StringRef thenFnName,
func::FuncOp *elseFn, StringRef elseFnName) {
IRRewriter rewriter(b);
Location loc = ifOp.getLoc();
FailureOr<func::FuncOp> outlinedFuncOpOrFailure;
if (thenFn && !ifOp.getThenRegion().empty()) {
outlinedFuncOpOrFailure = outlineSingleBlockRegion(
rewriter, loc, ifOp.getThenRegion(), thenFnName);
if (failed(outlinedFuncOpOrFailure))
return failure();
*thenFn = *outlinedFuncOpOrFailure;
}
if (elseFn && !ifOp.getElseRegion().empty()) {
outlinedFuncOpOrFailure = outlineSingleBlockRegion(
rewriter, loc, ifOp.getElseRegion(), elseFnName);
if (failed(outlinedFuncOpOrFailure))
return failure();
*elseFn = *outlinedFuncOpOrFailure;
}
return success();
}
bool mlir::getInnermostParallelLoops(Operation *rootOp,
SmallVectorImpl<scf::ParallelOp> &result) {
assert(rootOp != nullptr && "Root operation must not be a nullptr.");
bool rootEnclosesPloops = false;
for (Region &region : rootOp->getRegions()) {
for (Block &block : region.getBlocks()) {
for (Operation &op : block) {
bool enclosesPloops = getInnermostParallelLoops(&op, result);
rootEnclosesPloops |= enclosesPloops;
if (auto ploop = dyn_cast<scf::ParallelOp>(op)) {
rootEnclosesPloops = true;
// Collect parallel loop if it is an innermost one.
if (!enclosesPloops)
result.push_back(ploop);
}
}
}
}
return rootEnclosesPloops;
}
// Build the IR that performs ceil division of a positive value by a constant:
// ceildiv(a, B) = divis(a + (B-1), B)
// where divis is rounding-to-zero division.
static Value ceilDivPositive(OpBuilder &builder, Location loc, Value dividend,
int64_t divisor) {
assert(divisor > 0 && "expected positive divisor");
assert(dividend.getType().isIntOrIndex() &&
"expected integer or index-typed value");
Value divisorMinusOneCst = arith::ConstantOp::create(
builder, loc, builder.getIntegerAttr(dividend.getType(), divisor - 1));
Value divisorCst = arith::ConstantOp::create(
builder, loc, builder.getIntegerAttr(dividend.getType(), divisor));
Value sum = arith::AddIOp::create(builder, loc, dividend, divisorMinusOneCst);
return arith::DivUIOp::create(builder, loc, sum, divisorCst);
}
// Build the IR that performs ceil division of a positive value by another
// positive value:
// ceildiv(a, b) = divis(a + (b - 1), b)
// where divis is rounding-to-zero division.
static Value ceilDivPositive(OpBuilder &builder, Location loc, Value dividend,
Value divisor) {
assert(dividend.getType().isIntOrIndex() &&
"expected integer or index-typed value");
Value cstOne = arith::ConstantOp::create(
builder, loc, builder.getOneAttr(dividend.getType()));
Value divisorMinusOne = arith::SubIOp::create(builder, loc, divisor, cstOne);
Value sum = arith::AddIOp::create(builder, loc, dividend, divisorMinusOne);
return arith::DivUIOp::create(builder, loc, sum, divisor);
}
void mlir::generateUnrolledLoop(
Block *loopBodyBlock, Value iv, uint64_t unrollFactor,
function_ref<Value(unsigned, Value, OpBuilder)> ivRemapFn,
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
ValueRange iterArgs, ValueRange yieldedValues,
IRMapping *clonedToSrcOpsMap) {
// Check if the op was cloned from another source op, and return it if found
// (or the same op if not found)
auto findOriginalSrcOp =
[](Operation *op, const IRMapping &clonedToSrcOpsMap) -> Operation * {
Operation *srcOp = op;
// If the source op derives from another op: traverse the chain to find the
// original source op
while (srcOp && clonedToSrcOpsMap.contains(srcOp))
srcOp = clonedToSrcOpsMap.lookup(srcOp);
return srcOp;
};
// Builder to insert unrolled bodies just before the terminator of the body of
// the loop.
auto builder = OpBuilder::atBlockTerminator(loopBodyBlock);
static const auto noopAnnotateFn = [](unsigned, Operation *, OpBuilder) {};
if (!annotateFn)
annotateFn = noopAnnotateFn;
// Keep a pointer to the last non-terminator operation in the original block
// so that we know what to clone (since we are doing this in-place).
Block::iterator srcBlockEnd = std::prev(loopBodyBlock->end(), 2);
// Unroll the contents of the loop body (append unrollFactor - 1 additional
// copies).
SmallVector<Value, 4> lastYielded(yieldedValues);
for (unsigned i = 1; i < unrollFactor; i++) {
// Prepare operand map.
IRMapping operandMap;
operandMap.map(iterArgs, lastYielded);
// If the induction variable is used, create a remapping to the value for
// this unrolled instance.
if (!iv.use_empty()) {
Value ivUnroll = ivRemapFn(i, iv, builder);
operandMap.map(iv, ivUnroll);
}
// Clone the original body of 'forOp'.
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++) {
Operation *srcOp = &(*it);
Operation *clonedOp = builder.clone(*srcOp, operandMap);
annotateFn(i, clonedOp, builder);
if (clonedToSrcOpsMap)
clonedToSrcOpsMap->map(clonedOp,
findOriginalSrcOp(srcOp, *clonedToSrcOpsMap));
}
// Update yielded values.
for (unsigned i = 0, e = lastYielded.size(); i < e; i++)
lastYielded[i] = operandMap.lookupOrDefault(yieldedValues[i]);
}
// Make sure we annotate the Ops in the original body. We do this last so that
// any annotations are not copied into the cloned Ops above.
for (auto it = loopBodyBlock->begin(); it != std::next(srcBlockEnd); it++)
annotateFn(0, &*it, builder);
// Update operands of the yield statement.
loopBodyBlock->getTerminator()->setOperands(lastYielded);
}
/// Unrolls 'forOp' by 'unrollFactor', returns the unrolled main loop and the
/// epilogue loop, if the loop is unrolled.
FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
scf::ForOp forOp, uint64_t unrollFactor,
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
assert(unrollFactor > 0 && "expected positive unroll factor");
// Return if the loop body is empty.
if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
return UnrolledLoopInfo{forOp, std::nullopt};
// Compute tripCount = ceilDiv((upperBound - lowerBound), step) and populate
// 'upperBoundUnrolled' and 'stepUnrolled' for static and dynamic cases.
OpBuilder boundsBuilder(forOp);
IRRewriter rewriter(forOp.getContext());
auto loc = forOp.getLoc();
Value step = forOp.getStep();
Value upperBoundUnrolled;
Value stepUnrolled;
bool generateEpilogueLoop = true;
std::optional<APInt> constTripCount = forOp.getStaticTripCount();
if (constTripCount) {
// Constant loop bounds computation.
int64_t lbCst = getConstantIntValue(forOp.getLowerBound()).value();
int64_t ubCst = getConstantIntValue(forOp.getUpperBound()).value();
int64_t stepCst = getConstantIntValue(forOp.getStep()).value();
if (unrollFactor == 1) {
if (constTripCount->isOne() &&
failed(forOp.promoteIfSingleIteration(rewriter)))
return failure();
return UnrolledLoopInfo{forOp, std::nullopt};
}
uint64_t tripCount = constTripCount->getZExtValue();
uint64_t tripCountEvenMultiple = tripCount - tripCount % unrollFactor;
int64_t upperBoundUnrolledCst = lbCst + tripCountEvenMultiple * stepCst;
int64_t stepUnrolledCst = stepCst * unrollFactor;
// Create constant for 'upperBoundUnrolled' and set epilogue loop flag.
generateEpilogueLoop = upperBoundUnrolledCst < ubCst;
if (generateEpilogueLoop)
upperBoundUnrolled = arith::ConstantOp::create(
boundsBuilder, loc,
boundsBuilder.getIntegerAttr(forOp.getUpperBound().getType(),
upperBoundUnrolledCst));
else
upperBoundUnrolled = forOp.getUpperBound();
// Create constant for 'stepUnrolled'.
stepUnrolled =
stepCst == stepUnrolledCst
? step
: arith::ConstantOp::create(boundsBuilder, loc,
boundsBuilder.getIntegerAttr(
step.getType(), stepUnrolledCst));
} else {
// Dynamic loop bounds computation.
// TODO: Add dynamic asserts for negative lb/ub/step, or
// consider using ceilDiv from AffineApplyExpander.
auto lowerBound = forOp.getLowerBound();
auto upperBound = forOp.getUpperBound();
Value diff =
arith::SubIOp::create(boundsBuilder, loc, upperBound, lowerBound);
Value tripCount = ceilDivPositive(boundsBuilder, loc, diff, step);
Value unrollFactorCst = arith::ConstantOp::create(
boundsBuilder, loc,
boundsBuilder.getIntegerAttr(tripCount.getType(), unrollFactor));
Value tripCountRem =
arith::RemSIOp::create(boundsBuilder, loc, tripCount, unrollFactorCst);
// Compute tripCountEvenMultiple = tripCount - (tripCount % unrollFactor)
Value tripCountEvenMultiple =
arith::SubIOp::create(boundsBuilder, loc, tripCount, tripCountRem);
// Compute upperBoundUnrolled = lowerBound + tripCountEvenMultiple * step
upperBoundUnrolled = arith::AddIOp::create(
boundsBuilder, loc, lowerBound,
arith::MulIOp::create(boundsBuilder, loc, tripCountEvenMultiple, step));
// Scale 'step' by 'unrollFactor'.
stepUnrolled =
arith::MulIOp::create(boundsBuilder, loc, step, unrollFactorCst);
}
UnrolledLoopInfo resultLoops;
// Create epilogue clean up loop starting at 'upperBoundUnrolled'.
if (generateEpilogueLoop) {
OpBuilder epilogueBuilder(forOp->getContext());
epilogueBuilder.setInsertionPointAfter(forOp);
auto epilogueForOp = cast<scf::ForOp>(epilogueBuilder.clone(*forOp));
epilogueForOp.setLowerBound(upperBoundUnrolled);
// Update uses of loop results.
auto results = forOp.getResults();
auto epilogueResults = epilogueForOp.getResults();
for (auto e : llvm::zip(results, epilogueResults)) {
std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
}
epilogueForOp->setOperands(epilogueForOp.getNumControlOperands(),
epilogueForOp.getInitArgs().size(), results);
if (epilogueForOp.promoteIfSingleIteration(rewriter).failed())
resultLoops.epilogueLoopOp = epilogueForOp;
}
// Create unrolled loop.
forOp.setUpperBound(upperBoundUnrolled);
forOp.setStep(stepUnrolled);
auto iterArgs = ValueRange(forOp.getRegionIterArgs());
auto yieldedValues = forOp.getBody()->getTerminator()->getOperands();
generateUnrolledLoop(
forOp.getBody(), forOp.getInductionVar(), unrollFactor,
[&](unsigned i, Value iv, OpBuilder b) {
// iv' = iv + step * i;
auto stride = arith::MulIOp::create(
b, loc, step,
arith::ConstantOp::create(b, loc,
b.getIntegerAttr(iv.getType(), i)));
return arith::AddIOp::create(b, loc, iv, stride);
},
annotateFn, iterArgs, yieldedValues);
// Promote the loop body up if this has turned into a single iteration loop.
if (forOp.promoteIfSingleIteration(rewriter).failed())
resultLoops.mainLoopOp = forOp;
return resultLoops;
}
/// Unrolls this loop completely.
LogicalResult mlir::loopUnrollFull(scf::ForOp forOp) {
IRRewriter rewriter(forOp.getContext());
std::optional<APInt> mayBeConstantTripCount = forOp.getStaticTripCount();
if (!mayBeConstantTripCount.has_value())
return failure();
const APInt &tripCount = *mayBeConstantTripCount;
if (tripCount.isZero())
return success();
if (tripCount.isOne())
return forOp.promoteIfSingleIteration(rewriter);
return loopUnrollByFactor(forOp, tripCount.getZExtValue());
}
/// Check if bounds of all inner loops are defined outside of `forOp`
/// and return false if not.
static bool areInnerBoundsInvariant(scf::ForOp forOp) {
auto walkResult = forOp.walk([&](scf::ForOp innerForOp) {
if (!forOp.isDefinedOutsideOfLoop(innerForOp.getLowerBound()) ||
!forOp.isDefinedOutsideOfLoop(innerForOp.getUpperBound()) ||
!forOp.isDefinedOutsideOfLoop(innerForOp.getStep()))
return WalkResult::interrupt();
return WalkResult::advance();
});
return !walkResult.wasInterrupted();
}
/// Unrolls and jams this loop by the specified factor.
LogicalResult mlir::loopUnrollJamByFactor(scf::ForOp forOp,
uint64_t unrollJamFactor) {
assert(unrollJamFactor > 0 && "unroll jam factor should be positive");
if (unrollJamFactor == 1)
return success();
// If any control operand of any inner loop of `forOp` is defined within
// `forOp`, no unroll jam.
if (!areInnerBoundsInvariant(forOp)) {
LDBG() << "failed to unroll and jam: inner bounds are not invariant";
return failure();
}
// Currently, for operations with results are not supported.
if (forOp->getNumResults() > 0) {
LDBG() << "failed to unroll and jam: unsupported loop with results";
return failure();
}
// Currently, only constant trip count that divided by the unroll factor is
// supported.
std::optional<APInt> tripCount = forOp.getStaticTripCount();
if (!tripCount.has_value()) {
// If the trip count is dynamic, do not unroll & jam.
LDBG() << "failed to unroll and jam: trip count could not be determined";
return failure();
}
uint64_t tripCountValue = tripCount->getZExtValue();
if (tripCountValue == 0)
return success();
if (unrollJamFactor > tripCountValue) {
LDBG() << "unroll and jam factor is greater than trip count, set factor to "
"trip "
"count";
unrollJamFactor = tripCountValue;
} else if (tripCountValue % unrollJamFactor != 0) {
LDBG() << "failed to unroll and jam: unsupported trip count that is not a "
"multiple of unroll jam factor";
return failure();
}
// Nothing in the loop body other than the terminator.
if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
return success();
// Gather all sub-blocks to jam upon the loop being unrolled.
JamBlockGatherer<scf::ForOp> jbg;
jbg.walk(forOp);
auto &subBlocks = jbg.subBlocks;
// Collect inner loops.
SmallVector<scf::ForOp> innerLoops;
forOp.walk([&](scf::ForOp innerForOp) { innerLoops.push_back(innerForOp); });
// `operandMaps[i - 1]` carries old->new operand mapping for the ith unrolled
// iteration. There are (`unrollJamFactor` - 1) iterations.
SmallVector<IRMapping> operandMaps(unrollJamFactor - 1);
// For any loop with iter_args, replace it with a new loop that has
// `unrollJamFactor` copies of its iterOperands, iter_args and yield
// operands.
SmallVector<scf::ForOp> newInnerLoops;
IRRewriter rewriter(forOp.getContext());
for (scf::ForOp oldForOp : innerLoops) {
SmallVector<Value> dupIterOperands, dupYieldOperands;
ValueRange oldIterOperands = oldForOp.getInits();
ValueRange oldIterArgs = oldForOp.getRegionIterArgs();
ValueRange oldYieldOperands =
cast<scf::YieldOp>(oldForOp.getBody()->getTerminator()).getOperands();
// Get additional iterOperands, iterArgs, and yield operands. We will
// fix iterOperands and yield operands after cloning of sub-blocks.
for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
dupIterOperands.append(oldIterOperands.begin(), oldIterOperands.end());
dupYieldOperands.append(oldYieldOperands.begin(), oldYieldOperands.end());
}
// Create a new loop with additional iterOperands, iter_args and yield
// operands. This new loop will take the loop body of the original loop.
bool forOpReplaced = oldForOp == forOp;
scf::ForOp newForOp =
cast<scf::ForOp>(*oldForOp.replaceWithAdditionalYields(
rewriter, dupIterOperands, /*replaceInitOperandUsesInLoop=*/false,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBbArgs) {
return dupYieldOperands;
}));
newInnerLoops.push_back(newForOp);
// `forOp` has been replaced with a new loop.
if (forOpReplaced)
forOp = newForOp;
// Update `operandMaps` for `newForOp` iterArgs and results.
ValueRange newIterArgs = newForOp.getRegionIterArgs();
unsigned oldNumIterArgs = oldIterArgs.size();
ValueRange newResults = newForOp.getResults();
unsigned oldNumResults = newResults.size() / unrollJamFactor;
assert(oldNumIterArgs == oldNumResults &&
"oldNumIterArgs must be the same as oldNumResults");
for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
for (unsigned j = 0; j < oldNumIterArgs; ++j) {
// `newForOp` has `unrollJamFactor` - 1 new sets of iterArgs and
// results. Update `operandMaps[i - 1]` to map old iterArgs and results
// to those in the `i`th new set.
operandMaps[i - 1].map(newIterArgs[j],
newIterArgs[i * oldNumIterArgs + j]);
operandMaps[i - 1].map(newResults[j],
newResults[i * oldNumResults + j]);
}
}
}
// Scale the step of loop being unroll-jammed by the unroll-jam factor.
rewriter.setInsertionPoint(forOp);
int64_t step = forOp.getConstantStep()->getSExtValue();
auto newStep = rewriter.createOrFold<arith::MulIOp>(
forOp.getLoc(), forOp.getStep(),
rewriter.createOrFold<arith::ConstantOp>(
forOp.getLoc(), rewriter.getIndexAttr(unrollJamFactor)));
forOp.setStep(newStep);
auto forOpIV = forOp.getInductionVar();
// Unroll and jam (appends unrollJamFactor - 1 additional copies).
for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
for (auto &subBlock : subBlocks) {
// Builder to insert unroll-jammed bodies. Insert right at the end of
// sub-block.
OpBuilder builder(subBlock.first->getBlock(), std::next(subBlock.second));
// If the induction variable is used, create a remapping to the value for
// this unrolled instance.
if (!forOpIV.use_empty()) {
// iv' = iv + i * step, i = 1 to unrollJamFactor-1.
auto ivTag = builder.createOrFold<arith::ConstantOp>(
forOp.getLoc(), builder.getIndexAttr(step * i));
auto ivUnroll =
builder.createOrFold<arith::AddIOp>(forOp.getLoc(), forOpIV, ivTag);
operandMaps[i - 1].map(forOpIV, ivUnroll);
}
// Clone the sub-block being unroll-jammed.
for (auto it = subBlock.first; it != std::next(subBlock.second); ++it)
builder.clone(*it, operandMaps[i - 1]);
}
// Fix iterOperands and yield op operands of newly created loops.
for (auto newForOp : newInnerLoops) {
unsigned oldNumIterOperands =
newForOp.getNumRegionIterArgs() / unrollJamFactor;
unsigned numControlOperands = newForOp.getNumControlOperands();
auto yieldOp = cast<scf::YieldOp>(newForOp.getBody()->getTerminator());
unsigned oldNumYieldOperands = yieldOp.getNumOperands() / unrollJamFactor;
assert(oldNumIterOperands == oldNumYieldOperands &&
"oldNumIterOperands must be the same as oldNumYieldOperands");
for (unsigned j = 0; j < oldNumIterOperands; ++j) {
// The `i`th duplication of an old iterOperand or yield op operand
// needs to be replaced with a mapped value from `operandMaps[i - 1]`
// if such mapped value exists.
newForOp.setOperand(numControlOperands + i * oldNumIterOperands + j,
operandMaps[i - 1].lookupOrDefault(
newForOp.getOperand(numControlOperands + j)));
yieldOp.setOperand(
i * oldNumYieldOperands + j,
operandMaps[i - 1].lookupOrDefault(yieldOp.getOperand(j)));
}
}
}
// Promote the loop body up if this has turned into a single iteration loop.
(void)forOp.promoteIfSingleIteration(rewriter);
return success();
}
static Range emitNormalizedLoopBoundsForIndexType(RewriterBase &rewriter,
Location loc, OpFoldResult lb,
OpFoldResult ub,
OpFoldResult step) {
Range normalizedLoopBounds;
normalizedLoopBounds.offset = rewriter.getIndexAttr(0);
normalizedLoopBounds.stride = rewriter.getIndexAttr(1);
AffineExpr s0, s1, s2;
bindSymbols(rewriter.getContext(), s0, s1, s2);
AffineExpr e = (s1 - s0).ceilDiv(s2);
normalizedLoopBounds.size =
affine::makeComposedFoldedAffineApply(rewriter, loc, e, {lb, ub, step});
return normalizedLoopBounds;
}
Range mlir::emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
OpFoldResult lb, OpFoldResult ub,
OpFoldResult step) {
if (getType(lb).isIndex()) {
return emitNormalizedLoopBoundsForIndexType(rewriter, loc, lb, ub, step);
}
// For non-index types, generate `arith` instructions
// Check if the loop is already known to have a constant zero lower bound or
// a constant one step.
bool isZeroBased = false;
if (auto lbCst = getConstantIntValue(lb))
isZeroBased = lbCst.value() == 0;
bool isStepOne = false;
if (auto stepCst = getConstantIntValue(step))
isStepOne = stepCst.value() == 1;
Type rangeType = getType(lb);
assert(rangeType == getType(ub) && rangeType == getType(step) &&
"expected matching types");
// Compute the number of iterations the loop executes: ceildiv(ub - lb, step)
// assuming the step is strictly positive. Update the bounds and the step
// of the loop to go from 0 to the number of iterations, if necessary.
if (isZeroBased && isStepOne)
return {lb, ub, step};
OpFoldResult diff = ub;
if (!isZeroBased) {
diff = rewriter.createOrFold<arith::SubIOp>(
loc, getValueOrCreateConstantIntOp(rewriter, loc, ub),
getValueOrCreateConstantIntOp(rewriter, loc, lb));
}
OpFoldResult newUpperBound = diff;
if (!isStepOne) {
newUpperBound = rewriter.createOrFold<arith::CeilDivSIOp>(
loc, getValueOrCreateConstantIntOp(rewriter, loc, diff),
getValueOrCreateConstantIntOp(rewriter, loc, step));
}
OpFoldResult newLowerBound = rewriter.getZeroAttr(rangeType);
OpFoldResult newStep = rewriter.getOneAttr(rangeType);
return {newLowerBound, newUpperBound, newStep};
}
static void denormalizeInductionVariableForIndexType(RewriterBase &rewriter,
Location loc,
Value normalizedIv,
OpFoldResult origLb,
OpFoldResult origStep) {
AffineExpr d0, s0, s1;
bindSymbols(rewriter.getContext(), s0, s1);
bindDims(rewriter.getContext(), d0);
AffineExpr e = d0 * s1 + s0;
OpFoldResult denormalizedIv = affine::makeComposedFoldedAffineApply(
rewriter, loc, e, ArrayRef<OpFoldResult>{normalizedIv, origLb, origStep});
Value denormalizedIvVal =
getValueOrCreateConstantIndexOp(rewriter, loc, denormalizedIv);
SmallPtrSet<Operation *, 1> preservedUses;
// If an `affine.apply` operation is generated for denormalization, the use
// of `origLb` in those ops must not be replaced. These arent not generated
// when `origLb == 0` and `origStep == 1`.
if (!isZeroInteger(origLb) || !isOneInteger(origStep)) {
if (Operation *preservedUse = denormalizedIvVal.getDefiningOp()) {
preservedUses.insert(preservedUse);
}
}
rewriter.replaceAllUsesExcept(normalizedIv, denormalizedIvVal, preservedUses);
}
void mlir::denormalizeInductionVariable(RewriterBase &rewriter, Location loc,
Value normalizedIv, OpFoldResult origLb,
OpFoldResult origStep) {
if (getType(origLb).isIndex()) {
return denormalizeInductionVariableForIndexType(rewriter, loc, normalizedIv,
origLb, origStep);
}
Value denormalizedIv;
SmallPtrSet<Operation *, 2> preserve;
bool isStepOne = isOneInteger(origStep);
bool isZeroBased = isZeroInteger(origLb);
Value scaled = normalizedIv;
if (!isStepOne) {
Value origStepValue =
getValueOrCreateConstantIntOp(rewriter, loc, origStep);
scaled = arith::MulIOp::create(rewriter, loc, normalizedIv, origStepValue);
preserve.insert(scaled.getDefiningOp());
}
denormalizedIv = scaled;
if (!isZeroBased) {
Value origLbValue = getValueOrCreateConstantIntOp(rewriter, loc, origLb);
denormalizedIv = arith::AddIOp::create(rewriter, loc, scaled, origLbValue);
preserve.insert(denormalizedIv.getDefiningOp());
}
rewriter.replaceAllUsesExcept(normalizedIv, denormalizedIv, preserve);
}
static OpFoldResult getProductOfIndexes(RewriterBase &rewriter, Location loc,
ArrayRef<OpFoldResult> values) {
assert(!values.empty() && "unexecpted empty array");
AffineExpr s0, s1;
bindSymbols(rewriter.getContext(), s0, s1);
AffineExpr mul = s0 * s1;
OpFoldResult products = rewriter.getIndexAttr(1);
for (auto v : values) {
products = affine::makeComposedFoldedAffineApply(
rewriter, loc, mul, ArrayRef<OpFoldResult>{products, v});
}
return products;
}
/// Helper function to multiply a sequence of values.
static Value getProductOfIntsOrIndexes(RewriterBase &rewriter, Location loc,
ArrayRef<Value> values) {
assert(!values.empty() && "unexpected empty list");
if (getType(values.front()).isIndex()) {
SmallVector<OpFoldResult> ofrs = getAsOpFoldResult(values);
OpFoldResult product = getProductOfIndexes(rewriter, loc, ofrs);
return getValueOrCreateConstantIndexOp(rewriter, loc, product);
}
std::optional<Value> productOf;
for (auto v : values) {
auto vOne = getConstantIntValue(v);
if (vOne && vOne.value() == 1)
continue;
if (productOf)
productOf = arith::MulIOp::create(rewriter, loc, productOf.value(), v)
.getResult();
else
productOf = v;
}
if (!productOf) {
productOf = arith::ConstantOp::create(
rewriter, loc, rewriter.getOneAttr(getType(values.front())))
.getResult();
}
return productOf.value();
}
/// For each original loop, the value of the
/// induction variable can be obtained by dividing the induction variable of
/// the linearized loop by the total number of iterations of the loops nested
/// in it modulo the number of iterations in this loop (remove the values
/// related to the outer loops):
/// iv_i = floordiv(iv_linear, product-of-loop-ranges-until-i) mod range_i.
/// Compute these iteratively from the innermost loop by creating a "running
/// quotient" of division by the range.
static std::pair<SmallVector<Value>, SmallPtrSet<Operation *, 2>>
delinearizeInductionVariable(RewriterBase &rewriter, Location loc,
Value linearizedIv, ArrayRef<Value> ubs) {
if (linearizedIv.getType().isIndex()) {
Operation *delinearizedOp = affine::AffineDelinearizeIndexOp::create(
rewriter, loc, linearizedIv, ubs);
auto resultVals = llvm::map_to_vector(
delinearizedOp->getResults(), [](OpResult r) -> Value { return r; });
return {resultVals, SmallPtrSet<Operation *, 2>{delinearizedOp}};
}
SmallVector<Value> delinearizedIvs(ubs.size());
SmallPtrSet<Operation *, 2> preservedUsers;
llvm::BitVector isUbOne(ubs.size());
for (auto [index, ub] : llvm::enumerate(ubs)) {
auto ubCst = getConstantIntValue(ub);
if (ubCst && ubCst.value() == 1)
isUbOne.set(index);
}
// Prune the lead ubs that are all ones.
unsigned numLeadingOneUbs = 0;
for (auto [index, ub] : llvm::enumerate(ubs)) {
if (!isUbOne.test(index)) {
break;
}
delinearizedIvs[index] = arith::ConstantOp::create(
rewriter, loc, rewriter.getZeroAttr(ub.getType()));
numLeadingOneUbs++;
}
Value previous = linearizedIv;
for (unsigned i = numLeadingOneUbs, e = ubs.size(); i < e; ++i) {
unsigned idx = ubs.size() - (i - numLeadingOneUbs) - 1;
if (i != numLeadingOneUbs && !isUbOne.test(idx + 1)) {
previous = arith::DivSIOp::create(rewriter, loc, previous, ubs[idx + 1]);
preservedUsers.insert(previous.getDefiningOp());
}
Value iv = previous;
if (i != e - 1) {
if (!isUbOne.test(idx)) {
iv = arith::RemSIOp::create(rewriter, loc, previous, ubs[idx]);
preservedUsers.insert(iv.getDefiningOp());
} else {
iv = arith::ConstantOp::create(
rewriter, loc, rewriter.getZeroAttr(ubs[idx].getType()));
}
}
delinearizedIvs[idx] = iv;
}
return {delinearizedIvs, preservedUsers};
}
LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
MutableArrayRef<scf::ForOp> loops) {
if (loops.size() < 2)
return failure();
scf::ForOp innermost = loops.back();
scf::ForOp outermost = loops.front();
// Bail out if any loop has a known zero step, as normalization
// would result in a division by zero.
for (auto loop : loops) {
if (auto step = getConstantIntValue(loop.getStep())) {
if (step.value() == 0) {
return failure();
}
}
}
// 1. Make sure all loops iterate from 0 to upperBound with step 1. This
// allows the following code to assume upperBound is the number of iterations.
for (auto loop : loops) {
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(outermost);
Value lb = loop.getLowerBound();
Value ub = loop.getUpperBound();
Value step = loop.getStep();
auto newLoopRange =
emitNormalizedLoopBounds(rewriter, loop.getLoc(), lb, ub, step);
rewriter.modifyOpInPlace(loop, [&]() {
loop.setLowerBound(getValueOrCreateConstantIntOp(rewriter, loop.getLoc(),
newLoopRange.offset));
loop.setUpperBound(getValueOrCreateConstantIntOp(rewriter, loop.getLoc(),
newLoopRange.size));
loop.setStep(getValueOrCreateConstantIntOp(rewriter, loop.getLoc(),
newLoopRange.stride));
});
rewriter.setInsertionPointToStart(innermost.getBody());
denormalizeInductionVariable(rewriter, loop.getLoc(),
loop.getInductionVar(), lb, step);
}
// 2. Emit code computing the upper bound of the coalesced loop as product
// of the number of iterations of all loops.
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(outermost);
Location loc = outermost.getLoc();
SmallVector<Value> upperBounds = llvm::map_to_vector(
loops, [](auto loop) { return loop.getUpperBound(); });
Value upperBound = getProductOfIntsOrIndexes(rewriter, loc, upperBounds);
outermost.setUpperBound(upperBound);
// Insert delinearization at the start of the outermost loop body.
rewriter.setInsertionPointToStart(outermost.getBody());
auto [delinearizeIvs, preservedUsers] = delinearizeInductionVariable(
rewriter, loc, outermost.getInductionVar(), upperBounds);
rewriter.replaceAllUsesExcept(outermost.getInductionVar(), delinearizeIvs[0],
preservedUsers);
for (int i = loops.size() - 1; i > 0; --i) {
auto outerLoop = loops[i - 1];
auto innerLoop = loops[i];
Operation *innerTerminator = innerLoop.getBody()->getTerminator();
auto yieldedVals = llvm::to_vector(innerTerminator->getOperands());
assert(llvm::equal(outerLoop.getRegionIterArgs(), innerLoop.getInitArgs()));
for (Value &yieldedVal : yieldedVals) {
// The yielded value may be an iteration argument of the inner loop
// which is about to be inlined.
auto iter = llvm::find(innerLoop.getRegionIterArgs(), yieldedVal);
if (iter != innerLoop.getRegionIterArgs().end()) {
unsigned iterArgIndex = iter - innerLoop.getRegionIterArgs().begin();
// `outerLoop` iter args identical to the `innerLoop` init args.
assert(iterArgIndex < innerLoop.getInitArgs().size());
yieldedVal = innerLoop.getInitArgs()[iterArgIndex];
}
}
rewriter.eraseOp(innerTerminator);
SmallVector<Value> innerBlockArgs;
innerBlockArgs.push_back(delinearizeIvs[i]);
llvm::append_range(innerBlockArgs, outerLoop.getRegionIterArgs());
rewriter.inlineBlockBefore(innerLoop.getBody(), outerLoop.getBody(),
Block::iterator(innerLoop), innerBlockArgs);
rewriter.replaceOp(innerLoop, yieldedVals);
}
return success();
}
LogicalResult mlir::coalesceLoops(MutableArrayRef<scf::ForOp> loops) {
if (loops.empty()) {
return failure();
}
IRRewriter rewriter(loops.front().getContext());
return coalesceLoops(rewriter, loops);
}
LogicalResult mlir::coalescePerfectlyNestedSCFForLoops(scf::ForOp op) {
LogicalResult result(failure());
SmallVector<scf::ForOp> loops;
getPerfectlyNestedLoops(loops, op);
// Look for a band of loops that can be coalesced, i.e. perfectly nested
// loops with bounds defined above some loop.
// 1. For each loop, find above which parent loop its bounds operands are
// defined.
SmallVector<unsigned> operandsDefinedAbove(loops.size());
for (unsigned i = 0, e = loops.size(); i < e; ++i) {
operandsDefinedAbove[i] = i;
for (unsigned j = 0; j < i; ++j) {
SmallVector<Value> boundsOperands = {loops[i].getLowerBound(),
loops[i].getUpperBound(),
loops[i].getStep()};
if (areValuesDefinedAbove(boundsOperands, loops[j].getRegion())) {
operandsDefinedAbove[i] = j;
break;
}
}
}
// 2. For each inner loop check that the iter_args for the immediately outer
// loop are the init for the immediately inner loop and that the yields of the
// return of the inner loop is the yield for the immediately outer loop. Keep
// track of where the chain starts from for each loop.
SmallVector<unsigned> iterArgChainStart(loops.size());
iterArgChainStart[0] = 0;
for (unsigned i = 1, e = loops.size(); i < e; ++i) {
// By default set the start of the chain to itself.
iterArgChainStart[i] = i;
auto outerloop = loops[i - 1];
auto innerLoop = loops[i];
if (outerloop.getNumRegionIterArgs() != innerLoop.getNumRegionIterArgs()) {
continue;
}
if (!llvm::equal(outerloop.getRegionIterArgs(), innerLoop.getInitArgs())) {
continue;
}
auto outerloopTerminator = outerloop.getBody()->getTerminator();
if (!llvm::equal(outerloopTerminator->getOperands(),
innerLoop.getResults())) {
continue;
}
iterArgChainStart[i] = iterArgChainStart[i - 1];
}
// 3. Identify bands of loops such that the operands of all of them are
// defined above the first loop in the band. Traverse the nest bottom-up
// so that modifications don't invalidate the inner loops.
for (unsigned end = loops.size(); end > 0; --end) {
unsigned start = 0;
for (; start < end - 1; ++start) {
auto maxPos =
*std::max_element(std::next(operandsDefinedAbove.begin(), start),
std::next(operandsDefinedAbove.begin(), end));
if (maxPos > start)
continue;
if (iterArgChainStart[end - 1] > start)
continue;
auto band = llvm::MutableArrayRef(loops.data() + start, end - start);
if (succeeded(coalesceLoops(band)))
result = success();
break;
}
// If a band was found and transformed, keep looking at the loops above
// the outermost transformed loop.
if (start != end - 1)
end = start + 1;
}
return result;
}
void mlir::collapseParallelLoops(
RewriterBase &rewriter, scf::ParallelOp loops,
ArrayRef<std::vector<unsigned>> combinedDimensions) {
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(loops);
Location loc = loops.getLoc();
// Presort combined dimensions.
auto sortedDimensions = llvm::to_vector<3>(combinedDimensions);
for (auto &dims : sortedDimensions)
llvm::sort(dims);
// Normalize ParallelOp's iteration pattern.
SmallVector<Value, 3> normalizedUpperBounds;
for (unsigned i = 0, e = loops.getNumLoops(); i < e; ++i) {
OpBuilder::InsertionGuard g2(rewriter);
rewriter.setInsertionPoint(loops);
Value lb = loops.getLowerBound()[i];
Value ub = loops.getUpperBound()[i];
Value step = loops.getStep()[i];
auto newLoopRange = emitNormalizedLoopBounds(rewriter, loc, lb, ub, step);
normalizedUpperBounds.push_back(getValueOrCreateConstantIntOp(
rewriter, loops.getLoc(), newLoopRange.size));
rewriter.setInsertionPointToStart(loops.getBody());
denormalizeInductionVariable(rewriter, loc, loops.getInductionVars()[i], lb,
step);
}
// Combine iteration spaces.
SmallVector<Value, 3> lowerBounds, upperBounds, steps;
auto cst0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
auto cst1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
for (auto &sortedDimension : sortedDimensions) {
Value newUpperBound = arith::ConstantIndexOp::create(rewriter, loc, 1);
for (auto idx : sortedDimension) {
newUpperBound = arith::MulIOp::create(rewriter, loc, newUpperBound,
normalizedUpperBounds[idx]);
}
lowerBounds.push_back(cst0);
steps.push_back(cst1);
upperBounds.push_back(newUpperBound);
}
// Create new ParallelLoop with conversions to the original induction values.
// The loop below uses divisions to get the relevant range of values in the
// new induction value that represent each range of the original induction
// value. The remainders then determine based on that range, which iteration
// of the original induction value this represents. This is a normalized value
// that is un-normalized already by the previous logic.
auto newPloop = scf::ParallelOp::create(
rewriter, loc, lowerBounds, upperBounds, steps,
[&](OpBuilder &insideBuilder, Location, ValueRange ploopIVs) {
for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
Value previous = ploopIVs[i];
unsigned numberCombinedDimensions = combinedDimensions[i].size();
// Iterate over all except the last induction value.
for (unsigned j = numberCombinedDimensions - 1; j > 0; --j) {
unsigned idx = combinedDimensions[i][j];
// Determine the current induction value's current loop iteration
Value iv = arith::RemSIOp::create(insideBuilder, loc, previous,
normalizedUpperBounds[idx]);
replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx), iv,
loops.getRegion());
// Remove the effect of the current induction value to prepare for
// the next value.
previous = arith::DivSIOp::create(insideBuilder, loc, previous,
normalizedUpperBounds[idx]);
}
// The final induction value is just the remaining value.
unsigned idx = combinedDimensions[i][0];
replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx),
previous, loops.getRegion());
}
});
// Replace the old loop with the new loop.
loops.getBody()->back().erase();
newPloop.getBody()->getOperations().splice(
Block::iterator(newPloop.getBody()->back()),
loops.getBody()->getOperations());
loops.erase();
}
// Hoist the ops within `outer` that appear before `inner`.
// Such ops include the ops that have been introduced by parametric tiling.
// Ops that come from triangular loops (i.e. that belong to the program slice
// rooted at `outer`) and ops that have side effects cannot be hoisted.
// Return failure when any op fails to hoist.
static LogicalResult hoistOpsBetween(scf::ForOp outer, scf::ForOp inner) {
SetVector<Operation *> forwardSlice;
ForwardSliceOptions options;
options.filter = [&inner](Operation *op) {
return op != inner.getOperation();
};
getForwardSlice(outer.getInductionVar(), &forwardSlice, options);
LogicalResult status = success();
SmallVector<Operation *, 8> toHoist;
for (auto &op : outer.getBody()->without_terminator()) {
// Stop when encountering the inner loop.
if (&op == inner.getOperation())
break;
// Skip over non-hoistable ops.
if (forwardSlice.count(&op) > 0) {
status = failure();
continue;
}
// Skip intermediate scf::ForOp, these are not considered a failure.
if (isa<scf::ForOp>(op))
continue;
// Skip other ops with regions.
if (op.getNumRegions() > 0) {
status = failure();
continue;
}
// Skip if op has side effects.
// TODO: loads to immutable memory regions are ok.
if (!isMemoryEffectFree(&op)) {
status = failure();
continue;
}
toHoist.push_back(&op);
}
auto *outerForOp = outer.getOperation();
for (auto *op : toHoist)
op->moveBefore(outerForOp);
return status;
}
// Traverse the interTile and intraTile loops and try to hoist ops such that
// bands of perfectly nested loops are isolated.
// Return failure if either perfect interTile or perfect intraTile bands cannot
// be formed.
static LogicalResult tryIsolateBands(const TileLoops &tileLoops) {
LogicalResult status = success();
const Loops &interTile = tileLoops.first;
const Loops &intraTile = tileLoops.second;
auto size = interTile.size();
assert(size == intraTile.size());
if (size <= 1)
return success();
for (unsigned s = 1; s < size; ++s)
status = succeeded(status) ? hoistOpsBetween(intraTile[0], intraTile[s])
: failure();
for (unsigned s = 1; s < size; ++s)
status = succeeded(status) ? hoistOpsBetween(interTile[0], interTile[s])
: failure();
return status;
}
/// Collect perfectly nested loops starting from `rootForOps`. Loops are
/// perfectly nested if each loop is the first and only non-terminator operation
/// in the parent loop. Collect at most `maxLoops` loops and append them to
/// `forOps`.
template <typename T>
static void getPerfectlyNestedLoopsImpl(
SmallVectorImpl<T> &forOps, T rootForOp,
unsigned maxLoops = std::numeric_limits<unsigned>::max()) {
for (unsigned i = 0; i < maxLoops; ++i) {
forOps.push_back(rootForOp);
Block &body = rootForOp.getRegion().front();
if (body.begin() != std::prev(body.end(), 2))
return;
rootForOp = dyn_cast<T>(&body.front());
if (!rootForOp)
return;
}
}
static Loops stripmineSink(scf::ForOp forOp, Value factor,
ArrayRef<scf::ForOp> targets) {
assert(!forOp.getUnsignedCmp() && "unsigned loops are not supported");
auto originalStep = forOp.getStep();
auto iv = forOp.getInductionVar();
OpBuilder b(forOp);
forOp.setStep(arith::MulIOp::create(b, forOp.getLoc(), originalStep, factor));
Loops innerLoops;
for (auto t : targets) {
assert(!t.getUnsignedCmp() && "unsigned loops are not supported");
// Save information for splicing ops out of t when done
auto begin = t.getBody()->begin();
auto nOps = t.getBody()->getOperations().size();
// Insert newForOp before the terminator of `t`.
auto b = OpBuilder::atBlockTerminator((t.getBody()));
Value stepped = arith::AddIOp::create(b, t.getLoc(), iv, forOp.getStep());
Value ub =
arith::MinSIOp::create(b, t.getLoc(), forOp.getUpperBound(), stepped);
// Splice [begin, begin + nOps - 1) into `newForOp` and replace uses.
auto newForOp = scf::ForOp::create(b, t.getLoc(), iv, ub, originalStep);
newForOp.getBody()->getOperations().splice(
newForOp.getBody()->getOperations().begin(),
t.getBody()->getOperations(), begin, std::next(begin, nOps - 1));
replaceAllUsesInRegionWith(iv, newForOp.getInductionVar(),
newForOp.getRegion());
innerLoops.push_back(newForOp);
}
return innerLoops;
}
// Stripmines a `forOp` by `factor` and sinks it under a single `target`.
// Returns the new for operation, nested immediately under `target`.
template <typename SizeType>
static scf::ForOp stripmineSink(scf::ForOp forOp, SizeType factor,
scf::ForOp target) {
// TODO: Use cheap structural assertions that targets are nested under
// forOp and that targets are not nested under each other when DominanceInfo
// exposes the capability. It seems overkill to construct a whole function
// dominance tree at this point.
auto res = stripmineSink(forOp, factor, ArrayRef<scf::ForOp>(target));
assert(res.size() == 1 && "Expected 1 inner forOp");
return res[0];
}
SmallVector<Loops, 8> mlir::tile(ArrayRef<scf::ForOp> forOps,
ArrayRef<Value> sizes,
ArrayRef<scf::ForOp> targets) {
SmallVector<SmallVector<scf::ForOp, 8>, 8> res;
SmallVector<scf::ForOp, 8> currentTargets(targets);
for (auto it : llvm::zip(forOps, sizes)) {
auto step = stripmineSink(std::get<0>(it), std::get<1>(it), currentTargets);
res.push_back(step);
currentTargets = step;
}
return res;
}
Loops mlir::tile(ArrayRef<scf::ForOp> forOps, ArrayRef<Value> sizes,
scf::ForOp target) {
SmallVector<scf::ForOp, 8> res;
for (auto loops : tile(forOps, sizes, ArrayRef<scf::ForOp>(target)))
res.push_back(llvm::getSingleElement(loops));
return res;
}
Loops mlir::tilePerfectlyNested(scf::ForOp rootForOp, ArrayRef<Value> sizes) {
// Collect perfectly nested loops. If more size values provided than nested
// loops available, truncate `sizes`.
SmallVector<scf::ForOp, 4> forOps;
forOps.reserve(sizes.size());
getPerfectlyNestedLoopsImpl(forOps, rootForOp, sizes.size());
if (forOps.size() < sizes.size())
sizes = sizes.take_front(forOps.size());
return ::tile(forOps, sizes, forOps.back());
}
void mlir::getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
scf::ForOp root) {
getPerfectlyNestedLoopsImpl(nestedLoops, root);
}
TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
ArrayRef<int64_t> sizes) {
// Collect perfectly nested loops. If more size values provided than nested
// loops available, truncate `sizes`.
SmallVector<scf::ForOp, 4> forOps;
forOps.reserve(sizes.size());
getPerfectlyNestedLoopsImpl(forOps, rootForOp, sizes.size());
if (forOps.size() < sizes.size())
sizes = sizes.take_front(forOps.size());
// The strip-mining transformation splices loop bodies into a new inner loop
// without threading iter_args. If any of the collected loops carries
// iter_args, the splice would produce invalid IR (yielded values from the
// inner scope used in the outer terminator). Skip the transformation in
// that case.
if (llvm::any_of(forOps,
[](scf::ForOp op) { return !op.getInitArgs().empty(); }))
return {};
// Compute the tile sizes such that i-th outer loop executes size[i]
// iterations. Given that the loop current executes
// numIterations = ceildiv((upperBound - lowerBound), step)
// iterations, we need to tile with size ceildiv(numIterations, size[i]).
SmallVector<Value, 4> tileSizes;
tileSizes.reserve(sizes.size());
for (unsigned i = 0, e = sizes.size(); i < e; ++i) {
assert(sizes[i] > 0 && "expected strictly positive size for strip-mining");
auto forOp = forOps[i];
OpBuilder builder(forOp);
auto loc = forOp.getLoc();
Value diff = arith::SubIOp::create(builder, loc, forOp.getUpperBound(),
forOp.getLowerBound());
Value numIterations = ceilDivPositive(builder, loc, diff, forOp.getStep());
Value iterationsPerBlock =
ceilDivPositive(builder, loc, numIterations, sizes[i]);
tileSizes.push_back(iterationsPerBlock);
}
// Call parametric tiling with the given sizes.
auto intraTile = tile(forOps, tileSizes, forOps.back());
TileLoops tileLoops = std::make_pair(forOps, intraTile);
// TODO: for now we just ignore the result of band isolation.
// In the future, mapping decisions may be impacted by the ability to
// isolate perfectly nested bands.
(void)tryIsolateBands(tileLoops);
return tileLoops;
}
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp source,
RewriterBase &rewriter) {
unsigned numTargetOuts = target.getNumResults();
unsigned numSourceOuts = source.getNumResults();
// Create fused shared_outs.
SmallVector<Value> fusedOuts;
llvm::append_range(fusedOuts, target.getOutputs());
llvm::append_range(fusedOuts, source.getOutputs());
// Create a new scf.forall op after the source loop.
rewriter.setInsertionPointAfter(source);
scf::ForallOp fusedLoop = scf::ForallOp::create(
rewriter, source.getLoc(), source.getMixedLowerBound(),
source.getMixedUpperBound(), source.getMixedStep(), fusedOuts,
source.getMapping());
// Map control operands.
IRMapping mapping;
mapping.map(target.getInductionVars(), fusedLoop.getInductionVars());
mapping.map(source.getInductionVars(), fusedLoop.getInductionVars());
// Map shared outs.
mapping.map(target.getRegionIterArgs(),
fusedLoop.getRegionIterArgs().take_front(numTargetOuts));
mapping.map(source.getRegionIterArgs(),
fusedLoop.getRegionIterArgs().take_back(numSourceOuts));
// Append everything except the terminator into the fused operation.
rewriter.setInsertionPointToStart(fusedLoop.getBody());
for (Operation &op : target.getBody()->without_terminator())
rewriter.clone(op, mapping);
for (Operation &op : source.getBody()->without_terminator())
rewriter.clone(op, mapping);
// Fuse the old terminator in_parallel ops into the new one.
scf::InParallelOp targetTerm = target.getTerminator();
scf::InParallelOp sourceTerm = source.getTerminator();
scf::InParallelOp fusedTerm = fusedLoop.getTerminator();
rewriter.setInsertionPointToStart(fusedTerm.getBody());
for (Operation &op : targetTerm.getYieldingOps())
rewriter.clone(op, mapping);
for (Operation &op : sourceTerm.getYieldingOps())
rewriter.clone(op, mapping);
// Replace old loops by substituting their uses by results of the fused loop.
rewriter.replaceOp(target, fusedLoop.getResults().take_front(numTargetOuts));
rewriter.replaceOp(source, fusedLoop.getResults().take_back(numSourceOuts));
return fusedLoop;
}
scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ForOp source,
RewriterBase &rewriter) {
assert(source.getUnsignedCmp() == target.getUnsignedCmp() &&
"incompatible signedness");
unsigned numTargetOuts = target.getNumResults();
unsigned numSourceOuts = source.getNumResults();
// Create fused init_args, with target's init_args before source's init_args.
SmallVector<Value> fusedInitArgs;
llvm::append_range(fusedInitArgs, target.getInitArgs());
llvm::append_range(fusedInitArgs, source.getInitArgs());
// Create a new scf.for op after the source loop (with scf.yield terminator
// (without arguments) only in case its init_args is empty).
rewriter.setInsertionPointAfter(source);
scf::ForOp fusedLoop = scf::ForOp::create(
rewriter, source.getLoc(), source.getLowerBound(), source.getUpperBound(),
source.getStep(), fusedInitArgs, /*bodyBuilder=*/nullptr,
source.getUnsignedCmp());
// Map original induction variables and operands to those of the fused loop.
IRMapping mapping;
mapping.map(target.getInductionVar(), fusedLoop.getInductionVar());
mapping.map(target.getRegionIterArgs(),
fusedLoop.getRegionIterArgs().take_front(numTargetOuts));
mapping.map(source.getInductionVar(), fusedLoop.getInductionVar());
mapping.map(source.getRegionIterArgs(),
fusedLoop.getRegionIterArgs().take_back(numSourceOuts));
// Merge target's body into the new (fused) for loop and then source's body.
rewriter.setInsertionPointToStart(fusedLoop.getBody());
for (Operation &op : target.getBody()->without_terminator())
rewriter.clone(op, mapping);
for (Operation &op : source.getBody()->without_terminator())
rewriter.clone(op, mapping);
// Build fused yield results by appropriately mapping original yield operands.
SmallVector<Value> yieldResults;
for (Value operand : target.getBody()->getTerminator()->getOperands())
yieldResults.push_back(mapping.lookupOrDefault(operand));
for (Value operand : source.getBody()->getTerminator()->getOperands())
yieldResults.push_back(mapping.lookupOrDefault(operand));
if (!yieldResults.empty())
scf::YieldOp::create(rewriter, source.getLoc(), yieldResults);
// Replace old loops by substituting their uses by results of the fused loop.
rewriter.replaceOp(target, fusedLoop.getResults().take_front(numTargetOuts));
rewriter.replaceOp(source, fusedLoop.getResults().take_back(numSourceOuts));
return fusedLoop;
}
FailureOr<scf::ForallOp> mlir::normalizeForallOp(RewriterBase &rewriter,
scf::ForallOp forallOp) {
SmallVector<OpFoldResult> lbs = forallOp.getMixedLowerBound();
SmallVector<OpFoldResult> ubs = forallOp.getMixedUpperBound();
SmallVector<OpFoldResult> steps = forallOp.getMixedStep();
if (forallOp.isNormalized())
return forallOp;
OpBuilder::InsertionGuard g(rewriter);
auto loc = forallOp.getLoc();
rewriter.setInsertionPoint(forallOp);
SmallVector<OpFoldResult> newUbs;
for (auto [lb, ub, step] : llvm::zip_equal(lbs, ubs, steps)) {
Range normalizedLoopParams =
emitNormalizedLoopBounds(rewriter, loc, lb, ub, step);
newUbs.push_back(normalizedLoopParams.size);
}
(void)foldDynamicIndexList(newUbs);
// Use the normalized builder since the lower bounds are always 0 and the
// steps are always 1.
auto normalizedForallOp = scf::ForallOp::create(
rewriter, loc, newUbs, forallOp.getOutputs(), forallOp.getMapping(),
[](OpBuilder &, Location, ValueRange) {});
rewriter.inlineRegionBefore(forallOp.getBodyRegion(),
normalizedForallOp.getBodyRegion(),
normalizedForallOp.getBodyRegion().begin());
// Remove the original empty block in the new loop.
rewriter.eraseBlock(&normalizedForallOp.getBodyRegion().back());
rewriter.setInsertionPointToStart(normalizedForallOp.getBody());
// Update the users of the original loop variables.
for (auto [idx, iv] :
llvm::enumerate(normalizedForallOp.getInductionVars())) {
auto origLb = getValueOrCreateConstantIndexOp(rewriter, loc, lbs[idx]);
auto origStep = getValueOrCreateConstantIndexOp(rewriter, loc, steps[idx]);
denormalizeInductionVariable(rewriter, loc, iv, origLb, origStep);
}
rewriter.replaceOp(forallOp, normalizedForallOp);
return normalizedForallOp;
}
bool mlir::isPerfectlyNestedForLoops(
MutableArrayRef<LoopLikeOpInterface> loops) {
assert(!loops.empty() && "unexpected empty loop nest");
if (loops.size() == 1)
return isa_and_nonnull<scf::ForOp>(loops.front().getOperation());
for (auto [outerLoop, innerLoop] :
llvm::zip_equal(loops.drop_back(), loops.drop_front())) {
auto outerFor = dyn_cast_or_null<scf::ForOp>(outerLoop.getOperation());
auto innerFor = dyn_cast_or_null<scf::ForOp>(innerLoop.getOperation());
if (!outerFor || !innerFor)
return false;
auto outerBBArgs = outerFor.getRegionIterArgs();
auto innerIterArgs = innerFor.getInitArgs();
if (outerBBArgs.size() != innerIterArgs.size())
return false;
for (auto [outerBBArg, innerIterArg] :
llvm::zip_equal(outerBBArgs, innerIterArgs)) {
if (!llvm::hasSingleElement(outerBBArg.getUses()) ||
innerIterArg != outerBBArg)
return false;
}
ValueRange outerYields =
cast<scf::YieldOp>(outerFor.getBody()->getTerminator())->getOperands();
ValueRange innerResults = innerFor.getResults();
if (outerYields.size() != innerResults.size())
return false;
for (auto [outerYield, innerResult] :
llvm::zip_equal(outerYields, innerResults)) {
if (!llvm::hasSingleElement(innerResult.getUses()) ||
outerYield != innerResult)
return false;
}
}
return true;
}
llvm::SmallVector<std::tuple<int64_t, int64_t, int64_t>>
mlir::getConstLoopBounds(mlir::LoopLikeOpInterface loopOp) {
std::optional<SmallVector<OpFoldResult>> loBnds = loopOp.getLoopLowerBounds();
std::optional<SmallVector<OpFoldResult>> upBnds = loopOp.getLoopUpperBounds();
std::optional<SmallVector<OpFoldResult>> steps = loopOp.getLoopSteps();
if (!loBnds || !upBnds || !steps)
return {};
llvm::SmallVector<std::tuple<int64_t, int64_t, int64_t>> loopRanges;
for (auto [lb, ub, step] : llvm::zip(*loBnds, *upBnds, *steps)) {
auto lbCst = getConstantIntValue(lb);
auto ubCst = getConstantIntValue(ub);
auto stepCst = getConstantIntValue(step);
if (!lbCst || !ubCst || !stepCst)
return {};
loopRanges.emplace_back(*lbCst, *ubCst, *stepCst);
}
return loopRanges;
}
llvm::SmallVector<llvm::APInt>
mlir::getConstLoopTripCounts(mlir::LoopLikeOpInterface loopOp) {
std::optional<SmallVector<OpFoldResult>> loBnds = loopOp.getLoopLowerBounds();
std::optional<SmallVector<OpFoldResult>> upBnds = loopOp.getLoopUpperBounds();
std::optional<SmallVector<OpFoldResult>> steps = loopOp.getLoopSteps();
if (!loBnds || !upBnds || !steps)
return {};
llvm::SmallVector<llvm::APInt> tripCounts;
for (auto [lb, ub, step] : llvm::zip(*loBnds, *upBnds, *steps)) {
// TODO(#178506): Signedness is not handled correctly here.
std::optional<llvm::APInt> numIter = constantTripCount(
lb, ub, step, /*isSigned=*/true, scf::computeUbMinusLb);
if (!numIter)
return {};
tripCounts.push_back(*numIter);
}
return tripCounts;
}
FailureOr<scf::ParallelOp> mlir::parallelLoopUnrollByFactors(
scf::ParallelOp op, ArrayRef<uint64_t> unrollFactors,
RewriterBase &rewriter,
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
IRMapping *clonedToSrcOpsMap) {
const unsigned numLoops = op.getNumLoops();
assert(llvm::none_of(unrollFactors, [](uint64_t f) { return f == 0; }) &&
"Expected positive unroll factors");
assert((!unrollFactors.empty() && (unrollFactors.size() <= numLoops)) &&
"Expected non-empty unroll factors of size <= to the number of loops");
// Bail out if no valid unroll factors were provided
if (llvm::all_of(unrollFactors, [](uint64_t f) { return f == 1; }))
return rewriter.notifyMatchFailure(
op, "Unrolling not applied if all factors are 1");
// Return if the loop body is empty.
if (llvm::hasSingleElement(op.getBody()->getOperations()))
return rewriter.notifyMatchFailure(op, "Cannot unroll an empty loop body");
// If the provided unroll factors do not cover all the loop dims, they are
// applied to the inner loop dimensions.
const unsigned firstLoopDimIdx = numLoops - unrollFactors.size();
// Make sure that the unroll factors divide the iteration space evenly
// TODO: Support unrolling loops with dynamic iteration spaces.
const llvm::SmallVector<llvm::APInt> tripCounts = getConstLoopTripCounts(op);
if (tripCounts.empty())
return rewriter.notifyMatchFailure(
op, "Failed to compute constant trip counts for the loop. Note that "
"dynamic loop sizes are not supported.");
for (unsigned dimIdx = firstLoopDimIdx; dimIdx < numLoops; dimIdx++) {
const uint64_t unrollFactor = unrollFactors[dimIdx - firstLoopDimIdx];
if (tripCounts[dimIdx].urem(unrollFactor) != 0)
return rewriter.notifyMatchFailure(
op, "Unroll factors don't divide the iteration space evenly");
}
std::optional<SmallVector<OpFoldResult>> maybeFoldSteps = op.getLoopSteps();
if (!maybeFoldSteps)
return rewriter.notifyMatchFailure(op, "Failed to retrieve loop steps");
llvm::SmallVector<size_t> steps{};
for (auto step : *maybeFoldSteps)
steps.push_back(static_cast<size_t>(*getConstantIntValue(step)));
for (unsigned dimIdx = firstLoopDimIdx; dimIdx < numLoops; dimIdx++) {
const uint64_t unrollFactor = unrollFactors[dimIdx - firstLoopDimIdx];
if (unrollFactor == 1)
continue;
const size_t origStep = steps[dimIdx];
const int64_t newStep = origStep * unrollFactor;
IRMapping clonedToSrcOpsMap;
ValueRange iterArgs = ValueRange(op.getRegionIterArgs());
auto yieldedValues = op.getBody()->getTerminator()->getOperands();
generateUnrolledLoop(
op.getBody(), op.getInductionVars()[dimIdx], unrollFactor,
[&](unsigned i, Value iv, OpBuilder b) {
// iv' = iv + step * i;
const AffineExpr expr = b.getAffineDimExpr(0) + (origStep * i);
const auto map =
b.getDimIdentityMap().dropResult(0).insertResult(expr, 0);
return affine::AffineApplyOp::create(b, iv.getLoc(), map,
ValueRange{iv});
},
/*annotateFn*/ annotateFn, iterArgs, yieldedValues, &clonedToSrcOpsMap);
// Update loop step
auto prevInsertPoint = rewriter.saveInsertionPoint();
rewriter.setInsertionPoint(op);
op.getStepMutable()[dimIdx].assign(
arith::ConstantIndexOp::create(rewriter, op.getLoc(), newStep));
rewriter.restoreInsertionPoint(prevInsertPoint);
}
return op;
}
Reference in New Issue View Git Blame Copy Permalink