1679 lines · cpp
1//===- Transforms.cpp - Linalg transformations as patterns ----------------===//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 logic and helpers to expose Linalg transforms as rewrite10// patterns.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Linalg/Transforms/Transforms.h"15#include "mlir/Dialect/Affine/IR/AffineOps.h"16#include "mlir/Dialect/Arith/IR/Arith.h"17#include "mlir/Dialect/Func/IR/FuncOps.h"18#include "mlir/Dialect/Linalg/IR/Linalg.h"19#include "mlir/Dialect/Linalg/Utils/Utils.h"20#include "mlir/Dialect/SCF/Transforms/Transforms.h"21#include "mlir/Dialect/Tensor/IR/Tensor.h"22#include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"23#include "mlir/Dialect/Tensor/Utils/Utils.h"24#include "mlir/Dialect/Utils/IndexingUtils.h"25#include "mlir/Dialect/Utils/StaticValueUtils.h"26#include "mlir/Dialect/Utils/StructuredOpsUtils.h"27#include "mlir/Dialect/Vector/IR/VectorOps.h"28#include "mlir/IR/AffineExpr.h"29#include "mlir/Support/LLVM.h"30#include "llvm/ADT/TypeSwitch.h"31#include "llvm/Support/Debug.h"32#include "llvm/Support/DebugLog.h"33#include "llvm/Support/InterleavedRange.h"34#include "llvm/Support/raw_ostream.h"35#include <utility>36 37#define DEBUG_TYPE "linalg-transforms"38 39using namespace mlir;40using namespace mlir::linalg;41 42//===----------------------------------------------------------------------===//43// Transformations exposed as functional-style API calls.44//===----------------------------------------------------------------------===//45 46//===----------------------------------------------------------------------===//47// peelLoop transformation.48//===----------------------------------------------------------------------===//49 50/// Try to peel and canonicalize loop `op` and return the new result.51/// Also applies affine_min/max bounds simplification on the fly where relevant.52// TODO: Add support for scf.parallel and affine.for loops.53SmallVector<Value> mlir::linalg::peelLoop(RewriterBase &rewriter,54 Operation *op) {55 return llvm::TypeSwitch<Operation *, SmallVector<Value, 4>>(op)56 .Case<scf::ForOp>([&](scf::ForOp forOp) {57 scf::ForOp partialIteration;58 if (succeeded(scf::peelForLoopAndSimplifyBounds(rewriter, forOp,59 partialIteration)))60 return partialIteration->getResults();61 assert(!partialIteration && "expected that loop was not peeled");62 return forOp->getResults();63 })64 .Default([&](Operation *op) { return op->getResults(); });65}66 67/// Peel 'loops' and applies affine_min/max bounds simplification on the fly68/// where relevant.69void mlir::linalg::peelLoops(RewriterBase &rewriter,70 ArrayRef<scf::ForOp> loops) {71 for (auto loopOp : loops)72 peelLoop(rewriter, loopOp);73}74 75//===----------------------------------------------------------------------===//76// pack transformation.77//===----------------------------------------------------------------------===//78 79#ifndef NDEBUG80/// Return true if `map` has 0 or 1 result function of AffineDimExpr(dim).81static bool hasAtMostOneResultFunctionOfDim(AffineMap map, int64_t dim) {82 bool found = false;83 for (AffineExpr e : map.getResults()) {84 if (!e.isFunctionOfDim(dim))85 continue;86 if (found)87 return false;88 found = true;89 }90 return true;91}92#endif // NDEBUG93 94static std::string stringifyReassocIndices(ReassociationIndicesRef ri) {95 return llvm::interleaved(ri, ", ", /*Prefix=*/"|", /*Suffix=*/"");96}97 98/// Return the index of the first result of `map` that is a function of99/// AffineDimExpr(dim), std::nullopt otherwise.100static std::optional<int64_t> getFirstResultIndexFunctionOf(AffineMap map,101 int64_t dim) {102 for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {103 AffineExpr expr = map.getResult(i);104 if (!expr.isFunctionOfDim(dim))105 continue;106 return i;107 }108 return std::nullopt;109}110 111/// Perform one step of packing of a LinalgOp's metadata along `dim` into the112/// `newDim` at `iteratorTypes.size()` by:113/// 1. Appending `iteratorTypes[newDim]`, equal to `iteratorTypes[dim]`.114/// 2. Appending a `newDim` to the domain of every indexing map.115/// 3. For each operand (i.e. for each map in `indexingMaps`), perform packing116/// by potentially adding a `newDim` result to `map`.117/// The preserved invariant is that `iteratorTypes.size()` is always equal to118/// `map.getNumDims()` for every map in `indexingMaps`.119///120/// Update `indexingMaps` and `iteratorTypes` inplace as one step of the update.121/// Return a vector that records the optional packing for each operand.122/// Return failure if the packed indexing cannot be represented with a LinalgOp.123///124/// Further details:125/// ================126/// The current implementation of packing (i.e. data tiling) consists of127/// rewriting a linearized strip-mined form into a higher-dimensional access.128/// e.g. consider an access `A[I][f(j, k, l)]` and packing by 4; we rewrite129/// `I` into `4 * i + ii`, where `0 <= ii < 4`.130/// The access is further rewritten as `A[i][f(j, k, l)][ii]`.131///132/// This rewrite into higher dimensional access is not possible for general133/// AffineExpr in Linalg atm, it is restricted to an AffineDimExpr:134/// e.g. consider an access `A[I + J][f(j, k, l)]` and packing by 4; we135/// rewrite `I + J` into `4 * i + ii + J`, where `0 <= ii < 4`.136/// The rewrite of the access would be a form not representable in Linalg:137/// `A[i + (ii + J) / 4][f(j, k, l)][(ii + J) % 4]`.138/// Note however that as `J` and `ii` iterate, the accesses do not have a139/// particular alignment, so packing does not achieve alignment in this case140///141/// In the future, we may want to consider a mixed-form that allows some142/// alignment in the presence of multiple accesses:143/// `A[I][f(j, k, l)]` and `B[I + J][f(j, k, l)]`144/// And would rewrite accesses as:145/// `A[i][f(j, k, l)][ii]` and `B[4 * i + ii + J][f(j, k, l)]`146static FailureOr<SmallVector<std::optional<int64_t>>>147packLinalgMetadataOnce(SmallVectorImpl<AffineMap> &indexingMaps,148 SmallVectorImpl<utils::IteratorType> &iteratorTypes,149 int64_t dim) {150 int64_t newDim = iteratorTypes.size();151 iteratorTypes.push_back(iteratorTypes[dim]);152 153 SmallVector<std::optional<int64_t>> packedDimPerIndexingMap(154 indexingMaps.size(), std::nullopt);155 SmallVector<AffineMap> newMaps;156 for (int64_t operandIdx = 0, e = indexingMaps.size(); operandIdx < e;157 ++operandIdx) {158 AffineMap map = indexingMaps[operandIdx];159 160 // Add the `newDim` to map whatever the case.161 assert(map.getNumDims() == newDim && "num dims invariant violation");162 map = map.shiftDims(1, newDim);163 164 // Get the at-most-1 index of the result that is a function of `dim`.165 // If we can find one, we insert `AffineDimExpr(newDim)` to the map, which166 // logically chunks dimension `dim` into `K * dim + newDim`, where the167 // packing factor `K` is specified separately.168 assert(hasAtMostOneResultFunctionOfDim(map, dim) &&169 "num results invariant violation");170 auto maybeOperandDimensionToPack = getFirstResultIndexFunctionOf(map, dim);171 if (!maybeOperandDimensionToPack.has_value()) {172 newMaps.push_back(map);173 continue;174 }175 176 // We can only pack AffineDimExpr atm.177 if (!isa<AffineDimExpr>(map.getResult(maybeOperandDimensionToPack.value())))178 return failure();179 180 // Add `newDim` to the results of the map.181 map = map.insertResult(Builder(map.getContext()).getAffineDimExpr(newDim),182 map.getNumResults());183 newMaps.push_back(map);184 185 // Record the that `operandIdx` is packed.186 packedDimPerIndexingMap[operandIdx] = maybeOperandDimensionToPack;187 }188 indexingMaps = newMaps;189 190 return packedDimPerIndexingMap;191}192 193namespace {194 195/// Helper struct to encode packing along one dimension of a LinalgOp.196struct PackedOperandsDim {197 OpFoldResult packedSize;198 SmallVector<std::optional<int64_t>> packedDimForEachOperand;199};200 201/// Helper struct to encode packing along all dimensions of a LinalgOp.202struct PackedOperandsDimList {203 void pushBack(PackedOperandsDim &&packedOperandsDims) {204 spec.emplace_back(packedOperandsDims);205 }206 /// Return all the dims that have been packed for operand @ `operandPos`.207 SmallVector<int64_t> extractPackedDimsForOperand(int64_t operandPos);208 /// Return all the pack sizes by which an operand @ `operandPos` is packed.209 SmallVector<OpFoldResult> extractPackSizesForOperand(int64_t operandPos);210 211private:212 SmallVector<PackedOperandsDim> spec;213};214 215} // namespace216 217FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,218 linalg::PackOp packOp,219 bool lowerPadLikeWithInsertSlice) {220 // 1. Filter out NYI cases.221 auto packedTensorType =222 cast<RankedTensorType>(packOp->getResultTypes().front());223 if (llvm::any_of(packOp.getStaticInnerTiles(), ShapedType::isDynamic)) {224 return rewriter.notifyMatchFailure(225 packOp,226 "non-static shape NYI, needs a more powerful tensor.expand_shape op");227 }228 229 Location loc = packOp->getLoc();230 OpBuilder::InsertionGuard g(rewriter);231 rewriter.setInsertionPoint(packOp);232 233 // 2. Compute the permutation vector to shuffle packed shape into the shape234 // before any outer or inner permutations have been applied.235 PackingMetadata packingMetadata;236 SmallVector<int64_t> packedToStripMinedShapePerm =237 getPackInverseDestPerm(packOp, packingMetadata);238 239 // 3. Compute the stripMinedShape: this is the packed shape before any outer240 // or inner permutations have been applied.241 SmallVector<int64_t> stripMinedShape(packedTensorType.getShape());242 applyPermutationToVector(stripMinedShape, packedToStripMinedShapePerm);243 244 // 4. Pad the source of packOp to a shape we can expand into stripMinedShape.245 SmallVector<OpFoldResult> lows(packOp.getSourceRank(),246 rewriter.getIndexAttr(0));247 SmallVector<OpFoldResult> highs(packOp.getSourceRank(),248 rewriter.getIndexAttr(0));249 for (auto [pos, innerSize] :250 llvm::zip_equal(packOp.getInnerDimsPos(), packOp.getMixedTiles())) {251 int outerPos =252 packedToStripMinedShapePerm[packingMetadata.outerPositions[pos]];253 OpFoldResult origSize =254 tensor::getMixedSize(rewriter, loc, packOp.getSource(), pos);255 OpFoldResult outerSize =256 tensor::getMixedSize(rewriter, loc, packOp.getDest(), outerPos);257 AffineExpr s0, d0, d1;258 bindDims(rewriter.getContext(), d0, d1);259 bindSymbols(rewriter.getContext(), s0);260 auto map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/1, d0 * s0 - d1);261 highs[pos] = affine::makeComposedFoldedAffineApply(262 rewriter, loc, map, {outerSize, origSize, innerSize});263 }264 RankedTensorType collapsed = tensor::CollapseShapeOp::inferCollapsedType(265 RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape),266 packingMetadata.reassociations);267 Value paddingValue = packOp.getPaddingValue();268 if (!paddingValue) {269 paddingValue = arith::ConstantOp::create(270 rewriter, loc, rewriter.getZeroAttr(getElementTypeOrSelf(collapsed)));271 }272 auto padOp =273 tensor::PadOp::create(rewriter, loc, collapsed, packOp.getSource(), lows,274 highs, paddingValue, /*nofold=*/false);275 276 LDBG() << "insertPositions: "277 << llvm::interleaved(packingMetadata.insertPositions);278 LDBG() << "outerPositions: "279 << llvm::interleaved(packingMetadata.outerPositions);280 LDBG() << "packedShape: " << llvm::interleaved(packedTensorType.getShape());281 LDBG() << "packedToStripMinedShapePerm: "282 << llvm::interleaved(packedToStripMinedShapePerm);283 LDBG() << "reassociations: "284 << llvm::interleaved(llvm::map_range(packingMetadata.reassociations,285 stringifyReassocIndices));286 LDBG() << "stripMinedShape: " << llvm::interleaved(stripMinedShape);287 LDBG() << "collapsed type: " << collapsed;288 289 if (lowerPadLikeWithInsertSlice && packOp.isLikePad()) {290 // Pack ops which operate as simple pads may not produce legal291 // tensor.insert_slice operations when the packed type does not rank reduce292 // to the padded type.293 SliceVerificationResult rankReduces =294 isRankReducedType(packedTensorType, padOp.getResultType());295 296 if (rankReduces == SliceVerificationResult::Success) {297 // This pack is just a plain pad.298 // Just insert the pad in the higher ranked tensor.299 // Offsets.300 SmallVector<OpFoldResult> zeros(packOp.getDestRank(),301 rewriter.getIndexAttr(0));302 // Strides.303 SmallVector<OpFoldResult> ones(packOp.getDestRank(),304 rewriter.getIndexAttr(1));305 SmallVector<OpFoldResult> sizes =306 tensor::getMixedSizes(rewriter, loc, packOp.getDest());307 308 auto insertSliceOp = tensor::InsertSliceOp::create(309 rewriter, loc, /*source=*/padOp, /*dest=*/packOp.getDest(),310 /*offsets=*/zeros, sizes, /*strides=*/ones);311 312 LDBG() << "insert_slice op: " << insertSliceOp;313 314 rewriter.replaceOp(packOp, insertSliceOp->getResults());315 316 return LowerPackResult{padOp, /*reshapeOp=*/nullptr,317 /*transposeOp=*/nullptr};318 }319 }320 321 // 5. Expand from the padded result to the stripMinedShape.322 auto expandShapeResultType =323 RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape);324 auto reshapeOp = tensor::ExpandShapeOp::create(325 rewriter, loc, expandShapeResultType, padOp.getResult(),326 packingMetadata.reassociations);327 328 // 6. Transpose stripMinedShape to packedShape.329 SmallVector<int64_t> transpPerm =330 invertPermutationVector(packedToStripMinedShapePerm);331 auto transposeOp = linalg::TransposeOp::create(332 rewriter, loc, reshapeOp.getResult(), packOp.getDest(), transpPerm);333 334 LDBG() << "reshape op: " << reshapeOp;335 LDBG() << "transpPerm: " << llvm::interleaved(transpPerm);336 LDBG() << "transpose op: " << transposeOp;337 338 // 7. Replace packOp by transposeOp.339 rewriter.replaceOp(packOp, transposeOp->getResults());340 341 return LowerPackResult{padOp, reshapeOp, transposeOp};342}343 344FailureOr<LowerUnPackOpResult>345linalg::lowerUnPack(RewriterBase &rewriter, linalg::UnPackOp unPackOp,346 bool lowerUnpadLikeWithExtractSlice) {347 Location loc = unPackOp->getLoc();348 OpBuilder::InsertionGuard g(rewriter);349 rewriter.setInsertionPoint(unPackOp);350 351 RankedTensorType packedTensorType = unPackOp.getSourceType();352 int64_t packedRank = packedTensorType.getRank();353 354 OpFoldResult zero = rewriter.getIndexAttr(0), one = rewriter.getIndexAttr(1);355 auto destTensorType = cast<RankedTensorType>(unPackOp.getDest().getType());356 if (lowerUnpadLikeWithExtractSlice && unPackOp.isLikeUnPad()) {357 // This unpack is just a plain unpad.358 // Just extract the slice from the higher ranked tensor.359 ArrayRef<int64_t> destShape = destTensorType.getShape();360 // The inner dimensions stay the same as the destination tensor, but the361 // outer ones are additional 1s.362 SmallVector<OpFoldResult> sizes(packedRank - destShape.size(), one);363 sizes.append(tensor::getMixedSizes(rewriter, loc, unPackOp.getDest()));364 365 auto extractSliceOp = tensor::ExtractSliceOp::create(366 rewriter, loc, destTensorType, unPackOp.getSource(),367 SmallVector<OpFoldResult>(packedRank, zero), sizes,368 SmallVector<OpFoldResult>(packedRank, one));369 370 rewriter.replaceOp(unPackOp, extractSliceOp->getResults());371 372 return LowerUnPackOpResult{/*emptyOp=*/nullptr, /*transposeOp=*/nullptr,373 /*reshapeOp=*/nullptr, extractSliceOp};374 }375 376 // 1. Compute the permutation vector to shuffle packed shape into the shape377 // before any outer or inner permutations have been applied.378 PackingMetadata packingMetadata;379 SmallVector<int64_t> packedToStripMinedShapePerm =380 getUnPackInverseSrcPerm(unPackOp, packingMetadata);381 382 // 2. Compute the stripMinedShape: this is the packed shape without outer and383 // inner permutations.384 SmallVector<int64_t> stripMinedShape(packedTensorType.getShape());385 applyPermutationToVector(stripMinedShape, packedToStripMinedShapePerm);386 387 // 3. Transpose packedShape to stripMinedShape.388 RankedTensorType stripMinedTensorType =389 RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape);390 RankedTensorType collapsedType = tensor::CollapseShapeOp::inferCollapsedType(391 stripMinedTensorType, packingMetadata.reassociations);392 393 // Get dynamic dims from input tensor based on packedToStripMinedShapePerm394 // permutation.395 SmallVector<OpFoldResult, 4> dims =396 tensor::getMixedSizes(rewriter, loc, unPackOp.getSource());397 applyPermutationToVector(dims, packedToStripMinedShapePerm);398 auto emptyOp = tensor::EmptyOp::create(rewriter, loc, dims,399 stripMinedTensorType.getElementType());400 auto transposeOp =401 linalg::TransposeOp::create(rewriter, loc, unPackOp.getSource(), emptyOp,402 packedToStripMinedShapePerm);403 404 LDBG() << "insertPositions: "405 << llvm::interleaved(packingMetadata.insertPositions);406 LDBG() << "packedShape: " << llvm::interleaved(packedTensorType.getShape());407 LDBG() << "packedToStripMinedShapePerm: "408 << llvm::interleaved(packedToStripMinedShapePerm);409 LDBG() << "reassociations: "410 << llvm::interleaved(llvm::map_range(packingMetadata.reassociations,411 stringifyReassocIndices));412 LDBG() << "stripMinedShape: " << llvm::interleaved(stripMinedShape);413 LDBG() << "collapsed type: " << collapsedType;414 415 // 4. Collapse from the stripMinedShape to the padded result.416 auto reshapeOp = tensor::CollapseShapeOp::create(417 rewriter, loc, collapsedType, transposeOp->getResult(0),418 packingMetadata.reassociations);419 420 // 5. ExtractSlice.421 int64_t destRank = destTensorType.getRank();422 auto extractSliceOp = tensor::ExtractSliceOp::create(423 rewriter, loc, destTensorType, reshapeOp->getResult(0),424 SmallVector<OpFoldResult>(destRank, zero),425 tensor::getMixedSizes(rewriter, loc, unPackOp.getDest()),426 SmallVector<OpFoldResult>(destRank, one));427 428 // 6. Inject a copy to preserve DPS.429 auto copyOp = linalg::CopyOp::create(430 rewriter, loc, extractSliceOp->getResult(0), unPackOp.getDest());431 432 // 7. Replace unPackOp by copyOp.433 rewriter.replaceOp(unPackOp, copyOp->getResults());434 435 return LowerUnPackOpResult{emptyOp, transposeOp, reshapeOp, extractSliceOp};436}437 438SmallVector<int64_t>439PackedOperandsDimList::extractPackedDimsForOperand(int64_t operandPos) {440 SmallVector<int64_t> res;441 for (auto &i : spec) {442 if (!i.packedDimForEachOperand[operandPos].has_value())443 continue;444 res.push_back(i.packedDimForEachOperand[operandPos].value());445 }446 return res;447}448 449SmallVector<OpFoldResult>450PackedOperandsDimList::extractPackSizesForOperand(int64_t operandPos) {451 SmallVector<OpFoldResult> res;452 for (auto &i : spec) {453 if (!i.packedDimForEachOperand[operandPos].has_value())454 continue;455 res.push_back(i.packedSize);456 }457 return res;458}459 460/// Implement packing of a single LinalgOp by performing packing by461/// `packedSizes`. There must be one packedSizes entry per `linalgOp` iterator.462/// Return the packed Linalg op on success, failure otherwise.463FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,464 linalg::LinalgOp linalgOp,465 ArrayRef<OpFoldResult> packedSizes) {466 if (packedSizes.size() != linalgOp.getNumLoops()) {467 return rewriter.notifyMatchFailure(linalgOp,468 "incorrect number of pack sizes");469 }470 471 Location loc = linalgOp->getLoc();472 SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();473 SmallVector<utils::IteratorType> iteratorTypes =474 linalgOp.getIteratorTypesArray();475 LDBG() << "Start packing: " << linalgOp;476 LDBG() << "maps: " << llvm::interleaved(indexingMaps);477 LDBG() << "iterators: " << llvm::interleaved(iteratorTypes);478 479 SmallVector<linalg::PackOp> packOps;480 SmallVector<linalg::UnPackOp> unPackOps;481 // Step 1. Pack each dim of the LinalgOp metadata by packedSizes[i].482 PackedOperandsDimList listOfPackedOperandsDim;483 for (int64_t i = 0, e = packedSizes.size(); i < e; ++i) {484 std::optional<int64_t> maybeConstant = getConstantIntValue(packedSizes[i]);485 // Skip tile sizes explicitly set to 0.486 if (maybeConstant.has_value() && maybeConstant.value() == 0)487 continue;488 489 PackedOperandsDim packedOperandsDims;490 packedOperandsDims.packedSize = packedSizes[i];491 FailureOr<SmallVector<std::optional<int64_t>>>492 maybePackedDimForEachOperand =493 packLinalgMetadataOnce(indexingMaps, iteratorTypes, i);494 if (failed(maybePackedDimForEachOperand))495 return failure();496 packedOperandsDims.packedDimForEachOperand = *maybePackedDimForEachOperand;497 498 LDBG() << "++++ After pack size #" << i << ": " << packedSizes[i];499 LDBG() << "maps: " << llvm::interleaved(indexingMaps);500 LDBG() << "iterators: " << llvm::interleaved(iteratorTypes);501 LDBG() << "packedDimForEachOperand: "502 << llvm::interleaved(packedOperandsDims.packedDimForEachOperand);503 504 listOfPackedOperandsDim.pushBack(std::move(packedOperandsDims));505 }506 507 // Step 2. Propagate packing to all LinalgOp operands.508 SmallVector<Value> inputsAndInits, results;509 SmallVector<OpOperand *> initOperands =510 llvm::to_vector(llvm::make_pointer_range(linalgOp.getDpsInitsMutable()));511 SmallVector<OpOperand *> inputOperands = linalgOp.getDpsInputOperands();512 for (const auto &operandsList : {inputOperands, initOperands}) {513 for (OpOperand *opOperand : operandsList) {514 int64_t pos = opOperand->getOperandNumber();515 Value operand = opOperand->get();516 SmallVector<int64_t> innerPos =517 listOfPackedOperandsDim.extractPackedDimsForOperand(pos);518 SmallVector<OpFoldResult> innerPackSizes =519 listOfPackedOperandsDim.extractPackSizesForOperand(pos);520 LDBG() << "operand: " << operand;521 LDBG() << "innerPos: " << llvm::interleaved(innerPos);522 LDBG() << "innerPackSizes: " << llvm::interleaved(innerPackSizes);523 if (innerPackSizes.empty()) {524 inputsAndInits.push_back(operand);525 continue;526 }527 Value dest = linalg::PackOp::createDestinationTensor(528 rewriter, loc, operand, innerPackSizes, innerPos,529 /*outerDimsPerm=*/{});530 ShapedType operandType = cast<ShapedType>(operand.getType());531 bool areConstantTiles =532 llvm::all_of(innerPackSizes, [](OpFoldResult tile) {533 return getConstantIntValue(tile).has_value();534 });535 if (areConstantTiles && operandType.hasStaticShape() &&536 !linalg::PackOp::requirePaddingValue(537 operandType.getShape(), innerPos,538 cast<ShapedType>(dest.getType()).getShape(), {},539 innerPackSizes)) {540 packOps.push_back(linalg::PackOp::create(rewriter, loc, operand, dest,541 innerPos, innerPackSizes));542 } else {543 // TODO: value of the padding attribute should be determined by544 // consumers.545 auto zeroAttr =546 rewriter.getZeroAttr(getElementTypeOrSelf(dest.getType()));547 Value zero = arith::ConstantOp::create(rewriter, loc, zeroAttr);548 packOps.push_back(linalg::PackOp::create(549 rewriter, loc, operand, dest, innerPos, innerPackSizes, zero));550 }551 inputsAndInits.push_back(packOps.back());552 }553 }554 555 // Step 3. Build the packed op, use the type of `inits` as result types.556 ValueRange inputs =557 ValueRange{inputsAndInits}.take_front(linalgOp.getNumDpsInputs());558 ValueRange inits =559 ValueRange{inputsAndInits}.take_back(linalgOp.getNumDpsInits());560 auto packedLinalgOp =561 linalg::GenericOp::create(rewriter, linalgOp.getLoc(), inits.getTypes(),562 inputs, inits, indexingMaps, iteratorTypes);563 packedLinalgOp.getRegion().takeBody(linalgOp->getRegion(0));564 565 // Step 4. Propagate packing to all the op results.566 for (OpResult result : packedLinalgOp->getResults()) {567 int64_t resultNum = result.getResultNumber();568 linalg::PackOp maybePackedInit =569 inits[resultNum].getDefiningOp<linalg::PackOp>();570 if (!maybePackedInit) {571 results.push_back(result);572 continue;573 }574 // Build the symmetrical UnPackOp to the existing PackOp.575 unPackOps.push_back(linalg::UnPackOp::create(576 rewriter, packedLinalgOp->getLoc(), result, maybePackedInit.getSource(),577 maybePackedInit.getInnerDimsPos(), maybePackedInit.getMixedTiles()));578 results.push_back(unPackOps.back());579 }580 581 // Step 5. Replace `linalgOp`.582 rewriter.replaceOp(linalgOp, results);583 584 // Return packedLinalgOp.585 return PackResult{packOps,586 cast<linalg::LinalgOp>(packedLinalgOp.getOperation()),587 unPackOps};588}589 590//===----------------------------------------------------------------------===//591// packTranspose transformation.592//===----------------------------------------------------------------------===//593 594/// Return a copy of `tensorType` after permutation by `permutationVector`.595// Note: Should be a new method in of MemRef/RankedTensor/VectorType::Builder596// but this would introduce a dependence on Dialect in IR.597// TODO: Restructure.598static RankedTensorType permuteShape(RankedTensorType tensorType,599 ArrayRef<int64_t> permutationVector) {600 SmallVector<int64_t> shape(tensorType.getShape());601 applyPermutationToVector(shape, permutationVector);602 return RankedTensorType::Builder(tensorType).setShape(shape);603}604 605/// Return a new GenericOp obtained by transposing opOperand by the permutation606/// vector:607/// - the corresponding indexing map is transposed by `permutation`608/// - the corresponding operand value is replaced by `transposedValue`609/// `linalgOp` is replaced by the return op in the process.610/// Asserts that `transposedValue` is of the proper transposed ShapedType.611static LinalgOp transposeOneLinalgOperandAndReplace(612 RewriterBase &rewriter, LinalgOp linalgOp, OpOperand &opOperand,613 ArrayRef<int64_t> permutation, Value transposedValue) {614 // Sanity check the operand.615 assert(linalgOp == opOperand.getOwner() && "linalg op must own the operand");616 617 // Sanity check of the expected transposed tensor type.618 auto tensorType = permuteShape(619 cast<RankedTensorType>(opOperand.get().getType()), permutation);620 (void)tensorType;621 assert(tensorType == transposedValue.getType() &&622 "expected tensor type mismatch");623 624 // Compute the transposed indexing map.625 // Sigh unsigned pollution.626 SmallVector<unsigned> tmpTransposition = llvm::to_vector(627 llvm::map_range(permutation, [](int64_t i) -> unsigned { return i; }));628 AffineMap permutationMap =629 AffineMap::getPermutationMap(tmpTransposition, rewriter.getContext());630 AffineMap transposedMap =631 permutationMap.compose(linalgOp.getMatchingIndexingMap(&opOperand));632 633 // Set the transposed indexing map in the proper position.634 SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();635 indexingMaps[linalgOp.getIndexingMapIndex(&opOperand)] = transposedMap;636 // Set the transposedValue in the proper operand position.637 SmallVector<Value> operands = linalgOp->getOperands();638 operands[opOperand.getOperandNumber()] = transposedValue;639 640 ValueRange operandsRef(operands);641 auto transposedGenericOp = linalg::GenericOp::create(642 rewriter,643 /*location=*/linalgOp->getLoc(),644 /*resultTensorTypes=*/645 operandsRef.drop_front(linalgOp.getNumDpsInputs()).getTypes(),646 /*inputs=*/operandsRef.take_front(linalgOp.getNumDpsInputs()),647 /*outputs=*/operandsRef.drop_front(linalgOp.getNumDpsInputs()),648 /*indexingMaps=*/indexingMaps,649 /*iteratorTypes=*/linalgOp.getIteratorTypesArray());650 transposedGenericOp.getRegion().takeBody(linalgOp->getRegion(0));651 rewriter.replaceOp(linalgOp, transposedGenericOp->getResults());652 653 return cast<linalg::LinalgOp>(transposedGenericOp.getOperation());654}655 656FailureOr<PackTransposeResult>657linalg::packTranspose(RewriterBase &rewriter, linalg::PackOp packOp,658 linalg::LinalgOp linalgOp, linalg::UnPackOp maybeUnPackOp,659 ArrayRef<int64_t> outerPerm,660 ArrayRef<int64_t> innerPerm) {661 Location loc = linalgOp.getLoc();662 663 // Step 1. Transpose packOp.664 rewriter.setInsertionPoint(packOp);665 linalg::PackOp transposedPackOp =666 packOp.createTransposedClone(rewriter, loc, innerPerm, outerPerm);667 668 if (!packOp.getResult().hasOneUse())669 return rewriter.notifyMatchFailure(linalgOp, "expect single pack use");670 671 OpOperand &packUse = *packOp->getUses().begin();672 if (packUse.getOwner() != linalgOp) {673 return rewriter.notifyMatchFailure(674 linalgOp, "not a single use by the LinalgOp target");675 }676 if (maybeUnPackOp &&677 (!linalgOp.isDpsInit(&packUse) ||678 maybeUnPackOp.getSource() != linalgOp.getTiedOpResult(&packUse))) {679 return rewriter.notifyMatchFailure(linalgOp,680 "not produced by the LinalgOp target");681 }682 683 // Step 2. Transpose linalgOp.684 // transposedPackOp.getOuterDimsPerm() may be empty, in which case it is the685 // identity. Don't rely on it.686 int64_t numLeadingDims = packOp.getSourceRank();687 int64_t numTrailingDims = packOp.getInnerDimsPos().size();688 // Step 2.a. Compute the permutation on the whole operand.689 // Leading part just reuse the outerPerm.690 SmallVector<int64_t> permutation(outerPerm);691 if (permutation.empty())692 llvm::append_range(permutation, llvm::seq<int64_t>(0, numLeadingDims));693 // Trailing part needs to reindex positions by `numLeadingDims`.694 if (innerPerm.empty()) {695 llvm::append_range(696 permutation,697 llvm::seq<int64_t>(numLeadingDims, numLeadingDims + numTrailingDims));698 } else {699 llvm::append_range(permutation,700 llvm::map_range(innerPerm, [&](int64_t pos) {701 return numLeadingDims + pos;702 }));703 }704 if (!isPermutationVector(permutation))705 return rewriter.notifyMatchFailure(linalgOp, "invalid permutation");706 707 // Step 2.b. Save the transposedPackUse operand number in case we need to708 // get the tied OpResult after `linalgOp` has been replaced.709 int64_t packUseOperandNumber = packUse.getOperandNumber();710 // Step 2.c. Actually perform the transposition.711 rewriter.setInsertionPoint(linalgOp);712 linalg::LinalgOp transposedLinalgOp = transposeOneLinalgOperandAndReplace(713 rewriter, linalgOp, packUse, permutation, transposedPackOp.getResult());714 715 // Step 3. Maybe transpose unPackOp.716 linalg::UnPackOp transposedUnPackOp;717 if (maybeUnPackOp) {718 OpOperand &opOperand =719 transposedLinalgOp->getOpOperand(packUseOperandNumber);720 OpResult transposedResult = transposedLinalgOp.getTiedOpResult(&opOperand);721 rewriter.setInsertionPoint(maybeUnPackOp);722 transposedUnPackOp = maybeUnPackOp.createTransposedClone(723 rewriter, loc, transposedResult, innerPerm, outerPerm);724 725 rewriter.replaceOp(maybeUnPackOp, transposedUnPackOp->getResults());726 }727 728 // Step 4. Finally, replace packOp now that we don't need it anymore.729 rewriter.replaceOp(packOp, transposedPackOp->getResults());730 731 return PackTransposeResult{transposedPackOp, transposedLinalgOp,732 transposedUnPackOp};733}734 735//===----------------------------------------------------------------------===//736// packMatmulGreedily transformation.737//===----------------------------------------------------------------------===//738 739/// Pack a LinalgOp by greedily inferring matmul dimensions (m, n, k) where m740/// and n are proper parallel dimensions and k is a proper reduction741/// dimension. Packing occurs by rewriting the op as a linalg.generic and742/// calling linalg::pack by `mnkPackedSizes`. The order of the packed743/// dimensions is customizable: the `mnkOrder` is a permutation of {0, 1, 2}744/// to reorder {m, n, k} into one of the 8 possible forms. The outer745/// dimensions of the operands are not permuted at this time, this is left for746/// future work.747FailureOr<PackResult>748linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,749 ArrayRef<OpFoldResult> mnkPackedSizes,750 ArrayRef<int64_t> mnkPaddedSizesNextMultipleOf,751 ArrayRef<int64_t> mnkOrder) {752 assert(mnkPackedSizes.size() == 3 && "unexpected num of packing sizes");753 assert((mnkPaddedSizesNextMultipleOf.empty() ||754 mnkPaddedSizesNextMultipleOf.size() == 3) &&755 "num of packing sizes next multiple should be empty or of size 3");756 assert(mnkOrder.size() == 3 && "unexpected mnkOrder size");757 assert(isPermutationVector(mnkOrder) && "expected a permutation");758 759 int64_t numLoops = linalgOp.getNumLoops();760 if (numLoops <= 2) {761 LDBG() << "need 3+ loops to find a matmul to pack, got " << numLoops762 << " in: " << linalgOp;763 return rewriter.notifyMatchFailure(764 linalgOp, "need 3+ loops to find a matmul to pack");765 }766 767 // Locally adjust the desired iterator position of mnk and packing sizes.768 int64_t numPackedDims = mnkPackedSizes.size();769 SmallVector<int64_t> mmnnkkPos(numPackedDims);770 for (int64_t i = 0, e = numPackedDims; i < e; ++i)771 mmnnkkPos[i] = numLoops - numPackedDims + mnkOrder[i];772 SmallVector<OpFoldResult> packedSizes(numPackedDims);773 for (int64_t i = 0, e = numPackedDims; i < e; ++i)774 packedSizes[mnkOrder[i]] = mnkPackedSizes[i];775 SmallVector<int64_t> paddedSizesNextMultipleOf(numPackedDims);776 for (int64_t i = 0, e = numPackedDims; i < e; ++i) {777 paddedSizesNextMultipleOf[mnkOrder[i]] =778 mnkPaddedSizesNextMultipleOf.empty() ? 0779 : mnkPaddedSizesNextMultipleOf[i];780 }781 782 // 1. Infer dims that are important for matmul.783 FailureOr<ContractionDimensions> maybeDimensions =784 inferContractionDims(linalgOp);785 if (failed(maybeDimensions)) {786 LDBG() << "couldn't infer matmul iterators in: " << linalgOp;787 return rewriter.notifyMatchFailure(linalgOp,788 "couldn't infer matmul iterators");789 }790 791 // 2. Normalize linalgOp to an kmn-matmul-like with [red, par, par] most792 // minor iterators. In cases with multiple options for m, n, k bias towards793 // the most minor embedding.794 // If we wanted a different normalization order, this is where it would have795 // to plug a heuristic.796 int64_t mPos = maybeDimensions->m.back(), nPos = maybeDimensions->n.back(),797 kPos = maybeDimensions->k.back();798 LDBG() << "Start packing generic op greedily with (m@" << mPos << ", n@"799 << nPos << ", k@" << kPos << "): " << linalgOp;800 801 // 2.a. Rewrite as a generic.802 auto genericOp = dyn_cast<GenericOp>(linalgOp.getOperation());803 if (!genericOp) {804 FailureOr<GenericOp> generalizeResult =805 generalizeNamedOp(rewriter, linalgOp);806 assert(succeeded(generalizeResult) && "unexpected failure generalizing op");807 genericOp = *generalizeResult;808 }809 810 // 2.b. Interchange to move the dimensions (k, m, n) as most-minor811 // iterators. Note that this only normalized the iteration order and does812 // not change the indexings of any operand.813 SmallVector<int64_t> permutation =814 computePermutationVector(numLoops, {mPos, nPos, kPos}, mmnnkkPos);815 LDBG() << "perm: " << llvm::interleaved(permutation);816 // Sign .. unsigned pollution.817 SmallVector<unsigned> unsignedPerm(permutation.begin(), permutation.end());818 FailureOr<GenericOp> interchangeResult =819 interchangeGenericOp(rewriter, genericOp, unsignedPerm);820 assert(succeeded(interchangeResult) && "unexpected failure interchanging op");821 genericOp = *interchangeResult;822 LDBG() << "Generalized Op to pack: " << genericOp;823 824 // At this point, the op iterators are normalized to {leading, k, m, n}.825 // The layouts induced by packing will always be:826 // - LHS{leading_lhs, kk, mm}827 // - RHS{leading_rhs, kk, nn}828 // - RES{leading_res, mm, nn}829 // If we wanted to change the packed order, we would reorder (k, m, n) to830 // something else above.831 //832 // Additional permutations of the outer dims of the operands (i.e.833 // leading_lhs, leading_rhs and leading_res) could follow by computing the834 // desired outerPerm for each operand.835 // This is left for future work.836 837 // TODO: this creates too much IR, go use reifyResultShapes.838 SmallVector<Range, 4> loopRanges =839 cast<LinalgOp>(genericOp.getOperation())840 .createLoopRanges(rewriter, genericOp.getLoc());841 842 // Add leading zeros to match numLoops, we only pack the last 3 dimensions843 // post interchange.844 LDBG() << "paddedSizesNextMultipleOf: "845 << llvm::interleaved(paddedSizesNextMultipleOf);846 LDBG() << "loopRanges: "847 << llvm::interleaved(848 llvm::map_range(loopRanges, [](Range r) { return r.size; }));849 SmallVector<OpFoldResult> adjustedPackedSizes(numLoops - packedSizes.size(),850 rewriter.getIndexAttr(0));851 for (int64_t i = 0, e = numPackedDims; i < e; ++i) {852 if (paddedSizesNextMultipleOf[i] == 0) {853 adjustedPackedSizes.push_back(packedSizes[i]);854 continue;855 }856 AffineExpr d0, s0;857 bindDims(rewriter.getContext(), d0);858 bindSymbols(rewriter.getContext(), s0);859 adjustedPackedSizes.push_back(affine::makeComposedFoldedAffineApply(860 rewriter, genericOp->getLoc(), d0.ceilDiv(s0) * s0,861 {loopRanges[adjustedPackedSizes.size()].size,862 rewriter.getIndexAttr(paddedSizesNextMultipleOf[i])}));863 }864 LDBG() << "adjustedPackedSizes: " << llvm::interleaved(adjustedPackedSizes);865 866 // TODO: If we wanted to give the genericOp a name after packing, after867 // calling `pack` would be a good time. One would still need to check that868 // `containsMostMinorMatmul(packingRes->packedLinalgOp)` is true, since we869 // also allow degenerate matmul cases (i.e. matvec, dot).870 return pack(rewriter, genericOp, adjustedPackedSizes);871}872 873//===----------------------------------------------------------------------===//874// Transformations exposed as rewrite patterns.875//===----------------------------------------------------------------------===//876 877LinalgTilingOptions &878mlir::linalg::LinalgTilingOptions::setTileSizes(ArrayRef<int64_t> ts) {879 assert(!tileSizeComputationFunction && "tile sizes already set");880 SmallVector<int64_t, 4> tileSizes(ts);881 tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {882 OpBuilder::InsertionGuard guard(b);883 b.setInsertionPointToStart(884 &op->getParentOfType<func::FuncOp>().getBody().front());885 return llvm::to_vector<4>(map_range(tileSizes, [&](int64_t s) {886 Value v = arith::ConstantIndexOp::create(b, op->getLoc(), s);887 return v;888 }));889 };890 return *this;891}892 893LogicalResult mlir::linalg::CopyVectorizationPattern::matchAndRewrite(894 memref::CopyOp copyOp, PatternRewriter &rewriter) const {895 return vectorizeCopy(rewriter, copyOp);896}897 898/// Filling `dest` using FillOp constant padding value if possible.899/// Otherwise, generate a tensor::GenerateOp.900Value DecomposePadOpPattern::createFillOrGenerateOp(901 RewriterBase &rewriter, tensor::PadOp padOp, Value dest,902 const SmallVector<Value> &dynSizes) const {903 auto padValue = padOp.getConstantPaddingValue();904 if (padValue) {905 // Move the padding value defined inside the PadOp block to outside.906 if (padValue.getParentBlock() == &padOp.getRegion().front())907 rewriter.moveOpBefore(padValue.getDefiningOp(), padOp);908 return FillOp::create(rewriter, padOp.getLoc(), padValue, dest).result();909 }910 911 // Fill could not be optimized: Lower to tensor::GenerateOp with region.912 auto generateOp = tensor::GenerateOp::create(rewriter, padOp.getLoc(),913 padOp.getResultType(), dynSizes);914 // Copy region to new op.915 IRMapping bvm;916 padOp.getRegion().cloneInto(&generateOp.getRegion(), bvm);917 return generateOp;918}919 920LogicalResult921DecomposePadOpPattern::matchAndRewrite(tensor::PadOp padOp,922 PatternRewriter &rewriter) const {923 // Given an OpFoldResult, return an index-typed value.924 auto getIdxValue = [&](OpFoldResult ofr) {925 if (auto val = llvm::dyn_cast_if_present<Value>(ofr))926 return val;927 return arith::ConstantIndexOp::create(928 rewriter, padOp.getLoc(),929 cast<IntegerAttr>(cast<Attribute>(ofr)).getInt())930 .getResult();931 };932 933 auto resultType = padOp.getResultType();934 // Compute size of EmptyOp. Any combination of static/dynamic is supported.935 SmallVector<Value> dynSizes;936 SmallVector<int64_t> staticSizes;937 for (unsigned dim = 0; dim < resultType.getRank(); ++dim) {938 if (resultType.isDynamicDim(dim)) {939 auto srcSize = getIdxValue(tensor::getMixedSize(rewriter, padOp.getLoc(),940 padOp.getSource(), dim));941 // Add low and high padding value.942 auto plusLow = rewriter.createOrFold<arith::AddIOp>(943 padOp.getLoc(), srcSize, getIdxValue(padOp.getMixedLowPad()[dim]));944 auto plusHigh = rewriter.createOrFold<arith::AddIOp>(945 padOp.getLoc(), plusLow, getIdxValue(padOp.getMixedHighPad()[dim]));946 dynSizes.push_back(plusHigh);947 }948 staticSizes.push_back(resultType.getDimSize(dim));949 }950 951 // Init tensor and fill it with padding.952 Value emptyTensor =953 tensor::EmptyOp::create(rewriter, padOp.getLoc(), staticSizes,954 resultType.getElementType(), dynSizes);955 Value fill = createFillOrGenerateOp(rewriter, padOp, emptyTensor, dynSizes);956 957 // Generate a InsertSliceOp for copying the PadOp source.958 auto sourceType = padOp.getSourceType();959 // Compute size of source of tensor::PadOp.960 SmallVector<OpFoldResult> srcSizes =961 tensor::getMixedSizes(rewriter, padOp.getLoc(), padOp.getSource());962 // Strides of InsertSliceOp are all 1.963 SmallVector<OpFoldResult> strides(sourceType.getRank(),964 rewriter.getIndexAttr(1));965 rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(966 padOp, padOp.getSource(), fill, padOp.getMixedLowPad(), srcSizes,967 strides);968 969 return success();970}971 972LogicalResult ExtractSliceOfPadTensorSwapPattern::matchAndRewrite(973 tensor::ExtractSliceOp sliceOp, PatternRewriter &rewriter) const {974 if (!sliceOp.hasUnitStride())975 return failure();976 977 auto padOp = sliceOp.getSource().getDefiningOp<tensor::PadOp>();978 if (!padOp)979 return failure();980 981 bool zeroSliceGuard = true;982 if (controlFn) {983 if (std::optional<bool> control = controlFn(sliceOp))984 zeroSliceGuard = *control;985 else986 return failure();987 }988 989 FailureOr<TilingResult> tilingResult =990 tensor::bubbleUpPadSlice(rewriter, padOp, sliceOp.getMixedOffsets(),991 sliceOp.getMixedSizes(), zeroSliceGuard);992 if (failed(tilingResult))993 return failure();994 995 RankedTensorType sourceType = sliceOp.getSourceType();996 RankedTensorType resultType = sliceOp.getResultType();997 998 // If the extract_slice is not rank-reduced, all shapes are static and the999 // data source is actually used. Rewrite into pad(extract_slice(x)).1000 if (sourceType.getRank() == resultType.getRank()) {1001 rewriter.replaceOp(sliceOp, tilingResult->tiledValues);1002 return success();1003 }1004 1005 // Handle rank-reduced slice by creating another extract_slice op.1006 Value rankReduced = tensor::createCanonicalRankReducingExtractSliceOp(1007 rewriter, sliceOp.getLoc(), tilingResult->tiledValues[0], resultType);1008 1009 rewriter.replaceOp(sliceOp, rankReduced);1010 return success();1011}1012 1013/// If padding value is set, returns a tensor.pad Op for the source tensor,1014/// with the output shape matching the output of `packOp`. Otherwise, returns1015/// the source directly.1016///1017/// This method assumes that all outer dims for this pack Op are 1.1018static Value getPackOpSourceOrPaddedSource(OpBuilder &builder,1019 linalg::PackOp packOp) {1020 Value input = packOp.getSource();1021 if (!packOp.getPaddingValue()) {1022 return input;1023 }1024 1025 assert(llvm::all_of(packOp.getAllOuterDims(),1026 [](int64_t val) { return val == 1; }) &&1027 "some outer dims are != 1");1028 1029 Location loc = packOp.getLoc();1030 ShapedType inputType = packOp.getSourceType();1031 int64_t inputRank = inputType.getRank();1032 1033 DenseMap<int64_t, OpFoldResult> tileAndPosMapping =1034 packOp.getDimAndTileMapping();1035 1036 // The sizes of dynamic tiles1037 SmallVector<Value> dynamicTileSizes;1038 1039 // Collect dims for the padded shape.1040 SmallVector<int64_t> paddedShape;1041 for (int64_t dimIdx = 0; dimIdx < inputRank; ++dimIdx) {1042 // 1. Non-tiled outer dims.1043 // These dims should be 1 and we simply preserve them.1044 if (!tileAndPosMapping.count(dimIdx)) {1045 int64_t inputDimSize = inputType.getDimSize(dimIdx);1046 assert(inputDimSize == 1 &&1047 "with all outer dims == 1, this non-tiled input dim should be 1!");1048 paddedShape.push_back(inputDimSize);1049 continue;1050 }1051 1052 // 2. Tiled outer dims1053 // As all outer dims == 1, it is safe to use the tile size for the padded1054 // shape.1055 OpFoldResult tileSizeForDim = tileAndPosMapping.lookup(dimIdx);1056 1057 // 2.1 Static tile sizes1058 std::optional<int64_t> cstTileSize = getConstantIntValue(tileSizeForDim);1059 if (cstTileSize.has_value()) {1060 paddedShape.push_back(cstTileSize.value());1061 continue;1062 }1063 1064 // 2.2 Dynamic tile sizes1065 paddedShape.push_back(ShapedType::kDynamic);1066 1067 // Get the value that holds the dynamic size.1068 dynamicTileSizes.push_back(llvm::dyn_cast<Value>(tileSizeForDim));1069 }1070 auto resultType =1071 RankedTensorType::get(paddedShape, inputType.getElementType());1072 return tensor::createPadHighOp(resultType, input, packOp.getPaddingValue(),1073 /*nofold=*/false, loc, builder,1074 dynamicTileSizes);1075}1076 1077// Normalizes a permutation on a higher rank space to its actual size, e.g.1078// perm = [1, 4, 2]1079// becomes1080// norm = [0, 2, 1]1081static SmallVector<int64_t>1082getPackUnpackNormalizedPerm(int rank, ArrayRef<int64_t> perm) {1083 constexpr int64_t kNonTiledMarker = -1;1084 SmallVector<int64_t> vec(rank, kNonTiledMarker);1085 for (auto [index, value] : llvm::enumerate(perm))1086 vec[value] = index;1087 SmallVector<int64_t> normalizedPerm = llvm::filter_to_vector(1088 vec, [&](int64_t v) { return v != kNonTiledMarker; });1089 // This inverts the permutation in addition to normalizing so invert back.1090 return invertPermutationVector(normalizedPerm);1091}1092 1093// Gets the normalized permutation implied by innerDimsPos and outerDimsPerm1094// assuming rank reduction of unit outer dims.1095static SmallVector<int64_t>1096getPackUnpackRankReducedPerm(ArrayRef<int64_t> shape,1097 ArrayRef<int64_t> innerDimsPos,1098 ArrayRef<int64_t> outerDimsPerm) {1099 SmallVector<int64_t> rankReducedOuterDimsPerm;1100 SmallVector<int64_t> outerDims;1101 SmallVector<int64_t> innerDims;1102 int64_t dim = 0;1103 int64_t unpackedRank = shape.size();1104 for (auto i : llvm::seq<unsigned>(0, unpackedRank)) {1105 if (llvm::is_contained(innerDimsPos, i)) {1106 innerDims.push_back(dim++);1107 continue;1108 }1109 if (shape[i] == 1)1110 continue;1111 outerDims.push_back(dim++);1112 if (!outerDimsPerm.empty())1113 rankReducedOuterDimsPerm.push_back(outerDimsPerm[i]);1114 }1115 1116 // Get the position of the inner dims after permutation.1117 SmallVector<int64_t> innerPerm =1118 getPackUnpackNormalizedPerm(unpackedRank, innerDimsPos);1119 applyPermutationToVector<int64_t>(innerDims, innerPerm);1120 1121 // Ditto for the outer dims.1122 SmallVector<int64_t> perm = outerDims;1123 1124 rankReducedOuterDimsPerm =1125 getPackUnpackNormalizedPerm(unpackedRank, rankReducedOuterDimsPerm);1126 if (!rankReducedOuterDimsPerm.empty())1127 applyPermutationToVector<int64_t>(perm, rankReducedOuterDimsPerm);1128 1129 // The tile always ends up as the inner most dims after packing.1130 perm.append(innerDims);1131 1132 return perm;1133}1134 1135LogicalResult DecomposeOuterUnitDimsPackOpPattern::matchAndRewrite(1136 linalg::PackOp packOp, PatternRewriter &rewriter) const {1137 if (llvm::any_of(packOp.getTiledOuterDims(),1138 [](int64_t dim) { return dim != 1; })) {1139 return rewriter.notifyMatchFailure(1140 packOp, "not all outer dimensions of the result are 1s");1141 }1142 1143 ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();1144 auto outerDimsPerm = packOp.getOuterDimsPerm();1145 1146 // Verify that there are no:1147 // * non-unit + un-tiled-outer-dims,1148 // that are permuted. Supporting such cases would require refining the logic1149 // that generates the Transpose Op.1150 if (!llvm::all_of(outerDimsPerm, [&innerDimsPos, &packOp](int64_t dim) {1151 static int prev = 0;1152 // Skip tiled dims - these can be permuted.1153 if (llvm::is_contained(innerDimsPos, dim))1154 return true;1155 1156 // Check whether this dim has been permuted. Permuting unit dims is fine1157 // as that's effectively a no-op.1158 if (dim < prev && (packOp.getType().getShape()[prev] != 1 ||1159 packOp.getType().getShape()[dim] != 1))1160 return false;1161 1162 prev = dim;1163 return true;1164 })) {1165 return rewriter.notifyMatchFailure(1166 packOp, "At least one non-unit and un-tiled outer dim is permuted, "1167 "this is not supported ATM!");1168 }1169 1170 Location loc = packOp.getLoc();1171 1172 int64_t srcRank = packOp.getSourceRank();1173 1174 // 1. Get the input that is going to be packed. If the input requires padding,1175 // add a padding operation and return that as the input.1176 Value input = getPackOpSourceOrPaddedSource(rewriter, packOp);1177 1178 // 2. Transpose the input to match the inner tile order:1179 // %init = tensor.empty()1180 // %transposed_tile = linalg.transpose ins(%source_or_padded_source),1181 // outs(%init)1182 // Assumptions made:1183 // - All tiled outer dims are 1 - the corresponding transposition order1184 // doesn't matter, but requires all dim indices to be present.1185 // - Un-tiled outer dims remain un-permuted.1186 1187 // 2.1 Get the permutation for linalg.transpose:1188 // [ untiled-dims, inner-dims-pos ]1189 // Note, this logic assumes that the untiled dims are not permuted.1190 SmallVector<int64_t> srcPermForTranspose;1191 for (int64_t i = 0; i < srcRank; i++) {1192 // We assume the `k` dimensions of the inner dim position, where `k` is the1193 // rank of the inner tiling, correspond to the last `k` indices of the1194 // transpose permutation. This is done by adding the indices not contained1195 // in the inner dimension position in order from 0 to `n`. Where n is the1196 // rank of the source tensor. For example if we have a source tensor with1197 // indices [0, 1, 2, 3] and inner dim position of [3, 0], the remaining1198 // indices are [1, 2]. and the transpose will be [1, 2, 3, 0].1199 if (llvm::is_contained(innerDimsPos, i))1200 continue;1201 srcPermForTranspose.push_back(i);1202 }1203 srcPermForTranspose.append(innerDimsPos.begin(), innerDimsPos.end());1204 1205 // 2.2 Create the init tensor for linalg.transpose with the correct shape:1206 // [ untiled-dims, tiled-dims ]1207 ShapedType inputTy = cast<ShapedType>(input.getType());1208 SmallVector<OpFoldResult> shapeForEmptyOp;1209 for (int64_t i = 0; i < srcRank; i++) {1210 if (llvm::is_contained(innerDimsPos, i)) {1211 // The tiled dims are appended after this loop.1212 continue;1213 }1214 if (inputTy.isStaticDim(i))1215 shapeForEmptyOp.push_back(rewriter.getIndexAttr(inputTy.getShape()[i]));1216 else1217 shapeForEmptyOp.emplace_back(1218 tensor::DimOp::create(rewriter, loc, input, i).getResult());1219 }1220 shapeForEmptyOp.append(packOp.getMixedTiles());1221 1222 // getMixedTiles() may contain Values pointing to constant ops, not the1223 // constant attributes. Replace them with a true OpFoldResult.1224 llvm::transform(shapeForEmptyOp, shapeForEmptyOp.begin(),1225 [&](OpFoldResult ofr) {1226 if (auto val = llvm::dyn_cast<Value>(ofr))1227 return getAsOpFoldResult(val);1228 return ofr;1229 });1230 1231 LDBG() << "Pack permutation: " << packOp;1232 LDBG() << "perm: " << llvm::interleaved(srcPermForTranspose);1233 LDBG() << "Shape of empty tensor: " << llvm::interleaved(shapeForEmptyOp);1234 1235 Value empty = tensor::EmptyOp::create(1236 rewriter, loc, shapeForEmptyOp, packOp.getSourceType().getElementType());1237 1238 // 2.3 Create linalg.transpose1239 auto transposedOp = linalg::TransposeOp::create(rewriter, loc, input, empty,1240 srcPermForTranspose);1241 1242 // 3. Insert the inner tile into the destination tensor:1243 // %inserted_tile = tensor.insert_slice(%transposed_tile)1244 1245 // Compute the sizes attribute:1246 // [ outer-dims, tile-sizes ]1247 // Note that the output from the transpose Op excludes the tiled outer dims.1248 // However, given the assumption that:1249 // * all tiled outer dims == 1,1250 // we can just use a rank-expanding tensor.insert_slice.1251 SmallVector<OpFoldResult> writeSizes;1252 for (auto size : packOp.getAllOuterDims()) {1253 writeSizes.push_back(rewriter.getIndexAttr(size));1254 }1255 1256 for (auto tileSize : packOp.getMixedTiles()) {1257 auto [_, tileSizeOfr] =1258 getSimplifiedOfrAndStaticSizePair(tileSize, rewriter);1259 writeSizes.push_back(tileSizeOfr);1260 }1261 1262 auto insert = tensor::InsertSliceOp::create(1263 rewriter, loc, transposedOp.getResult()[0], packOp.getDest(), writeSizes);1264 1265 // 4. Replace tensor.packOp with tensor.insert_slice created above1266 rewriter.replaceOp(packOp, insert.getResult());1267 1268 return success();1269}1270 1271LogicalResult DecomposeOuterUnitDimsUnPackOpPattern::matchAndRewrite(1272 linalg::UnPackOp unpackOp, PatternRewriter &rewriter) const {1273 int64_t destRank = unpackOp.getDestRank();1274 ArrayRef<int64_t> srcShape = unpackOp.getSourceType().getShape();1275 ArrayRef<int64_t> innerDimsPos = unpackOp.getInnerDimsPos();1276 if (llvm::any_of(unpackOp.getTiledOuterDims(),1277 [](int64_t dim) { return dim != 1; })) {1278 return rewriter.notifyMatchFailure(1279 unpackOp,1280 "require the tiled outer dimensions of the result are all 1s");1281 }1282 1283 // 1. Use rank-reduced tensor.extract_slice op to extract the tile:1284 // %extracted_tile = tensor.extract_slice(%unpack_op_input)1285 Location loc = unpackOp.getLoc();1286 Value source = unpackOp.getSource();1287 DenseMap<int64_t, OpFoldResult> dimAndTileMapping =1288 unpackOp.getDimAndTileMapping();1289 Attribute oneIdxAttr = rewriter.getIndexAttr(1);1290 1291 // The shape for ExtractSliceOp. Note that this will consist of 3 blocks of1292 // dims:1293 // [ outer-untiled-dims, outer-tiled-dims, tile-sizes ]1294 SmallVector<int64_t> readShapeForExtractSlice;1295 // The sizes attribute for ExtractSliceOp. Due to rank-reducing (and1296 // outer-tiled-dims being all 1), this will be1297 // [ outer-untiled-dims, tile-sizes ]1298 SmallVector<OpFoldResult> extractSliceSizes;1299 1300 // Shape for EmptyOp that's used as the init value for TransposeOp below.1301 // This should be:1302 // [ outer-untiled-dims, tile-sizes ]1303 // However, skip unit dims - TransposeOp (below) applies rank-reduced1304 // permutation.1305 SmallVector<OpFoldResult> shapeForEmptyOp;1306 1307 for (auto i : llvm::seq<unsigned>(0, destRank)) {1308 // Compute sizes attribute for ExtractSliceOp - outer-tiled-dims.1309 //1310 // As all outer tiled dims are 1, so the corresponding1311 // slice size to read will also 1. As this will be rank-reducing "extract1312 // slice" (i.e. the unit dims will be "collapsed"), there's no need to1313 // update:1314 // * the output shape for ExtractSliceOp, nor1315 // * the shape for EmptyOp.1316 if (dimAndTileMapping.count(i)) {1317 extractSliceSizes.push_back(oneIdxAttr);1318 continue;1319 }1320 1321 // Compute sizes attribute for ExtractSliceOp + EmptyOp -1322 // outer-untiled-dims1323 if (ShapedType::isDynamic(srcShape[i])) {1324 OpFoldResult dynamicDim =1325 tensor::DimOp::create(rewriter, loc, source, i).getResult();1326 extractSliceSizes.push_back(dynamicDim);1327 shapeForEmptyOp.push_back(dynamicDim);1328 } else {1329 extractSliceSizes.push_back(rewriter.getIndexAttr(srcShape[i]));1330 if (srcShape[i] != 1)1331 shapeForEmptyOp.push_back(rewriter.getIndexAttr(srcShape[i]));1332 }1333 // Compute the output shape for ExtractSliceOp - outer-untiled-dims (take1334 // into account rank-reducing)1335 if (srcShape[i] != 1) {1336 readShapeForExtractSlice.push_back(srcShape[i]);1337 }1338 }1339 // Append the tile sizes to "sizes attribute" for ExtractSliceOp and the1340 // shape for EmptyOp.1341 auto mixedTiles = unpackOp.getMixedTiles();1342 extractSliceSizes.append(mixedTiles.begin(), mixedTiles.end());1343 shapeForEmptyOp.append(mixedTiles.begin(), mixedTiles.end());1344 1345 // Explicitly create the type for extract_slice op because the inner tile1346 // size could be 1. We want to represent the whole inner tile in this case.1347 auto tileShape = srcShape.drop_front(destRank);1348 // Append the inner tile shape to the permuted and rank-reduced outer shape.1349 readShapeForExtractSlice.append(tileShape.begin(), tileShape.end());1350 Type elemType = unpackOp.getSourceType().getElementType();1351 auto readType = RankedTensorType::get(readShapeForExtractSlice, elemType);1352 Value innerTile = tensor::ExtractSliceOp::create(1353 rewriter, loc, readType, unpackOp.getSource(), extractSliceSizes);1354 1355 // 2. Transpose the tile to match the outer corresponding tile order.1356 SmallVector<int64_t> perm = getPackUnpackRankReducedPerm(1357 srcShape.take_front(destRank), innerDimsPos, unpackOp.getOuterDimsPerm());1358 // Unpack is a transition out of packed space so we invert the permutation.1359 perm = invertPermutationVector(perm);1360 applyPermutationToVector<OpFoldResult>(shapeForEmptyOp, perm);1361 1362 Value empty =1363 tensor::EmptyOp::create(rewriter, loc, shapeForEmptyOp, elemType);1364 auto transposedOp =1365 linalg::TransposeOp::create(rewriter, loc, innerTile, empty, perm);1366 1367 // 3. Handle in-complete tiles if needed. It truncates trailing data from the1368 // transposed tile.1369 SmallVector<OpFoldResult> tileSizes;1370 ArrayRef<int64_t> destShape = unpackOp.getDestType().getShape();1371 for (auto i : llvm::seq<unsigned>(0, destRank)) {1372 if (dimAndTileMapping.count(i) || destShape[i] != 1)1373 tileSizes.push_back(1374 tensor::getMixedSize(rewriter, loc, unpackOp.getDest(), i));1375 }1376 1377 auto partialTile =1378 tensor::ExtractSliceOp::create(rewriter, loc, RankedTensorType(),1379 transposedOp.getResult()[0], tileSizes);1380 1381 // 4. Insert the result to the destination tensor.1382 SmallVector<OpFoldResult> writeSizes;1383 for (int i = 0, idx = 0; i < destRank; ++i) {1384 if (dimAndTileMapping.count(i) || destShape[i] != 1)1385 writeSizes.push_back(tileSizes[idx++]);1386 else1387 writeSizes.push_back(oneIdxAttr);1388 }1389 auto insert = tensor::InsertSliceOp::create(rewriter, loc, partialTile,1390 unpackOp.getDest(), writeSizes);1391 rewriter.replaceOp(unpackOp, insert.getResult());1392 1393 return success();1394}1395 1396// The following are patterns for downscaling convolution ops with size-11397// window dimensions.1398//1399// Note that we'd eventually want to write such transformations in a generic1400// way, e.g., converting to linalg.generic, removing the size-1 dimensions,1401// and then turning back to named ops. But for now it's fine to have a few1402// patterns matching special ops to get started.1403 1404template <typename Conv2DOp, typename Conv1DOp>1405FailureOr<Conv1DOp> DownscaleSizeOneWindowed2DConvolution<Conv2DOp, Conv1DOp>::1406 returningMatchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const {1407 if (convOp.hasPureBufferSemantics())1408 return failure(); // To be implemented.1409 1410 Value input = convOp.getInputs().front();1411 Value kernel = convOp.getInputs().back();1412 Value output = convOp.getOutputs().front();1413 1414 auto inputType = dyn_cast<RankedTensorType>(input.getType());1415 auto kernelType = dyn_cast<RankedTensorType>(kernel.getType());1416 auto outputType = dyn_cast<RankedTensorType>(output.getType());1417 1418 auto kernelShape = kernelType.getShape();1419 auto outputShape = outputType.getShape();1420 1421 // Get domain indices based on conv2D layout.1422 auto [khIndex, kwIndex, ohIndex, owIndex] =1423 TypeSwitch<Operation *, std::tuple<int64_t, int64_t, int64_t, int64_t>>(1424 convOp)1425 .Case([&](linalg::Conv2DNhwcHwcfOp op) {1426 return std::make_tuple(0, 1, 1, 2);1427 })1428 .Case([&](linalg::Conv2DNchwFchwOp op) {1429 return std::make_tuple(2, 3, 2, 3);1430 })1431 .Case([&](linalg::PoolingNhwcSumOp op) {1432 return std::make_tuple(0, 1, 1, 2);1433 })1434 .Case([&](linalg::PoolingNchwSumOp op) {1435 return std::make_tuple(0, 1, 2, 3);1436 })1437 .Case([&](linalg::PoolingNhwcMaxOp op) {1438 return std::make_tuple(0, 1, 1, 2);1439 })1440 .Case([&](linalg::PoolingNhwcMaxUnsignedOp op) {1441 return std::make_tuple(0, 1, 1, 2);1442 })1443 .Case([&](linalg::PoolingNhwcMinOp op) {1444 return std::make_tuple(0, 1, 1, 2);1445 })1446 .Case([&](linalg::PoolingNhwcMinUnsignedOp op) {1447 return std::make_tuple(0, 1, 1, 2);1448 })1449 .Case([&](linalg::PoolingNchwMaxOp op) {1450 return std::make_tuple(0, 1, 2, 3);1451 })1452 .DefaultUnreachable("unexpected conv2d/pool2d operation.");1453 1454 // Only handle the case where at least one of the window dimensions is1455 // of size 1. Other cases can rely on tiling to reduce to such cases.1456 int64_t khSize = kernelShape[khIndex], kwSize = kernelShape[kwIndex];1457 int64_t ohSize = outputShape[ohIndex], owSize = outputShape[owIndex];1458 bool removeH = (khSize == 1 && ohSize == 1);1459 bool removeW = (kwSize == 1 && owSize == 1);1460 if (!removeH && !removeW)1461 return failure();1462 1463 // Get new shapes and types for all operands by removing the size-11464 // dimension.1465 using RTTBuilder = RankedTensorType::Builder;1466 RankedTensorType newInputType =1467 RTTBuilder(inputType).dropDim((removeH ? ohIndex : owIndex));1468 RankedTensorType newKernelType =1469 RTTBuilder(kernelType).dropDim((removeH ? khIndex : kwIndex));1470 RankedTensorType newOutputType =1471 RTTBuilder(outputType).dropDim((removeH ? ohIndex : owIndex));1472 1473 // Rank-reduce operands.1474 Location loc = convOp.getLoc();1475 Value newInput = tensor::createCanonicalRankReducingExtractSliceOp(1476 rewriter, loc, input, newInputType);1477 Value newKernel = tensor::createCanonicalRankReducingExtractSliceOp(1478 rewriter, loc, kernel, newKernelType);1479 Value newOutput = tensor::createCanonicalRankReducingExtractSliceOp(1480 rewriter, loc, output, newOutputType);1481 1482 // Rank-reduce strides and dilations too.1483 // TODO: dropDim 1-liner helper.1484 auto strides =1485 llvm::to_vector<4>(convOp.getStrides().template getValues<int64_t>());1486 strides.erase(strides.begin() + (removeH ? 0 : 1));1487 auto stridesAttr = rewriter.getI64VectorAttr(strides);1488 1489 auto dilations =1490 llvm::to_vector<4>(convOp.getDilations().template getValues<int64_t>());1491 dilations.erase(dilations.begin() + (removeH ? 0 : 1));1492 auto dilationsAttr = rewriter.getI64VectorAttr(dilations);1493 1494 auto conv1DOp = Conv1DOp::create(1495 rewriter, loc, newOutputType, ValueRange{newInput, newKernel},1496 ValueRange{newOutput}, stridesAttr, dilationsAttr);1497 1498 // Insert back.1499 Value inserted = tensor::createCanonicalRankReducingInsertSliceOp(1500 rewriter, loc, conv1DOp.getResult(0), output);1501 rewriter.replaceOp(convOp, inserted);1502 1503 return conv1DOp;1504}1505 1506template struct linalg::DownscaleSizeOneWindowed2DConvolution<Conv2DNhwcHwcfOp,1507 Conv1DNwcWcfOp>;1508template struct linalg::DownscaleSizeOneWindowed2DConvolution<Conv2DNchwFchwOp,1509 Conv1DNcwFcwOp>;1510template struct linalg::DownscaleSizeOneWindowed2DConvolution<PoolingNhwcSumOp,1511 PoolingNwcSumOp>;1512template struct linalg::DownscaleSizeOneWindowed2DConvolution<PoolingNchwSumOp,1513 PoolingNcwSumOp>;1514template struct linalg::DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMaxOp,1515 PoolingNwcMaxOp>;1516template struct linalg::DownscaleSizeOneWindowed2DConvolution<1517 PoolingNhwcMaxUnsignedOp, PoolingNwcMaxUnsignedOp>;1518template struct linalg::DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMinOp,1519 PoolingNwcMinOp>;1520template struct linalg::DownscaleSizeOneWindowed2DConvolution<1521 PoolingNhwcMinUnsignedOp, PoolingNwcMinUnsignedOp>;1522template struct linalg::DownscaleSizeOneWindowed2DConvolution<PoolingNchwMaxOp,1523 PoolingNcwMaxOp>;1524 1525FailureOr<DepthwiseConv1DNwcWcOp>1526DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite(1527 DepthwiseConv2DNhwcHwcOp convOp, PatternRewriter &rewriter) const {1528 if (convOp.hasPureBufferSemantics())1529 return failure(); // To be implemented.1530 1531 Value input = convOp.getInputs().front();1532 Value kernel = convOp.getInputs().back();1533 Value output = convOp.getOutputs().front();1534 1535 auto inputType = dyn_cast<RankedTensorType>(input.getType());1536 auto kernelType = dyn_cast<RankedTensorType>(kernel.getType());1537 auto outputType = dyn_cast<RankedTensorType>(output.getType());1538 1539 auto kernelShape = kernelType.getShape();1540 auto outputShape = outputType.getShape();1541 1542 // Only handle the case where at least one of the window dimensions is1543 // of size 1. Other cases can rely on tiling to reduce to such cases.1544 int64_t khSize = kernelShape[0], kwSize = kernelShape[1];1545 int64_t ohSize = outputShape[1], owSize = outputShape[2];1546 bool removeH = (khSize == 1 && ohSize == 1);1547 bool removeW = (kwSize == 1 && owSize == 1);1548 if (!removeH && !removeW)1549 return failure();1550 1551 // Get new shapes and types for all operands by removing the size-11552 // dimension.1553 using RTTBuilder = RankedTensorType::Builder;1554 RankedTensorType newInputType =1555 RTTBuilder(inputType).dropDim((removeH ? 1 : 2));1556 RankedTensorType newKernelType =1557 RTTBuilder(kernelType).dropDim((removeH ? 0 : 1));1558 RankedTensorType newOutputType =1559 RTTBuilder(outputType).dropDim(removeH ? 1 : 2);1560 1561 // Rank-reduce operands.1562 Location loc = convOp.getLoc();1563 Value newInput = tensor::createCanonicalRankReducingExtractSliceOp(1564 rewriter, loc, input, newInputType);1565 Value newKernel = tensor::createCanonicalRankReducingExtractSliceOp(1566 rewriter, loc, kernel, newKernelType);1567 Value newOutput = tensor::createCanonicalRankReducingExtractSliceOp(1568 rewriter, loc, output, newOutputType);1569 1570 // Rank-reduce strides and dilations too.1571 // TODO: dropDim 1-liner helper.1572 auto strides = llvm::to_vector<4>(convOp.getStrides().getValues<int64_t>());1573 strides.erase(strides.begin() + (removeH ? 0 : 1));1574 auto stridesAttr = rewriter.getI64VectorAttr(strides);1575 1576 auto dilations =1577 llvm::to_vector<4>(convOp.getDilations().getValues<int64_t>());1578 dilations.erase(dilations.begin() + (removeH ? 0 : 1));1579 auto dilationsAttr = rewriter.getI64VectorAttr(dilations);1580 1581 auto conv1DOp = DepthwiseConv1DNwcWcOp::create(1582 rewriter, loc, newOutputType, ValueRange{newInput, newKernel},1583 ValueRange{newOutput}, stridesAttr, dilationsAttr);1584 1585 // Insert back.1586 Value inserted = tensor::createCanonicalRankReducingInsertSliceOp(1587 rewriter, loc, conv1DOp.getResult(0), output);1588 rewriter.replaceOp(convOp, inserted);1589 1590 return conv1DOp;1591}1592 1593FailureOr<Conv1DOp>1594DownscaleConv2DOp::returningMatchAndRewrite(Conv2DOp convOp,1595 PatternRewriter &rewriter) const {1596 if (convOp.hasPureBufferSemantics())1597 return failure(); // To be implemented.1598 1599 Value input = convOp.getInputs().front();1600 Value kernel = convOp.getInputs().back();1601 Value output = convOp.getOutputs().front();1602 1603 auto inputType = dyn_cast<RankedTensorType>(input.getType());1604 auto kernelType = dyn_cast<RankedTensorType>(kernel.getType());1605 auto outputType = dyn_cast<RankedTensorType>(output.getType());1606 1607 auto kernelShape = kernelType.getShape();1608 auto outputShape = outputType.getShape();1609 1610 // Only handle the case where at least one of the window dimensions is1611 // of size 1. Other cases can rely on tiling to reduce to such cases.1612 int64_t khSize = kernelShape[0], kwSize = kernelShape[1];1613 int64_t ohSize = outputShape[0], owSize = outputShape[1];1614 bool removeH = (khSize == 1 && ohSize == 1);1615 bool removeW = (kwSize == 1 && owSize == 1);1616 if (!removeH && !removeW)1617 return failure();1618 1619 // Get new shapes and types for all operands by removing the size-11620 // dimension.1621 using RTTBuilder = RankedTensorType::Builder;1622 RankedTensorType newInputType =1623 RTTBuilder(inputType).dropDim((removeH ? 0 : 1));1624 RankedTensorType newKernelType =1625 RTTBuilder(kernelType).dropDim((removeH ? 0 : 1));1626 RankedTensorType newOutputType =1627 RTTBuilder(outputType).dropDim(removeH ? 0 : 1);1628 1629 // Rank-reduce operands.1630 Location loc = convOp.getLoc();1631 Value newInput = tensor::createCanonicalRankReducingExtractSliceOp(1632 rewriter, loc, input, newInputType);1633 Value newKernel = tensor::createCanonicalRankReducingExtractSliceOp(1634 rewriter, loc, kernel, newKernelType);1635 Value newOutput = tensor::createCanonicalRankReducingExtractSliceOp(1636 rewriter, loc, output, newOutputType);1637 1638 auto conv1DOp =1639 Conv1DOp::create(rewriter, loc, newOutputType,1640 ValueRange{newInput, newKernel}, ValueRange{newOutput});1641 1642 // Insert back.1643 Value inserted = tensor::createCanonicalRankReducingInsertSliceOp(1644 rewriter, loc, conv1DOp.getResult(0), output);1645 rewriter.replaceOp(convOp, inserted);1646 1647 return conv1DOp;1648}1649 1650void linalg::populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,1651 PatternBenefit benefit) {1652 patterns.add<DownscaleSizeOneWindowed2DConvolution<linalg::Conv2DNhwcHwcfOp,1653 Conv1DNwcWcfOp>,1654 DownscaleSizeOneWindowed2DConvolution<linalg::Conv2DNchwFchwOp,1655 Conv1DNcwFcwOp>,1656 DownscaleDepthwiseConv2DNhwcHwcOp, DownscaleConv2DOp>(1657 patterns.getContext(), benefit);1658 patterns.add<1659 DownscaleSizeOneWindowed2DConvolution<PoolingNhwcSumOp, PoolingNwcSumOp>,1660 DownscaleSizeOneWindowed2DConvolution<PoolingNchwSumOp, PoolingNcwSumOp>,1661 DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMaxOp, PoolingNwcMaxOp>,1662 DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMaxUnsignedOp,1663 PoolingNwcMaxUnsignedOp>,1664 DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMinOp, PoolingNwcMinOp>,1665 DownscaleSizeOneWindowed2DConvolution<PoolingNhwcMinUnsignedOp,1666 PoolingNwcMinUnsignedOp>,1667 DownscaleSizeOneWindowed2DConvolution<PoolingNchwMaxOp, PoolingNcwMaxOp>>(1668 patterns.getContext(), benefit);1669}1670 1671void linalg::populateDecomposePackUnpackPatterns(RewritePatternSet &patterns) {1672 patterns.add<DecomposeOuterUnitDimsPackOpPattern>(patterns.getContext());1673 patterns.add<DecomposeOuterUnitDimsUnPackOpPattern>(patterns.getContext());1674}1675 1676void linalg::populateDecomposePadPatterns(RewritePatternSet &patterns) {1677 patterns.add<DecomposePadOpPattern>(patterns.getContext());1678}1679