1125 lines · cpp
1//===- DropUnitDims.cpp - Pass to drop use of unit-extent for broadcasting ===//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 patterns/pass to remove usage of unit-extent dimensions10// to specify broadcasting in favor of more canonical representation of the11// computation12//13//===----------------------------------------------------------------------===//14 15#include "mlir/Dialect/Linalg/Passes.h"16 17#include "mlir/Dialect/Affine/IR/AffineOps.h"18#include "mlir/Dialect/Arith/IR/Arith.h"19#include "mlir/Dialect/Linalg/IR/Linalg.h"20#include "mlir/Dialect/Linalg/Transforms/Transforms.h"21#include "mlir/Dialect/Linalg/Utils/Utils.h"22#include "mlir/Dialect/MemRef/Transforms/Transforms.h"23#include "mlir/Dialect/Tensor/IR/Tensor.h"24#include "mlir/Dialect/Tensor/Transforms/Transforms.h"25#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"26#include "mlir/IR/AffineExpr.h"27#include "mlir/IR/AffineMap.h"28#include "mlir/IR/BuiltinTypes.h"29#include "mlir/Transforms/FoldUtils.h"30#include "mlir/Transforms/GreedyPatternRewriteDriver.h"31#include "llvm/Support/Debug.h"32 33namespace mlir {34#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS35#include "mlir/Dialect/Linalg/Passes.h.inc"36} // namespace mlir37 38#define DEBUG_TYPE "linalg-drop-unit-dims"39 40using namespace mlir;41using namespace mlir::linalg;42 43namespace {44/// Pattern to move init operands to ins when all the loops are parallel and45/// blockArgument corresponding to init is used in the region. This is a fix-up46/// when unit reduction dimensions are all folded away. In this context, it47/// becomes a elementwise generic op. E.g., it converts48///49/// %0 = tensor.empty() : tensor<1x1xf32>50/// %1 = linalg.fill51/// ins(%cst : f32)52/// outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>53/// %2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,54/// affine_map<(d0) -> (0, d0)>],55/// iterator_types = ["parallel"]}56/// ins(%arg0 : tensor<1x?x1x1xf32>)57/// outs(%1 : tensor<1x1xf32>) {58/// ^bb0(%in: f32, %out: f32):59/// %3 = arith.addf %in, %out : f3260/// linalg.yield %3 : f3261/// } -> tensor<1x1xf32>62///63/// into64///65/// %0 = tensor.empty() : tensor<1x1xf32>66/// %1 = linalg.fill67/// ins(%cst : f32)68/// outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>69/// %2 = tensor.empty() : tensor<1x1xf32>70/// %3 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,71/// affine_map<(d0) -> (0, d0)>,72/// affine_map<(d0) -> (0, d0)>],73/// iterator_types = ["parallel"]}74/// ins(%arg0, %1 : tensor<1x?x1x1xf32>, tensor<1x1xf32>)75/// outs(%2 : tensor<1x1xf32>) {76/// ^bb0(%in: f32, %in_0: f32, %out: f32):77/// %4 = arith.addf %in, %in_0 : f3278/// linalg.yield %4 : f3279/// } -> tensor<1x1xf32>80struct MoveInitOperandsToInput : public OpRewritePattern<GenericOp> {81 using OpRewritePattern<GenericOp>::OpRewritePattern;82 LogicalResult matchAndRewrite(GenericOp genericOp,83 PatternRewriter &rewriter) const override {84 if (!genericOp.hasPureTensorSemantics())85 return failure();86 if (genericOp.getNumParallelLoops() != genericOp.getNumLoops())87 return failure();88 89 auto outputOperands = genericOp.getDpsInitsMutable();90 SetVector<OpOperand *> candidates;91 for (OpOperand &op : outputOperands) {92 if (genericOp.getMatchingBlockArgument(&op).use_empty())93 continue;94 candidates.insert(&op);95 }96 97 if (candidates.empty())98 return failure();99 100 // Compute the modified indexing maps.101 int64_t origNumInput = genericOp.getNumDpsInputs();102 SmallVector<Value> newInputOperands = genericOp.getDpsInputs();103 SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();104 SmallVector<AffineMap> newIndexingMaps;105 newIndexingMaps.append(indexingMaps.begin(),106 std::next(indexingMaps.begin(), origNumInput));107 for (OpOperand *op : candidates) {108 newInputOperands.push_back(op->get());109 newIndexingMaps.push_back(genericOp.getMatchingIndexingMap(op));110 }111 newIndexingMaps.append(std::next(indexingMaps.begin(), origNumInput),112 indexingMaps.end());113 114 Location loc = genericOp.getLoc();115 SmallVector<Value> newOutputOperands =116 llvm::to_vector(genericOp.getDpsInits());117 for (OpOperand *op : candidates) {118 OpBuilder::InsertionGuard guard(rewriter);119 rewriter.setInsertionPointAfterValue(op->get());120 auto elemType = cast<ShapedType>(op->get().getType()).getElementType();121 auto empty = tensor::EmptyOp::create(122 rewriter, loc, tensor::getMixedSizes(rewriter, loc, op->get()),123 elemType);124 125 unsigned start = genericOp.getDpsInits().getBeginOperandIndex();126 newOutputOperands[op->getOperandNumber() - start] = empty.getResult();127 }128 129 auto newOp = GenericOp::create(130 rewriter, loc, genericOp.getResultTypes(), newInputOperands,131 newOutputOperands, newIndexingMaps, genericOp.getIteratorTypesArray(),132 /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp));133 134 OpBuilder::InsertionGuard guard(rewriter);135 Region ®ion = newOp.getRegion();136 Block *block = rewriter.createBlock(®ion);137 IRMapping mapper;138 for (auto bbarg : genericOp.getRegionInputArgs())139 mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));140 141 for (OpOperand *op : candidates) {142 BlockArgument bbarg = genericOp.getMatchingBlockArgument(op);143 mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));144 }145 146 for (OpOperand &op : outputOperands) {147 BlockArgument bbarg = genericOp.getMatchingBlockArgument(&op);148 if (candidates.count(&op))149 block->addArgument(bbarg.getType(), loc);150 else151 mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));152 }153 154 for (auto &op : genericOp.getBody()->getOperations()) {155 rewriter.clone(op, mapper);156 }157 rewriter.replaceOp(genericOp, newOp.getResults());158 159 return success();160 }161};162} // namespace163 164//===---------------------------------------------------------------------===//165// Drop loops that are unit-extents within Linalg operations.166//===---------------------------------------------------------------------===//167 168/// Implements a pass that canonicalizes the uses of unit-extent dimensions for169/// broadcasting. For example,170///171/// ```mlir172/// #accesses = [173/// affine_map<(d0, d1) -> (0, d1)>,174/// affine_map<(d0, d1) -> (d0, 0)>,175/// affine_map<(d0, d1) -> (d0, d1)>176/// ]177///178/// #trait = {179/// indexing_maps = #accesses,180/// iterator_types = ["parallel", "parallel"],181/// library_call = "some_external_fn"182/// }183///184/// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->185/// tensor<5x5xf32>186/// {187/// %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] :188/// tensor<5xf32> into tensor<1x5xf32>189/// %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] :190/// tensor<5xf32> into tensor<5x1xf32>191/// %2 = linalg.generic #trait %0, %1 {192/// ^bb0(%arg2: f32, %arg3: f32):193/// %3 = arith.addf %arg2, %arg3 : f32194/// linalg.yield %3 : f32195/// } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32>196/// return %2 : tensor<5x5xf32>197/// }198///199/// would canonicalize to200///201/// ```mlir202/// #accesses = [203/// affine_map<(d0, d1) -> (d1)>,204/// affine_map<(d0, d1) -> (d0)>,205/// affine_map<(d0, d1) -> (d0, d1)>206/// ]207///208/// #trait = {209/// indexing_maps = #accesses,210/// iterator_types = ["parallel", "parallel"],211/// library_call = "some_external_fn"212/// }213///214/// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->215/// tensor<5x5xf32>216/// {217/// %0 = linalg.generic #trait %arg0, %arg1 {218/// ^bb0(%arg2: f32, %arg3: f32):219/// %3 = arith.addf %arg2, %arg3 : f32220/// linalg.yield %3 : f32221/// } : tensor<5xf32>, tensor<5xf32> -> tensor<5x5xf32>222/// return %0 : tensor<5x5xf32>223/// }224 225/// Update the index accesses of linalg operations having index semantics.226static void227replaceUnitDimIndexOps(GenericOp genericOp,228 const llvm::SmallDenseSet<unsigned> &unitDims,229 RewriterBase &rewriter) {230 for (IndexOp indexOp :231 llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {232 OpBuilder::InsertionGuard guard(rewriter);233 rewriter.setInsertionPoint(indexOp);234 if (unitDims.count(indexOp.getDim()) != 0) {235 rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(indexOp, 0);236 } else {237 // Update the dimension of the index operation if needed.238 unsigned droppedDims = llvm::count_if(239 unitDims, [&](unsigned dim) { return dim < indexOp.getDim(); });240 if (droppedDims != 0)241 rewriter.replaceOpWithNewOp<IndexOp>(indexOp,242 indexOp.getDim() - droppedDims);243 }244 }245}246 247/// Expand the given `value` so that the type matches the type of `origDest`.248/// The `reassociation` is used when `rankReductionStrategy` is set to249/// `RankReductionStrategy::ReassociativeReshape`.250static Value251expandValue(RewriterBase &rewriter, Location loc, Value result, Value origDest,252 ArrayRef<ReassociationIndices> reassociation,253 ControlDropUnitDims::RankReductionStrategy rankReductionStrategy) {254 // There are no results for memref outputs.255 auto origResultType = cast<RankedTensorType>(origDest.getType());256 if (rankReductionStrategy ==257 ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {258 unsigned rank = origResultType.getRank();259 SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));260 SmallVector<OpFoldResult> sizes =261 tensor::getMixedSizes(rewriter, loc, origDest);262 SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));263 return rewriter.createOrFold<tensor::InsertSliceOp>(264 loc, result, origDest, offsets, sizes, strides);265 }266 267 assert(rankReductionStrategy ==268 ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape &&269 "unknown rank reduction strategy");270 return tensor::ExpandShapeOp::create(rewriter, loc, origResultType, result,271 reassociation)272 .getResult();273}274 275/// Collapse the given `value` so that the type matches the type of276/// `origOutput`. The `reassociation` is used when `rankReductionStrategy` is277/// set to `RankReductionStrategy::ReassociativeReshape`.278static Value collapseValue(279 RewriterBase &rewriter, Location loc, Value operand,280 ArrayRef<int64_t> targetShape, ArrayRef<ReassociationIndices> reassociation,281 ControlDropUnitDims::RankReductionStrategy rankReductionStrategy) {282 if (auto memrefType = dyn_cast<MemRefType>(operand.getType())) {283 if (rankReductionStrategy ==284 ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {285 FailureOr<Value> rankReducingExtract =286 memref::SubViewOp::rankReduceIfNeeded(rewriter, loc, operand,287 targetShape);288 assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");289 return *rankReducingExtract;290 }291 292 assert(293 rankReductionStrategy ==294 ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape &&295 "unknown rank reduction strategy");296 MemRefLayoutAttrInterface layout;297 auto targetType = MemRefType::get(targetShape, memrefType.getElementType(),298 layout, memrefType.getMemorySpace());299 return memref::CollapseShapeOp::create(rewriter, loc, targetType, operand,300 reassociation);301 }302 if (auto tensorType = dyn_cast<RankedTensorType>(operand.getType())) {303 if (rankReductionStrategy ==304 ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {305 FailureOr<Value> rankReducingExtract =306 tensor::ExtractSliceOp::rankReduceIfNeeded(rewriter, loc, operand,307 targetShape);308 assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");309 return *rankReducingExtract;310 }311 312 assert(313 rankReductionStrategy ==314 ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape &&315 "unknown rank reduction strategy");316 auto targetType =317 RankedTensorType::get(targetShape, tensorType.getElementType());318 return tensor::CollapseShapeOp::create(rewriter, loc, targetType, operand,319 reassociation);320 }321 llvm_unreachable("unsupported operand type");322}323 324/// Compute the modified metadata for an operands of operation325/// whose unit dims are being dropped. Return the new indexing map326/// to use, the shape of the operand in the replacement op327/// and the `reassocation` to use to go from original operand shape328/// to modified operand shape.329struct UnitExtentReplacementInfo {330 AffineMap indexMap;331 SmallVector<ReassociationIndices> reassociation;332 SmallVector<int64_t> targetShape;333};334static UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata(335 MLIRContext *context, IndexingMapOpInterface op, OpOperand *opOperand,336 llvm::SmallDenseMap<unsigned, unsigned> &oldDimsToNewDimsMap,337 ArrayRef<AffineExpr> dimReplacements) {338 UnitExtentReplacementInfo info;339 ReassociationIndices reassociationGroup;340 SmallVector<AffineExpr> newIndexExprs;341 AffineMap indexingMap = op.getMatchingIndexingMap(opOperand);342 SmallVector<int64_t> operandShape = op.getStaticOperandShape(opOperand);343 ArrayRef<AffineExpr> exprs = indexingMap.getResults();344 345 auto isUnitDim = [&](unsigned dim) {346 if (auto dimExpr = dyn_cast<AffineDimExpr>(exprs[dim])) {347 unsigned oldPosition = dimExpr.getPosition();348 return !oldDimsToNewDimsMap.count(oldPosition) &&349 (operandShape[dim] == 1);350 }351 // Handle the other case where the shape is 1, and is accessed using a352 // constant 0.353 if (operandShape[dim] == 1) {354 auto constAffineExpr = dyn_cast<AffineConstantExpr>(exprs[dim]);355 return constAffineExpr && constAffineExpr.getValue() == 0;356 }357 return false;358 };359 360 unsigned dim = 0;361 while (dim < operandShape.size() && isUnitDim(dim))362 reassociationGroup.push_back(dim++);363 while (dim < operandShape.size()) {364 assert(!isUnitDim(dim) && "expected non unit-extent");365 reassociationGroup.push_back(dim);366 AffineExpr newExpr = exprs[dim].replaceDims(dimReplacements);367 newIndexExprs.push_back(newExpr);368 info.targetShape.push_back(operandShape[dim]);369 ++dim;370 // Fold all following dimensions that are unit-extent.371 while (dim < operandShape.size() && isUnitDim(dim)) {372 reassociationGroup.push_back(dim++);373 }374 info.reassociation.push_back(reassociationGroup);375 reassociationGroup.clear();376 }377 info.indexMap =378 AffineMap::get(oldDimsToNewDimsMap.size(), indexingMap.getNumSymbols(),379 newIndexExprs, context);380 return info;381}382 383FailureOr<DropUnitDimsResult>384linalg::dropUnitDims(RewriterBase &rewriter, IndexingMapOpInterface op,385 const DroppedUnitDimsBuilder &droppedUnitDimsBuilder,386 const ControlDropUnitDims &options) {387 auto dpsOp = dyn_cast<DestinationStyleOpInterface>(op.getOperation());388 if (!dpsOp) {389 return rewriter.notifyMatchFailure(390 op, "op should implement DestinationStyleOpInterface");391 }392 393 SmallVector<AffineMap> indexingMaps = op.getIndexingMapsArray();394 if (indexingMaps.empty())395 return failure();396 397 // 1. Check if any of the iteration dimensions are unit-trip count. They will398 // end up being unit-trip count if they are used to index into a unit-dim399 // tensor/memref.400 AffineMap invertedMap =401 inversePermutation(concatAffineMaps(indexingMaps, rewriter.getContext()));402 if (!invertedMap) {403 return rewriter.notifyMatchFailure(op,404 "invalid indexing maps for operation");405 }406 407 SmallVector<int64_t> allShapesSizes;408 for (OpOperand &opOperand : op->getOpOperands())409 llvm::append_range(allShapesSizes, op.getStaticOperandShape(&opOperand));410 411 // 1a. Get the allowed list of dimensions to drop from the `options`.412 SmallVector<unsigned> allowedUnitDims = options.controlFn(op);413 if (allowedUnitDims.empty()) {414 return rewriter.notifyMatchFailure(415 op, "control function returns no allowed unit dims to prune");416 }417 llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),418 allowedUnitDims.end());419 llvm::SmallDenseSet<unsigned> unitDims;420 for (const auto &expr : enumerate(invertedMap.getResults())) {421 if (AffineDimExpr dimExpr = dyn_cast<AffineDimExpr>(expr.value())) {422 if (allShapesSizes[dimExpr.getPosition()] == 1 &&423 unitDimsFilter.count(expr.index()))424 unitDims.insert(expr.index());425 }426 }427 428 // 2. Compute the new loops of the modified op by dropping the one-trip429 // count loops.430 llvm::SmallDenseMap<unsigned, unsigned> oldDimToNewDimMap;431 SmallVector<AffineExpr> dimReplacements;432 unsigned newDims = 0;433 for (auto index : llvm::seq<int64_t>(op.getStaticLoopRanges().size())) {434 if (unitDims.count(index)) {435 dimReplacements.push_back(436 getAffineConstantExpr(0, rewriter.getContext()));437 } else {438 oldDimToNewDimMap[index] = newDims;439 dimReplacements.push_back(440 getAffineDimExpr(newDims, rewriter.getContext()));441 newDims++;442 }443 }444 445 // 3. For each of the operands, find the446 // - modified affine map to use.447 // - shape of the operands after the unit-dims are dropped.448 // - the reassociation indices used to convert from the original449 // operand type to modified operand (needed only when using reshapes450 // for rank reduction strategy)451 // Note that the indexing maps might need changing even if there are no452 // unit dimensions that are dropped to handle cases where `0` is used to453 // access a unit-extent tensor. Consider moving this out of this specific454 // transformation as a stand-alone transformation. Kept here right now due455 // to legacy.456 SmallVector<AffineMap> newIndexingMaps;457 SmallVector<SmallVector<ReassociationIndices>> reassociations;458 SmallVector<SmallVector<int64_t>> targetShapes;459 SmallVector<bool> collapsed;460 auto hasCollapsibleType = [](OpOperand &operand) {461 Type operandType = operand.get().getType();462 if (auto memrefOperandType = dyn_cast_or_null<MemRefType>(operandType)) {463 return memrefOperandType.getLayout().isIdentity();464 }465 if (auto tensorOperandType = dyn_cast<RankedTensorType>(operandType)) {466 return tensorOperandType.getEncoding() == nullptr;467 }468 return false;469 };470 for (OpOperand &opOperand : op->getOpOperands()) {471 auto indexingMap = op.getMatchingIndexingMap(&opOperand);472 SmallVector<int64_t> shape = op.getStaticOperandShape(&opOperand);473 if (!hasCollapsibleType(opOperand)) {474 AffineMap newIndexingMap = indexingMap.replaceDimsAndSymbols(475 dimReplacements, ArrayRef<AffineExpr>{}, oldDimToNewDimMap.size(), 0);476 newIndexingMaps.push_back(newIndexingMap);477 targetShapes.push_back(llvm::to_vector(shape));478 collapsed.push_back(false);479 reassociations.push_back({});480 continue;481 }482 auto replacementInfo =483 dropUnitExtentFromOperandMetadata(rewriter.getContext(), op, &opOperand,484 oldDimToNewDimMap, dimReplacements);485 reassociations.push_back(replacementInfo.reassociation);486 newIndexingMaps.push_back(replacementInfo.indexMap);487 targetShapes.push_back(replacementInfo.targetShape);488 collapsed.push_back(!(replacementInfo.indexMap.getNumResults() ==489 indexingMap.getNumResults()));490 }491 492 // Abort if the indexing maps of the result operation are not invertible493 // (i.e. not legal) or if no dimension was reduced.494 if (newIndexingMaps == indexingMaps ||495 !inversePermutation(496 concatAffineMaps(newIndexingMaps, rewriter.getContext())))497 return failure();498 499 Location loc = op.getLoc();500 // 4. For each of the operands, collapse the operand to convert501 // from original shape to shape in the modified operation if needed,502 // either through use of reshapes or rank-reducing slices as503 // specified in `options`.504 SmallVector<Value> newOperands;505 for (OpOperand &opOperand : op->getOpOperands()) {506 int64_t idx = opOperand.getOperandNumber();507 if (!collapsed[idx]) {508 newOperands.push_back(opOperand.get());509 continue;510 }511 newOperands.push_back(collapseValue(rewriter, loc, opOperand.get(),512 targetShapes[idx], reassociations[idx],513 options.rankReductionStrategy));514 }515 516 IndexingMapOpInterface replacementOp = droppedUnitDimsBuilder(517 loc, rewriter, op, newOperands, newIndexingMaps, unitDims);518 519 // 6. If any result type changes, insert a reshape/slice to convert from the520 // original type to the new type.521 SmallVector<Value> resultReplacements;522 for (auto [index, result] : llvm::enumerate(replacementOp->getResults())) {523 unsigned opOperandIndex = index + dpsOp.getNumDpsInputs();524 Value origDest = dpsOp.getDpsInitOperand(index)->get();525 if (!collapsed[opOperandIndex]) {526 resultReplacements.push_back(result);527 continue;528 }529 Value expandedValue = expandValue(rewriter, loc, result, origDest,530 reassociations[opOperandIndex],531 options.rankReductionStrategy);532 resultReplacements.push_back(expandedValue);533 }534 535 return DropUnitDimsResult{replacementOp, resultReplacements};536}537 538FailureOr<DropUnitDimsResult>539linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,540 const ControlDropUnitDims &options) {541 542 DroppedUnitDimsBuilder build =543 [](Location loc, OpBuilder &b, IndexingMapOpInterface op,544 ArrayRef<Value> newOperands, ArrayRef<AffineMap> newIndexingMaps,545 const llvm::SmallDenseSet<unsigned> &droppedDims)546 -> IndexingMapOpInterface {547 auto genericOp = cast<GenericOp>(op);548 // Compute the iterator types of the modified op by dropping the one-trip549 // count loops.550 SmallVector<utils::IteratorType> newIteratorTypes;551 for (auto [index, attr] :552 llvm::enumerate(genericOp.getIteratorTypesArray())) {553 if (!droppedDims.count(index))554 newIteratorTypes.push_back(attr);555 }556 557 // Create the `linalg.generic` operation with the new operands,558 // indexing maps, iterator types and result types.559 ArrayRef<Value> newInputs =560 ArrayRef<Value>(newOperands).take_front(genericOp.getNumDpsInputs());561 ArrayRef<Value> newOutputs =562 ArrayRef<Value>(newOperands).take_back(genericOp.getNumDpsInits());563 SmallVector<Type> resultTypes;564 resultTypes.reserve(genericOp.getNumResults());565 for (unsigned i : llvm::seq<unsigned>(0, genericOp.getNumResults()))566 resultTypes.push_back(newOutputs[i].getType());567 GenericOp replacementOp =568 GenericOp::create(b, loc, resultTypes, newInputs, newOutputs,569 newIndexingMaps, newIteratorTypes);570 b.cloneRegionBefore(genericOp.getRegion(), replacementOp.getRegion(),571 replacementOp.getRegion().begin());572 // 5a. Replace `linalg.index` operations that refer to the dropped unit573 // dimensions.574 IRRewriter rewriter(b);575 replaceUnitDimIndexOps(replacementOp, droppedDims, rewriter);576 577 return replacementOp;578 };579 580 return dropUnitDims(rewriter, genericOp, build, options);581}582 583namespace {584struct DropUnitDims : public OpRewritePattern<GenericOp> {585 DropUnitDims(MLIRContext *context, ControlDropUnitDims options = {},586 PatternBenefit benefit = 1)587 : OpRewritePattern(context, benefit), options(std::move(options)) {}588 589 LogicalResult matchAndRewrite(GenericOp genericOp,590 PatternRewriter &rewriter) const override {591 FailureOr<DropUnitDimsResult> result =592 dropUnitDims(rewriter, genericOp, options);593 if (failed(result)) {594 return failure();595 }596 rewriter.replaceOp(genericOp, result->replacements);597 return success();598 }599 600private:601 ControlDropUnitDims options;602};603} // namespace604 605//===---------------------------------------------------------------------===//606// Drop dimensions that are unit-extents within tensor operations.607//===---------------------------------------------------------------------===//608 609namespace {610struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {611 DropPadUnitDims(MLIRContext *context, ControlDropUnitDims options = {},612 PatternBenefit benefit = 1)613 : OpRewritePattern(context, benefit), options(std::move(options)) {}614 615 LogicalResult matchAndRewrite(tensor::PadOp padOp,616 PatternRewriter &rewriter) const override {617 // 1a. Get the allowed list of dimensions to drop from the `options`.618 SmallVector<unsigned> allowedUnitDims = options.controlFn(padOp);619 if (allowedUnitDims.empty()) {620 return rewriter.notifyMatchFailure(621 padOp, "control function returns no allowed unit dims to prune");622 }623 624 if (padOp.getSourceType().getEncoding()) {625 return rewriter.notifyMatchFailure(626 padOp, "cannot collapse dims of tensor with encoding");627 }628 629 // Fail for non-constant padding values. The body of the pad could630 // depend on the padding indices and/or properties of the padded631 // tensor so for now we fail.632 // TODO: Support non-constant padding values.633 Value paddingVal = padOp.getConstantPaddingValue();634 if (!paddingVal) {635 return rewriter.notifyMatchFailure(636 padOp, "unimplemented: non-constant padding value");637 }638 639 ArrayRef<int64_t> sourceShape = padOp.getSourceType().getShape();640 ArrayRef<int64_t> resultShape = padOp.getResultType().getShape();641 int64_t padRank = sourceShape.size();642 643 auto isStaticZero = [](OpFoldResult f) {644 return getConstantIntValue(f) == 0;645 };646 647 llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),648 allowedUnitDims.end());649 llvm::SmallDenseSet<unsigned> unitDims;650 SmallVector<int64_t> newShape;651 SmallVector<int64_t> newResultShape;652 SmallVector<OpFoldResult> newLowPad;653 SmallVector<OpFoldResult> newHighPad;654 for (const auto [dim, size, outSize, low, high] : zip_equal(655 llvm::seq(static_cast<int64_t>(0), padRank), sourceShape,656 resultShape, padOp.getMixedLowPad(), padOp.getMixedHighPad())) {657 if (unitDimsFilter.contains(dim) && size == 1 && isStaticZero(low) &&658 isStaticZero(high)) {659 unitDims.insert(dim);660 } else {661 newShape.push_back(size);662 newResultShape.push_back(outSize);663 newLowPad.push_back(low);664 newHighPad.push_back(high);665 }666 }667 668 if (unitDims.empty()) {669 return rewriter.notifyMatchFailure(padOp, "no unit dims to collapse");670 }671 672 ReassociationIndices reassociationGroup;673 SmallVector<ReassociationIndices> reassociationMap;674 int64_t dim = 0;675 while (dim < padRank && unitDims.contains(dim))676 reassociationGroup.push_back(dim++);677 while (dim < padRank) {678 assert(!unitDims.contains(dim) && "expected non unit-extent");679 reassociationGroup.push_back(dim);680 dim++;681 // Fold all following dimensions that are unit-extent.682 while (dim < padRank && unitDims.contains(dim))683 reassociationGroup.push_back(dim++);684 reassociationMap.push_back(reassociationGroup);685 reassociationGroup.clear();686 }687 688 Value collapsedSource =689 collapseValue(rewriter, padOp.getLoc(), padOp.getSource(), newShape,690 reassociationMap, options.rankReductionStrategy);691 692 auto newResultType = RankedTensorType::get(693 newResultShape, padOp.getResultType().getElementType());694 auto newPadOp = tensor::PadOp::create(695 rewriter, padOp.getLoc(), /*result=*/newResultType, collapsedSource,696 newLowPad, newHighPad, paddingVal, padOp.getNofold());697 698 Value dest = padOp.getResult();699 if (options.rankReductionStrategy ==700 ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {701 SmallVector<OpFoldResult> expandedSizes;702 int64_t numUnitDims = 0;703 for (auto dim : llvm::seq(static_cast<int64_t>(0), padRank)) {704 if (unitDims.contains(dim)) {705 expandedSizes.push_back(rewriter.getIndexAttr(1));706 numUnitDims++;707 continue;708 }709 expandedSizes.push_back(tensor::getMixedSize(710 rewriter, padOp.getLoc(), newPadOp, dim - numUnitDims));711 }712 dest = tensor::EmptyOp::create(rewriter, padOp.getLoc(), expandedSizes,713 padOp.getResultType().getElementType());714 }715 716 Value expandedValue =717 expandValue(rewriter, padOp.getLoc(), newPadOp.getResult(), dest,718 reassociationMap, options.rankReductionStrategy);719 rewriter.replaceOp(padOp, expandedValue);720 return success();721 }722 723private:724 ControlDropUnitDims options;725};726} // namespace727 728namespace {729/// Convert `extract_slice` operations to rank-reduced versions.730struct RankReducedExtractSliceOp731 : public OpRewritePattern<tensor::ExtractSliceOp> {732 using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;733 734 LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,735 PatternRewriter &rewriter) const override {736 RankedTensorType resultType = sliceOp.getType();737 SmallVector<OpFoldResult> targetShape;738 for (auto size : resultType.getShape())739 targetShape.push_back(rewriter.getIndexAttr(size));740 auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);741 if (!reassociation ||742 reassociation->size() == static_cast<size_t>(resultType.getRank()))743 return failure();744 745 SmallVector<OpFoldResult> offsets = sliceOp.getMixedOffsets();746 SmallVector<OpFoldResult> strides = sliceOp.getMixedStrides();747 SmallVector<OpFoldResult> sizes = sliceOp.getMixedSizes();748 auto rankReducedType = cast<RankedTensorType>(749 tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(750 reassociation->size(), sliceOp.getSourceType(), sizes));751 752 Location loc = sliceOp.getLoc();753 Value newSlice = tensor::ExtractSliceOp::create(754 rewriter, loc, rankReducedType, sliceOp.getSource(), offsets, sizes,755 strides);756 rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(757 sliceOp, resultType, newSlice, *reassociation);758 return success();759 }760};761 762/// Convert `insert_slice` operations to rank-reduced versions.763/// This patterns works with both InsertSliceOp and ParallelInsertSliceOp.764template <typename InsertOpTy>765struct RankReducedInsertSliceOp : public OpRewritePattern<InsertOpTy> {766 using OpRewritePattern<InsertOpTy>::OpRewritePattern;767 768 LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,769 PatternRewriter &rewriter) const override {770 RankedTensorType sourceType = insertSliceOp.getSourceType();771 SmallVector<OpFoldResult> targetShape;772 for (auto size : sourceType.getShape())773 targetShape.push_back(rewriter.getIndexAttr(size));774 auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);775 if (!reassociation ||776 reassociation->size() == static_cast<size_t>(sourceType.getRank()))777 return failure();778 779 Location loc = insertSliceOp.getLoc();780 tensor::CollapseShapeOp reshapedSource;781 {782 OpBuilder::InsertionGuard g(rewriter);783 // The only difference between InsertSliceOp and ParallelInsertSliceOp784 // is the insertion point is just before the ParallelCombiningOp in the785 // parallel case.786 if (std::is_same<InsertOpTy, tensor::ParallelInsertSliceOp>::value)787 rewriter.setInsertionPoint(insertSliceOp->getParentOp());788 reshapedSource = tensor::CollapseShapeOp::create(789 rewriter, loc, insertSliceOp.getSource(), *reassociation);790 }791 rewriter.replaceOpWithNewOp<InsertOpTy>(792 insertSliceOp, reshapedSource, insertSliceOp.getDest(),793 insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),794 insertSliceOp.getMixedStrides());795 return success();796 }797};798} // namespace799 800/// Patterns that are used to canonicalize the use of unit-extent dims for801/// broadcasting.802static void803populateFoldUnitExtentDimsViaReshapesPatterns(RewritePatternSet &patterns,804 ControlDropUnitDims &options) {805 auto *context = patterns.getContext();806 patterns.add<DropUnitDims>(context, options);807 patterns.add<DropPadUnitDims>(context, options);808 // TODO: Patterns unrelated to unit dim folding should be factored out.809 patterns.add<RankReducedExtractSliceOp,810 RankReducedInsertSliceOp<tensor::InsertSliceOp>,811 RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(812 context);813 linalg::FillOp::getCanonicalizationPatterns(patterns, context);814 tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);815 tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);816 tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);817 tensor::populateFoldTensorEmptyPatterns(patterns);818 memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);819 memref::populateResolveShapedTypeResultDimsPatterns(patterns);820}821 822static void823populateFoldUnitExtentDimsViaSlicesPatterns(RewritePatternSet &patterns,824 ControlDropUnitDims &options) {825 auto *context = patterns.getContext();826 patterns.add<DropUnitDims>(context, options);827 patterns.add<DropPadUnitDims>(context, options);828 // TODO: Patterns unrelated to unit dim folding should be factored out.829 linalg::FillOp::getCanonicalizationPatterns(patterns, context);830 tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);831 tensor::populateFoldTensorEmptyPatterns(patterns);832 memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);833 memref::populateResolveShapedTypeResultDimsPatterns(patterns);834}835 836void mlir::linalg::populateFoldUnitExtentDimsPatterns(837 RewritePatternSet &patterns, linalg::ControlDropUnitDims &options) {838 if (options.rankReductionStrategy ==839 linalg::ControlDropUnitDims::RankReductionStrategy::ExtractInsertSlice) {840 populateFoldUnitExtentDimsViaSlicesPatterns(patterns, options);841 } else if (options.rankReductionStrategy ==842 linalg::ControlDropUnitDims::RankReductionStrategy::843 ReassociativeReshape) {844 populateFoldUnitExtentDimsViaReshapesPatterns(patterns, options);845 }846}847 848void mlir::linalg::populateMoveInitOperandsToInputPattern(849 RewritePatternSet &patterns) {850 patterns.add<MoveInitOperandsToInput>(patterns.getContext());851}852 853namespace {854/// Pass that removes unit-extent dims within generic ops.855struct LinalgFoldUnitExtentDimsPass856 : public impl::LinalgFoldUnitExtentDimsPassBase<857 LinalgFoldUnitExtentDimsPass> {858 using impl::LinalgFoldUnitExtentDimsPassBase<859 LinalgFoldUnitExtentDimsPass>::LinalgFoldUnitExtentDimsPassBase;860 void runOnOperation() override {861 Operation *op = getOperation();862 MLIRContext *context = op->getContext();863 RewritePatternSet patterns(context);864 ControlDropUnitDims options;865 if (useRankReducingSlices) {866 options.rankReductionStrategy = linalg::ControlDropUnitDims::867 RankReductionStrategy::ExtractInsertSlice;868 }869 linalg::populateFoldUnitExtentDimsPatterns(patterns, options);870 populateMoveInitOperandsToInputPattern(patterns);871 (void)applyPatternsGreedily(op, std::move(patterns));872 }873};874 875} // namespace876 877namespace {878 879/// Returns reassociation indices for collapsing/expanding a880/// tensor of rank `rank` at position `pos`.881static SmallVector<ReassociationIndices>882getReassociationForReshapeAtDim(int64_t rank, int64_t pos) {883 SmallVector<ReassociationIndices> reassociation(rank - 1, {0, 1});884 bool lastDim = pos == rank - 1;885 if (rank > 2) {886 for (int64_t i = 0; i < rank - 1; i++) {887 if (i == pos || (lastDim && i == pos - 1))888 reassociation[i] = ReassociationIndices{i, i + 1};889 else if (i < pos)890 reassociation[i] = ReassociationIndices{i};891 else892 reassociation[i] = ReassociationIndices{i + 1};893 }894 }895 return reassociation;896}897 898/// Returns a collapsed `val` where the collapsing occurs at dim `pos`.899/// If `pos < 0`, then don't collapse.900static Value collapseSingletonDimAt(PatternRewriter &rewriter, Value val,901 int64_t pos) {902 if (pos < 0)903 return val;904 auto valType = cast<ShapedType>(val.getType());905 SmallVector<int64_t> collapsedShape(valType.getShape());906 collapsedShape.erase(collapsedShape.begin() + pos);907 return collapseValue(908 rewriter, val.getLoc(), val, collapsedShape,909 getReassociationForReshapeAtDim(valType.getRank(), pos),910 ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape);911}912 913/// Base class for all rank reduction patterns for contraction ops914/// with unit dimensions. All patterns should convert one named op915/// to another named op. Intended to reduce only one iteration space dim916/// at a time.917/// Reducing multiple dims will happen with recusive application of918/// pattern rewrites.919template <typename FromOpTy, typename ToOpTy>920struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {921 using OpRewritePattern<FromOpTy>::OpRewritePattern;922 923 /// Collapse all collapsable operands.924 SmallVector<Value>925 collapseOperands(PatternRewriter &rewriter, ArrayRef<Value> operands,926 ArrayRef<int64_t> operandCollapseDims) const {927 assert(operandCollapseDims.size() == 3 && operands.size() == 3 &&928 "expected 3 operands and dims");929 return llvm::map_to_vector(930 llvm::zip(operands, operandCollapseDims), [&](auto pair) {931 return collapseSingletonDimAt(rewriter, std::get<0>(pair),932 std::get<1>(pair));933 });934 }935 936 /// Expand result tensor.937 Value expandResult(PatternRewriter &rewriter, Value result,938 RankedTensorType expandedType, int64_t dim) const {939 return tensor::ExpandShapeOp::create(940 rewriter, result.getLoc(), expandedType, result,941 getReassociationForReshapeAtDim(expandedType.getRank(), dim));942 }943 944 LogicalResult matchAndRewrite(FromOpTy contractionOp,945 PatternRewriter &rewriter) const override {946 if (contractionOp.hasUserDefinedMaps()) {947 return rewriter.notifyMatchFailure(948 contractionOp, "ops with user-defined maps are not supported");949 }950 951 auto loc = contractionOp.getLoc();952 auto inputs = contractionOp.getDpsInputs();953 auto inits = contractionOp.getDpsInits();954 if (inputs.size() != 2 || inits.size() != 1)955 return rewriter.notifyMatchFailure(contractionOp,956 "expected 2 inputs and 1 init");957 auto lhs = inputs[0];958 auto rhs = inputs[1];959 auto init = inits[0];960 SmallVector<Value> operands{lhs, rhs, init};961 962 SmallVector<int64_t> operandUnitDims;963 if (failed(getOperandUnitDims(contractionOp, operandUnitDims)))964 return rewriter.notifyMatchFailure(contractionOp,965 "no reducable dims found");966 967 SmallVector<Value> collapsedOperands =968 collapseOperands(rewriter, operands, operandUnitDims);969 Value collapsedLhs = collapsedOperands[0];970 Value collapsedRhs = collapsedOperands[1];971 Value collapsedInit = collapsedOperands[2];972 SmallVector<Type, 1> collapsedResultTy;973 if (isa<RankedTensorType>(collapsedInit.getType()))974 collapsedResultTy.push_back(collapsedInit.getType());975 auto collapsedOp = ToOpTy::create(rewriter, loc, collapsedResultTy,976 ValueRange{collapsedLhs, collapsedRhs},977 ValueRange{collapsedInit});978 for (auto attr : contractionOp->getAttrs()) {979 if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName ||980 attr.getName() == "indexing_maps")981 continue;982 collapsedOp->setAttr(attr.getName(), attr.getValue());983 }984 985 auto results = contractionOp.getResults();986 assert(results.size() < 2 && "expected at most one result");987 if (results.empty()) {988 rewriter.replaceOp(contractionOp, collapsedOp);989 } else {990 rewriter.replaceOp(991 contractionOp,992 expandResult(rewriter, collapsedOp.getResultTensors()[0],993 cast<RankedTensorType>(results[0].getType()),994 operandUnitDims[2]));995 }996 997 return success();998 }999 1000 /// Populate `operandUnitDims` with 3 indices indicating the unit dim1001 /// for each operand that should be collapsed in this pattern. If an1002 /// operand shouldn't be collapsed, the index should be negative.1003 virtual LogicalResult1004 getOperandUnitDims(LinalgOp op,1005 SmallVectorImpl<int64_t> &operandUnitDims) const = 0;1006};1007 1008/// Patterns for unbatching batched contraction ops1009template <typename FromOpTy, typename ToOpTy>1010struct RankReduceToUnBatched : RankReduceContractionOps<FromOpTy, ToOpTy> {1011 using RankReduceContractionOps<FromOpTy, ToOpTy>::RankReduceContractionOps;1012 1013 /// Look for unit batch dims to collapse.1014 LogicalResult1015 getOperandUnitDims(LinalgOp op,1016 SmallVectorImpl<int64_t> &operandUnitDims) const override {1017 FailureOr<ContractionDimensions> maybeContractionDims =1018 inferContractionDims(op);1019 if (failed(maybeContractionDims)) {1020 LLVM_DEBUG(llvm::dbgs() << "could not infer contraction dims");1021 return failure();1022 }1023 ContractionDimensions contractionDims = maybeContractionDims.value();1024 1025 if (contractionDims.batch.size() != 1)1026 return failure();1027 auto batchDim = contractionDims.batch[0];1028 SmallVector<std::pair<Value, unsigned>, 3> bOperands;1029 op.mapIterationSpaceDimToAllOperandDims(batchDim, bOperands);1030 if (bOperands.size() != 3 || llvm::any_of(bOperands, [](auto pair) {1031 return cast<ShapedType>(std::get<0>(pair).getType())1032 .getShape()[std::get<1>(pair)] != 1;1033 })) {1034 LLVM_DEBUG(llvm::dbgs() << "specified unit dims not found");1035 return failure();1036 }1037 1038 operandUnitDims = SmallVector<int64_t>{std::get<1>(bOperands[0]),1039 std::get<1>(bOperands[1]),1040 std::get<1>(bOperands[2])};1041 return success();1042 }1043};1044 1045/// Patterns for reducing non-batch dimensions1046template <typename FromOpTy, typename ToOpTy>1047struct RankReduceMatmul : RankReduceContractionOps<FromOpTy, ToOpTy> {1048 using RankReduceContractionOps<FromOpTy, ToOpTy>::RankReduceContractionOps;1049 1050 /// Helper for determining whether the lhs/init or rhs/init are reduced.1051 static bool constexpr reduceLeft =1052 (std::is_same_v<FromOpTy, BatchMatmulOp> &&1053 std::is_same_v<ToOpTy, BatchVecmatOp>) ||1054 (std::is_same_v<FromOpTy, MatmulOp> &&1055 std::is_same_v<ToOpTy, VecmatOp>) ||1056 (std::is_same_v<FromOpTy, MatvecOp> && std::is_same_v<ToOpTy, DotOp>);1057 1058 /// Look for non-batch spatial dims to collapse.1059 LogicalResult1060 getOperandUnitDims(LinalgOp op,1061 SmallVectorImpl<int64_t> &operandUnitDims) const override {1062 FailureOr<ContractionDimensions> maybeContractionDims =1063 inferContractionDims(op);1064 if (failed(maybeContractionDims)) {1065 LLVM_DEBUG(llvm::dbgs() << "could not infer contraction dims");1066 return failure();1067 }1068 ContractionDimensions contractionDims = maybeContractionDims.value();1069 1070 if constexpr (reduceLeft) {1071 auto m = contractionDims.m[0];1072 SmallVector<std::pair<Value, unsigned>, 2> mOperands;1073 op.mapIterationSpaceDimToAllOperandDims(m, mOperands);1074 if (mOperands.size() != 2)1075 return failure();1076 if (llvm::all_of(mOperands, [](auto pair) {1077 return cast<ShapedType>(std::get<0>(pair).getType())1078 .getShape()[std::get<1>(pair)] == 1;1079 })) {1080 operandUnitDims = SmallVector<int64_t>{std::get<1>(mOperands[0]), -1,1081 std::get<1>(mOperands[1])};1082 return success();1083 }1084 } else {1085 auto n = contractionDims.n[0];1086 SmallVector<std::pair<Value, unsigned>, 2> nOperands;1087 op.mapIterationSpaceDimToAllOperandDims(n, nOperands);1088 if (nOperands.size() != 2)1089 return failure();1090 if (llvm::all_of(nOperands, [](auto pair) {1091 return cast<ShapedType>(std::get<0>(pair).getType())1092 .getShape()[std::get<1>(pair)] == 1;1093 })) {1094 operandUnitDims = SmallVector<int64_t>{-1, std::get<1>(nOperands[0]),1095 std::get<1>(nOperands[1])};1096 return success();1097 }1098 }1099 LLVM_DEBUG(llvm::dbgs() << "specified unit dims not found");1100 return failure();1101 }1102};1103 1104} // namespace1105 1106void mlir::linalg::populateContractionOpRankReducingPatterns(1107 RewritePatternSet &patterns) {1108 MLIRContext *context = patterns.getContext();1109 // Unbatching patterns for unit batch size1110 patterns.add<RankReduceToUnBatched<BatchMatmulOp, MatmulOp>>(context);1111 patterns.add<RankReduceToUnBatched<BatchMatvecOp, MatvecOp>>(context);1112 patterns.add<RankReduceToUnBatched<BatchVecmatOp, VecmatOp>>(context);1113 1114 // Non-batch rank 1 reducing patterns1115 patterns.add<RankReduceMatmul<MatmulOp, VecmatOp>>(context);1116 patterns.add<RankReduceMatmul<MatmulOp, MatvecOp>>(context);1117 // Batch rank 1 reducing patterns1118 patterns.add<RankReduceMatmul<BatchMatmulOp, BatchVecmatOp>>(context);1119 patterns.add<RankReduceMatmul<BatchMatmulOp, BatchMatvecOp>>(context);1120 1121 // Non-batch rank 0 reducing patterns1122 patterns.add<RankReduceMatmul<MatvecOp, DotOp>>(context);1123 patterns.add<RankReduceMatmul<VecmatOp, DotOp>>(context);1124}1125