1661 lines · cpp
1//===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements utilities for the Linalg dialect.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Linalg/Utils/Utils.h"14 15#include "mlir/Analysis/SliceAnalysis.h"16#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"17#include "mlir/Dialect/Affine/IR/AffineOps.h"18#include "mlir/Dialect/Affine/LoopUtils.h"19#include "mlir/Dialect/Arith/IR/Arith.h"20#include "mlir/Dialect/Arith/Utils/Utils.h"21#include "mlir/Dialect/Func/IR/FuncOps.h"22#include "mlir/Dialect/Linalg/IR/Linalg.h"23#include "mlir/Dialect/MemRef/IR/MemRef.h"24#include "mlir/Dialect/SCF/IR/SCF.h"25#include "mlir/Dialect/Tensor/IR/Tensor.h"26#include "mlir/Dialect/Tensor/Utils/Utils.h"27#include "mlir/Dialect/Utils/IndexingUtils.h"28#include "mlir/Dialect/Utils/StaticValueUtils.h"29#include "mlir/IR/AffineExpr.h"30#include "mlir/IR/AffineExprVisitor.h"31#include "mlir/IR/AffineMap.h"32#include "mlir/IR/Matchers.h"33#include "llvm/ADT/TypeSwitch.h"34#include "llvm/Support/Debug.h"35#include <optional>36 37#define DEBUG_TYPE "linalg-utils"38 39using namespace mlir;40using namespace presburger;41using namespace mlir::affine;42using namespace mlir::linalg;43using namespace mlir::scf;44 45namespace {46 47// Helper visitor to determine whether an AffineExpr is tiled.48// This is achieved by traversing every AffineDimExpr with position `pos` and49// checking whether the corresponding `tileSizes[pos]` is non-zero.50// This also enforces only positive coefficients occur in multiplications.51//52// Example:53// `d0 + 2 * d1 + d3` is tiled by [0, 0, 0, 2] but not by [0, 0, 2, 0]54//55struct TileCheck : public AffineExprVisitor<TileCheck> {56 TileCheck(ArrayRef<OpFoldResult> tileSizes) : tileSizes(tileSizes) {}57 58 void visitDimExpr(AffineDimExpr expr) {59 isTiled |= !isZeroInteger(tileSizes[expr.getPosition()]);60 }61 void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {62 visit(expr.getLHS());63 visit(expr.getRHS());64 if (expr.getKind() == mlir::AffineExprKind::Mul)65 assert(cast<AffineConstantExpr>(expr.getRHS()).getValue() > 0 &&66 "nonpositive multiplying coefficient");67 }68 bool isTiled = false;69 ArrayRef<OpFoldResult> tileSizes;70};71 72} // namespace73 74static bool isTiled(AffineExpr expr, ArrayRef<OpFoldResult> tileSizes) {75 if (!expr)76 return false;77 TileCheck t(tileSizes);78 t.visit(expr);79 return t.isTiled;80}81 82// Checks whether the `map varies with respect to a non-zero `tileSize`.83static bool isTiled(AffineMap map, ArrayRef<OpFoldResult> tileSizes) {84 if (!map)85 return false;86 for (unsigned r = 0; r < map.getNumResults(); ++r)87 if (isTiled(map.getResult(r), tileSizes))88 return true;89 return false;90}91 92std::optional<RegionMatcher::BinaryOpKind>93RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {94 auto ®ion = op.getRegion();95 if (!region.hasOneBlock())96 return std::nullopt;97 98 Block &block = region.front();99 if (block.getNumArguments() != 2 ||100 !block.getArgument(0).getType().isSignlessIntOrFloat() ||101 !block.getArgument(1).getType().isSignlessIntOrFloat())102 return std::nullopt;103 104 auto &ops = block.getOperations();105 if (!llvm::hasSingleElement(block.without_terminator()))106 return std::nullopt;107 108 using mlir::matchers::m_Val;109 auto a = m_Val(block.getArgument(0));110 auto b = m_Val(block.getArgument(1));111 112 auto addPattern = m_Op<linalg::YieldOp>(m_Op<arith::AddIOp>(a, b));113 if (addPattern.match(&ops.back()))114 return BinaryOpKind::IAdd;115 116 return std::nullopt;117}118 119/// Explicit instantiation of loop nest generator for different loop types.120template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;121template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;122template struct mlir::linalg::GenerateLoopNest<AffineForOp>;123 124/// Given a list of subview ranges, extract individual values for lower, upper125/// bounds and steps and put them into the corresponding vectors.126static void unpackRanges(OpBuilder &builder, Location loc,127 ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,128 SmallVectorImpl<Value> &ubs,129 SmallVectorImpl<Value> &steps) {130 for (Range range : ranges) {131 lbs.emplace_back(132 getValueOrCreateConstantIndexOp(builder, loc, range.offset));133 ubs.emplace_back(getValueOrCreateConstantIndexOp(builder, loc, range.size));134 steps.emplace_back(135 getValueOrCreateConstantIndexOp(builder, loc, range.stride));136 }137}138 139//===----------------------------------------------------------------------===//140// General utilities141//===----------------------------------------------------------------------===//142//143/// The permutation can be obtained from two permutations:144/// a) Compute the permutation vector to move the last `numPackedDims` into145/// the `innerPosDims` of a shape of rank `rank`.146/// b) Compute the permutation vector to move outer dims if the147/// `outerPerm` parameter is not empty.148/// Apply (b) permutation on (a) permutation to get the final permutation.149static SmallVector<int64_t>150computePackUnPackPerm(int64_t rank, ArrayRef<int64_t> &innerDimsPos,151 ArrayRef<int64_t> &outerPerm,152 PackingMetadata &packingMetadata) {153 int64_t numPackedDims = innerDimsPos.size();154 auto lastDims =155 llvm::to_vector(llvm::seq<int64_t>(rank - numPackedDims, rank));156 packingMetadata = computePackingMetadata(rank, innerDimsPos);157 SmallVector<int64_t> innerPositionsPerm =158 computePermutationVector(rank, lastDims, packingMetadata.insertPositions);159 160 SmallVector<int64_t> outerPos = packingMetadata.outerPositions;161 if (!outerPerm.empty())162 applyPermutationToVector(outerPos, outerPerm);163 SmallVector<int64_t> outerPositionPerm =164 computePermutationVector(rank, packingMetadata.outerPositions, outerPos);165 166 SmallVector<int64_t> packInverseDestPermutation = innerPositionsPerm;167 applyPermutationToVector(packInverseDestPermutation, outerPositionPerm);168 return packInverseDestPermutation;169}170 171namespace mlir {172namespace linalg {173 174SmallVector<int64_t> getPackInverseDestPerm(PackOp packOp,175 PackingMetadata &metadata) {176 177 int64_t packedRank = packOp.getDestType().getRank();178 ArrayRef<int64_t> innerDimPos = packOp.getInnerDimsPos();179 ArrayRef<int64_t> outerPerm = packOp.getOuterDimsPerm();180 SmallVector<int64_t> packInvDestPerm =181 computePackUnPackPerm(packedRank, innerDimPos, outerPerm, metadata);182 return packInvDestPerm;183}184 185SmallVector<int64_t> getUnPackInverseSrcPerm(UnPackOp unpackOp,186 PackingMetadata &metadata) {187 int64_t packedRank = unpackOp.getSourceType().getRank();188 ArrayRef<int64_t> innerDimPos = unpackOp.getInnerDimsPos();189 ArrayRef<int64_t> outerPerm = unpackOp.getOuterDimsPerm();190 SmallVector<int64_t> unpackInvSrcPerm =191 computePackUnPackPerm(packedRank, innerDimPos, outerPerm, metadata);192 return unpackInvSrcPerm;193}194 195bool allIndexingsAreProjectedPermutation(LinalgOp op) {196 return llvm::all_of(op.getIndexingMapsArray(), [](AffineMap m) {197 return m.isProjectedPermutation(/*allowZeroInResults=*/true);198 });199}200 201bool hasOnlyScalarElementwiseOp(Region &r) {202 if (!r.hasOneBlock())203 return false;204 for (Operation &op : r.front()) {205 if (!(isa<arith::ConstantOp, func::ConstantOp, tensor::ExtractOp,206 linalg::YieldOp, linalg::IndexOp, AffineApplyOp>(op) ||207 OpTrait::hasElementwiseMappableTraits(&op)) ||208 llvm::any_of(op.getResultTypes(),209 [](Type type) { return !type.isIntOrIndexOrFloat(); }))210 return false;211 }212 return true;213}214 215bool isElementwise(LinalgOp op) {216 if (op.getNumLoops() != op.getNumParallelLoops())217 return false;218 219 if (!allIndexingsAreProjectedPermutation(op))220 return false;221 222 // TODO: relax the restrictions on indexing map.223 for (OpOperand &opOperand : op.getDpsInitsMutable()) {224 if (!op.getMatchingIndexingMap(&opOperand).isPermutation())225 return false;226 }227 return hasOnlyScalarElementwiseOp(op->getRegion(0));228}229 230bool isParallelIterator(utils::IteratorType iteratorType) {231 return iteratorType == utils::IteratorType::parallel;232}233 234bool isReductionIterator(utils::IteratorType iteratorType) {235 return iteratorType == utils::IteratorType::reduction;236}237 238//===----------------------------------------------------------------------===//239// Convolution matcher utilities240//===----------------------------------------------------------------------===//241 242/// Returns the BlockArgument that leads to `val`, if any. Traverses optional243/// ext* ops.244static BlockArgument getBlockArgumentWithOptionalExtOps(Value val) {245 BlockArgument blockArg = dyn_cast<BlockArgument>(val);246 if ((blockArg))247 return blockArg;248 249 Operation *defOp = val.getDefiningOp();250 if (!dyn_cast_if_present<arith::ExtFOp>(defOp) &&251 !dyn_cast_if_present<arith::ExtSIOp>(defOp) &&252 !dyn_cast_if_present<arith::ExtUIOp>(defOp)) {253 return nullptr;254 }255 return dyn_cast<BlockArgument>(defOp->getOperand(0));256}257 258/// Utility to match block body for convolution ops.259/// The body is thus expected to yield :-260/// %out + (%lhs * %rhs)261/// where: %lhs, %rhs and %out are block arguments and262/// %lhs and %rhs can have optional upcast operation.263static bool bodyMatcherForConvolutionOps(Value yieldVal, Block *body) {264 Operation *addOp = yieldVal.getDefiningOp();265 if (!isa_and_present<arith::AddIOp, arith::AddFOp>(addOp))266 return false;267 268 Operation *mulOp = addOp->getOperand(1).getDefiningOp();269 if (!isa_and_present<arith::MulIOp, arith::MulFOp>(mulOp))270 return false;271 272 BlockArgument lhsBlockArg =273 getBlockArgumentWithOptionalExtOps(mulOp->getOperand(0));274 BlockArgument rhsBlockArg =275 getBlockArgumentWithOptionalExtOps(mulOp->getOperand(1));276 BlockArgument outBlockArg =277 getBlockArgumentWithOptionalExtOps(addOp->getOperand(0));278 if (!lhsBlockArg || !rhsBlockArg || !outBlockArg ||279 lhsBlockArg.getOwner() != body || rhsBlockArg.getOwner() != body ||280 outBlockArg.getOwner() != body || lhsBlockArg.getArgNumber() != 0 ||281 rhsBlockArg.getArgNumber() != 1 || outBlockArg.getArgNumber() != 2)282 return false;283 return true;284}285 286/// Utility to match block body for linalg.pool* ops.287template <typename... OpTypes>288static bool bodyMatcherForPoolOps(Value yieldVal, Block *body) {289 Operation *defOp = yieldVal.getDefiningOp();290 if (!(isa_and_present<OpTypes>(defOp) || ...))291 return false;292 293 BlockArgument lhsArg =294 getBlockArgumentWithOptionalExtOps(defOp->getOperand(0));295 BlockArgument rhsArg =296 getBlockArgumentWithOptionalExtOps(defOp->getOperand(1));297 if (!lhsArg || !rhsArg || lhsArg.getOwner() != body ||298 rhsArg.getOwner() != body || lhsArg.getArgNumber() != 2 ||299 rhsArg.getArgNumber() != 0)300 return false;301 return true;302}303 304static bool bodyMatcherForMaxSignedPoolOps(Value yieldVal, Block *body) {305 return bodyMatcherForPoolOps<arith::MaximumFOp, arith::MaxSIOp>(yieldVal,306 body);307}308 309// max_unsigned ops should not allow float data type.310// TODO(#164800): Retire OPDSL logic.311static bool bodyMatcherForMaxUnsignedPoolOps(Value yieldVal, Block *body) {312 return bodyMatcherForPoolOps<arith::MaximumFOp, arith::MaxUIOp>(yieldVal,313 body);314}315 316static bool bodyMatcherForMinSignedPoolOps(Value yieldVal, Block *body) {317 return bodyMatcherForPoolOps<arith::MinimumFOp, arith::MinSIOp>(yieldVal,318 body);319}320 321// min_unsigned ops should not allow float data type.322// TODO(#164800): Retire OPDSL logic.323static bool bodyMatcherForMinUnsignedPoolOps(Value yieldVal, Block *body) {324 return bodyMatcherForPoolOps<arith::MinimumFOp, arith::MinUIOp>(yieldVal,325 body);326}327 328static bool bodyMatcherForSumPoolOps(Value yieldVal, Block *body) {329 return bodyMatcherForPoolOps<arith::AddIOp, arith::AddFOp>(yieldVal, body);330}331 332static AffineExpr getAffineMapDim(ArrayAttr indexingMaps, uint32_t mapIndex,333 uint32_t dimIndex) {334 auto affineMap = cast<AffineMapAttr>(indexingMaps[mapIndex]).getValue();335 if (dimIndex < affineMap.getNumResults())336 return affineMap.getResult(dimIndex);337 return nullptr;338}339 340/// Check if `expr` is either:341/// - a dimension expr alone (implying multiplication by 1), or342/// - a multiplication of dimension expr by any positive constant != 1343/// In both cases we will capture the dimension expression into `dim` and344/// return the constant multiplier. Returns -1 in case of a match failure.345static int64_t isDimTimesConstantOrDimOnly(AffineExpr expr, AffineExpr &dim) {346 if ((dim = dyn_cast<AffineDimExpr>(expr)))347 return 1;348 349 auto mulExpr = dyn_cast<AffineBinaryOpExpr>(expr);350 if (!mulExpr || mulExpr.getKind() != AffineExprKind::Mul)351 return -1;352 353 AffineExpr lhs = mulExpr.getLHS();354 AffineExpr rhs = mulExpr.getRHS();355 356 AffineConstantExpr cst = nullptr;357 if (((dim = dyn_cast<AffineDimExpr>(lhs)) &&358 (cst = dyn_cast<AffineConstantExpr>(rhs))) ||359 ((dim = dyn_cast<AffineDimExpr>(rhs)) &&360 (cst = dyn_cast<AffineConstantExpr>(lhs))))361 return cst.getValue();362 return -1;363}364 365/// Given an array of AffineMaps `indexingMaps` verify the following366/// commutatively:-367/// indexingMaps[0].getResult(iDim) ==368/// indexingMaps[1].getResult(fDim) * <c0> +369/// indexingMaps[n-1].getResult(oDim) * <c1>370/// where,371/// - c0 and c1 can be any constant,372/// - n is the size of the indexingMaps' array,373/// - 0, 1 and n-1 are input, filter and output map indices respectively,374/// - iDim, fDim and oDim are the input, filter and output dimension375/// indices in their respective indexing maps376/// Example:377/// #inputMap = affine_map<(d0, d1, d2, d3, d4, d5, d6)378/// -> (d0, d1 * 2 + d4 * 3, d2 + d5, d6)>379/// #filterMap = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>380/// #outputMap = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>381///382/// Here,383/// #inputMap[1] = #outputMap[1] * 2 + #filterMap[0] * 3384/// Therefore,385/// matchConvDimAddExprPattern(indexingMaps, 1, 0, 1, dilation, stride)386/// would return true and update dilation = 3 and stride = 2387static bool matchConvDimAddExprPattern(ArrayAttr indexingMaps, unsigned iDim,388 unsigned fDim, unsigned oDim,389 int64_t &dilation, int64_t &stride) {390 unsigned inputMapIdx = 0, filterMapIdx = 1,391 outputMapIdx = indexingMaps.size() - 1;392 AffineExpr inpExpr = getAffineMapDim(indexingMaps, inputMapIdx, iDim);393 auto addExpr = dyn_cast_or_null<AffineBinaryOpExpr>(inpExpr);394 if (!addExpr || addExpr.getKind() != AffineExprKind::Add)395 return false;396 397 AffineExpr dim0, dim1;398 int64_t c0 = isDimTimesConstantOrDimOnly(addExpr.getLHS(), dim0);399 int64_t c1 = isDimTimesConstantOrDimOnly(addExpr.getRHS(), dim1);400 401 if (c0 == -1 || c1 == -1)402 return false;403 // Pattern matched with dims and constants extracted.404 AffineExpr fExpr = getAffineMapDim(indexingMaps, filterMapIdx, fDim);405 AffineExpr oExpr = getAffineMapDim(indexingMaps, outputMapIdx, oDim);406 if (dim0 == fExpr && dim1 == oExpr) {407 dilation = c0;408 stride = c1;409 return true;410 }411 if (dim1 == fExpr && dim0 == oExpr) {412 dilation = c1;413 stride = c0;414 return true;415 }416 return false;417}418 419/// Returns true if the given indexing maps matches with the expected indexing420/// maps.421static bool convLayoutMatches(ArrayRef<ArrayRef<AffineExpr>> mapListExpected,422 ArrayAttr indexingMaps, MLIRContext *context) {423 SmallVector<AffineMap, 4> expectedIndexingMaps =424 AffineMap::inferFromExprList(mapListExpected, context);425 return indexingMaps ==426 ArrayAttr::get(427 context, llvm::to_vector<4>(llvm::map_range(428 expectedIndexingMaps, [&](AffineMap m) -> Attribute {429 return AffineMapAttr::get(m);430 })));431}432 433/// Enum representing pooling operation types used by ConvMatcherBuilder.434enum class PoolingType {435 None,436 MaxSigned,437 MaxUnsigned,438 MinSigned,439 MinUnsigned,440 Sum441};442 443/// Helper class for building convolution op matchers with minimal boilerplate.444/// Reduces repetitive code across Conv1D/2D/3D and Depthwise variants as well445/// as Pooling ops.446///447/// Usage: Create an instance with the op, spatial rank, and output pointers for448/// extracted dilations/strides. Then chain matchStride() calls for each spatial449/// dimension, followed by matchMaps() to verify indexing maps, and finally450/// matchBody() to verify the operation body pattern.451///452/// The `matched` flag starts as `true` and is set to `false` if any match step453/// fails. This allows chaining multiple match calls; once any match fails, all454/// subsequent calls become no-ops and the final result is `false`.455///456/// The `dilations` and `strides` pointers are output parameters that get457/// populated with the extracted dilation and stride values from the operation's458/// indexing maps during matchStride() calls. These values are initially set to459/// 1 for each spatial dimension and updated as patterns are matched.460class ConvMatcherBuilder {461 LinalgOp op;462 MLIRContext *ctx;463 SmallVector<int64_t> *dilations, *strides;464 ArrayAttr indexingMaps;465 PoolingType poolingType;466 bool matched = true;467 468public:469 ConvMatcherBuilder(LinalgOp op, unsigned spatialRank, SmallVector<int64_t> *d,470 SmallVector<int64_t> *s,471 PoolingType poolingType = PoolingType::None)472 : op(op), ctx(op->getContext()), dilations(d), strides(s),473 indexingMaps(op.getIndexingMaps()), poolingType(poolingType) {474 *dilations = SmallVector<int64_t>(spatialRank, 1);475 *strides = SmallVector<int64_t>(spatialRank, 1);476 }477 478 /// Get affine dimension expression for dimension `i`.479 AffineExpr dim(unsigned i) { return getAffineDimExpr(i, ctx); }480 481 /// Build strided expression: base * stride[idx] + kernel * dilation[idx].482 AffineExpr strided(AffineExpr base, AffineExpr kernel, unsigned idx) {483 return base * (*strides)[idx] + kernel * (*dilations)[idx];484 }485 486 /// Match stride/dilation pattern for a spatial dimension.487 /// Returns *this for method chaining.488 ConvMatcherBuilder &matchStride(unsigned iDim, unsigned fDim, unsigned oDim,489 unsigned idx) {490 if (matched) {491 matched &= matchConvDimAddExprPattern(indexingMaps, iDim, fDim, oDim,492 (*dilations)[idx], (*strides)[idx]);493 }494 return *this;495 }496 497 /// Match expected indexing maps layout. Returns *this for method chaining.498 ConvMatcherBuilder &matchMaps(ArrayRef<ArrayRef<AffineExpr>> maps) {499 if (matched)500 matched &= convLayoutMatches(maps, indexingMaps, ctx);501 return *this;502 }503 504 /// Match body pattern. This should be called last.505 bool matchBody() {506 if (!matched)507 return false;508 Block *body = op.getBlock();509 auto yieldOp = cast<linalg::YieldOp>(body->getTerminator());510 switch (poolingType) {511 case PoolingType::None:512 return bodyMatcherForConvolutionOps(yieldOp.getOperand(0), body);513 case PoolingType::MaxSigned:514 return bodyMatcherForMaxSignedPoolOps(yieldOp.getOperand(0), body);515 case PoolingType::MaxUnsigned:516 return bodyMatcherForMaxUnsignedPoolOps(yieldOp.getOperand(0), body);517 case PoolingType::MinSigned:518 return bodyMatcherForMinSignedPoolOps(yieldOp.getOperand(0), body);519 case PoolingType::MinUnsigned:520 return bodyMatcherForMinUnsignedPoolOps(yieldOp.getOperand(0), body);521 case PoolingType::Sum:522 return bodyMatcherForSumPoolOps(yieldOp.getOperand(0), body);523 }524 return false;525 }526};527 528//===----------------------------------------------------------------------===//529// Matchers for specific convolution operation.530//===----------------------------------------------------------------------===//531 532// #inputMap = affine_map<(W, w) -> (W + w)>533// #filterMap = affine_map<(W, w) -> (w)>534// #outputMap = affine_map<(W, w) -> (W)>535template <>536bool isaConvolutionOpOfType<linalg::Conv1DOp>(LinalgOp op,537 SmallVector<int64_t> *dilations,538 SmallVector<int64_t> *strides) {539 if (isa<linalg::Conv1DOp>(op))540 return true;541 542 assert(isaConvolutionOpInterface(op) &&543 "expected op to implement ConvolutionOpInterface");544 545 ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);546 AffineExpr W = m.dim(0);547 AffineExpr w = m.dim(1);548 549 return m.matchStride(/*iDim=*/0, /*fDim=*/0, /*oDim=*/0, /*idx=*/0)550 .matchMaps({/*inputMap=*/{m.strided(W, w, 0)},551 /*filterMap=*/{w},552 /*outputMap=*/{W}})553 .matchBody();554}555 556// #inputMap = affine_map<(N, W, F, w, c) -> (N, W + w, c)>557// #filterMap = affine_map<(N, W, F, w, c) -> (w, c, F)>558// #outputMap = affine_map<(N, W, F, w, c) -> (N, W, F)>559template <>560bool isaConvolutionOpOfType<linalg::Conv1DNwcWcfOp>(561 LinalgOp op, SmallVector<int64_t> *dilations,562 SmallVector<int64_t> *strides) {563 if (isa<linalg::Conv1DNwcWcfOp>(op))564 return true;565 566 assert(isaConvolutionOpInterface(op) &&567 "expected op to implement ConvolutionOpInterface");568 569 ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);570 AffineExpr N = m.dim(0);571 AffineExpr W = m.dim(1);572 AffineExpr F = m.dim(2);573 AffineExpr w = m.dim(3);574 AffineExpr c = m.dim(4);575 576 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)577 .matchMaps({/*inputMap=*/{N, m.strided(W, w, 0), c},578 /*filterMap=*/{w, c, F},579 /*outputMap=*/{N, W, F}})580 .matchBody();581}582 583// #inputMap = affine_map<(N, F, W, c, w) -> (N, c, W + w)>584// #filterMap = affine_map<(N, F, W, c, w) -> (F, c, w)>585// #outputMap = affine_map<(N, F, W, c, w) -> (N, F, W)>586template <>587bool isaConvolutionOpOfType<linalg::Conv1DNcwFcwOp>(588 LinalgOp op, SmallVector<int64_t> *dilations,589 SmallVector<int64_t> *strides) {590 if (isa<linalg::Conv1DNcwFcwOp>(op))591 return true;592 593 assert(isaConvolutionOpInterface(op) &&594 "expected op to implement ConvolutionOpInterface");595 596 ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);597 AffineExpr N = m.dim(0);598 AffineExpr F = m.dim(1);599 AffineExpr W = m.dim(2);600 AffineExpr c = m.dim(3);601 AffineExpr w = m.dim(4);602 603 return m.matchStride(/*iDim=*/2, /*fDim=*/2, /*oDim=*/2, /*idx=*/0)604 .matchMaps({/*inputMap=*/{N, c, m.strided(W, w, 0)},605 /*filterMap=*/{F, c, w},606 /*outputMap=*/{N, F, W}})607 .matchBody();608}609 610// #inputMap = affine_map<(H, W, h, w) -> (H + h, W + w)>611// #filterMap = affine_map<(H, W, h, w) -> (h, w)>612// #outputMap = affine_map<(H, W, h, w) -> (H, W)>613template <>614bool isaConvolutionOpOfType<linalg::Conv2DOp>(LinalgOp op,615 SmallVector<int64_t> *dilations,616 SmallVector<int64_t> *strides) {617 if (isa<linalg::Conv2DOp>(op))618 return true;619 620 assert(isaConvolutionOpInterface(op) &&621 "expected op to implement ConvolutionOpInterface");622 623 ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);624 AffineExpr H = m.dim(0);625 AffineExpr W = m.dim(1);626 AffineExpr h = m.dim(2);627 AffineExpr w = m.dim(3);628 629 return m.matchStride(/*iDim=*/0, /*fDim=*/0, /*oDim=*/0, /*idx=*/0)630 .matchStride(/*iDim=*/1, /*fDim=*/1, /*oDim=*/1, /*idx=*/1)631 .matchMaps({/*inputMap=*/{m.strided(H, h, 0), m.strided(W, w, 1)},632 /*filterMap=*/{h, w},633 /*outputMap=*/{H, W}})634 .matchBody();635}636 637// #inputMap = affine_map<(D, H, W, d, h, w) -> (D + d, H + h, W + w)>638// #filterMap = affine_map<(D, H, W, d, h, w) -> (d, h, w)>639// #outputMap = affine_map<(D, H, W, d, h, w) -> (D, H, W)>640template <>641bool isaConvolutionOpOfType<linalg::Conv3DOp>(LinalgOp op,642 SmallVector<int64_t> *dilations,643 SmallVector<int64_t> *strides) {644 if (isa<linalg::Conv3DOp>(op))645 return true;646 647 assert(isaConvolutionOpInterface(op) &&648 "expected op to implement ConvolutionOpInterface");649 650 ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);651 AffineExpr D = m.dim(0);652 AffineExpr H = m.dim(1);653 AffineExpr W = m.dim(2);654 AffineExpr d = m.dim(3);655 AffineExpr h = m.dim(4);656 AffineExpr w = m.dim(5);657 658 return m.matchStride(/*iDim=*/0, /*fDim=*/0, /*oDim=*/0, /*idx=*/0)659 .matchStride(/*iDim=*/1, /*fDim=*/1, /*oDim=*/1, /*idx=*/1)660 .matchStride(/*iDim=*/2, /*fDim=*/2, /*oDim=*/2, /*idx=*/2)661 .matchMaps({/*inputMap=*/{m.strided(D, d, 0), m.strided(H, h, 1),662 m.strided(W, w, 2)},663 /*filterMap=*/{d, h, w},664 /*outputMap=*/{D, H, W}})665 .matchBody();666}667 668// #inputMap = affine_map<(N, W, C, w) -> (N, C, W + w)>669// #filterMap = affine_map<(N, W, C, w) -> (C, w)>670// #outputMap = affine_map<(N, W, C, w) -> (N, C, W)>671template <>672bool isaConvolutionOpOfType<linalg::DepthwiseConv1DNcwCwOp>(673 LinalgOp op, SmallVector<int64_t> *dilations,674 SmallVector<int64_t> *strides) {675 if (isa<linalg::DepthwiseConv1DNcwCwOp>(op))676 return true;677 678 assert(isaConvolutionOpInterface(op) &&679 "expected op to implement ConvolutionOpInterface");680 681 ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);682 AffineExpr N = m.dim(0);683 AffineExpr W = m.dim(1);684 AffineExpr C = m.dim(2);685 AffineExpr w = m.dim(3);686 687 return m.matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/0)688 .matchMaps({/*inputMap=*/{N, C, m.strided(W, w, 0)},689 /*filterMap=*/{C, w},690 /*outputMap=*/{N, C, W}})691 .matchBody();692}693 694// #inputMap = affine_map<(N, W, C, w) -> (N, W + w, C)>695// #filterMap = affine_map<(N, W, C, w) -> (w, C)>696// #outputMap = affine_map<(N, W, C, w) -> (N, W, C)>697template <>698bool isaConvolutionOpOfType<linalg::DepthwiseConv1DNwcWcOp>(699 LinalgOp op, SmallVector<int64_t> *dilations,700 SmallVector<int64_t> *strides) {701 if (isa<linalg::DepthwiseConv1DNwcWcOp>(op))702 return true;703 704 assert(isaConvolutionOpInterface(op) &&705 "expected op to implement ConvolutionOpInterface");706 707 ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);708 AffineExpr N = m.dim(0);709 AffineExpr W = m.dim(1);710 AffineExpr C = m.dim(2);711 AffineExpr w = m.dim(3);712 713 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)714 .matchMaps({/*inputMap=*/{N, m.strided(W, w, 0), C},715 /*filterMap=*/{w, C},716 /*outputMap=*/{N, W, C}})717 .matchBody();718}719 720// #inputMap = affine_map<(N, W, C, CM, w) -> (N, W + w, C)>721// #filterMap = affine_map<(N, W, C, CM, w) -> (w, C, CM)>722// #outputMap = affine_map<(N, W, C, CM, w) -> (N, W, C, CM)>723template <>724bool isaConvolutionOpOfType<linalg::DepthwiseConv1DNwcWcmOp>(725 LinalgOp op, SmallVector<int64_t> *dilations,726 SmallVector<int64_t> *strides) {727 if (isa<linalg::DepthwiseConv1DNwcWcmOp>(op))728 return true;729 730 assert(isaConvolutionOpInterface(op) &&731 "expected op to implement ConvolutionOpInterface");732 733 ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);734 AffineExpr N = m.dim(0);735 AffineExpr W = m.dim(1);736 AffineExpr C = m.dim(2);737 AffineExpr CM = m.dim(3);738 AffineExpr w = m.dim(4);739 740 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)741 .matchMaps({/*inputMap=*/{N, m.strided(W, w, 0), C},742 /*filterMap=*/{w, C, CM},743 /*outputMap=*/{N, W, C, CM}})744 .matchBody();745}746 747// #inputMap = affine_map<(N, H, W, C, h, w) -> (N, C, H + h, W + w)>748// #filterMap = affine_map<(N, H, W, C, h, w) -> (C, h, w)>749// #outputMap = affine_map<(N, H, W, C, h, w) -> (N, C, H, W)>750template <>751bool isaConvolutionOpOfType<linalg::DepthwiseConv2DNchwChwOp>(752 LinalgOp op, SmallVector<int64_t> *dilations,753 SmallVector<int64_t> *strides) {754 if (isa<linalg::DepthwiseConv2DNchwChwOp>(op))755 return true;756 757 assert(isaConvolutionOpInterface(op) &&758 "expected op to implement ConvolutionOpInterface");759 760 ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);761 AffineExpr N = m.dim(0);762 AffineExpr H = m.dim(1);763 AffineExpr W = m.dim(2);764 AffineExpr C = m.dim(3);765 AffineExpr h = m.dim(4);766 AffineExpr w = m.dim(5);767 768 return m.matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/0)769 .matchStride(/*iDim=*/3, /*fDim=*/2, /*oDim=*/3, /*idx=*/1)770 .matchMaps({/*inputMap=*/{N, C, m.strided(H, h, 0), m.strided(W, w, 1)},771 /*filterMap=*/{C, h, w},772 /*outputMap=*/{N, C, H, W}})773 .matchBody();774}775 776// #inputMap = affine_map<(N, D, H, W, CM, d, h, w, C)777// -> (N, D + d, H + h, W + w, C)>778// #filterMap = affine_map<(N, D, H, W, CM, d, h, w, C)779// -> (d, h, w, C, CM)>780// #outputMap = affine_map<(N, D, H, W, CM, d, h, w, C)781// -> (N, D, H, W, C, CM)>782template <>783bool isaConvolutionOpOfType<linalg::DepthwiseConv3DNdhwcDhwcmOp>(784 LinalgOp op, SmallVector<int64_t> *dilations,785 SmallVector<int64_t> *strides) {786 if (isa<linalg::DepthwiseConv3DNdhwcDhwcmOp>(op))787 return true;788 789 assert(isaConvolutionOpInterface(op) &&790 "expected op to implement ConvolutionOpInterface");791 792 ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);793 AffineExpr N = m.dim(0);794 AffineExpr D = m.dim(1);795 AffineExpr H = m.dim(2);796 AffineExpr W = m.dim(3);797 AffineExpr CM = m.dim(4);798 AffineExpr d = m.dim(5);799 AffineExpr h = m.dim(6);800 AffineExpr w = m.dim(7);801 AffineExpr C = m.dim(8);802 803 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)804 .matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/1)805 .matchStride(/*iDim=*/3, /*fDim=*/2, /*oDim=*/3, /*idx=*/2)806 .matchMaps({/*inputMap=*/{N, m.strided(D, d, 0), m.strided(H, h, 1),807 m.strided(W, w, 2), C},808 /*filterMap=*/{d, h, w, C, CM},809 /*outputMap=*/{N, D, H, W, C, CM}})810 .matchBody();811}812 813// #inputMap = affine_map<(N, H, W, C, h, w) -> (N, H + h, W + w, C)>814// #filterMap = affine_map<(N, H, W, C, h, w) -> (h, w)>815// #outputMap = affine_map<(N, H, W, C, h, w) -> (N, H, W, C)>816template <>817bool isaConvolutionOpOfType<linalg::PoolingNhwcMaxOp>(818 LinalgOp op, SmallVector<int64_t> *dilations,819 SmallVector<int64_t> *strides) {820 if (isa<linalg::PoolingNhwcMaxOp>(op))821 return true;822 823 assert(isaConvolutionOpInterface(op) &&824 "expected op to implement ConvolutionOpInterface");825 826 ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,827 PoolingType::MaxSigned);828 AffineExpr N = m.dim(0);829 AffineExpr H = m.dim(1);830 AffineExpr W = m.dim(2);831 AffineExpr C = m.dim(3);832 AffineExpr h = m.dim(4);833 AffineExpr w = m.dim(5);834 835 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)836 .matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/1)837 .matchMaps({/*inputMap=*/{N, m.strided(H, h, 0), m.strided(W, w, 1), C},838 /*filterMap=*/{h, w},839 /*outputMap=*/{N, H, W, C}})840 .matchBody();841}842 843// #inputMap = affine_map<(N, H, W, C, h, w) -> (N, H + h, W + w, C)>844// #filterMap = affine_map<(N, H, W, C, h, w) -> (h, w)>845// #outputMap = affine_map<(N, H, W, C, h, w) -> (N, H, W, C)>846template <>847bool isaConvolutionOpOfType<linalg::PoolingNhwcMinOp>(848 LinalgOp op, SmallVector<int64_t> *dilations,849 SmallVector<int64_t> *strides) {850 if (isa<linalg::PoolingNhwcMinOp>(op))851 return true;852 853 assert(isaConvolutionOpInterface(op) &&854 "expected op to implement ConvolutionOpInterface");855 856 ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,857 PoolingType::MinSigned);858 AffineExpr N = m.dim(0);859 AffineExpr H = m.dim(1);860 AffineExpr W = m.dim(2);861 AffineExpr C = m.dim(3);862 AffineExpr h = m.dim(4);863 AffineExpr w = m.dim(5);864 865 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)866 .matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/1)867 .matchMaps({/*inputMap=*/{N, m.strided(H, h, 0), m.strided(W, w, 1), C},868 /*filterMap=*/{h, w},869 /*outputMap=*/{N, H, W, C}})870 .matchBody();871}872 873// #inputMap = affine_map<(N, H, W, C, h, w) -> (N, H + h, W + w, C)>874// #filterMap = affine_map<(N, H, W, C, h, w) -> (h, w)>875// #outputMap = affine_map<(N, H, W, C, h, w) -> (N, H, W, C)>876template <>877bool isaConvolutionOpOfType<linalg::PoolingNhwcSumOp>(878 LinalgOp op, SmallVector<int64_t> *dilations,879 SmallVector<int64_t> *strides) {880 if (isa<linalg::PoolingNhwcSumOp>(op))881 return true;882 883 assert(isaConvolutionOpInterface(op) &&884 "expected op to implement ConvolutionOpInterface");885 886 ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,887 PoolingType::Sum);888 AffineExpr N = m.dim(0);889 AffineExpr H = m.dim(1);890 AffineExpr W = m.dim(2);891 AffineExpr C = m.dim(3);892 AffineExpr h = m.dim(4);893 AffineExpr w = m.dim(5);894 895 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)896 .matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/1)897 .matchMaps({/*inputMap=*/{N, m.strided(H, h, 0), m.strided(W, w, 1), C},898 /*filterMap=*/{h, w},899 /*outputMap=*/{N, H, W, C}})900 .matchBody();901}902 903// #inputMap = affine_map<(N, H, W, C, h, w) -> (N, H + h, W + w, C)>904// #filterMap = affine_map<(N, H, W, C, h, w) -> (h, w)>905// #outputMap = affine_map<(N, H, W, C, h, w) -> (N, H, W, C)>906template <>907bool isaConvolutionOpOfType<linalg::PoolingNhwcMaxUnsignedOp>(908 LinalgOp op, SmallVector<int64_t> *dilations,909 SmallVector<int64_t> *strides) {910 if (isa<linalg::PoolingNhwcMaxUnsignedOp>(op))911 return true;912 913 assert(isaConvolutionOpInterface(op) &&914 "expected op to implement ConvolutionOpInterface");915 916 ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,917 PoolingType::MaxUnsigned);918 AffineExpr N = m.dim(0);919 AffineExpr H = m.dim(1);920 AffineExpr W = m.dim(2);921 AffineExpr C = m.dim(3);922 AffineExpr h = m.dim(4);923 AffineExpr w = m.dim(5);924 925 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)926 .matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/1)927 .matchMaps({/*inputMap=*/{N, m.strided(H, h, 0), m.strided(W, w, 1), C},928 /*filterMap=*/{h, w},929 /*outputMap=*/{N, H, W, C}})930 .matchBody();931}932 933// #inputMap = affine_map<(N, H, W, C, h, w) -> (N, H + h, W + w, C)>934// #filterMap = affine_map<(N, H, W, C, h, w) -> (h, w)>935// #outputMap = affine_map<(N, H, W, C, h, w) -> (N, H, W, C)>936template <>937bool isaConvolutionOpOfType<linalg::PoolingNhwcMinUnsignedOp>(938 LinalgOp op, SmallVector<int64_t> *dilations,939 SmallVector<int64_t> *strides) {940 if (isa<linalg::PoolingNhwcMinUnsignedOp>(op))941 return true;942 943 assert(isaConvolutionOpInterface(op) &&944 "expected op to implement ConvolutionOpInterface");945 946 ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,947 PoolingType::MinUnsigned);948 AffineExpr N = m.dim(0);949 AffineExpr H = m.dim(1);950 AffineExpr W = m.dim(2);951 AffineExpr C = m.dim(3);952 AffineExpr h = m.dim(4);953 AffineExpr w = m.dim(5);954 955 return m.matchStride(/*iDim=*/1, /*fDim=*/0, /*oDim=*/1, /*idx=*/0)956 .matchStride(/*iDim=*/2, /*fDim=*/1, /*oDim=*/2, /*idx=*/1)957 .matchMaps({/*inputMap=*/{N, m.strided(H, h, 0), m.strided(W, w, 1), C},958 /*filterMap=*/{h, w},959 /*outputMap=*/{N, H, W, C}})960 .matchBody();961}962 963Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,964 Value source, Value pad, bool nofold,965 ValueRange typeDynDims) {966 // Exit if `source` is not defined by an ExtractSliceOp.967 auto sliceOp = source.getDefiningOp<tensor::ExtractSliceOp>();968 if (!sliceOp)969 return tensor::createPadHighOp(type, source, pad, nofold, loc, b,970 typeDynDims);971 972 // Search the `source` use-def chain for padded LinalgOps.973 Value current = sliceOp.getSource();974 while (current) {975 auto linalgOp = current.getDefiningOp<LinalgOp>();976 if (!linalgOp)977 break;978 OpResult opResult = cast<OpResult>(current);979 current = linalgOp.getDpsInitOperand(opResult.getResultNumber())->get();980 }981 auto padOp = current ? current.getDefiningOp<tensor::PadOp>() : nullptr;982 983 // Exit if the search fails to match a tensor::PadOp at the end of the matched984 // LinalgOp sequence.985 if (!padOp)986 return tensor::createPadHighOp(type, source, pad, nofold, loc, b,987 typeDynDims);988 989 // Exit if the padded result type does not match.990 if (sliceOp.getSource().getType() != type)991 return tensor::createPadHighOp(type, source, pad, nofold, loc, b,992 typeDynDims);993 994 // Exit if the LinalgOps are not high padded.995 if (llvm::any_of(padOp.getMixedLowPad(), [](OpFoldResult ofr) {996 return getConstantIntValue(ofr) != static_cast<int64_t>(0);997 }))998 return tensor::createPadHighOp(type, source, pad, nofold, loc, b,999 typeDynDims);1000 1001 // Exit if `padOpSliceOp`, which defines the slice used by1002 // `padOp`, is rank-reducing.1003 auto padOpSliceOp = padOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();1004 if (!padOpSliceOp ||1005 sliceOp.getMixedSizes().size() != padOpSliceOp.getMixedSizes().size())1006 return tensor::createPadHighOp(type, source, pad, nofold, loc, b,1007 typeDynDims);1008 1009 // Exit if the sizes of the dynamic sizes of `sliceOp` do not match the size1010 // of the slice padded by `padOp`.1011 if (llvm::any_of(1012 llvm::zip(sliceOp.getMixedSizes(), padOpSliceOp.getMixedSizes()),1013 [](std::tuple<OpFoldResult, OpFoldResult> it) {1014 return !isEqualConstantIntOrValue(std::get<0>(it), std::get<1>(it));1015 }))1016 return tensor::createPadHighOp(type, source, pad, nofold, loc, b,1017 typeDynDims);1018 1019 // Exit if the padding values do not match.1020 Attribute padOpPadAttr, padAttr;1021 Value padOpPad = padOp.getConstantPaddingValue();1022 if (!padOpPad || !matchPattern(padOpPad, m_Constant(&padOpPadAttr)) ||1023 !matchPattern(pad, m_Constant(&padAttr)) || padOpPadAttr != padAttr)1024 return tensor::createPadHighOp(type, source, pad, nofold, loc, b,1025 typeDynDims);1026 1027 // Return the padded result if the padding values and sizes match.1028 return sliceOp.getSource();1029}1030 1031GenericOp makeMemRefCopyOp(OpBuilder &b, Location loc, Value from, Value to) {1032 auto memrefTypeTo = cast<MemRefType>(to.getType());1033#ifndef NDEBUG1034 auto memrefTypeFrom = cast<MemRefType>(from.getType());1035 assert(memrefTypeFrom.getRank() == memrefTypeTo.getRank() &&1036 "`from` and `to` memref must have the same rank");1037#endif // NDEBUG1038 1039 AffineMap id =1040 AffineMap::getMultiDimIdentityMap(memrefTypeTo.getRank(), b.getContext());1041 SmallVector<utils::IteratorType> iteratorTypes(memrefTypeTo.getRank(),1042 utils::IteratorType::parallel);1043 return linalg::GenericOp::create(1044 b, loc,1045 /*inputs=*/from,1046 /*outputs=*/to,1047 /*indexingMaps=*/llvm::ArrayRef({id, id}),1048 /*iteratorTypes=*/iteratorTypes,1049 [](OpBuilder &b, Location loc, ValueRange args) {1050 linalg::YieldOp::create(b, loc, args.front());1051 });1052}1053 1054/// Specialization to build an scf "for" nest.1055template <>1056void GenerateLoopNest<scf::ForOp>::doit(1057 OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,1058 ArrayRef<utils::IteratorType> iteratorTypes,1059 function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,1060 ValueRange)>1061 bodyBuilderFn,1062 ArrayRef<linalg::ProcInfo> procInfo) {1063 assert((procInfo.empty() || (procInfo.size() == loopRanges.size())) &&1064 "expected as many entries for proc info as number of loops, even if "1065 "they are null entries");1066 SmallVector<Value> iterArgInitValues;1067 if (!linalgOp.hasPureBufferSemantics())1068 llvm::append_range(iterArgInitValues, linalgOp.getDpsInits());1069 SmallVector<Value, 4> lbs, ubs, steps;1070 unpackRanges(b, loc, loopRanges, lbs, ubs, steps);1071 LoopNest loopNest = mlir::scf::buildLoopNest(1072 b, loc, lbs, ubs, steps, iterArgInitValues,1073 [&](OpBuilder &b, Location loc, ValueRange ivs, ValueRange iterArgs) {1074 assert(iterArgs.size() == iterArgInitValues.size() &&1075 "expect the number of output tensors and iter args to match");1076 SmallVector<Value> operandValuesToUse = linalgOp->getOperands();1077 if (!iterArgs.empty()) {1078 operandValuesToUse = linalgOp.getDpsInputs();1079 operandValuesToUse.append(iterArgs.begin(), iterArgs.end());1080 }1081 return bodyBuilderFn(b, loc, ivs, operandValuesToUse);1082 });1083 1084 if (loopNest.loops.empty() || procInfo.empty())1085 return;1086 1087 // Filter out scf.for loops that were created out of parallel dimensions.1088 for (const auto &loop : llvm::enumerate(loopNest.loops)) {1089 if (procInfo[loop.index()].distributionMethod ==1090 DistributionMethod::Cyclic) {1091 mapLoopToProcessorIds(loop.value(), procInfo[loop.index()].procId,1092 procInfo[loop.index()].nprocs);1093 }1094 }1095}1096 1097/// Specialization to build affine "for" nest.1098template <>1099void GenerateLoopNest<AffineForOp>::doit(1100 OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,1101 ArrayRef<utils::IteratorType> iteratorTypes,1102 function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,1103 ValueRange)>1104 bodyBuilderFn,1105 ArrayRef<linalg::ProcInfo> /*procInfo*/) {1106 SmallVector<Value> iterArgInitValues;1107 if (!linalgOp.hasPureBufferSemantics())1108 llvm::append_range(iterArgInitValues, linalgOp.getDpsInits());1109 assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");1110 SmallVector<Value, 4> lbs, ubs, steps;1111 unpackRanges(b, loc, loopRanges, lbs, ubs, steps);1112 1113 // Affine loops require constant steps.1114 SmallVector<int64_t, 4> constantSteps;1115 constantSteps.reserve(steps.size());1116 for (Value v : steps) {1117 auto constVal = getConstantIntValue(v);1118 assert(constVal.has_value() && "Affine loops require constant steps");1119 constantSteps.push_back(constVal.value());1120 }1121 1122 affine::buildAffineLoopNest(b, loc, lbs, ubs, constantSteps,1123 [&](OpBuilder &b, Location loc, ValueRange ivs) {1124 bodyBuilderFn(b, loc, ivs,1125 linalgOp->getOperands());1126 });1127}1128 1129/// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.1130void updateBoundsForCyclicDistribution(OpBuilder &b, Location loc, Value procId,1131 Value nprocs, Value &lb, Value &ub,1132 Value &step) {1133 AffineExpr d0, d1;1134 bindDims(b.getContext(), d0, d1);1135 AffineExpr s0 = getAffineSymbolExpr(0, b.getContext());1136 lb =1137 affine::makeComposedAffineApply(b, loc, d0 + d1 * s0, {lb, procId, step});1138 step = affine::makeComposedAffineApply(b, loc, d0 * s0, {nprocs, step});1139}1140 1141/// Generates a loop nest consisting of scf.parallel and scf.for, depending1142/// on the `iteratorTypes.` Consecutive parallel loops create a single1143/// scf.parallel operation; each sequential loop creates a new scf.for1144/// operation. The body of the innermost loop is populated by1145/// `bodyBuilderFn` that accepts a range of induction variables for all1146/// loops. `ivStorage` is used to store the partial list of induction1147/// variables.1148// TODO: this function can be made iterative instead. However, it1149// will have at most as many recursive calls as nested loops, which rarely1150// exceeds 10.1151static void generateParallelLoopNest(1152 OpBuilder &b, Location loc, ValueRange lbs, ValueRange ubs,1153 ValueRange steps, ArrayRef<utils::IteratorType> iteratorTypes,1154 ArrayRef<linalg::ProcInfo> procInfo,1155 function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn,1156 SmallVectorImpl<Value> &ivStorage) {1157 assert(lbs.size() == ubs.size());1158 assert(lbs.size() == steps.size());1159 assert(lbs.size() == iteratorTypes.size());1160 assert(procInfo.empty() || (lbs.size() == procInfo.size()));1161 1162 // If there are no (more) loops to be generated, generate the body and be1163 // done with it.1164 if (iteratorTypes.empty()) {1165 bodyBuilderFn(b, loc, ivStorage);1166 return;1167 }1168 1169 // If there are no outer parallel loops, generate one sequential loop and1170 // recurse.1171 if (!isParallelIterator(iteratorTypes.front())) {1172 LoopNest singleLoop = buildLoopNest(1173 b, loc, lbs.take_front(), ubs.take_front(), steps.take_front(),1174 [&](OpBuilder &b, Location loc, ValueRange ivs) {1175 ivStorage.append(ivs.begin(), ivs.end());1176 generateParallelLoopNest(1177 b, loc, lbs.drop_front(), ubs.drop_front(), steps.drop_front(),1178 iteratorTypes.drop_front(),1179 procInfo.empty() ? procInfo : procInfo.drop_front(),1180 bodyBuilderFn, ivStorage);1181 });1182 return;1183 }1184 1185 unsigned nLoops = iteratorTypes.size();1186 unsigned numProcessed = 0;1187 DistributionMethod distributionMethod = DistributionMethod::None;1188 if (procInfo.empty()) {1189 numProcessed = nLoops - iteratorTypes.drop_while(isParallelIterator).size();1190 } else {1191 distributionMethod = procInfo.front().distributionMethod;1192 numProcessed =1193 nLoops - procInfo1194 .drop_while([&](linalg::ProcInfo p) {1195 return p.distributionMethod == distributionMethod;1196 })1197 .size();1198 }1199 1200 auto remainderProcInfo =1201 procInfo.empty() ? procInfo : procInfo.drop_front(numProcessed);1202 switch (distributionMethod) {1203 case DistributionMethod::None: {1204 // Generate a single parallel loop-nest operation for all outermost1205 // parallel loops and recurse.1206 scf::ParallelOp::create(1207 b, loc, lbs.take_front(numProcessed), ubs.take_front(numProcessed),1208 steps.take_front(numProcessed),1209 [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {1210 ivStorage.append(localIvs.begin(), localIvs.end());1211 generateParallelLoopNest(1212 nestedBuilder, nestedLoc, lbs.drop_front(numProcessed),1213 ubs.drop_front(numProcessed), steps.drop_front(numProcessed),1214 iteratorTypes.drop_front(numProcessed), remainderProcInfo,1215 bodyBuilderFn, ivStorage);1216 });1217 return;1218 }1219 case DistributionMethod::Cyclic: {1220 // Generate a single parallel loop-nest operation for all outermost1221 // parallel loops and recurse.1222 scf::ParallelOp::create(1223 b, loc, lbs.take_front(numProcessed), ubs.take_front(numProcessed),1224 steps.take_front(numProcessed),1225 [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {1226 ivStorage.append(localIvs.begin(), localIvs.end());1227 generateParallelLoopNest(1228 nestedBuilder, nestedLoc, lbs.drop_front(numProcessed),1229 ubs.drop_front(numProcessed), steps.drop_front(numProcessed),1230 iteratorTypes.drop_front(numProcessed), remainderProcInfo,1231 bodyBuilderFn, ivStorage);1232 });1233 return;1234 }1235 case DistributionMethod::CyclicNumProcsGeNumIters: {1236 // Check (for the processed loops) that the iteration is in-bounds.1237 ArithBuilder ab(b, loc);1238 Value cond = ab.slt(lbs[0], ubs[0]);1239 for (unsigned i = 1; i < numProcessed; ++i)1240 cond = ab._and(cond, ab.slt(lbs[i], ubs[i]));1241 ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));1242 scf::IfOp::create(b, loc, cond, [&](OpBuilder &b, Location loc) {1243 generateParallelLoopNest(b, loc, lbs.drop_front(numProcessed),1244 ubs.drop_front(numProcessed),1245 steps.drop_front(numProcessed),1246 iteratorTypes.drop_front(numProcessed),1247 remainderProcInfo, bodyBuilderFn, ivStorage);1248 scf::YieldOp::create(b, loc, ValueRange{});1249 });1250 return;1251 }1252 case DistributionMethod::CyclicNumProcsEqNumIters:1253 // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed1254 // with inner loop generation.1255 ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));1256 generateParallelLoopNest(1257 b, loc, lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),1258 steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),1259 remainderProcInfo, bodyBuilderFn, ivStorage);1260 return;1261 }1262}1263 1264/// Specialization for generating a mix of parallel and sequential scf loops.1265template <>1266void GenerateLoopNest<scf::ParallelOp>::doit(1267 OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,1268 ArrayRef<utils::IteratorType> iteratorTypes,1269 function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,1270 ValueRange)>1271 bodyBuilderFn,1272 ArrayRef<linalg::ProcInfo> procInfo) {1273 SmallVector<Value> iterArgInitValues;1274 if (!linalgOp.hasPureBufferSemantics())1275 llvm::append_range(iterArgInitValues, linalgOp.getDpsInits());1276 assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");1277 // This function may be passed more iterator types than ranges.1278 assert(iteratorTypes.size() >= loopRanges.size() &&1279 "expected iterator type for all ranges");1280 assert((procInfo.empty() || (procInfo.size() == loopRanges.size())) &&1281 "expected proc information for all loops when present");1282 iteratorTypes = iteratorTypes.take_front(loopRanges.size());1283 SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;1284 unsigned numLoops = iteratorTypes.size();1285 ivs.reserve(numLoops);1286 lbsStorage.reserve(numLoops);1287 ubsStorage.reserve(numLoops);1288 stepsStorage.reserve(numLoops);1289 1290 // Get the loop lb, ub, and step.1291 unpackRanges(b, loc, loopRanges, lbsStorage, ubsStorage, stepsStorage);1292 1293 // Modify the lb, ub, and step based on the distribution options.1294 for (const auto &it : llvm::enumerate(procInfo)) {1295 if (it.value().distributionMethod != linalg::DistributionMethod::None) {1296 updateBoundsForCyclicDistribution(1297 b, loc, it.value().procId, it.value().nprocs, lbsStorage[it.index()],1298 ubsStorage[it.index()], stepsStorage[it.index()]);1299 }1300 }1301 ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);1302 generateParallelLoopNest(1303 b, loc, lbs, ubs, steps, iteratorTypes, procInfo,1304 [&](OpBuilder &b, Location loc, ValueRange ivs) {1305 bodyBuilderFn(b, loc, ivs, linalgOp->getOperands());1306 },1307 ivs);1308 1309 assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");1310}1311 1312static Operation *materializeTiledShape(OpBuilder &builder, Location loc,1313 Value valueToTile,1314 const SliceParameters &sliceParams) {1315 auto shapedType = dyn_cast<ShapedType>(valueToTile.getType());1316 auto *sliceOp = TypeSwitch<ShapedType, Operation *>(shapedType)1317 .Case([&](MemRefType) {1318 return memref::SubViewOp::create(1319 builder, loc, valueToTile, sliceParams.offsets,1320 sliceParams.sizes, sliceParams.strides);1321 })1322 .Case([&](RankedTensorType) {1323 return tensor::ExtractSliceOp::create(1324 builder, loc, valueToTile, sliceParams.offsets,1325 sliceParams.sizes, sliceParams.strides);1326 })1327 .DefaultUnreachable("Unexpected shaped type");1328 return sliceOp;1329}1330 1331Operation *makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,1332 ArrayRef<OpFoldResult> tileSizes, AffineMap map,1333 ArrayRef<OpFoldResult> lbs,1334 ArrayRef<OpFoldResult> ubs,1335 ArrayRef<OpFoldResult> subShapeSizes,1336 bool omitPartialTileCheck) {1337 SliceParameters sliceParams =1338 computeSliceParameters(builder, loc, valueToTile, tileSizes, map, lbs,1339 ubs, subShapeSizes, omitPartialTileCheck);1340 return materializeTiledShape(builder, loc, valueToTile, sliceParams);1341}1342 1343SliceParameters1344computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,1345 ArrayRef<OpFoldResult> tileSizes, AffineMap map,1346 ArrayRef<OpFoldResult> lbs, ArrayRef<OpFoldResult> ubs,1347 ArrayRef<OpFoldResult> subShapeSizes,1348 bool omitPartialTileCheck) {1349 auto shapedType = dyn_cast<ShapedType>(valueToTile.getType());1350 assert(shapedType && "only shaped types can be tiled");1351 ArrayRef<int64_t> shape = shapedType.getShape();1352 int64_t rank = shapedType.getRank();1353 1354 // Compute offsets/sizes/strides for the tile.1355 SliceParameters sliceParams;1356 sliceParams.offsets.reserve(rank);1357 sliceParams.sizes.reserve(rank);1358 sliceParams.strides.reserve(rank);1359 for (unsigned r = 0; r < rank; ++r) {1360 LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: for dim#" << r);1361 if (!isTiled(map.getSubMap({r}), tileSizes)) {1362 sliceParams.offsets.push_back(builder.getIndexAttr(0));1363 OpFoldResult dim = createFoldedDimOp(builder, loc, valueToTile, r);1364 sliceParams.sizes.push_back(dim);1365 sliceParams.strides.push_back(builder.getIndexAttr(1));1366 LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");1367 continue;1368 }1369 LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");1370 1371 // Tiling creates a new slice at the proper index, the slice step is 11372 // (i.e. the op does not subsample, stepping occurs in the loop).1373 auto m = map.getSubMap({r});1374 LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: submap: " << m << "\n");1375 IRRewriter rewriter(builder);1376 // The offset of the slice is m(lbs) - m(0).1377 SmallVector<Attribute> zeros(lbs.size(), rewriter.getIndexAttr(0));1378 SmallVector<Attribute> mAtZero;1379 [[maybe_unused]] auto res = m.constantFold(zeros, mAtZero);1380 assert(succeeded(res) && "affine_map must be evaluatable (not symbols)");1381 int64_t mAtZeroInt =1382 cast<IntegerAttr>(mAtZero[0]).getValue().getSExtValue();1383 OpFoldResult offset = makeComposedFoldedAffineApply(1384 rewriter, loc, m.getResult(0) - mAtZeroInt, lbs);1385 sliceParams.offsets.push_back(offset);1386 1387 OpFoldResult closedIntSize =1388 makeComposedFoldedAffineApply(rewriter, loc, m, subShapeSizes);1389 // Resulting size needs to be made half open interval again.1390 AffineExpr s0 = getAffineSymbolExpr(0, builder.getContext());1391 OpFoldResult size =1392 makeComposedFoldedAffineApply(rewriter, loc, s0 + 1, closedIntSize);1393 LLVM_DEBUG(llvm::dbgs()1394 << "computeSliceParameters: raw size: " << size << "\n");1395 LLVM_DEBUG(llvm::dbgs()1396 << "computeSliceParameters: new offset: " << offset << "\n");1397 sliceParams.strides.push_back(builder.getIndexAttr(1));1398 1399 if (omitPartialTileCheck) {1400 // We statically know that the partial/boundary tile condition is1401 // unnecessary.1402 LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n");1403 sliceParams.sizes.push_back(size);1404 continue;1405 }1406 1407 // The size of the subview / extract_slice should be trimmed to avoid1408 // out-of-bounds accesses, unless:1409 // a. We statically know the subshape size divides the shape size evenly.1410 // b. The subshape size is 1. According to the way the loops are set up,1411 // tensors with "0" dimensions would never be constructed.1412 int64_t shapeSize = shape[r];1413 std::optional<int64_t> sizeCst = getConstantIntValue(size);1414 auto hasTileSizeOne = sizeCst == 1;1415 auto dividesEvenly = sizeCst && ShapedType::isStatic(shapeSize) &&1416 ((shapeSize % *sizeCst) == 0);1417 if (!hasTileSizeOne && !dividesEvenly) {1418 LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: shapeSize=" << shapeSize1419 << ", size: " << size1420 << ": make sure in bound with affine.min\n");1421 1422 AffineExpr dim0, dim1, dim2;1423 MLIRContext *context = builder.getContext();1424 bindDims(context, dim0, dim1, dim2);1425 1426 // Get the dimension size for this dimension. We need to first calculate1427 // the max index and then plus one. This is important because for1428 // convolution ops, we have its input window dimension's affine map of the1429 // form `(d0 * s0 + d1)`, where `d0`/`d1 is an output/filter window1430 // dimension and `s0` is stride. Directly use the dimension size of1431 // output/filer window dimensions will cause incorrect calculation.1432 AffineMap minusOneMap = AffineMap::inferFromExprList(1433 {ArrayRef<AffineExpr>{dim0 - 1}}, context)1434 .front();1435 AffineMap plusOneMap = AffineMap::inferFromExprList(1436 {ArrayRef<AffineExpr>{dim0 + 1}}, context)1437 .front();1438 SmallVector<OpFoldResult> maxIndices =1439 llvm::to_vector(llvm::map_range(ubs, [&](OpFoldResult ub) {1440 return makeComposedFoldedAffineApply(rewriter, loc, minusOneMap,1441 {ub});1442 }));1443 OpFoldResult maxIndex =1444 makeComposedFoldedAffineApply(rewriter, loc, m, maxIndices);1445 OpFoldResult d =1446 makeComposedFoldedAffineApply(rewriter, loc, plusOneMap, {maxIndex});1447 1448 // Compute min(dim - offset, size) to avoid out-of-bounds accesses.1449 AffineMap minMap = AffineMap::inferFromExprList(1450 {ArrayRef<AffineExpr>{dim1 - dim2, dim0}}, context)1451 .front();1452 size =1453 makeComposedFoldedAffineMin(rewriter, loc, minMap, {size, d, offset});1454 }1455 LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: new size: " << size << "\n");1456 sliceParams.sizes.push_back(size);1457 }1458 return sliceParams;1459}1460 1461SmallVector<OpFoldResult> computeTileOffsets(OpBuilder &b, Location loc,1462 ArrayRef<OpFoldResult> ivs,1463 ArrayRef<OpFoldResult> tileSizes) {1464 SmallVector<OpFoldResult> offsets;1465 for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {1466 LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");1467 bool isTiled = !isZeroInteger(tileSizes[idx]);1468 offsets.push_back(isTiled ? ivs[idxIvs++] : b.getIndexAttr(0));1469 LLVM_DEBUG(llvm::dbgs()1470 << "computeTileOffsets: " << offsets.back() << "\n");1471 }1472 return offsets;1473}1474 1475SmallVector<OpFoldResult> computeTileSizes(OpBuilder &b, Location loc,1476 ArrayRef<OpFoldResult> tileSizes,1477 ArrayRef<OpFoldResult> sizeBounds) {1478 SmallVector<OpFoldResult> sizes;1479 for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx) {1480 bool isTiled = !isZeroInteger(tileSizes[idx]);1481 // Before composing, we need to make range a closed interval.1482 OpFoldResult size = isTiled ? tileSizes[idx] : sizeBounds[idx];1483 AffineExpr d0 = getAffineDimExpr(0, b.getContext());1484 IRRewriter rewriter(b);1485 sizes.push_back(makeComposedFoldedAffineApply(rewriter, loc, d0 - 1, size));1486 LLVM_DEBUG(llvm::dbgs() << "computeTileSizes: " << sizes.back() << "\n");1487 }1488 return sizes;1489}1490 1491SmallVector<Type> getTensorOutputTypes(LinalgOp op, ValueRange operands) {1492 if (op.hasPureBufferSemantics())1493 return {};1494 return llvm::to_vector(1495 llvm::map_range(op.getDpsInitsMutable(), [&](OpOperand &opOperand) {1496 return operands[opOperand.getOperandNumber()].getType();1497 }));1498}1499 1500SmallVector<Value> insertSlicesBack(OpBuilder &builder, Location loc,1501 LinalgOp op, ValueRange operands,1502 ValueRange results) {1503 if (op.hasPureBufferSemantics())1504 return {};1505 SmallVector<Value> tensorResults;1506 tensorResults.reserve(results.size());1507 // Insert a insert_slice for each output tensor.1508 unsigned resultIdx = 0;1509 for (OpOperand &opOperand : op.getDpsInitsMutable()) {1510 // TODO: use an interface/adaptor to avoid leaking position in1511 // `tiledOperands`.1512 Value outputTensor = operands[opOperand.getOperandNumber()];1513 if (auto sliceOp = outputTensor.getDefiningOp<tensor::ExtractSliceOp>()) {1514 Value inserted = tensor::InsertSliceOp::create(1515 builder, loc, sliceOp.getSource().getType(), results[resultIdx],1516 sliceOp.getSource(), sliceOp.getOffsets(), sliceOp.getSizes(),1517 sliceOp.getStrides(), sliceOp.getStaticOffsets(),1518 sliceOp.getStaticSizes(), sliceOp.getStaticStrides());1519 tensorResults.push_back(inserted);1520 } else {1521 tensorResults.push_back(results[resultIdx]);1522 }1523 ++resultIdx;1524 }1525 return tensorResults;1526}1527 1528SmallVector<std::optional<SliceParameters>>1529computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp,1530 ValueRange valuesToTile, ArrayRef<OpFoldResult> ivs,1531 ArrayRef<OpFoldResult> tileSizes,1532 ArrayRef<OpFoldResult> sizeBounds,1533 bool omitPartialTileCheck) {1534 assert(ivs.size() == static_cast<size_t>(llvm::count_if(1535 llvm::make_range(tileSizes.begin(), tileSizes.end()),1536 [](OpFoldResult v) { return !isZeroInteger(v); })) &&1537 "expected as many ivs as non-zero sizes");1538 1539 // Construct (potentially temporary) mins and maxes on which to apply maps1540 // that define tile subshapes.1541 SmallVector<OpFoldResult> lbs =1542 computeTileOffsets(builder, loc, ivs, tileSizes);1543 SmallVector<OpFoldResult> subShapeSizes =1544 computeTileSizes(builder, loc, tileSizes, sizeBounds);1545 1546 assert(static_cast<int64_t>(valuesToTile.size()) <=1547 linalgOp->getNumOperands() &&1548 "more value to tile than operands.");1549 SmallVector<std::optional<SliceParameters>> allSliceParams;1550 allSliceParams.reserve(valuesToTile.size());1551 for (auto [opOperand, val] :1552 llvm::zip(linalgOp->getOpOperands(), valuesToTile)) {1553 Value shapedOp = val;1554 LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);1555 AffineMap map = linalgOp.getMatchingIndexingMap(&opOperand);1556 // Use `opOperand` as is if it is not tiled and not an output tensor. Having1557 // an extract/insert slice pair for all output tensors simplifies follow up1558 // transformations such as padding and bufferization since the1559 // extract/insert slice pairs make the accessed iteration argument1560 // subdomains explicit.1561 1562 Type operandType = opOperand.get().getType();1563 if (!isTiled(map, tileSizes) && !(isa<RankedTensorType>(operandType) &&1564 linalgOp.isDpsInit(&opOperand))) {1565 allSliceParams.push_back(std::nullopt);1566 LLVM_DEBUG(llvm::dbgs()1567 << ": not tiled: use shape: " << operandType << "\n");1568 continue;1569 }1570 LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");1571 1572 allSliceParams.push_back(computeSliceParameters(1573 builder, loc, shapedOp, tileSizes, map, lbs, sizeBounds, subShapeSizes,1574 omitPartialTileCheck));1575 }1576 1577 return allSliceParams;1578}1579 1580SmallVector<Value> makeTiledShapes(OpBuilder &builder, Location loc,1581 LinalgOp linalgOp, ValueRange valuesToTile,1582 ArrayRef<OpFoldResult> ivs,1583 ArrayRef<OpFoldResult> tileSizes,1584 ArrayRef<OpFoldResult> sizeBounds,1585 bool omitPartialTileCheck) {1586 SmallVector<std::optional<SliceParameters>> allSliceParameter =1587 computeAllSliceParameters(builder, loc, linalgOp, valuesToTile, ivs,1588 tileSizes, sizeBounds, omitPartialTileCheck);1589 SmallVector<Value> tiledShapes;1590 for (auto item : llvm::zip(valuesToTile, allSliceParameter)) {1591 Value valueToTile = std::get<0>(item);1592 std::optional<SliceParameters> sliceParams = std::get<1>(item);1593 tiledShapes.push_back(1594 sliceParams.has_value()1595 ? materializeTiledShape(builder, loc, valueToTile, *sliceParams)1596 ->getResult(0)1597 : valueToTile);1598 }1599 return tiledShapes;1600}1601 1602void offsetIndices(OpBuilder &b, LinalgOp linalgOp,1603 ArrayRef<OpFoldResult> offsets) {1604 IRRewriter rewriter(b);1605 offsetIndices(rewriter, linalgOp, offsets);1606}1607 1608void offsetIndices(RewriterBase &b, LinalgOp linalgOp,1609 ArrayRef<OpFoldResult> offsets) {1610 if (!linalgOp.hasIndexSemantics())1611 return;1612 1613 for (IndexOp indexOp : linalgOp.getBlock()->getOps<IndexOp>()) {1614 if (indexOp.getDim() >= offsets.size() || !offsets[indexOp.getDim()])1615 continue;1616 OpBuilder::InsertionGuard guard(b);1617 b.setInsertionPointAfter(indexOp);1618 AffineExpr index, offset;1619 bindDims(b.getContext(), index, offset);1620 OpFoldResult applied = makeComposedFoldedAffineApply(1621 b, indexOp.getLoc(), index + offset,1622 {getAsOpFoldResult(indexOp.getResult()), offsets[indexOp.getDim()]});1623 Value materialized =1624 getValueOrCreateConstantIndexOp(b, indexOp.getLoc(), applied);1625 b.replaceUsesWithIf(indexOp, materialized, [&](OpOperand &use) {1626 return use.getOwner() != materialized.getDefiningOp();1627 });1628 }1629}1630 1631/// Get the reassociation maps to fold the result of a extract_slice (or source1632/// of a insert_slice) operation with given offsets, and sizes to its1633/// rank-reduced version. This is only done for the cases where the size is 11634/// and offset is 0. Strictly speaking the offset 0 is not required in general,1635/// but non-zero offsets are not handled by SPIR-V backend at this point (and1636/// potentially cannot be handled).1637std::optional<SmallVector<ReassociationIndices>>1638getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes) {1639 SmallVector<ReassociationIndices> reassociation;1640 ReassociationIndices curr;1641 for (const auto &it : llvm::enumerate(mixedSizes)) {1642 auto dim = it.index();1643 auto size = it.value();1644 curr.push_back(dim);1645 auto attr = llvm::dyn_cast_if_present<Attribute>(size);1646 if (attr && cast<IntegerAttr>(attr).getInt() == 1)1647 continue;1648 reassociation.emplace_back(ReassociationIndices{});1649 std::swap(reassociation.back(), curr);1650 }1651 // When the reassociations are not empty, then fold the remaining1652 // unit-dimensions into the last dimension. If the reassociations so far is1653 // empty, then leave it emtpy. This will fold everything to a rank-0 tensor.1654 if (!curr.empty() && !reassociation.empty())1655 reassociation.back().append(curr.begin(), curr.end());1656 return reassociation;1657}1658 1659} // namespace linalg1660} // namespace mlir1661