diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -2206,11 +2206,60 @@ return success(); } }; + +/// Pattern to rewrite an InsertStridedSliceOp(ExtractStridedSliceOp(dst), dst) +/// to dst. +class FoldInsertStridedSliceOfExtract final + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(InsertStridedSliceOp insertStridedSliceOp, + PatternRewriter &rewriter) const override { + auto extractStridedSliceOp = + insertStridedSliceOp.getSource() + .getDefiningOp(); + + if (!extractStridedSliceOp) + return failure(); + + if (extractStridedSliceOp.getOperand() != insertStridedSliceOp.getDest()) + return failure(); + + auto getElement = [](ArrayAttr array, int idx) { + return array[idx].cast().getInt(); + }; + + auto extractStrides = extractStridedSliceOp.getStrides(); + auto extractOffsets = extractStridedSliceOp.getOffsets(); + auto insertStrides = insertStridedSliceOp.getStrides(); + auto insertOffsets = insertStridedSliceOp.getOffsets(); + + // Check if have the same strides. + if (extractStrides.size() != insertStrides.size()) + return failure(); + for (size_t i = 0; i < extractStrides.size(); ++i) + if (getElement(extractStrides, i) != getElement(insertStrides, i)) + return failure(); + + // Check if have the same offsets. + if (extractOffsets.size() != insertOffsets.size()) + return failure(); + for (size_t i = 0; i < extractOffsets.size(); ++i) + if (getElement(extractOffsets, i) != getElement(insertOffsets, i)) + return failure(); + + rewriter.replaceOp(insertStridedSliceOp, insertStridedSliceOp.getDest()); + return success(); + } +}; + } // namespace void vector::InsertStridedSliceOp::getCanonicalizationPatterns( RewritePatternSet &results, MLIRContext *context) { - results.add(context); + results.add( + context); } OpFoldResult InsertStridedSliceOp::fold(ArrayRef operands) { diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir --- a/mlir/test/Dialect/Vector/canonicalize.mlir +++ b/mlir/test/Dialect/Vector/canonicalize.mlir @@ -1641,3 +1641,17 @@ : vector<4x4xf32> into vector<8x16xf32> return %0 : vector<8x16xf32> } + + +// ----- + +// CHECK-LABEL: @insert_extract_strided_slice +// CHECK-SAME: (%[[ARG:.*]]: vector<8x16xf32>) +// CHECK-NEXT: return %[[ARG]] : vector<8x16xf32> +func.func @insert_extract_strided_slice(%x: vector<8x16xf32>) -> (vector<8x16xf32>) { + %0 = vector.extract_strided_slice %x {offsets = [0, 8], sizes = [2, 4], strides = [1, 1]} + : vector<8x16xf32> to vector<2x4xf32> + %1 = vector.insert_strided_slice %0, %x {offsets = [0, 8], strides = [1, 1]} + : vector<2x4xf32> into vector<8x16xf32> + return %1 : vector<8x16xf32> +}