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);
129 Operation *tiledOp =
clone(b, linalgOp, resultTensorTypes, tiledOperands);
144 unsigned numLoops = linalgOp.getNumLoops();
145 auto tilingInterfaceOp = cast<TilingInterface>(linalgOp.getOperation());
146 mappedOffsets.resize(numLoops);
147 mappedSizes.resize(numLoops);
150 tilingInterfaceOp.getIterationDomain(b);
152 mappedOffsets[index] = value.offset;
153 mappedSizes[index] = value.size;
156 for (
const auto &&[index, value] :
158 unsigned dimPosition = cast<AffineDimExpr>(value).getPosition();
159 mappedOffsets[dimPosition] = offsets[index];
160 mappedSizes[dimPosition] = sizes[index];
166 LogicalResult getIterationDomainTileFromOperandTile(
171 auto linalgOp = cast<LinalgOp>(op);
178 linalgOp.getMatchingIndexingMap(&op->
getOpOperand(operandNumber));
181 <<
"unhandled get iter domain position when operand is not "
182 "accessed using a permuted projection";
185 getMappedOffsetAndSize(linalgOp, b, indexingMap, offsets, sizes,
186 iterDomainOffsets, iterDomainSizes);
199 LinalgOp linalgOp = cast<LinalgOp>(op);
204 llvm::to_vector(llvm::map_range(sizes, [&](
OpFoldResult ofr) {
208 OpOperand *outOperand = linalgOp.getDpsInitOperand(resultNumber);
210 b, loc, outOperand->
get(), sizes,
211 linalgOp.getMatchingIndexingMap(outOperand), offsets,
212 {}, subShapeSizes,
true);
213 resultOffsets = sliceParams.
offsets;
214 resultSizes = sliceParams.
sizes;
218 LogicalResult getIterationDomainTileFromResultTile(
223 auto linalgOp = cast<LinalgOp>(op);
230 linalgOp.getIndexingMapMatchingResult(op->
getResult(resultNumber));
233 "unhandled tiled implementation generation when result is not "
234 "accessed using a permuted projection");
237 getMappedOffsetAndSize(linalgOp, b, indexingMap, offsets, sizes,
238 iterDomainOffsets, iterDomainSizes);
242 FailureOr<TilingResult>
247 if (failed(getIterationDomainTileFromResultTile(
248 op, b, resultNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
251 auto tilingInterfaceOp = cast<TilingInterface>(op);
252 FailureOr<TilingResult> tilingResult =
253 tilingInterfaceOp.getTiledImplementation(b, mappedOffsets, mappedSizes);
255 if (failed(tilingResult))
258 if (tilingResult->tiledOps.size() != 1)
259 return op->
emitOpError(
"failed to generate tiled implementation");
268 FailureOr<TilingResult> getTiledImplementationFromOperandTile(
272 if (failed(getIterationDomainTileFromOperandTile(
273 op, b, operandNumber, offsets, sizes, mappedOffsets,
277 return getTiledImplementation(op, b, mappedOffsets, mappedSizes);
283 auto linalgOp = cast<LinalgOp>(op);
284 if (!linalgOp.hasPureBufferSemantics())
285 return op->
emitOpError(
"expected operation to have buffer semantics");
288 indexedValues.reserve(linalgOp->getNumOperands());
292 for (
OpOperand &operand : linalgOp->getOpOperands()) {
293 if (!linalgOp.payloadUsesValueFromOperand(&operand)) {
294 indexedValues.push_back(
nullptr);
297 if (linalgOp.isScalar(&operand)) {
298 indexedValues.push_back(operand.get());
302 builder, linalgOpLoc, linalgOp.getMatchingIndexingMap(&operand), ivs);
304 builder.
create<memref::LoadOp>(linalgOpLoc, operand.get(), indices);
305 indexedValues.push_back(load);
318 template <
typename LinalgOpTy>
319 struct LinalgOpPartialReductionInterface
320 :
public PartialReductionOpInterface::ExternalModel<
321 LinalgOpPartialReductionInterface<LinalgOpTy>, LinalgOpTy> {
322 FailureOr<SmallVector<Value>> generateInitialTensorForPartialReduction(
325 auto linalgOp = cast<LinalgOp>(op);
328 if (linalgOp.hasPureBufferSemantics())
329 return op->
emitOpError(
"expected operation to have tensor semantics");
332 for (
int initIdx = 0, e = linalgOp.getNumDpsInits(); initIdx < e;
339 combinerOps.size() != 1)
340 return op->
emitOpError(
"Failed to anaysis the reduction operation.");
344 if (!identity.has_value())
346 "Failed to get an identity value for the reduction operation.");
349 linalgOp.getShape(linalgOp.getDpsInitOperand(initIdx));
355 int64_t currReductionDims = 0;
357 reductionDims.end());
359 llvm::seq<int64_t>(0, oldShape.size() + reductionDims.size())) {
360 if (reductionDimsSet.contains(idx)) {
365 int64_t oldIdx = idx - currReductionDims;
366 int64_t dim = oldShape[oldIdx];
367 newOutputShape.push_back(dim);
368 if (ShapedType::isDynamic(dim))
369 dynamicDims.push_back(b.
create<tensor::DimOp>(
370 loc, linalgOp.getDpsInitOperand(initIdx)->get(), oldIdx));
374 linalgOp.getRegionOutputArgs()[initIdx].getType(), dynamicDims);
375 Value constantOp = b.
create<arith::ConstantOp>(loc, *identity);
376 auto identityTensor =
377 b.
create<linalg::FillOp>(loc, constantOp, emptyTensor);
378 inits.push_back(identityTensor.getResult(0));
384 FailureOr<TilingResult>
390 auto linalgOp = cast<LinalgOp>(op);
395 newInitMaps.reserve(linalgOp.getNumDpsInits());
396 for (
int idx : llvm::seq<int>(0, linalgOp.getNumDpsInits())) {
400 linalgOp.getMatchingIndexingMap(linalgOp.getDpsInitOperand(idx));
401 for (
int redPos : reductionDims) {
405 newInitMaps.push_back(newMap);
410 b, loc, linalgOp, linalgOp.getDpsInputs(), offsets, sizes, {},
true);
414 for (
auto [valueMap, valueToTile] : llvm::zip_equal(newInitMaps, init)) {
415 int64_t initRank = valueMap.getNumResults();
419 for (
AffineExpr dimExpr : valueMap.getResults()) {
420 auto dim = cast<AffineDimExpr>(dimExpr);
421 initSizes.push_back(sizes[dim.getPosition()]);
424 auto extractSlice = b.
create<tensor::ExtractSliceOp>(
425 loc, valueToTile, initOffset, initSizes, initStride);
426 tiledInits.push_back(extractSlice);
432 for (
int idx : llvm::seq<int>(0, linalgOp.getNumDpsInits())) {
435 OpOperand *initOperand = linalgOp.getDpsInitOperand(idx);
436 int64_t mapIdx = linalgOp.getIndexingMapIndex(initOperand);
437 newMaps[mapIdx] = newInitMaps[idx];
442 linalgOp.getIteratorTypesArray();
443 for (
int dim : reductionDims)
444 newIteratorTypes[dim] = utils::IteratorType::parallel;
449 tiledInits, newMaps, newIteratorTypes);
452 genericOp.getRegion().begin(), mapping);
454 {genericOp.getOperation()},
455 llvm::map_to_vector(genericOp->getResults(),
462 auto linalgOp = cast<LinalgOp>(op);
464 reductionDims.end());
465 auto reduction = b.
create<linalg::ReduceOp>(
466 loc, partialReduce, linalgOp.getDpsInits(), reductionDimsInt64,
468 int64_t numInits = linalgOp.getNumDpsInits();
470 for (
int idx : llvm::seq<int>(0, numInits)) {
476 clonedReductionOp->
setOperand(0, inputs[idx]);
477 clonedReductionOp->
setOperand(1, inputs[numInits + idx]);
479 yieldedValues.push_back(clonedReductionOp->
getResult(0));
481 b.create<linalg::YieldOp>(loc, yieldedValues);
484 {reduction.getOperation()},
485 llvm::map_to_vector(reduction->getResults(),
492 template <
typename OpType>
494 OpType::template attachInterface<LinalgOpTilingInterface<OpType>>(*ctx);
495 OpType::template attachInterface<LinalgOpPartialReductionInterface<OpType>>(
500 template <
typename... OpTypes>
502 (registerOne<OpTypes>(ctx), ...);
510 registerOne<linalg::GenericOp>(ctx);
512 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
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)
AffineExpr getAffineDimExpr(unsigned position)
MLIRContext * getContext() const
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
void addExtension(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.
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...
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.
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...
void dispatchIndexOpFoldResults(ArrayRef< OpFoldResult > ofrs, SmallVectorImpl< Value > &dynamicVec, SmallVectorImpl< int64_t > &staticVec)
Helper function to dispatch multiple OpFoldResults according to the behavior of dispatchIndexOpFoldRe...
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
Container for the result of merge operation of tiling.
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