33 #include "llvm/ADT/STLExtras.h"
34 #include "llvm/Support/CommandLine.h"
38 #define GEN_PASS_DEF_LINALGTILINGPASS
39 #include "mlir/Dialect/Linalg/Passes.h.inc"
47 #define DEBUG_TYPE "linalg-tiling"
61 for (
int idx = 0, e = tileSizes.size(), zerosCount = 0; idx < e; ++idx) {
63 static_cast<int64_t
>(0)) {
64 shapeSizes.erase(shapeSizes.begin() + idx - zerosCount);
65 tileSizes.erase(tileSizes.begin() + idx - zerosCount);
69 loopIndexToRangeIndex[idx] = idx - zerosCount;
74 for (
unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx)
76 return std::make_tuple(res, loopIndexToRangeIndex);
84 auto rangeIndex = loopIndexToRangeIndex.find(en.index());
85 if (rangeIndex == loopIndexToRangeIndex.end())
87 en.value() = ivs[rangeIndex->second];
97 if (
auto attr = llvm::dyn_cast_if_present<Attribute>(value)) {
98 assert(cast<IntegerAttr>(attr).getValue().isStrictlyPositive() &&
99 "expected strictly positive tile size and divisor");
104 Value condition = b.
create<arith::CmpIOp>(arith::CmpIPredicate::sgt,
105 cast<Value>(value), zero);
108 b.
getStringAttr(
"expected strictly positive tile size and divisor"));
111 FailureOr<StaticContinuousTileSizeSpecification>
113 unsigned targetSize) {
115 assert(!op.hasDynamicShape() &&
116 "cannot compute static multi-tile sizes for an op with dynamic shape");
117 assert(targetSize > 0 &&
"target size must be non-negative");
118 assert(dimension < op.getNumLoops() &&
"dimension overflow");
121 int64_t loopRange = op.getStaticLoopRanges()[dimension];
122 int64_t tripCount = loopRange / targetSize;
124 unsigned tileSize = targetSize;
129 int64_t remainderChunk = loopRange % targetSize;
131 while (tileSize > 1 && remainderChunk != 0) {
133 uint64_t maxPower = llvm::bit_floor(tileSize);
134 tileSize = maxPower == tileSize ? maxPower >> 1 : maxPower;
136 tripCount = remainderChunk / tileSize;
143 remainderChunk = remainderChunk % tileSize;
148 int64_t range) ->
bool {
149 int64_t computedRange = 0;
150 for (
auto [tileSize, tripCount] : llvm::zip(tileSizes, tripCounts))
151 computedRange += tileSize * tripCount;
152 return range == computedRange;
161 FailureOr<ContinuousTileSizeSpecification>
165 bool emitAssertions) {
168 unsigned numLoops = loopRanges.size();
171 if (dimension >= numLoops)
177 if (emitAssertions) {
180 Value targetSizeValue =
196 Value tripCountValue = apply(s0.
floorDiv(s1), {loopRange, targetSizeValue});
197 Value remainderChunkValue = apply(s0 % s1, {loopRange, targetSizeValue});
206 assert(tileSizeInt > 0 &&
"target size must be non-negative");
208 spec.
tileSizes.push_back(targetSizeValue);
211 while (tileSizeInt > 1) {
212 uint64_t maxPower = llvm::bit_floor(tileSizeInt);
213 tileSizeInt = maxPower == tileSizeInt ? maxPower >> 1 : maxPower;
216 tripCountValue = apply(s0.
floorDiv(s1), {remainderChunkValue, constStepOp});
219 b, b.
getLoc(), s0.
floorDiv(s1), {remainderChunkValue, constStepOp});
222 if (
Attribute attr = llvm::dyn_cast_if_present<Attribute>(tripCountSize)) {
223 auto intAttr = cast<IntegerAttr>(attr);
224 bool isTripCountZero = intAttr.getValue().isZero();
226 if (!isTripCountZero) {
235 remainderChunkValue = apply(s0 % s1, {remainderChunkValue, constStepOp});
241 FailureOr<StaticMultiSizeSpecification>
243 int64_t targetSize, int64_t divisor) {
244 assert(!op.hasDynamicShape() &&
245 "cannot compute static multi-tile sizes for an op with dynamic shape");
246 assert(targetSize > 0 &&
"target size must be non-negative");
247 assert(divisor > 0 &&
"divisor must be non-negative");
248 assert(dimension < op.getNumLoops() &&
"dimension overflow");
251 int64_t tripCount = op.getStaticLoopRanges()[dimension];
252 int64_t a = tripCount / divisor;
253 int64_t t = (targetSize + divisor - 1) / divisor;
254 int64_t totalTripCount = (a + t - 1) / t;
267 FailureOr<MultiSizeSpecification>
272 if (dimension >= op.getNumLoops())
278 if (emitAssertions) {
282 Value targetSizeValue =
289 op.createFlatListOfOperandDims(b, b.
getLoc());
290 AffineMap shapesToLoops = op.getShapesToLoopsMap();
305 Value t = apply((s0 + s1 - 1).floorDiv(s1), {targetSizeValue, divisorValue});
306 Value d = apply((s0 + s1 - 1).floorDiv(s1), {a, t});
308 Value v = apply(s0 % s1, {a, d});
309 Value u = apply(s0 - s1, {d, v});
313 spec.highTileSize = apply(s0 + s1, {s, divisorValue});
314 spec.lowTripCount = u;
315 spec.highTripCount = v;
321 if (emitAssertions) {
324 apply(s0 * s1 + s2 * s3, {spec.lowTileSize, spec.lowTripCount,
325 spec.highTileSize, spec.highTripCount});
326 Value equals = b.
create<arith::CmpIOp>(arith::CmpIPredicate::eq,
327 coveredSize, tripCount);
330 "could not compute dynamic multi-size tile shapes"));
344 if (!tileSizeConst || !numThreadsConst || !iterSizeConst)
346 return *tileSizeConst * (*numThreadsConst - 1) < *iterSizeConst;
370 bool omitTileOffsetBoundsCheck,
380 int64_t nLoops = loopRanges.size();
381 tiledOffsets.reserve(nLoops);
382 tiledSizes.reserve(nLoops);
383 for (
unsigned loopIdx = 0, threadIdIdx = 0; loopIdx < nLoops; ++loopIdx) {
384 bool overflow = loopIdx >= numThreads.size();
385 bool isZero = !overflow &&
isZeroInteger(numThreads[loopIdx]);
387 if (overflow || isZero) {
388 tiledOffsets.push_back(loopRanges[loopIdx].offset);
389 tiledSizes.push_back(loopRanges[loopIdx].size);
403 nominalTileSizes.has_value()
404 ? (*nominalTileSizes)[loopIdx]
411 b, loc, i +
j * m, {offset, threadId, tileSizePerThread});
414 b, loc, i +
j * m - n,
415 {offset, nonZeroNumThreads[threadIdIdx], tileSizePerThread, size});
418 b, loc, -i + m, {offsetPerThread, size});
420 buildMin(b, loc, {sizeMinusOffsetPerThread, tileSizePerThread});
423 tiledOffsets.push_back(offsetPerThread);
425 if (!omitTileOffsetBoundsCheck &&
427 nonZeroNumThreads[threadIdIdx], size))
431 tiledSizes.push_back(tileSizePerThread);
436 template <
typename LoopTy>
437 static FailureOr<TiledLinalgOp>
442 auto nLoops = op.getNumLoops();
444 tileSizes = tileSizes.take_front(nLoops);
450 tiledOp.
op = cast<LinalgOp>(b.
clone(*op.getOperation()));
452 tiledOp.
op->result_end());
458 op.createFlatListOfOperandDims(b, op.getLoc());
459 AffineMap shapeSizesToLoopsMap = op.getShapesToLoopsMap();
460 if (!shapeSizesToLoopsMap)
464 b, op.getLoc(), shapeSizesToLoopsMap, allShapeSizes, tileSizes);
467 for (
const auto &attr :
enumerate(op.getIteratorTypesArray())) {
468 if (loopIndexToRangeIndex.count(attr.index()))
469 iteratorTypes.push_back(attr.value());
473 auto invPermutationMap =
475 if (!
options.interchangeVector.empty()) {
479 interchangeVector.reserve(
options.interchangeVector.size());
480 for (
auto pos :
options.interchangeVector) {
481 auto it = loopIndexToRangeIndex.find(pos);
482 if (it == loopIndexToRangeIndex.end())
484 interchangeVector.push_back(it->second);
490 assert(invPermutationMap);
492 interchangeVector.end());
502 iteratorTypes.size(),
509 parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
511 auto returnedProcInfo =
512 options.distribution->procInfo(b, op.getLoc(), parallelLoopRanges);
513 unsigned procIdIdx = 0;
518 procInfo[iteratorType.index()] = returnedProcInfo[procIdIdx++];
525 auto tiledLoopBodyBuilder =
528 ivs.assign(localIvs.begin(), localIvs.end());
534 if (!
options.interchangeVector.empty()) {
535 for (
AffineExpr result : invPermutationMap.getResults())
536 interchangedIvs.push_back(
537 ivs[cast<AffineDimExpr>(result).getPosition()]);
539 interchangedIvs.assign(ivs.begin(), ivs.end());
544 assert(operandValuesToUse.size() ==
545 static_cast<size_t>(op->getNumOperands()) &&
546 "expect the number of operands and inputs and outputs to match");
558 res =
clone(b, op, resultTensorTypes, tiledOperands);
564 tiledLoopBodyBuilder, procInfo);
571 loops.reserve(ivs.size());
572 for (
auto iv : ivs) {
573 if (isa<BlockArgument>(iv)) {
574 loops.push_back(cast<BlockArgument>(iv).getOwner()->getParentOp());
575 assert(loops.back() &&
"no owner found for induction variable!");
579 loops.push_back(
nullptr);
587 if ((outermostLoop = loop))
591 res, loops, outermostLoop ? outermostLoop->
getResults() : tensorResults};
597 std::optional<ArrayAttr> mapping) {
604 auto tilingInterfaceOp = cast<TilingInterface>(op.getOperation());
610 auto destinationStyleOp =
611 dyn_cast<DestinationStyleOpInterface>(op.getOperation());
612 if (!destinationStyleOp)
616 auto linalgOp = dyn_cast<linalg::LinalgOp>(op.getOperation());
621 if (op->getNumResults() != 1)
623 op,
"don't support ops with multiple results for now");
626 tilingInterfaceOp.getLoopIteratorTypes();
628 linalgOp.getReductionDims(redDims);
629 if (redDims.size() != 1)
631 op,
"only support ops with one reduction dimension.");
632 if (!tileSizes.empty() && tileSizes.size() != numThreads.size())
634 "many elements as number of threads");
636 if (redDims.front() >= numThreads.size())
638 op,
"reduction dimension must be mapped to threads");
641 unsigned reductionDim = redDims.front();
643 reductionDims.insert(reductionDim);
644 FailureOr<SmallVector<Value>> maybeInitTensors =
645 op.generateInitialTensorForPartialReduction(b, loc, numThreads,
647 if (failed(maybeInitTensors))
649 op,
"Failed to create inital tensors for partial reduction");
665 scf::ForallOp forallOp = b.
create<scf::ForallOp>(
674 std::nullopt, tiledOffsets,
687 for (
Value initOperand : destinationStyleOp.getDpsInits()) {
688 auto *it = llvm::find(dest, initOperand);
689 assert(it != dest.end() &&
"dest operand not found in dest");
690 unsigned destNum = std::distance(dest.begin(), it);
696 outOffsets[reductionDim] = forallOp.getInductionVars()[0];
698 tiledDpsInitOperands.push_back(b.
create<tensor::ExtractSliceOp>(
699 loc, cast<RankedTensorType>(initOperand.getType()),
700 destBbArgs[destNum], outOffsets, sizes, strides));
708 for (
auto [initOperandPtr, tiledInitValue] : llvm::zip_equal(
709 cast<DestinationStyleOpInterface>(clonedOp).getDpsInitsMutable(),
710 tiledDpsInitOperands)) {
711 initOperandPtr.set(tiledInitValue);
716 if (tileSizes.empty()) {
717 FailureOr<TilingResult> tilingResult =
718 cast<TilingInterface>(clonedOp).getTiledImplementation(
719 b, tiledOffsets, tiledSizes);
720 if (failed(tilingResult))
721 return clonedOp->
emitError(
"Failed to tile op: ");
722 if (tilingResult->tiledOps.size() != 1) {
723 return clonedOp->
emitError(
"expected a single produced tiled op, got ")
724 << tilingResult->tiledOps.size();
726 tiledOp = tilingResult->tiledOps.front();
727 tilingResults = tilingResult->tiledValues;
730 FailureOr<TiledLinalgOp> maybeTiled = tileLinalgOpImpl<scf::ForOp>(
731 b, cast<LinalgOp>(clonedOp), tileSizes,
options);
732 if (failed(maybeTiled))
737 materializedNonZeroNumThreads);
738 if (maybeTiled->loops.size() != 1) {
739 return clonedOp->
emitError(
"expected a single produced loop");
741 tiledOp = maybeTiled->op;
742 tilingResults = maybeTiled->loops.front()->getResults();
749 for (
auto [index, result, bbArg] : llvm::zip(
750 llvm::seq<unsigned>(0, dest.size()), tilingResults, destBbArgs)) {
756 if (failed(tilingInterfaceOp.getResultTilePosition(
757 b, index, tiledOffsets, tiledSizes, resultOffsets, resultSizes)))
758 return op->emitOpError(
"output offsets couldn't be calculated");
762 for (int64_t i = 0, e = numThreads.size(); i < e; ++i) {
763 if (i == reductionDim) {
764 resultOffsetsRank.push_back(forallOp.getInductionVars()[0]);
768 resultOffsetsRank.push_back(resultOffsets[offIdx++]);
769 resultSizesRank.push_back(resultSizes[sizeIdx++]);
777 b.
create<tensor::ParallelInsertSliceOp>(
778 loc, result, bbArg, resultOffsetsRank, resultSizesRank, strides);
783 FailureOr<MergeResult> mergeResult =
784 op.mergeReductions(b, loc, forallOp->getResults(), reductionDims);
785 if (failed(mergeResult)) {
788 b.
replaceOp(op, mergeResult->replacements);
793 results.
loops = forallOp;
795 results.
mergeOps.append(mergeResult->mergeOps);
799 template <
typename LoopTy>
805 if (!
options.tileSizeComputationFunction)
811 auto nLoops = op.getNumLoops();
814 if (tileSizeVector.size() < nLoops) {
815 tileSizeVector.append(nLoops - tileSizeVector.size(), b.
getIndexAttr(0));
818 return tileLinalgOpImpl<LoopTy>(b, op, tileSizeVector,
options);
821 FailureOr<TiledLinalgOp>
826 return tileLinalgOpImpl<scf::ForOp>(b, op,
options);
827 case LinalgTilingLoopType::ParallelLoops:
828 return tileLinalgOpImpl<scf::ParallelOp>(b, op,
options);
836 template <
typename... OpTypes>
837 class CanonicalizationPatternList;
840 class CanonicalizationPatternList<> {
845 template <
typename OpTy,
typename... OpTypes>
846 class CanonicalizationPatternList<OpTy, OpTypes...> {
850 CanonicalizationPatternList<OpTypes...>::insert(
patterns);
865 affine::AffineApplyOp::getCanonicalizationPatterns(
patterns, ctx);
866 affine::AffineForOp::getCanonicalizationPatterns(
patterns, ctx);
867 affine::AffineMinOp::getCanonicalizationPatterns(
patterns, ctx);
868 affine::AffineMaxOp::getCanonicalizationPatterns(
patterns, ctx);
869 arith::ConstantIndexOp::getCanonicalizationPatterns(
patterns, ctx);
871 memref::SubViewOp::getCanonicalizationPatterns(
patterns, ctx);
872 memref::ViewOp::getCanonicalizationPatterns(
patterns, ctx);
874 scf::ForOp::getCanonicalizationPatterns(
patterns, ctx);
875 scf::ParallelOp::getCanonicalizationPatterns(
patterns, ctx);
877 tensor::CastOp::getCanonicalizationPatterns(
patterns, ctx);
878 tensor::EmptyOp::getCanonicalizationPatterns(
patterns, ctx);
879 tensor::ExtractSliceOp::getCanonicalizationPatterns(
patterns, ctx);
880 tensor::InsertSliceOp::getCanonicalizationPatterns(
patterns, ctx);
881 tensor::PadOp::getCanonicalizationPatterns(
patterns, ctx);
882 ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(
patterns);
884 CanonicalizationPatternList<
886 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
static llvm::ManagedStatic< PassManagerOptions > options
static bool canOmitTileOffsetInBoundsCheck(OpFoldResult tileSize, OpFoldResult numThreads, OpFoldResult iterationSize)
Returns true if the maximum tile offset tileSize * numThreads-1 is less than iterationSize.
static void emitIsPositiveIndexAssertion(ImplicitLocOpBuilder &b, OpFoldResult value)
Asserts that the given index-typed value is strictly positive.
static OpFoldResult buildMax(OpBuilder &b, Location loc, ArrayRef< OpFoldResult > vals)
Build an affine_max of all the vals.
static void calculateTileOffsetsAndSizes(RewriterBase &b, Location loc, scf::ForallOp forallOp, ArrayRef< OpFoldResult > numThreads, SmallVector< Range > loopRanges, bool omitTileOffsetBoundsCheck, std::optional< ArrayRef< OpFoldResult >> nominalTileSizes, SmallVector< OpFoldResult > &tiledOffsets, SmallVector< OpFoldResult > &tiledSizes)
Fill out the tiledOffsets and tiledSizes to be used to tile to a given number of threads.
static FailureOr< TiledLinalgOp > tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ArrayRef< OpFoldResult > tileSizes, const LinalgTilingOptions &options)
static OpFoldResult buildMin(OpBuilder &b, Location loc, ArrayRef< OpFoldResult > vals)
Build an affine_min of all the vals.
Base type for affine expression.
AffineExpr floorDiv(uint64_t v) const
AffineExpr ceilDiv(uint64_t v) const
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
static AffineMap getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with 'numDims' identity result dim exprs.
unsigned getNumResults() const
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Attributes are known-constant values of operations.
IntegerAttr getIndexAttr(int64_t value)
AffineExpr getAffineSymbolExpr(unsigned position)
StringAttr getStringAttr(const Twine &bytes)
MLIRContext * getContext() const
ImplicitLocOpBuilder maintains a 'current location', allowing use of the create<> method without spec...
Location getLoc() const
Accessors for the implied location.
OpTy create(Args &&...args)
Create an operation of specific op type at the current insertion point and location.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext * getContext() const
Return the context this location is uniqued in.
MLIRContext is the top-level object for a collection of MLIR operations.
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
Operation * clone(Operation &op, IRMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
void setInsertionPointToEnd(Block *block)
Sets the insertion point to the end of the specified block.
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...
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
This class represents a single result from folding an operation.
Operation is the basic unit of execution within MLIR.
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
result_range getResults()
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Specialization of arith.constant op that returns an integer of index type.
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...
SmallVector< OpFoldResult > makeComposedFoldedMultiResultAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
Variant of makeComposedFoldedAffineApply suitable for multi-result maps.
OpFoldResult makeComposedFoldedAffineMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineMinOp that computes a maximum across the results of applying map to operands,...
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
OpFoldResult makeComposedFoldedAffineMin(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineMinOp that computes a minimum across the results of applying map to operands,...
void mapLoopToProcessorIds(scf::ForOp forOp, ArrayRef< Value > processorId, ArrayRef< Value > numProcessors)
Maps forOp for execution on a parallel grid of virtual processorIds of size given by numProcessors.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
SmallVector< Value > makeTiledShapes(OpBuilder &builder, Location loc, LinalgOp linalgOp, ValueRange valuesToTile, ArrayRef< OpFoldResult > ivs, ArrayRef< OpFoldResult > tileSizes, ArrayRef< OpFoldResult > sizeBounds, bool omitPartialTileCheck)
Creates extract_slice/subview ops for all valuesToTile of the given linalgOp with builder,...
void transformIndexOps(RewriterBase &b, LinalgOp op, SmallVectorImpl< Value > &ivs, const LoopIndexToRangeIndexMap &loopIndexToRangeIndex)
All indices returned by IndexOp should be invariant with respect to tiling.
bool isParallelIterator(utils::IteratorType iteratorType)
Check if iterator type has "parallel" semantics.
void populateLinalgTilingCanonicalizationPatterns(RewritePatternSet &patterns)
SmallVector< Value > insertSlicesBack(OpBuilder &builder, Location loc, LinalgOp op, ValueRange operands, ValueRange results)
Creates insert_slice ops that insert results back into larger tensors they were originally extracted ...
std::tuple< SmallVector< Range, 4 >, LoopIndexToRangeIndexMap > makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > allShapeSizes, ArrayRef< OpFoldResult > allTileSizes)
void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef< OpFoldResult > offests)
Add the specified offsets to any linalg.index ops contained in the given linalgOp.
FailureOr< StaticMultiSizeSpecification > computeStaticMultiTileSizes(LinalgOp op, unsigned dimension, int64_t targetSize, int64_t divisor)
FailureOr< ContinuousTileSizeSpecification > computeContinuousTileSizes(OpBuilder &builder, TilingInterface op, unsigned dimension, OpFoldResult targetSize, bool emitAssertions)
FailureOr< StaticContinuousTileSizeSpecification > computeStaticContinuousTileSizes(LinalgOp op, unsigned dimension, unsigned targetSize)
FailureOr< ForallReductionTilingResult > tileReductionUsingForall(RewriterBase &b, PartialReductionOpInterface op, ArrayRef< OpFoldResult > numThreads, ArrayRef< OpFoldResult > tileSizes={}, std::optional< ArrayAttr > mapping=std::nullopt)
Method to tile a reduction to parallel iterations computing partial reductions.
FailureOr< TiledLinalgOp > tileLinalgOp(RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options)
RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx)
Canonicalization patterns relevant to apply after tiling patterns.
SmallVector< Type > getTensorOutputTypes(LinalgOp op, ValueRange operands)
Returns the list of tensor output types produced when the given structured operation op is applied to...
FailureOr< MultiSizeSpecification > computeMultiTileSizes(OpBuilder &builder, LinalgOp op, unsigned dimension, OpFoldResult targetSize, OpFoldResult divisor, bool emitAssertions=true)
Emits the IR computing the multi-sized tiling specification with two tile sizes not exceeding targetS...
SmallVector< Value > ValueVector
An owning vector of values, handy to return from functions.
LogicalResult getOrCreateDestinations(OpBuilder &b, Location loc, Operation *op, SmallVector< Value > &result)
This is a helper function for DestinationStyleOpInterface.
Include the generated interface declarations.
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
const FrozenRewritePatternSet & patterns
bool isZeroInteger(OpFoldResult v)
Return true if v is an IntegerAttr with value 0.
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
OpFoldResult getAsOpFoldResult(Value val)
Given a value, try to extract a constant Attribute.
SmallVector< scf::ForOp, 8 > Loops
Tile a nest of standard for loops rooted at rootForOp by finding such parametric tile sizes that the ...
void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)
Apply the permutation defined by permutation to inVec.
Represents a range (offset, size, and stride) where each element of the triple may be dynamic or stat...
Transformation information returned after reduction tiling.
SmallVector< Operation * > mergeOps
The final reduction operation merging all the partial reductions.
SmallVector< Value > initialValues
Initial values used for partial reductions.
scf::ForallOp loops
The scf.forall operation that iterate over the tiles.
SmallVector< Operation * > parallelTiledOps
The partial reduction tiled op generated.
A description of a multi-size tiling comprising tile sizes and numbers of tiles, expressed as Values ...
Callback function type used to get processor ID, and number of processors used for distribution for a...
Perform standalone tiling of a single LinalgOp by tileSizes.
SmallVector< Value, 4 > tensorResults
SmallVector< T > tileSizes
Tile sizes.
SmallVector< T > tripCounts
Number of tiles associated with each size.
T lowTripCount
Number of tiles associated with each size.
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.