27 #include "llvm/ADT/DenseMap.h" 31 using scf::ParallelOp;
39 constantIndices.reserve(op.getUpperBound().size());
40 for (
auto bound : op.getUpperBound()) {
41 auto minOp = bound.getDefiningOp<AffineMinOp>();
45 for (
AffineExpr expr : minOp.map().getResults()) {
51 constantIndices.push_back(minConstant);
57 for (
auto bound : llvm::zip(op.getUpperBound(), constantIndices)) {
60 Value cmp = b.
create<arith::CmpIOp>(op.getLoc(), arith::CmpIPredicate::eq,
61 std::get<0>(bound), constant);
62 cond = cond ? b.
create<arith::AndIOp>(op.getLoc(), cond, cmp) : cmp;
63 map.
map(std::get<0>(bound), constant);
65 auto ifOp = b.
create<scf::IfOp>(op.getLoc(), cond,
true);
66 ifOp.getThenBodyBuilder().
clone(*op.getOperation(), map);
67 ifOp.getElseBodyBuilder().
clone(*op.getOperation());
76 auto bound = op.getUpperBound();
77 auto minOp = bound.getDefiningOp<AffineMinOp>();
81 for (
AffineExpr expr : minOp.map().getResults()) {
91 Value cond = b.
create<arith::CmpIOp>(op.getLoc(), arith::CmpIPredicate::eq,
93 map.
map(bound, constant);
94 auto ifOp = b.
create<scf::IfOp>(op.getLoc(), cond,
true);
95 ifOp.getThenBodyBuilder().
clone(*op.getOperation(), map);
96 ifOp.getElseBodyBuilder().
clone(*op.getOperation());
112 ForOp &partialIteration,
Value &splitBound) {
113 RewriterBase::InsertionGuard guard(b);
119 if (lbInt && ubInt && stepInt && (*ubInt - *lbInt) % *stepInt == 0)
122 if (stepInt == static_cast<int64_t>(1))
125 auto loc = forOp.getLoc();
129 auto modMap =
AffineMap::get(0, 3, {sym1 - ((sym1 - sym0) % sym2)});
133 forOp.getUpperBound(),
138 partialIteration = cast<ForOp>(b.
clone(*forOp.getOperation()));
139 partialIteration.getLowerBoundMutable().assign(splitBound);
140 forOp.replaceAllUsesWith(partialIteration->getResults());
141 partialIteration.getInitArgsMutable().assign(forOp->getResults());
145 forOp, [&]() { forOp.getUpperBoundMutable().assign(splitBound); });
150 template <
typename OpTy,
bool IsMin>
152 ForOp partialIteration,
154 Value mainIv = forOp.getInductionVar();
155 Value partialIv = partialIteration.getInductionVar();
156 assert(forOp.getStep() == partialIteration.getStep() &&
157 "expected same step in main and partial loop");
158 Value step = forOp.getStep();
160 forOp.walk([&](OpTy affineOp) {
163 affineOp.operands(), IsMin, mainIv,
167 partialIteration.walk([&](OpTy affineOp) {
170 affineOp.operands(), IsMin, partialIv,
171 previousUb, step,
false);
177 ForOp &partialIteration) {
178 Value previousUb = forOp.getUpperBound();
185 rewriter, forOp, partialIteration, previousUb);
187 rewriter, forOp, partialIteration, previousUb);
197 ForLoopPeelingPattern(
MLIRContext *ctx,
bool skipPartial)
215 scf::ForOp partialIteration;
237 struct ParallelLoopSpecialization
238 :
public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {
239 void runOnOperation()
override {
245 struct ForLoopSpecialization
246 :
public SCFForLoopSpecializationBase<ForLoopSpecialization> {
247 void runOnOperation()
override {
252 struct ForLoopPeeling :
public SCFForLoopPeelingBase<ForLoopPeeling> {
253 void runOnOperation()
override {
254 func::FuncOp funcOp = getOperation();
257 patterns.
add<ForLoopPeelingPattern>(ctx, skipPartial);
270 return std::make_unique<ParallelLoopSpecialization>();
274 return std::make_unique<ForLoopSpecialization>();
278 return std::make_unique<ForLoopPeeling>();
Include the generated interface declarations.
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
static void rewriteAffineOpAfterPeeling(RewriterBase &rewriter, ForOp forOp, ForOp partialIteration, Value previousUb)
std::unique_ptr< Pass > createParallelLoopSpecializationPass()
Creates a pass that specializes parallel loop for unrolling and vectorization.
MLIRContext * getContext() const
void createOrFold(SmallVectorImpl< Value > &results, Location location, Args &&...args)
Create an operation of specific op type at the current insertion point, and immediately try to fold i...
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Operation is a basic unit of execution within MLIR.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Operation * clone(Operation &op, BlockAndValueMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Value constantIndex(OpBuilder &builder, Location loc, int64_t i)
Generates a constant of index type.
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
An integer constant appearing in affine expression.
void erase()
Remove this operation from its parent block and delete it.
void map(Block *from, Block *to)
Inserts a new mapping for 'from' to 'to'.
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents an efficient way to signal success or failure.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
bool hasAttr(StringAttr name)
Return true if the operation has an attribute with the provided name, false otherwise.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
Operation * clone(BlockAndValueMapping &mapper, CloneOptions options=CloneOptions::all())
Create a deep copy of this operation, remapping any operands that use values outside of the operation...
Attribute removeAttr(StringAttr name)
Remove the attribute with the specified name if it exists.
static constexpr char kPartialIterationLabel[]
Base type for affine expression.
static void specializeParallelLoopForUnrolling(ParallelOp op)
Rewrite a parallel loop with bounds defined by an affine.min with a constant into 2 loops after check...
std::unique_ptr< Pass > createForLoopSpecializationPass()
Creates a pass that specializes for loop for unrolling and vectorization.
void updateRootInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around a root update of an operation.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued...
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
static constexpr char kPeeledLoopLabel[]
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
LogicalResult peelAndCanonicalizeForLoop(RewriterBase &rewriter, ForOp forOp, scf::ForOp &partialIteration)
Rewrite a for loop with bounds/step that potentially do not divide evenly into a for loop where the s...
std::unique_ptr< Pass > createForLoopPeelingPass()
Creates a pass that peels for loops at their upper bounds for better vectorization.
static void specializeForLoopForUnrolling(ForOp op)
Rewrite a for loop with bounds defined by an affine.min with a constant into 2 loops after checking i...
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&... args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments...
Specialization of arith.constant op that returns an integer of index type.
MLIRContext is the top-level object for a collection of MLIR operations.
LogicalResult rewritePeeledMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map, ValueRange operands, bool isMin, Value iv, Value ub, Value step, bool insideLoop)
Try to simplify a min/max operation op after loop peeling.
Optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
LogicalResult applyPatternsAndFoldGreedily(MutableArrayRef< Region > regions, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig())
Rewrite the regions of the specified operation, which must be isolated from above, by repeatedly applying the highest benefit patterns in a greedy work-list driven manner.
This class helps build Operations.
This class provides an abstraction over the different types of ranges over Values.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp, ForOp &partialIteration, Value &splitBound)
Rewrite a for loop with bounds/step that potentially do not divide evenly into a for loop where the s...