40 Value v = b.
create<affine::AffineApplyOp>(loc, m, ivs);
50 Block *body = linalgOp.getBlock();
54 if (
auto indexOp = dyn_cast<IndexOp>(&op)) {
55 map.
map(indexOp.getResult(), ivs[indexOp.getDim()]);
65 OpOperand *storeInto = linalgOp.getDpsInitOperand(operand.index());
67 b, loc, linalgOp.getMatchingIndexingMap(storeInto), ivs);
69 loc, toStore, linalgOp.getDpsInitOperand(operand.index())->get(),
85 template <
typename LinalgOpTy>
86 struct LinalgOpTilingInterface
87 :
public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
91 LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
92 return concreteOp.getIteratorTypesArray();
100 LinalgOp linalgOp = cast<LinalgOp>(op);
102 linalgOp.createFlatListOfOperandDims(b, loc);
103 AffineMap map = linalgOp.getShapesToLoopsMap();
105 return llvm::to_vector(
107 OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
108 b, loc, loopExpr, allShapesSizes);
109 return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
114 FailureOr<TilingResult>
121 LinalgOp linalgOp = cast<LinalgOp>(op);
124 b, loc, linalgOp, valuesToTile, offsets, sizes, {},
true);
126 llvm::make_filter_range(
128 [](
Value v) ->
bool {
129 return isa_and_nonnull<tensor::ExtractSliceOp, memref::SubViewOp>(
137 Operation *tiledOp =
clone(b, linalgOp, resultTensorTypes, tiledOperands);
153 unsigned numLoops = linalgOp.getNumLoops();
154 auto tilingInterfaceOp = cast<TilingInterface>(linalgOp.getOperation());
155 mappedOffsets.resize(numLoops);
156 mappedSizes.resize(numLoops);
159 tilingInterfaceOp.getIterationDomain(b);
161 mappedOffsets[index] = value.offset;
162 mappedSizes[index] = value.size;
165 for (
const auto &&[index, value] :
167 unsigned dimPosition = cast<AffineDimExpr>(value).getPosition();
168 mappedOffsets[dimPosition] = offsets[index];
169 mappedSizes[dimPosition] = sizes[index];
175 LogicalResult getIterationDomainTileFromOperandTile(
180 auto linalgOp = cast<LinalgOp>(op);
187 linalgOp.getMatchingIndexingMap(&op->
getOpOperand(operandNumber));
190 <<
"unhandled get iter domain position when operand is not "
191 "accessed using a permuted projection";
194 getMappedOffsetAndSize(linalgOp, b, indexingMap, offsets, sizes,
195 iterDomainOffsets, iterDomainSizes);
208 LinalgOp linalgOp = cast<LinalgOp>(op);
213 llvm::to_vector(llvm::map_range(sizes, [&](
OpFoldResult ofr) {
217 OpOperand *outOperand = linalgOp.getDpsInitOperand(resultNumber);
219 b, loc, outOperand->
get(), sizes,
220 linalgOp.getMatchingIndexingMap(outOperand), offsets,
221 {}, subShapeSizes,
true);
222 resultOffsets = sliceParams.
offsets;
223 resultSizes = sliceParams.
sizes;
227 LogicalResult getIterationDomainTileFromResultTile(
232 auto linalgOp = cast<LinalgOp>(op);
239 linalgOp.getIndexingMapMatchingResult(op->
getResult(resultNumber));
242 "unhandled tiled implementation generation when result is not "
243 "accessed using a permuted projection");
246 getMappedOffsetAndSize(linalgOp, b, indexingMap, offsets, sizes,
247 iterDomainOffsets, iterDomainSizes);
251 FailureOr<TilingResult>
256 if (failed(getIterationDomainTileFromResultTile(
257 op, b, resultNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
260 auto tilingInterfaceOp = cast<TilingInterface>(op);
261 FailureOr<TilingResult> tilingResult =
262 tilingInterfaceOp.getTiledImplementation(b, mappedOffsets, mappedSizes);
264 if (failed(tilingResult))
267 if (tilingResult->tiledOps.size() != 1)
268 return op->
emitOpError(
"failed to generate tiled implementation");
273 tilingResult->generatedSlices};
278 FailureOr<TilingResult> getTiledImplementationFromOperandTile(
282 if (failed(getIterationDomainTileFromOperandTile(
283 op, b, operandNumber, offsets, sizes, mappedOffsets,
293 auto linalgOp = cast<LinalgOp>(op);
294 if (!linalgOp.hasPureBufferSemantics())
295 return op->
emitOpError(
"expected operation to have buffer semantics");
298 indexedValues.reserve(linalgOp->getNumOperands());
302 for (
OpOperand &operand : linalgOp->getOpOperands()) {
303 if (!linalgOp.payloadUsesValueFromOperand(&operand)) {
304 indexedValues.push_back(
nullptr);
307 if (linalgOp.isScalar(&operand)) {
308 indexedValues.push_back(operand.get());
312 builder, linalgOpLoc, linalgOp.getMatchingIndexingMap(&operand), ivs);
314 builder.
create<memref::LoadOp>(linalgOpLoc, operand.get(), indices);
315 indexedValues.push_back(load);
334 static AffineMap getPartialResultAffineMap(LinalgOp linalgOp,
336 unsigned resultNumber) {
338 linalgOp.getMatchingIndexingMap(linalgOp.getDpsInitOperand(resultNumber));
339 for (
int redPos : reductionDims) {
348 template <
typename LinalgOpTy>
349 struct LinalgOpPartialReductionInterface
350 :
public PartialReductionOpInterface::ExternalModel<
351 LinalgOpPartialReductionInterface<LinalgOpTy>, LinalgOpTy> {
352 FailureOr<SmallVector<Value>> generateInitialTensorForPartialReduction(
355 auto linalgOp = cast<LinalgOp>(op);
358 if (linalgOp.hasPureBufferSemantics())
359 return op->
emitOpError(
"expected operation to have tensor semantics");
362 auto tilingInterfaceOp = cast<TilingInterface>(linalgOp.getOperation());
364 llvm::map_to_vector(tilingInterfaceOp.getIterationDomain(b),
365 [](
Range x) { return x.size; });
368 for (
auto [tileSize, dimSize] : llvm::zip_equal(sizes, shape)) {
370 tiledShape.push_back(dimSize);
372 tiledShape.push_back(tileSize);
377 for (
int initIdx = 0, e = linalgOp.getNumDpsInits(); initIdx < e;
382 combinerOps.size() != 1)
383 return op->
emitOpError(
"Failed to anaysis the reduction operation.");
387 if (!identity.has_value())
389 "Failed to get an identity value for the reduction operation.");
393 getPartialResultAffineMap(linalgOp, reductionDims, initIdx);
396 auto dim = cast<AffineDimExpr>(dimExpr);
397 partialResultShape.push_back(tiledShape[dim.getPosition()]);
403 b.
create<tensor::EmptyOp>(loc, partialResultShape, elType);
404 Value constantOp = b.
create<arith::ConstantOp>(loc, *identity);
405 auto identityTensor =
406 b.
create<linalg::FillOp>(loc, constantOp, emptyTensor);
407 inits.push_back(identityTensor.getResult(0));
413 FailureOr<TilingResult>
419 auto linalgOp = cast<LinalgOp>(op);
424 newInitMaps.reserve(linalgOp.getNumDpsInits());
425 for (
int idx : llvm::seq<int>(0, linalgOp.getNumDpsInits())) {
429 getPartialResultAffineMap(linalgOp, reductionDims, idx);
430 newInitMaps.push_back(newMap);
435 b, loc, linalgOp, linalgOp.getDpsInputs(), offsets, sizes, {},
true);
437 llvm::make_filter_range(
443 for (
auto [valueMap, valueToTile] : llvm::zip_equal(newInitMaps, init)) {
444 int64_t initRank = valueMap.getNumResults();
448 for (
AffineExpr dimExpr : valueMap.getResults()) {
449 auto dim = cast<AffineDimExpr>(dimExpr);
450 initSizes.push_back(sizes[dim.getPosition()]);
453 auto extractSlice = b.
create<tensor::ExtractSliceOp>(
454 loc, valueToTile, initOffset, initSizes, initStride);
455 tiledInits.push_back(extractSlice);
456 generatedSlices.push_back(extractSlice);
462 for (
int idx : llvm::seq<int>(0, linalgOp.getNumDpsInits())) {
465 OpOperand *initOperand = linalgOp.getDpsInitOperand(idx);
466 int64_t mapIdx = linalgOp.getIndexingMapIndex(initOperand);
467 newMaps[mapIdx] = newInitMaps[idx];
472 linalgOp.getIteratorTypesArray();
473 for (
int dim : reductionDims)
474 newIteratorTypes[dim] = utils::IteratorType::parallel;
479 tiledInits, newMaps, newIteratorTypes);
482 genericOp.getRegion().begin(), mapping);
484 {genericOp.getOperation()},
485 llvm::map_to_vector(genericOp->getResults(),
493 auto linalgOp = cast<LinalgOp>(op);
497 int64_t numInits = linalgOp.getNumDpsInits();
500 for (
int idx : llvm::seq(numInits)) {
506 getPartialResultAffineMap(linalgOp, reductionDims, idx);
508 for (
auto [resultNum, dimExpr] :
510 unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
511 if (llvm::find(reductionDims, dim) != reductionDims.end()) {
512 partialReductionDims.push_back(resultNum);
516 Value partialResult = partialReduce[idx];
517 Value init = linalgOp.getDpsInits()[idx];
519 auto reduction = b.
create<linalg::ReduceOp>(
520 loc, partialResult, init, partialReductionDims,
529 b.create<linalg::YieldOp>(loc, clonedReductionOp->
getResult(0));
532 mergeOperations.push_back(reduction);
533 replacements.push_back(reduction->getResult(0));
539 LogicalResult getPartialResultTilePosition(
545 auto linalgOp = cast<LinalgOp>(op);
548 getPartialResultAffineMap(linalgOp, reductionDims, resultNumber);
550 unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
551 resultSizes.push_back(sizes[dim]);
553 if (llvm::find(reductionDims, dim) != reductionDims.end()) {
558 resultOffsets.push_back(offsets[dim]);
568 template <
typename OpType>
570 OpType::template attachInterface<LinalgOpTilingInterface<OpType>>(*ctx);
571 OpType::template attachInterface<LinalgOpPartialReductionInterface<OpType>>(
576 template <
typename... OpTypes>
578 (registerOne<OpTypes>(ctx), ...);
586 registerOne<linalg::GenericOp>(ctx);
588 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
static LogicalResult getResultTilePosition(RewriterBase &rewriter, int64_t index, Value tiledResult, TilingInterface op, ArrayRef< OpFoldResult > offsets, ArrayRef< OpFoldResult > sizes, SmallVector< OpFoldResult > &resultOffset, SmallVector< OpFoldResult > &resultSize, const scf::SCFTilingOptions &options)
static FailureOr< TilingResult > getTiledImplementation(RewriterBase &rewriter, TilingInterface op, ValueRange regionIterArg, ArrayRef< OpFoldResult > offsets, ArrayRef< OpFoldResult > sizes, const scf::SCFTilingOptions &options)
static SmallVector< Value > getIndicesForAccess(OpBuilder &b, Location loc, AffineMap indexingMap, ValueRange ivs)
Return the SSA values that represent the data point accessed using a given indexingMap for a given po...
static LogicalResult inlinePayload(OpBuilder &b, LinalgOp linalgOp, ValueRange ivs, ValueRange argValues)
Method to inline the payload of a linalgOp given the iteration space point and values for the argumen...
static void registerAll(MLIRContext *ctx)
Variadic helper function.
static void registerOne(MLIRContext *ctx)
Base type for affine expression.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
AffineMap insertResult(AffineExpr expr, unsigned pos) const
Returns a new AffineMap with the same number of dims and symbols and an extra result inserted at pos.
bool isProjectedPermutation(bool allowZeroInResults=false) const
Returns true if the AffineMap represents a subset (i.e.
unsigned getNumSymbols() const
unsigned getNumDims() const
ArrayRef< AffineExpr > getResults() const
unsigned getNumResults() const
bool isPermutation() const
Returns true if the AffineMap represents a symbol-less permutation map.
Block represents an ordered list of Operations.
Operation * getTerminator()
Get the terminator operation of this block.
BlockArgListType getArguments()
iterator_range< iterator > without_terminator()
Return an iterator range over the operation within this block excluding the terminator operation at t...
IntegerAttr getIndexAttr(int64_t value)
MLIRContext * getContext() const
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
bool addExtension(TypeID extensionID, std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
This is a utility class for mapping one set of IR entities to another.
auto lookupOrDefault(T from) const
Lookup a mapped value within the map.
void map(Value from, Value to)
Inserts a new mapping for 'from' to 'to'.
IRValueT get() const
Return the current value being used by this operand.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
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 setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents a single result from folding an operation.
This class represents an operand of an operation.
This is a value defined by a result of an operation.
Operation is the basic unit of execution within MLIR.
OpOperand & getOpOperand(unsigned idx)
void setOperand(unsigned idx, Value value)
Operation * clone(IRMapping &mapper, CloneOptions options=CloneOptions::all())
Create a deep copy of this operation, remapping any operands that use values outside of the operation...
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Location getLoc()
The source location the operation was defined or derived from.
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Region & getRegion(unsigned index)
Returns the region held by this operation at position 'index'.
operand_range getOperands()
Returns an iterator on the underlying Value's.
result_range getResults()
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
void cloneInto(Region *dest, IRMapping &mapper)
Clone the internal blocks from this region into dest.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class provides an abstraction over the different types of ranges over Values.
type_range getTypes() const
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
std::optional< TypedAttr > getNeutralElement(Operation *op)
Return the identity numeric value associated to the give op.
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 offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef< OpFoldResult > offests)
Add the specified offsets to any linalg.index ops contained in the given linalgOp.
void registerTilingInterfaceExternalModels(DialectRegistry ®istry)
SmallVector< Type > getTensorOutputTypes(LinalgOp op, ValueRange operands)
Returns the list of tensor output types produced when the given structured operation op is applied to...
SliceParameters computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile, ArrayRef< OpFoldResult > tileSizes, AffineMap map, ArrayRef< OpFoldResult > lbs, ArrayRef< OpFoldResult > ubs, ArrayRef< OpFoldResult > subShapeSizes, bool omitPartialTileCheck)
Computes SliceParameters for a single valueToTile assuming that its user is being tiled with the give...
Include the generated interface declarations.
bool isZeroIndex(OpFoldResult v)
Return true if v is an IntegerAttr with value 0 of a ConstantIndexOp with attribute with value 0.
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Value matchReduction(ArrayRef< BlockArgument > iterCarriedArgs, unsigned redPos, SmallVectorImpl< Operation * > &combinerOps)
Utility to match a generic reduction given a list of iteration-carried arguments, iterCarriedArgs and...
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Container for the result of merge operation of tiling.
Represents a range (offset, size, and stride) where each element of the triple may be dynamic or stat...
Container for result values of tiling.
SmallVector< Operation * > tiledOps
A struct containg offsets-sizes-strides arguments of the tiled shape.
SmallVector< OpFoldResult > sizes
SmallVector< OpFoldResult > offsets