2018-09-28 12:17:26 -07:00
|
|
|
//===- LoopUtils.cpp ---- Misc utilities for loop transformation ----------===//
|
2018-09-07 14:47:21 -07:00
|
|
|
//
|
|
|
|
|
// Copyright 2019 The MLIR Authors.
|
|
|
|
|
//
|
|
|
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
|
// you may not use this file except in compliance with the License.
|
|
|
|
|
// You may obtain a copy of the License at
|
|
|
|
|
//
|
|
|
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
|
//
|
|
|
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
|
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
|
// See the License for the specific language governing permissions and
|
|
|
|
|
// limitations under the License.
|
|
|
|
|
// =============================================================================
|
|
|
|
|
//
|
2018-09-28 12:17:26 -07:00
|
|
|
// This file implements miscellaneous loop transformation routines.
|
2018-09-07 14:47:21 -07:00
|
|
|
//
|
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
|
|
2018-09-18 10:22:03 -07:00
|
|
|
#include "mlir/Transforms/LoopUtils.h"
|
2018-09-12 10:21:23 -07:00
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
#include "mlir/AffineOps/AffineOps.h"
|
2019-02-15 09:32:18 -08:00
|
|
|
#include "mlir/Analysis/AffineAnalysis.h"
|
2019-02-22 16:51:08 -08:00
|
|
|
#include "mlir/Analysis/AffineStructures.h"
|
2018-09-12 10:21:23 -07:00
|
|
|
#include "mlir/Analysis/LoopAnalysis.h"
|
2018-09-18 10:22:03 -07:00
|
|
|
#include "mlir/IR/AffineExpr.h"
|
|
|
|
|
#include "mlir/IR/AffineMap.h"
|
2019-01-24 12:25:30 -08:00
|
|
|
#include "mlir/IR/BlockAndValueMapping.h"
|
2018-09-07 14:47:21 -07:00
|
|
|
#include "mlir/IR/Builders.h"
|
2019-02-03 09:49:39 -08:00
|
|
|
#include "mlir/IR/Instruction.h"
|
2019-03-01 13:48:24 -08:00
|
|
|
#include "mlir/StandardOps/Ops.h"
|
2018-09-28 12:17:26 -07:00
|
|
|
#include "llvm/ADT/DenseMap.h"
|
2018-10-22 13:44:31 -07:00
|
|
|
#include "llvm/Support/Debug.h"
|
Add a stripmineSink and imperfectly nested tiling primitives.
This CL adds a primitive to perform stripmining of a loop by a given factor and
sinking it under multiple target loops.
In turn this is used to implement imperfectly nested loop tiling (with interchange) by repeatedly calling the stripmineSink primitive.
The API returns the point loops and allows repeated invocations of tiling to achieve declarative, multi-level, imperfectly-nested tiling.
Note that this CL is only concerned with the mechanical aspects and does not worry about analysis and legality.
The API is demonstrated in an example which creates an EDSC block, emits the corresponding MLIR and applies imperfectly-nested tiling:
```cpp
auto block = edsc::block({
For(ArrayRef<edsc::Expr>{i, j}, {zero, zero}, {M, N}, {one, one}, {
For(k1, zero, O, one, {
C({i, j, k1}) = A({i, j, k1}) + B({i, j, k1})
}),
For(k2, zero, O, one, {
C({i, j, k2}) = A({i, j, k2}) + B({i, j, k2})
}),
}),
});
// clang-format on
emitter.emitStmts(block.getBody());
auto l_i = emitter.getAffineForOp(i), l_j = emitter.getAffineForOp(j),
l_k1 = emitter.getAffineForOp(k1), l_k2 = emitter.getAffineForOp(k2);
auto indicesL1 = mlir::tile({l_i, l_j}, {512, 1024}, {l_k1, l_k2});
auto l_ii1 = indicesL1[0][0], l_jj1 = indicesL1[1][0];
mlir::tile({l_jj1, l_ii1}, {32, 16}, l_jj1);
```
The edsc::Expr for the induction variables (i, j, k_1, k_2) provide the programmatic hooks from which tiling can be applied declaratively.
PiperOrigin-RevId: 235548228
2019-02-25 09:53:05 -08:00
|
|
|
|
2018-10-22 13:44:31 -07:00
|
|
|
#define DEBUG_TYPE "LoopUtils"
|
2018-09-07 14:47:21 -07:00
|
|
|
|
2018-09-18 10:22:03 -07:00
|
|
|
using namespace mlir;
|
|
|
|
|
|
|
|
|
|
/// Returns the upper bound of an unrolled loop with lower bound 'lb' and with
|
|
|
|
|
/// the specified trip count, stride, and unroll factor. Returns nullptr when
|
|
|
|
|
/// the trip count can't be expressed as an affine expression.
|
2019-02-01 16:42:18 -08:00
|
|
|
AffineMap mlir::getUnrolledLoopUpperBound(ConstOpPointer<AffineForOp> forOp,
|
2018-10-09 16:39:24 -07:00
|
|
|
unsigned unrollFactor,
|
2018-12-27 15:06:22 -08:00
|
|
|
FuncBuilder *builder) {
|
2019-02-01 16:42:18 -08:00
|
|
|
auto lbMap = forOp->getLowerBoundMap();
|
2018-09-18 10:22:03 -07:00
|
|
|
|
|
|
|
|
// Single result lower bound map only.
|
2018-10-09 16:39:24 -07:00
|
|
|
if (lbMap.getNumResults() != 1)
|
2019-01-26 10:41:17 -08:00
|
|
|
return AffineMap();
|
2018-09-18 10:22:03 -07:00
|
|
|
|
|
|
|
|
// Sometimes, the trip count cannot be expressed as an affine expression.
|
2019-02-01 16:42:18 -08:00
|
|
|
auto tripCount = getTripCountExpr(forOp);
|
2018-10-03 15:34:57 -07:00
|
|
|
if (!tripCount)
|
2019-01-26 10:41:17 -08:00
|
|
|
return AffineMap();
|
2018-09-18 10:22:03 -07:00
|
|
|
|
2018-10-09 16:39:24 -07:00
|
|
|
AffineExpr lb(lbMap.getResult(0));
|
2019-02-01 16:42:18 -08:00
|
|
|
unsigned step = forOp->getStep();
|
2018-10-03 15:36:53 -07:00
|
|
|
auto newUb = lb + (tripCount - tripCount % unrollFactor - 1) * step;
|
2018-10-03 15:34:57 -07:00
|
|
|
|
2018-10-09 16:39:24 -07:00
|
|
|
return builder->getAffineMap(lbMap.getNumDims(), lbMap.getNumSymbols(),
|
2018-10-03 15:34:57 -07:00
|
|
|
{newUb}, {});
|
2018-09-18 10:22:03 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Returns the lower bound of the cleanup loop when unrolling a loop with lower
|
|
|
|
|
/// bound 'lb' and with the specified trip count, stride, and unroll factor.
|
2018-10-09 16:39:24 -07:00
|
|
|
/// Returns an AffinMap with nullptr storage (that evaluates to false)
|
|
|
|
|
/// when the trip count can't be expressed as an affine expression.
|
2019-02-01 16:42:18 -08:00
|
|
|
AffineMap mlir::getCleanupLoopLowerBound(ConstOpPointer<AffineForOp> forOp,
|
2018-10-09 16:39:24 -07:00
|
|
|
unsigned unrollFactor,
|
2018-12-27 15:06:22 -08:00
|
|
|
FuncBuilder *builder) {
|
2019-02-01 16:42:18 -08:00
|
|
|
auto lbMap = forOp->getLowerBoundMap();
|
2018-09-18 10:22:03 -07:00
|
|
|
|
|
|
|
|
// Single result lower bound map only.
|
2018-10-09 16:39:24 -07:00
|
|
|
if (lbMap.getNumResults() != 1)
|
2019-01-26 10:41:17 -08:00
|
|
|
return AffineMap();
|
2018-09-18 10:22:03 -07:00
|
|
|
|
|
|
|
|
// Sometimes the trip count cannot be expressed as an affine expression.
|
2019-02-01 16:42:18 -08:00
|
|
|
AffineExpr tripCount(getTripCountExpr(forOp));
|
2018-10-03 15:34:57 -07:00
|
|
|
if (!tripCount)
|
2019-01-26 10:41:17 -08:00
|
|
|
return AffineMap();
|
2018-09-18 10:22:03 -07:00
|
|
|
|
2018-10-09 16:39:24 -07:00
|
|
|
AffineExpr lb(lbMap.getResult(0));
|
2019-02-01 16:42:18 -08:00
|
|
|
unsigned step = forOp->getStep();
|
2018-10-03 15:34:57 -07:00
|
|
|
auto newLb = lb + (tripCount - tripCount % unrollFactor) * step;
|
2018-10-09 16:39:24 -07:00
|
|
|
return builder->getAffineMap(lbMap.getNumDims(), lbMap.getNumSymbols(),
|
2018-10-03 15:34:57 -07:00
|
|
|
{newLb}, {});
|
2018-09-18 10:22:03 -07:00
|
|
|
}
|
|
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
/// Promotes the loop body of a forOp to its containing block if the forOp
|
2018-09-07 14:47:21 -07:00
|
|
|
/// was known to have a single iteration. Returns false otherwise.
|
2018-09-12 10:21:23 -07:00
|
|
|
// TODO(bondhugula): extend this for arbitrary affine bounds.
|
2019-02-01 16:42:18 -08:00
|
|
|
bool mlir::promoteIfSingleIteration(OpPointer<AffineForOp> forOp) {
|
|
|
|
|
Optional<uint64_t> tripCount = getConstantTripCount(forOp);
|
2018-09-18 10:22:03 -07:00
|
|
|
if (!tripCount.hasValue() || tripCount.getValue() != 1)
|
2018-09-07 14:47:21 -07:00
|
|
|
return false;
|
|
|
|
|
|
2018-09-18 10:22:03 -07:00
|
|
|
// TODO(mlir-team): there is no builder for a max.
|
2019-02-01 16:42:18 -08:00
|
|
|
if (forOp->getLowerBoundMap().getNumResults() != 1)
|
2018-09-07 14:47:21 -07:00
|
|
|
return false;
|
|
|
|
|
|
|
|
|
|
// Replaces all IV uses to its single iteration value.
|
2019-02-01 16:42:18 -08:00
|
|
|
auto *iv = forOp->getInductionVar();
|
2019-02-04 10:38:47 -08:00
|
|
|
Instruction *forInst = forOp->getInstruction();
|
2019-01-26 12:40:12 -08:00
|
|
|
if (!iv->use_empty()) {
|
2019-02-01 16:42:18 -08:00
|
|
|
if (forOp->hasConstantLowerBound()) {
|
2018-12-28 16:05:35 -08:00
|
|
|
auto *mlFunc = forInst->getFunction();
|
2018-12-29 15:33:43 -08:00
|
|
|
FuncBuilder topBuilder(mlFunc);
|
2018-10-06 17:21:53 -07:00
|
|
|
auto constOp = topBuilder.create<ConstantIndexOp>(
|
2019-02-01 16:42:18 -08:00
|
|
|
forOp->getLoc(), forOp->getConstantLowerBound());
|
2019-01-26 12:40:12 -08:00
|
|
|
iv->replaceAllUsesWith(constOp);
|
2018-09-18 10:22:03 -07:00
|
|
|
} else {
|
2019-02-01 16:42:18 -08:00
|
|
|
const AffineBound lb = forOp->getLowerBound();
|
2018-12-27 14:35:10 -08:00
|
|
|
SmallVector<Value *, 4> lbOperands(lb.operand_begin(), lb.operand_end());
|
2018-12-28 16:05:35 -08:00
|
|
|
FuncBuilder builder(forInst->getBlock(), Block::iterator(forInst));
|
2019-01-22 13:58:52 -08:00
|
|
|
if (lb.getMap() == builder.getDimIdentityMap()) {
|
2019-02-06 11:08:18 -08:00
|
|
|
// No need of generating an affine.apply.
|
2019-01-26 12:40:12 -08:00
|
|
|
iv->replaceAllUsesWith(lbOperands[0]);
|
2019-01-22 13:58:52 -08:00
|
|
|
} else {
|
|
|
|
|
auto affineApplyOp = builder.create<AffineApplyOp>(
|
|
|
|
|
forInst->getLoc(), lb.getMap(), lbOperands);
|
2019-01-27 09:33:19 -08:00
|
|
|
iv->replaceAllUsesWith(affineApplyOp);
|
2019-01-22 13:58:52 -08:00
|
|
|
}
|
2018-09-18 10:22:03 -07:00
|
|
|
}
|
|
|
|
|
}
|
2018-12-28 16:05:35 -08:00
|
|
|
// Move the loop body instructions to the loop's containing block.
|
|
|
|
|
auto *block = forInst->getBlock();
|
|
|
|
|
block->getInstructions().splice(Block::iterator(forInst),
|
2019-02-01 16:42:18 -08:00
|
|
|
forOp->getBody()->getInstructions());
|
|
|
|
|
forOp->erase();
|
2018-09-07 14:47:21 -07:00
|
|
|
return true;
|
|
|
|
|
}
|
|
|
|
|
|
2018-12-28 16:05:35 -08:00
|
|
|
/// Promotes all single iteration for inst's in the Function, i.e., moves
|
2018-12-28 13:07:39 -08:00
|
|
|
/// their body into the containing Block.
|
2018-12-28 08:48:09 -08:00
|
|
|
void mlir::promoteSingleIterationLoops(Function *f) {
|
2018-09-07 14:47:21 -07:00
|
|
|
// Gathers all innermost loops through a post order pruned walk.
|
2019-02-04 16:24:44 -08:00
|
|
|
f->walkPostOrder<AffineForOp>(
|
|
|
|
|
[](OpPointer<AffineForOp> forOp) { promoteIfSingleIteration(forOp); });
|
2018-09-07 14:47:21 -07:00
|
|
|
}
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2019-02-06 21:54:18 -08:00
|
|
|
/// Generates a 'for' inst with the specified lower and upper bounds while
|
|
|
|
|
/// generating the right IV remappings for the shifted instructions. The
|
2018-12-28 16:05:35 -08:00
|
|
|
/// instruction blocks that go into the loop are specified in instGroupQueue
|
2018-09-28 12:17:26 -07:00
|
|
|
/// starting from the specified offset, and in that order; the first element of
|
2018-12-28 16:05:35 -08:00
|
|
|
/// the pair specifies the shift applied to that group of instructions; note
|
|
|
|
|
/// that the shift is multiplied by the loop step before being applied. Returns
|
2018-09-28 12:17:26 -07:00
|
|
|
/// nullptr if the generated loop simplifies to a single iteration one.
|
2019-02-01 16:42:18 -08:00
|
|
|
static OpPointer<AffineForOp>
|
2018-12-10 15:17:25 -08:00
|
|
|
generateLoop(AffineMap lbMap, AffineMap ubMap,
|
2018-12-28 16:05:35 -08:00
|
|
|
const std::vector<std::pair<uint64_t, ArrayRef<Instruction *>>>
|
|
|
|
|
&instGroupQueue,
|
2019-02-01 16:42:18 -08:00
|
|
|
unsigned offset, OpPointer<AffineForOp> srcForInst,
|
|
|
|
|
FuncBuilder *b) {
|
2018-12-28 16:05:35 -08:00
|
|
|
SmallVector<Value *, 4> lbOperands(srcForInst->getLowerBoundOperands());
|
|
|
|
|
SmallVector<Value *, 4> ubOperands(srcForInst->getUpperBoundOperands());
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2018-12-10 15:17:25 -08:00
|
|
|
assert(lbMap.getNumInputs() == lbOperands.size());
|
|
|
|
|
assert(ubMap.getNumInputs() == ubOperands.size());
|
|
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
auto loopChunk =
|
|
|
|
|
b->create<AffineForOp>(srcForInst->getLoc(), lbOperands, lbMap,
|
|
|
|
|
ubOperands, ubMap, srcForInst->getStep());
|
|
|
|
|
loopChunk->createBody();
|
2019-01-26 12:40:12 -08:00
|
|
|
auto *loopChunkIV = loopChunk->getInductionVar();
|
|
|
|
|
auto *srcIV = srcForInst->getInductionVar();
|
2018-12-10 15:17:25 -08:00
|
|
|
|
2019-01-24 12:25:30 -08:00
|
|
|
BlockAndValueMapping operandMap;
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2018-12-28 16:05:35 -08:00
|
|
|
for (auto it = instGroupQueue.begin() + offset, e = instGroupQueue.end();
|
2018-09-28 12:17:26 -07:00
|
|
|
it != e; ++it) {
|
2018-12-10 15:17:25 -08:00
|
|
|
uint64_t shift = it->first;
|
2018-12-28 16:05:35 -08:00
|
|
|
auto insts = it->second;
|
|
|
|
|
// All 'same shift' instructions get added with their operands being
|
|
|
|
|
// remapped to results of cloned instructions, and their IV used remapped.
|
2018-12-10 15:17:25 -08:00
|
|
|
// Generate the remapping if the shift is not zero: remappedIV = newIV -
|
|
|
|
|
// shift.
|
2019-01-26 12:40:12 -08:00
|
|
|
if (!srcIV->use_empty() && shift != 0) {
|
2019-02-01 16:42:18 -08:00
|
|
|
FuncBuilder b(loopChunk->getBody());
|
2019-01-27 09:33:19 -08:00
|
|
|
auto ivRemap = b.create<AffineApplyOp>(
|
|
|
|
|
srcForInst->getLoc(),
|
|
|
|
|
b.getSingleDimShiftAffineMap(
|
|
|
|
|
-static_cast<int64_t>(srcForInst->getStep() * shift)),
|
|
|
|
|
loopChunkIV);
|
2019-01-26 12:40:12 -08:00
|
|
|
operandMap.map(srcIV, ivRemap);
|
2018-09-28 12:17:26 -07:00
|
|
|
} else {
|
2019-01-26 12:40:12 -08:00
|
|
|
operandMap.map(srcIV, loopChunkIV);
|
2018-09-28 12:17:26 -07:00
|
|
|
}
|
2018-12-28 16:05:35 -08:00
|
|
|
for (auto *inst : insts) {
|
|
|
|
|
loopChunk->getBody()->push_back(inst->clone(operandMap, b->getContext()));
|
2018-09-28 12:17:26 -07:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
if (promoteIfSingleIteration(loopChunk))
|
2019-02-01 16:42:18 -08:00
|
|
|
return OpPointer<AffineForOp>();
|
2018-09-28 12:17:26 -07:00
|
|
|
return loopChunk;
|
|
|
|
|
}
|
|
|
|
|
|
2019-02-06 21:54:18 -08:00
|
|
|
/// Skew the instructions in the body of a 'for' instruction with the specified
|
|
|
|
|
/// instruction-wise shifts. The shifts are with respect to the original
|
|
|
|
|
/// execution order, and are multiplied by the loop 'step' before being applied.
|
|
|
|
|
/// A shift of zero for each instruction will lead to no change.
|
2018-12-28 16:05:35 -08:00
|
|
|
// The skewing of instructions with respect to one another can be used for
|
|
|
|
|
// example to allow overlap of asynchronous operations (such as DMA
|
|
|
|
|
// communication) with computation, or just relative shifting of instructions
|
|
|
|
|
// for better register reuse, locality or parallelism. As such, the shifts are
|
|
|
|
|
// typically expected to be at most of the order of the number of instructions.
|
|
|
|
|
// This method should not be used as a substitute for loop distribution/fission.
|
|
|
|
|
// This method uses an algorithm// in time linear in the number of instructions
|
|
|
|
|
// in the body of the for loop - (using the 'sweep line' paradigm). This method
|
2018-09-28 12:17:26 -07:00
|
|
|
// asserts preservation of SSA dominance. A check for that as well as that for
|
|
|
|
|
// memory-based depedence preservation check rests with the users of this
|
|
|
|
|
// method.
|
2019-02-01 16:42:18 -08:00
|
|
|
UtilResult mlir::instBodySkew(OpPointer<AffineForOp> forOp,
|
|
|
|
|
ArrayRef<uint64_t> shifts,
|
2018-09-28 12:17:26 -07:00
|
|
|
bool unrollPrologueEpilogue) {
|
2019-02-01 16:42:18 -08:00
|
|
|
if (forOp->getBody()->empty())
|
2018-09-28 12:17:26 -07:00
|
|
|
return UtilResult::Success;
|
|
|
|
|
|
|
|
|
|
// If the trip counts aren't constant, we would need versioning and
|
|
|
|
|
// conditional guards (or context information to prevent such versioning). The
|
|
|
|
|
// better way to pipeline for such loops is to first tile them and extract
|
|
|
|
|
// constant trip count "full tiles" before applying this.
|
2019-02-01 16:42:18 -08:00
|
|
|
auto mayBeConstTripCount = getConstantTripCount(forOp);
|
2018-10-22 13:44:31 -07:00
|
|
|
if (!mayBeConstTripCount.hasValue()) {
|
2019-02-28 12:07:12 -08:00
|
|
|
LLVM_DEBUG(forOp->emitNote("non-constant trip count loop not handled"));
|
2018-10-22 13:44:31 -07:00
|
|
|
return UtilResult::Success;
|
|
|
|
|
}
|
2018-09-28 12:17:26 -07:00
|
|
|
uint64_t tripCount = mayBeConstTripCount.getValue();
|
|
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
assert(isInstwiseShiftValid(forOp, shifts) &&
|
2018-10-22 13:44:31 -07:00
|
|
|
"shifts will lead to an invalid transformation\n");
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
int64_t step = forOp->getStep();
|
2018-12-10 15:17:25 -08:00
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
unsigned numChildInsts = forOp->getBody()->getInstructions().size();
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2018-12-10 15:17:25 -08:00
|
|
|
// Do a linear time (counting) sort for the shifts.
|
|
|
|
|
uint64_t maxShift = 0;
|
2018-12-28 16:05:35 -08:00
|
|
|
for (unsigned i = 0; i < numChildInsts; i++) {
|
2018-12-10 15:17:25 -08:00
|
|
|
maxShift = std::max(maxShift, shifts[i]);
|
2018-09-28 12:17:26 -07:00
|
|
|
}
|
2018-12-10 15:17:25 -08:00
|
|
|
// Such large shifts are not the typical use case.
|
2018-12-28 16:05:35 -08:00
|
|
|
if (maxShift >= numChildInsts) {
|
2019-02-28 12:07:12 -08:00
|
|
|
forOp->emitWarning("not shifting because shifts are unrealistically large");
|
2018-10-22 13:44:31 -07:00
|
|
|
return UtilResult::Success;
|
|
|
|
|
}
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2018-12-28 16:05:35 -08:00
|
|
|
// An array of instruction groups sorted by shift amount; each group has all
|
|
|
|
|
// instructions with the same shift in the order in which they appear in the
|
2019-02-06 21:54:18 -08:00
|
|
|
// body of the 'for' inst.
|
2018-12-28 16:05:35 -08:00
|
|
|
std::vector<std::vector<Instruction *>> sortedInstGroups(maxShift + 1);
|
2018-09-28 12:17:26 -07:00
|
|
|
unsigned pos = 0;
|
2019-02-01 16:42:18 -08:00
|
|
|
for (auto &inst : *forOp->getBody()) {
|
2018-12-10 15:17:25 -08:00
|
|
|
auto shift = shifts[pos++];
|
2018-12-28 16:05:35 -08:00
|
|
|
sortedInstGroups[shift].push_back(&inst);
|
2018-09-28 12:17:26 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Unless the shifts have a specific pattern (which actually would be the
|
|
|
|
|
// common use case), prologue and epilogue are not meaningfully defined.
|
|
|
|
|
// Nevertheless, if 'unrollPrologueEpilogue' is set, we will treat the first
|
|
|
|
|
// loop generated as the prologue and the last as epilogue and unroll these
|
|
|
|
|
// fully.
|
2019-02-01 16:42:18 -08:00
|
|
|
OpPointer<AffineForOp> prologue;
|
|
|
|
|
OpPointer<AffineForOp> epilogue;
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2018-12-10 15:17:25 -08:00
|
|
|
// Do a sweep over the sorted shifts while storing open groups in a
|
2018-09-28 12:17:26 -07:00
|
|
|
// vector, and generating loop portions as necessary during the sweep. A block
|
2018-12-28 16:05:35 -08:00
|
|
|
// of instructions is paired with its shift.
|
|
|
|
|
std::vector<std::pair<uint64_t, ArrayRef<Instruction *>>> instGroupQueue;
|
2018-09-28 12:17:26 -07:00
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
auto origLbMap = forOp->getLowerBoundMap();
|
2018-12-10 15:17:25 -08:00
|
|
|
uint64_t lbShift = 0;
|
2019-02-01 16:42:18 -08:00
|
|
|
FuncBuilder b(forOp->getInstruction());
|
2018-12-28 16:05:35 -08:00
|
|
|
for (uint64_t d = 0, e = sortedInstGroups.size(); d < e; ++d) {
|
2018-12-10 15:17:25 -08:00
|
|
|
// If nothing is shifted by d, continue.
|
2018-12-28 16:05:35 -08:00
|
|
|
if (sortedInstGroups[d].empty())
|
2018-09-28 12:17:26 -07:00
|
|
|
continue;
|
2018-12-28 16:05:35 -08:00
|
|
|
if (!instGroupQueue.empty()) {
|
2018-09-28 12:17:26 -07:00
|
|
|
assert(d >= 1 &&
|
|
|
|
|
"Queue expected to be empty when the first block is found");
|
|
|
|
|
// The interval for which the loop needs to be generated here is:
|
2018-12-10 15:17:25 -08:00
|
|
|
// [lbShift, min(lbShift + tripCount, d)) and the body of the
|
2018-12-28 16:05:35 -08:00
|
|
|
// loop needs to have all instructions in instQueue in that order.
|
2019-02-01 16:42:18 -08:00
|
|
|
OpPointer<AffineForOp> res;
|
2018-12-10 15:17:25 -08:00
|
|
|
if (lbShift + tripCount * step < d * step) {
|
|
|
|
|
res = generateLoop(
|
|
|
|
|
b.getShiftedAffineMap(origLbMap, lbShift),
|
|
|
|
|
b.getShiftedAffineMap(origLbMap, lbShift + tripCount * step),
|
2019-02-01 16:42:18 -08:00
|
|
|
instGroupQueue, 0, forOp, &b);
|
2018-12-28 16:05:35 -08:00
|
|
|
// Entire loop for the queued inst groups generated, empty it.
|
|
|
|
|
instGroupQueue.clear();
|
2018-12-10 15:17:25 -08:00
|
|
|
lbShift += tripCount * step;
|
2018-09-28 12:17:26 -07:00
|
|
|
} else {
|
2018-12-10 15:17:25 -08:00
|
|
|
res = generateLoop(b.getShiftedAffineMap(origLbMap, lbShift),
|
2018-12-28 16:05:35 -08:00
|
|
|
b.getShiftedAffineMap(origLbMap, d), instGroupQueue,
|
2019-02-01 16:42:18 -08:00
|
|
|
0, forOp, &b);
|
2018-12-10 15:17:25 -08:00
|
|
|
lbShift = d * step;
|
2018-09-28 12:17:26 -07:00
|
|
|
}
|
|
|
|
|
if (!prologue && res)
|
|
|
|
|
prologue = res;
|
|
|
|
|
epilogue = res;
|
|
|
|
|
} else {
|
|
|
|
|
// Start of first interval.
|
2018-12-10 15:17:25 -08:00
|
|
|
lbShift = d * step;
|
2018-09-28 12:17:26 -07:00
|
|
|
}
|
2018-12-28 16:05:35 -08:00
|
|
|
// Augment the list of instructions that get into the current open interval.
|
|
|
|
|
instGroupQueue.push_back({d, sortedInstGroups[d]});
|
2018-09-28 12:17:26 -07:00
|
|
|
}
|
|
|
|
|
|
2018-12-28 16:05:35 -08:00
|
|
|
// Those instructions groups left in the queue now need to be processed (FIFO)
|
2018-09-28 12:17:26 -07:00
|
|
|
// and their loops completed.
|
2018-12-28 16:05:35 -08:00
|
|
|
for (unsigned i = 0, e = instGroupQueue.size(); i < e; ++i) {
|
|
|
|
|
uint64_t ubShift = (instGroupQueue[i].first + tripCount) * step;
|
2018-12-10 15:17:25 -08:00
|
|
|
epilogue = generateLoop(b.getShiftedAffineMap(origLbMap, lbShift),
|
|
|
|
|
b.getShiftedAffineMap(origLbMap, ubShift),
|
2019-02-01 16:42:18 -08:00
|
|
|
instGroupQueue, i, forOp, &b);
|
2018-12-10 15:17:25 -08:00
|
|
|
lbShift = ubShift;
|
2018-09-28 12:17:26 -07:00
|
|
|
if (!prologue)
|
|
|
|
|
prologue = epilogue;
|
|
|
|
|
}
|
|
|
|
|
|
2018-12-28 16:05:35 -08:00
|
|
|
// Erase the original for inst.
|
2019-02-01 16:42:18 -08:00
|
|
|
forOp->erase();
|
2018-09-28 12:17:26 -07:00
|
|
|
|
|
|
|
|
if (unrollPrologueEpilogue && prologue)
|
|
|
|
|
loopUnrollFull(prologue);
|
2019-02-01 16:42:18 -08:00
|
|
|
if (unrollPrologueEpilogue && !epilogue &&
|
|
|
|
|
epilogue->getInstruction() != prologue->getInstruction())
|
2018-09-28 12:17:26 -07:00
|
|
|
loopUnrollFull(epilogue);
|
|
|
|
|
|
|
|
|
|
return UtilResult::Success;
|
|
|
|
|
}
|
2018-11-14 14:15:48 -08:00
|
|
|
|
|
|
|
|
/// Unrolls this loop completely.
|
2019-02-01 16:42:18 -08:00
|
|
|
bool mlir::loopUnrollFull(OpPointer<AffineForOp> forOp) {
|
|
|
|
|
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
|
2018-11-14 14:15:48 -08:00
|
|
|
if (mayBeConstantTripCount.hasValue()) {
|
|
|
|
|
uint64_t tripCount = mayBeConstantTripCount.getValue();
|
|
|
|
|
if (tripCount == 1) {
|
2019-02-01 16:42:18 -08:00
|
|
|
return promoteIfSingleIteration(forOp);
|
2018-11-14 14:15:48 -08:00
|
|
|
}
|
2019-02-01 16:42:18 -08:00
|
|
|
return loopUnrollByFactor(forOp, tripCount);
|
2018-11-14 14:15:48 -08:00
|
|
|
}
|
|
|
|
|
return false;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Unrolls and jams this loop by the specified factor or by the trip count (if
|
|
|
|
|
/// constant) whichever is lower.
|
2019-02-01 16:42:18 -08:00
|
|
|
bool mlir::loopUnrollUpToFactor(OpPointer<AffineForOp> forOp,
|
|
|
|
|
uint64_t unrollFactor) {
|
|
|
|
|
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
|
2018-11-14 14:15:48 -08:00
|
|
|
|
|
|
|
|
if (mayBeConstantTripCount.hasValue() &&
|
|
|
|
|
mayBeConstantTripCount.getValue() < unrollFactor)
|
2019-02-01 16:42:18 -08:00
|
|
|
return loopUnrollByFactor(forOp, mayBeConstantTripCount.getValue());
|
|
|
|
|
return loopUnrollByFactor(forOp, unrollFactor);
|
2018-11-14 14:15:48 -08:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Unrolls this loop by the specified factor. Returns true if the loop
|
|
|
|
|
/// is successfully unrolled.
|
2019-02-01 16:42:18 -08:00
|
|
|
bool mlir::loopUnrollByFactor(OpPointer<AffineForOp> forOp,
|
|
|
|
|
uint64_t unrollFactor) {
|
2018-11-14 14:15:48 -08:00
|
|
|
assert(unrollFactor >= 1 && "unroll factor should be >= 1");
|
|
|
|
|
|
2019-01-22 15:48:07 -08:00
|
|
|
if (unrollFactor == 1)
|
2019-02-01 16:42:18 -08:00
|
|
|
return promoteIfSingleIteration(forOp);
|
2019-01-22 15:48:07 -08:00
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
if (forOp->getBody()->empty())
|
2018-11-14 14:15:48 -08:00
|
|
|
return false;
|
|
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
auto lbMap = forOp->getLowerBoundMap();
|
|
|
|
|
auto ubMap = forOp->getUpperBoundMap();
|
2018-11-14 14:15:48 -08:00
|
|
|
|
|
|
|
|
// Loops with max/min expressions won't be unrolled here (the output can't be
|
2018-12-28 08:48:09 -08:00
|
|
|
// expressed as a Function in the general case). However, the right way to
|
|
|
|
|
// do such unrolling for a Function would be to specialize the loop for the
|
2018-11-14 14:15:48 -08:00
|
|
|
// 'hotspot' case and unroll that hotspot.
|
|
|
|
|
if (lbMap.getNumResults() != 1 || ubMap.getNumResults() != 1)
|
|
|
|
|
return false;
|
|
|
|
|
|
|
|
|
|
// Same operand list for lower and upper bound for now.
|
|
|
|
|
// TODO(bondhugula): handle bounds with different operand lists.
|
2019-02-01 16:42:18 -08:00
|
|
|
if (!forOp->matchingBoundOperandList())
|
2018-11-14 14:15:48 -08:00
|
|
|
return false;
|
|
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
|
2018-11-14 14:15:48 -08:00
|
|
|
|
|
|
|
|
// If the trip count is lower than the unroll factor, no unrolled body.
|
|
|
|
|
// TODO(bondhugula): option to specify cleanup loop unrolling.
|
|
|
|
|
if (mayBeConstantTripCount.hasValue() &&
|
|
|
|
|
mayBeConstantTripCount.getValue() < unrollFactor)
|
|
|
|
|
return false;
|
|
|
|
|
|
|
|
|
|
// Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
|
2019-02-04 10:38:47 -08:00
|
|
|
Instruction *forInst = forOp->getInstruction();
|
2019-02-01 16:42:18 -08:00
|
|
|
if (getLargestDivisorOfTripCount(forOp) % unrollFactor != 0) {
|
2018-12-28 16:05:35 -08:00
|
|
|
FuncBuilder builder(forInst->getBlock(), ++Block::iterator(forInst));
|
2019-02-04 10:38:47 -08:00
|
|
|
auto cleanupForInst = builder.clone(*forInst)->cast<AffineForOp>();
|
2019-02-01 16:42:18 -08:00
|
|
|
auto clLbMap = getCleanupLoopLowerBound(forOp, unrollFactor, &builder);
|
2018-11-14 14:15:48 -08:00
|
|
|
assert(clLbMap &&
|
|
|
|
|
"cleanup loop lower bound map for single result bound maps can "
|
|
|
|
|
"always be determined");
|
2018-12-28 16:05:35 -08:00
|
|
|
cleanupForInst->setLowerBoundMap(clLbMap);
|
2018-11-14 14:15:48 -08:00
|
|
|
// Promote the loop body up if this has turned into a single iteration loop.
|
2018-12-28 16:05:35 -08:00
|
|
|
promoteIfSingleIteration(cleanupForInst);
|
2018-11-14 14:15:48 -08:00
|
|
|
|
|
|
|
|
// Adjust upper bound.
|
|
|
|
|
auto unrolledUbMap =
|
2019-02-01 16:42:18 -08:00
|
|
|
getUnrolledLoopUpperBound(forOp, unrollFactor, &builder);
|
2018-11-14 14:15:48 -08:00
|
|
|
assert(unrolledUbMap &&
|
|
|
|
|
"upper bound map can alwayys be determined for an unrolled loop "
|
|
|
|
|
"with single result bounds");
|
2019-02-01 16:42:18 -08:00
|
|
|
forOp->setUpperBoundMap(unrolledUbMap);
|
2018-11-14 14:15:48 -08:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Scale the step of loop being unrolled by unroll factor.
|
2019-02-01 16:42:18 -08:00
|
|
|
int64_t step = forOp->getStep();
|
|
|
|
|
forOp->setStep(step * unrollFactor);
|
2018-11-14 14:15:48 -08:00
|
|
|
|
2018-12-28 16:05:35 -08:00
|
|
|
// Builder to insert unrolled bodies right after the last instruction in the
|
2019-02-01 16:42:18 -08:00
|
|
|
// body of 'forOp'.
|
|
|
|
|
FuncBuilder builder(forOp->getBody(), forOp->getBody()->end());
|
2018-11-14 14:15:48 -08:00
|
|
|
|
2018-12-28 16:05:35 -08:00
|
|
|
// Keep a pointer to the last instruction in the original block so that we
|
|
|
|
|
// know what to clone (since we are doing this in-place).
|
2019-02-01 16:42:18 -08:00
|
|
|
Block::iterator srcBlockEnd = std::prev(forOp->getBody()->end());
|
2018-11-14 14:15:48 -08:00
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
// Unroll the contents of 'forOp' (append unrollFactor-1 additional copies).
|
|
|
|
|
auto *forOpIV = forOp->getInductionVar();
|
2018-11-14 14:15:48 -08:00
|
|
|
for (unsigned i = 1; i < unrollFactor; i++) {
|
2019-01-24 12:25:30 -08:00
|
|
|
BlockAndValueMapping operandMap;
|
2018-11-14 14:15:48 -08:00
|
|
|
|
|
|
|
|
// If the induction variable is used, create a remapping to the value for
|
|
|
|
|
// this unrolled instance.
|
2019-02-01 16:42:18 -08:00
|
|
|
if (!forOpIV->use_empty()) {
|
2018-11-14 14:15:48 -08:00
|
|
|
// iv' = iv + 1/2/3...unrollFactor-1;
|
|
|
|
|
auto d0 = builder.getAffineDimExpr(0);
|
|
|
|
|
auto bumpMap = builder.getAffineMap(1, 0, {d0 + i * step}, {});
|
2019-01-27 09:33:19 -08:00
|
|
|
auto ivUnroll =
|
2019-02-01 16:42:18 -08:00
|
|
|
builder.create<AffineApplyOp>(forOp->getLoc(), bumpMap, forOpIV);
|
|
|
|
|
operandMap.map(forOpIV, ivUnroll);
|
2018-11-14 14:15:48 -08:00
|
|
|
}
|
|
|
|
|
|
2019-02-01 16:42:18 -08:00
|
|
|
// Clone the original body of 'forOp'.
|
|
|
|
|
for (auto it = forOp->getBody()->begin(); it != std::next(srcBlockEnd);
|
2018-12-23 08:17:48 -08:00
|
|
|
it++) {
|
2018-11-14 14:15:48 -08:00
|
|
|
builder.clone(*it, operandMap);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Promote the loop body up if this has turned into a single iteration loop.
|
2019-02-01 16:42:18 -08:00
|
|
|
promoteIfSingleIteration(forOp);
|
2018-11-14 14:15:48 -08:00
|
|
|
|
|
|
|
|
return true;
|
|
|
|
|
}
|
2019-02-15 09:32:18 -08:00
|
|
|
|
|
|
|
|
/// Performs loop interchange on 'forOpA' and 'forOpB', where 'forOpB' is
|
|
|
|
|
/// nested within 'forOpA' as the only instruction in its block.
|
|
|
|
|
void mlir::interchangeLoops(OpPointer<AffineForOp> forOpA,
|
|
|
|
|
OpPointer<AffineForOp> forOpB) {
|
|
|
|
|
auto *forOpAInst = forOpA->getInstruction();
|
|
|
|
|
// 1) Slice forOpA's instruction list (which is just forOpB) just before
|
|
|
|
|
// forOpA (in forOpA's parent's block) this should leave 'forOpA's
|
|
|
|
|
// instruction list empty (because its perfectly nested).
|
|
|
|
|
assert(&*forOpA->getBody()->begin() == forOpB->getInstruction());
|
|
|
|
|
forOpAInst->getBlock()->getInstructions().splice(
|
|
|
|
|
Block::iterator(forOpAInst), forOpA->getBody()->getInstructions());
|
|
|
|
|
// 2) Slice forOpB's instruction list into forOpA's instruction list (this
|
|
|
|
|
// leaves forOpB's instruction list empty).
|
|
|
|
|
forOpA->getBody()->getInstructions().splice(
|
|
|
|
|
forOpA->getBody()->begin(), forOpB->getBody()->getInstructions());
|
|
|
|
|
// 3) Slice forOpA into forOpB's instruction list.
|
|
|
|
|
forOpB->getBody()->getInstructions().splice(
|
|
|
|
|
forOpB->getBody()->begin(), forOpAInst->getBlock()->getInstructions(),
|
|
|
|
|
Block::iterator(forOpAInst));
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Performs a series of loop interchanges to sink 'forOp' 'loopDepth' levels
|
|
|
|
|
/// deeper in the loop nest.
|
|
|
|
|
void mlir::sinkLoop(OpPointer<AffineForOp> forOp, unsigned loopDepth) {
|
|
|
|
|
for (unsigned i = 0; i < loopDepth; ++i) {
|
|
|
|
|
assert(forOp->getBody()->front().isa<AffineForOp>());
|
|
|
|
|
OpPointer<AffineForOp> nextForOp =
|
|
|
|
|
forOp->getBody()->front().cast<AffineForOp>();
|
|
|
|
|
interchangeLoops(forOp, nextForOp);
|
|
|
|
|
}
|
|
|
|
|
}
|
Add a stripmineSink and imperfectly nested tiling primitives.
This CL adds a primitive to perform stripmining of a loop by a given factor and
sinking it under multiple target loops.
In turn this is used to implement imperfectly nested loop tiling (with interchange) by repeatedly calling the stripmineSink primitive.
The API returns the point loops and allows repeated invocations of tiling to achieve declarative, multi-level, imperfectly-nested tiling.
Note that this CL is only concerned with the mechanical aspects and does not worry about analysis and legality.
The API is demonstrated in an example which creates an EDSC block, emits the corresponding MLIR and applies imperfectly-nested tiling:
```cpp
auto block = edsc::block({
For(ArrayRef<edsc::Expr>{i, j}, {zero, zero}, {M, N}, {one, one}, {
For(k1, zero, O, one, {
C({i, j, k1}) = A({i, j, k1}) + B({i, j, k1})
}),
For(k2, zero, O, one, {
C({i, j, k2}) = A({i, j, k2}) + B({i, j, k2})
}),
}),
});
// clang-format on
emitter.emitStmts(block.getBody());
auto l_i = emitter.getAffineForOp(i), l_j = emitter.getAffineForOp(j),
l_k1 = emitter.getAffineForOp(k1), l_k2 = emitter.getAffineForOp(k2);
auto indicesL1 = mlir::tile({l_i, l_j}, {512, 1024}, {l_k1, l_k2});
auto l_ii1 = indicesL1[0][0], l_jj1 = indicesL1[1][0];
mlir::tile({l_jj1, l_ii1}, {32, 16}, l_jj1);
```
The edsc::Expr for the induction variables (i, j, k_1, k_2) provide the programmatic hooks from which tiling can be applied declaratively.
PiperOrigin-RevId: 235548228
2019-02-25 09:53:05 -08:00
|
|
|
|
|
|
|
|
// Factors out common behavior to add max(`iv`, ...), min(`iv` + `offset`, ...)
|
|
|
|
|
// to loop bounds.
|
|
|
|
|
static void augmentMapAndBounds(FuncBuilder *b, Value *iv, AffineMap *map,
|
|
|
|
|
SmallVector<Value *, 4> *operands,
|
|
|
|
|
int64_t offset = 0) {
|
|
|
|
|
auto bounds = llvm::to_vector<4>(map->getResults());
|
|
|
|
|
operands->push_back(iv);
|
|
|
|
|
auto numOperands = operands->size();
|
|
|
|
|
bounds.push_back(b->getAffineDimExpr(numOperands - 1) + offset);
|
|
|
|
|
*map = b->getAffineMap(numOperands, map->getNumSymbols(), bounds, {});
|
|
|
|
|
canonicalizeMapAndOperands(map, operands);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Stripmines `forOp` by `factor` and sinks it under each of the `targets`.
|
|
|
|
|
// Stripmine-sink is a primitive building block for generalized tiling of
|
|
|
|
|
// imperfectly nested loops.
|
|
|
|
|
// This transformation is purely mechanical and does not check legality,
|
|
|
|
|
// profitability or even structural correctness. It is the user's
|
|
|
|
|
// responsibility to specify `targets` that are dominated by `forOp`.
|
|
|
|
|
// Returns the new AffineForOps, one per `targets`, nested immediately under
|
|
|
|
|
// each of the `targets`.
|
|
|
|
|
static SmallVector<OpPointer<AffineForOp>, 8>
|
|
|
|
|
stripmineSink(OpPointer<AffineForOp> forOp, uint64_t factor,
|
|
|
|
|
ArrayRef<OpPointer<AffineForOp>> targets) {
|
|
|
|
|
// TODO(ntv): 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 originalStep = forOp->getStep();
|
|
|
|
|
auto scaledStep = originalStep * factor;
|
|
|
|
|
forOp->setStep(scaledStep);
|
|
|
|
|
|
|
|
|
|
auto *forInst = forOp->getInstruction();
|
|
|
|
|
FuncBuilder b(forInst->getBlock(), ++Block::iterator(forInst));
|
|
|
|
|
|
|
|
|
|
// Lower-bound map creation.
|
|
|
|
|
auto lbMap = forOp->getLowerBoundMap();
|
|
|
|
|
SmallVector<Value *, 4> lbOperands(forOp->getLowerBoundOperands());
|
|
|
|
|
augmentMapAndBounds(&b, forOp->getInductionVar(), &lbMap, &lbOperands);
|
|
|
|
|
|
|
|
|
|
// Upper-bound map creation.
|
|
|
|
|
auto ubMap = forOp->getLowerBoundMap();
|
|
|
|
|
SmallVector<Value *, 4> ubOperands(forOp->getUpperBoundOperands());
|
|
|
|
|
augmentMapAndBounds(&b, forOp->getInductionVar(), &ubMap, &ubOperands,
|
|
|
|
|
/*offset=*/scaledStep);
|
|
|
|
|
|
|
|
|
|
SmallVector<OpPointer<AffineForOp>, 8> innerLoops;
|
|
|
|
|
for (auto t : targets) {
|
|
|
|
|
// Insert forOp just before the first instruction in the body.
|
|
|
|
|
auto *body = t->getBody();
|
|
|
|
|
auto &inst = body->getInstructions().front();
|
|
|
|
|
FuncBuilder b(&inst);
|
|
|
|
|
auto newLoop = b.create<AffineForOp>(t->getLoc(), lbOperands, lbMap,
|
|
|
|
|
ubOperands, ubMap, originalStep);
|
|
|
|
|
newLoop->createBody()->getInstructions().splice(
|
|
|
|
|
newLoop->getBody()->end(), body->getInstructions(), ++body->begin(),
|
|
|
|
|
body->end());
|
|
|
|
|
innerLoops.push_back(newLoop);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
return innerLoops;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Stripmines a `forOp` by `factor` and sinks it under a single `target`.
|
|
|
|
|
// Returns the new AffineForOps, nested immediately under `target`.
|
|
|
|
|
OpPointer<AffineForOp> stripmineSink(OpPointer<AffineForOp> forOp,
|
|
|
|
|
uint64_t factor,
|
|
|
|
|
OpPointer<AffineForOp> target) {
|
|
|
|
|
auto res =
|
|
|
|
|
stripmineSink(forOp, factor, ArrayRef<OpPointer<AffineForOp>>{target});
|
|
|
|
|
assert(res.size() == 1 && "Expected 1 inner forOp");
|
|
|
|
|
return res[0];
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
SmallVector<SmallVector<OpPointer<AffineForOp>, 8>, 8>
|
|
|
|
|
mlir::tile(ArrayRef<OpPointer<AffineForOp>> forOps, ArrayRef<uint64_t> sizes,
|
|
|
|
|
ArrayRef<OpPointer<AffineForOp>> targets) {
|
|
|
|
|
SmallVector<SmallVector<OpPointer<AffineForOp>, 8>, 8> res;
|
|
|
|
|
SmallVector<OpPointer<AffineForOp>, 8> currentTargets(targets.begin(),
|
|
|
|
|
targets.end());
|
|
|
|
|
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;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
SmallVector<OpPointer<AffineForOp>, 8>
|
|
|
|
|
mlir::tile(ArrayRef<OpPointer<AffineForOp>> forOps, ArrayRef<uint64_t> sizes,
|
|
|
|
|
OpPointer<AffineForOp> target) {
|
|
|
|
|
return tile(forOps, sizes, ArrayRef<OpPointer<AffineForOp>>{target})[0];
|
|
|
|
|
}
|