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/IR/BuiltinTypes.h" 13 #include "mlir/IR/OpDefinition.h" 14 #include "mlir/IR/PatternMatch.h" 15 16 using namespace mlir; 17 using namespace mlir::tensor; 18 19 namespace { 20 /// Merges consecutive tensor.extract_slice ops into one. 21 // TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps. 22 struct MergeConsecutiveExtractSlice : public OpRewritePattern<ExtractSliceOp> { 23 using OpRewritePattern::OpRewritePattern; 24 25 LogicalResult matchAndRewrite(ExtractSliceOp nextOp, 26 PatternRewriter &rewriter) const override { 27 auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>(); 28 if (!prevOp) 29 return failure(); 30 31 SmallVector<OpFoldResult> newOffsets, newSizes, newStrides; 32 if (failed(affine::mergeOffsetsSizesAndStrides( 33 rewriter, nextOp.getLoc(), prevOp, nextOp, prevOp.getDroppedDims(), 34 newOffsets, newSizes, newStrides))) 35 return failure(); 36 37 rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(), 38 prevOp.getSource(), newOffsets, 39 newSizes, newStrides); 40 return success(); 41 } 42 }; 43 44 /// Merges consecutive tensor.insert_slice ops into one. 45 // TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps. 46 template <typename OpTy> 47 struct MergeConsecutiveInsertSlice : public OpRewritePattern<OpTy> { 48 using OpRewritePattern<OpTy>::OpRewritePattern; 49 50 LogicalResult matchAndRewrite(OpTy nextOp, 51 PatternRewriter &rewriter) const override { 52 auto prevOp = nextOp.getSource().template getDefiningOp<InsertSliceOp>(); 53 if (!prevOp) 54 return failure(); 55 56 if (!prevOp.hasUnitStride() || !nextOp.hasUnitStride()) 57 return failure(); 58 59 // The first insert_slice op should be rank reducing to make sure we cover 60 // the full source tensor to be inserted in the second insert_slice op. 61 SliceVerificationResult result = 62 isRankReducedType(prevOp.getDestType(), prevOp.getSourceType()); 63 if (result != SliceVerificationResult::Success) 64 return failure(); 65 66 // Dynamic dimensions can pass rank reducing check in the above, e.g, 67 // inserting <?xf32> into <1x?x1xf32>. For such cases we cannot be certain 68 // the dynamic size covers the full tensor. 69 if (!prevOp.getSourceType().hasStaticShape() || 70 !prevOp.getDestType().hasStaticShape()) 71 return failure(); 72 73 rewriter.replaceOpWithNewOp<OpTy>( 74 nextOp, prevOp.getSource(), nextOp.getDest(), nextOp.getMixedOffsets(), 75 nextOp.getMixedSizes(), nextOp.getMixedStrides()); 76 return success(); 77 } 78 }; 79 } // namespace 80 81 void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns( 82 RewritePatternSet &patterns) { 83 patterns.add<MergeConsecutiveExtractSlice, 84 MergeConsecutiveInsertSlice<InsertSliceOp>, 85 MergeConsecutiveInsertSlice<ParallelInsertSliceOp>>( 86 patterns.getContext()); 87 } 88