1 //===- MergeConsecutiveInsertExtractSlicePatterns.cpp ---------------------===// 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-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h" 10 #include "mlir/Dialect/Tensor/IR/Tensor.h" 11 #include "mlir/Dialect/Tensor/Transforms/Transforms.h" 12 #include "mlir/Dialect/Tensor/Utils/Utils.h" 13 #include "mlir/IR/BuiltinTypes.h" 14 #include "mlir/IR/OpDefinition.h" 15 #include "mlir/IR/PatternMatch.h" 16 17 using namespace mlir; 18 using namespace mlir::tensor; 19 20 namespace { 21 /// Merges consecutive tensor.extract_slice ops into one. 22 // TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps. 23 struct MergeConsecutiveExtractSlice : public OpRewritePattern<ExtractSliceOp> { 24 using OpRewritePattern::OpRewritePattern; 25 26 LogicalResult matchAndRewrite(ExtractSliceOp nextOp, 27 PatternRewriter &rewriter) const override { 28 auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>(); 29 if (!prevOp) 30 return failure(); 31 32 SmallVector<OpFoldResult> newOffsets, newSizes, newStrides; 33 if (failed(affine::mergeOffsetsSizesAndStrides( 34 rewriter, nextOp.getLoc(), prevOp, nextOp, prevOp.getDroppedDims(), 35 newOffsets, newSizes, newStrides))) 36 return failure(); 37 38 rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(), 39 prevOp.getSource(), newOffsets, 40 newSizes, newStrides); 41 return success(); 42 } 43 }; 44 45 /// Merges consecutive tensor.insert_slice ops into one. 46 // TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps. 47 template <typename OpTy> 48 struct MergeConsecutiveInsertSlice : public OpRewritePattern<OpTy> { 49 using OpRewritePattern<OpTy>::OpRewritePattern; 50 51 LogicalResult matchAndRewrite(OpTy nextOp, 52 PatternRewriter &rewriter) const override { 53 auto prevOp = nextOp.getSource().template getDefiningOp<InsertSliceOp>(); 54 if (!prevOp) 55 return failure(); 56 57 if (!prevOp.hasUnitStride() || !nextOp.hasUnitStride()) 58 return failure(); 59 60 // The first insert_slice op should be rank reducing to make sure we cover 61 // the full source tensor to be inserted in the second insert_slice op. 62 SliceVerificationResult result = 63 isRankReducedType(prevOp.getDestType(), prevOp.getSourceType()); 64 if (result != SliceVerificationResult::Success) 65 return failure(); 66 67 // Dynamic dimensions can pass rank reducing check in the above, e.g, 68 // inserting <?xf32> into <1x?x1xf32>. For such cases we cannot be certain 69 // the dynamic size covers the full tensor. 70 if (!prevOp.getSourceType().hasStaticShape() || 71 !prevOp.getDestType().hasStaticShape()) 72 return failure(); 73 74 rewriter.replaceOpWithNewOp<OpTy>( 75 nextOp, prevOp.getSource(), nextOp.getDest(), nextOp.getMixedOffsets(), 76 nextOp.getMixedSizes(), nextOp.getMixedStrides()); 77 return success(); 78 } 79 }; 80 81 /// Drop redundant rank expansion. I.e., rank expansions that are directly 82 /// followed by rank reductions. E.g.: 83 /// %0 = tensor.insert_slice ... : tensor<5x10xf32> into tensor<1x1x5x10xf32> 84 /// %1 = tensor.extract_slice %0[0, 0, 2, 3] [1, 1, 2, 2] [1, 1, 1, 1] 85 /// : tensor<1x1x5x10xf32> to tensor<2x2xf32> 86 struct DropRedundantInsertSliceRankExpansion 87 : public OpRewritePattern<ExtractSliceOp> { 88 using OpRewritePattern::OpRewritePattern; 89 90 LogicalResult matchAndRewrite(ExtractSliceOp extractSliceOp, 91 PatternRewriter &rewriter) const override { 92 // Nothing to do if no dims are dropped. 93 llvm::SmallBitVector droppedDims = extractSliceOp.getDroppedDims(); 94 if (droppedDims.none()) 95 return failure(); 96 97 // Look for tensor.insert_slice op that has an inverse rank expansion. 98 auto insertSliceOp = 99 extractSliceOp.getSource().getDefiningOp<InsertSliceOp>(); 100 if (!insertSliceOp) 101 return failure(); 102 llvm::SmallBitVector expandedDims = insertSliceOp.getDroppedDims(); 103 104 // TODO: This could be extended to support cases where the dropped dims are 105 // a subset of the expanded dims. 106 if (expandedDims != droppedDims) 107 return failure(); 108 109 // The tensor.insert_slice may not be redundant if it has multiple users. 110 if (!insertSliceOp->hasOneUse()) 111 return failure(); 112 113 // Only consider tensor.insert_slice ops that are pure rank-reductions. 114 // I.e., no elements are taken from the destination. 115 if (!isCastLikeInsertSliceOp(insertSliceOp)) 116 return failure(); 117 118 // Extract directly from the source. 119 OpBuilder::InsertionGuard g(rewriter); 120 rewriter.setInsertionPoint(extractSliceOp); 121 SmallVector<OpFoldResult> newOffsets, newSizes, newStrides; 122 for (int64_t i = 0, e = extractSliceOp.getSourceType().getRank(); i < e; 123 ++i) { 124 if (droppedDims.test(i)) 125 continue; 126 newOffsets.push_back(extractSliceOp.getMixedOffsets()[i]); 127 newSizes.push_back(extractSliceOp.getMixedSizes()[i]); 128 newStrides.push_back(extractSliceOp.getMixedStrides()[i]); 129 } 130 rewriter.replaceOpWithNewOp<ExtractSliceOp>( 131 extractSliceOp, /*source=*/insertSliceOp.getSource(), newOffsets, 132 newSizes, newStrides); 133 rewriter.eraseOp(insertSliceOp); 134 return success(); 135 } 136 }; 137 } // namespace 138 139 void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns( 140 RewritePatternSet &patterns) { 141 patterns.add<MergeConsecutiveExtractSlice, 142 MergeConsecutiveInsertSlice<InsertSliceOp>, 143 MergeConsecutiveInsertSlice<ParallelInsertSliceOp>>( 144 patterns.getContext()); 145 } 146 147 void mlir::tensor::populateDropRedundantInsertSliceRankExpansionPatterns( 148 RewritePatternSet &patterns) { 149 patterns.add<DropRedundantInsertSliceRankExpansion>(patterns.getContext()); 150 } 151