[mlir][Vector] Fold InsertStridedSliceOp of two splat with the same input to splat.

[jacquesguan]

Summary

This patch folds InsertStridedSliceOp(SplatOp(X):src_type, SplatOp(X):dst_type) to SplatOp(X):dst_type.

[Mogball]

LGTM

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

2189

[jacquesguan]

use replaceOp with Dest rather than create a new SplatOp.

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

bool hasNonUnitStrides() {
       return llvm::any_of(getStrides(), [](Attribute attr) {
         return attr.cast<IntegerAttr>().getInt() != 1;
       });
     }
   }];
 
   let hasFolder = 1;
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Vector_OuterProductOp :
   Vector_Op<"outerproduct", [NoSideEffect,
     PredOpTrait<"lhs operand and result have same element type",
                 TCresVTEtIsSameAsOpBase<0, 0>>,
     PredOpTrait<"rhs operand and result have same element type",
                 TCresVTEtIsSameAsOpBase<0, 1>>]>,

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

if (failed(isIntegerArrayAttrConfinedToShape(*this, offsets, destShape,
           makeI64ArrayAttr(sourceShapeAsDestShape, getContext()), destShape,
           offName, "source vector shape",
           /*halfOpen=*/false, /*min=*/1)))
     return failure();
 
   return success();
 }
 
+namespace {
+/// Pattern to rewrite an InsertStridedSliceOp(SplatOp(X):src_type,
+/// SplatOp(X):dst_type) to SplatOp(X):dst_type.
+class FoldInsertStridedSliceSplat final
+    : public OpRewritePattern<InsertStridedSliceOp> {
+public:
+  using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(InsertStridedSliceOp insertStridedSliceOp,
+                                PatternRewriter &rewriter) const override {
+    auto srcSplatOp =
+        insertStridedSliceOp.getSource().getDefiningOp<vector::SplatOp>();
+    auto destSplatOp =
+        insertStridedSliceOp.getDest().getDefiningOp<vector::SplatOp>();
+
+    if (!srcSplatOp || !destSplatOp)
+      return failure();
+
+    if (srcSplatOp.getInput() != destSplatOp.getInput())
+      return failure();
+
+    rewriter.replaceOp(insertStridedSliceOp, insertStridedSliceOp.getDest());
+    return success();
+  }
+};
+} // namespace
+
+void vector::InsertStridedSliceOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  results.add<FoldInsertStridedSliceSplat>(context);
+}
+
 OpFoldResult InsertStridedSliceOp::fold(ArrayRef<Attribute> operands) {
   if (getSourceVectorType() == getDestVectorType())
     return getSource();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // OuterProductOp

mlir/test/Dialect/Vector/canonicalize.mlir

 // CHECK-LABEL: func @bitcast(
 //  CHECK-SAME:               %[[ARG:.*]]: vector<4x8xf32>) -> vector<4x16xi16> {
 //       CHECK: vector.bitcast %[[ARG:.*]] : vector<4x8xf32> to vector<4x16xi16>
 func.func @bitcast(%a: vector<4x8xf32>) -> vector<4x16xi16> {
   %0 = vector.bitcast %a : vector<4x8xf32> to vector<4x8xi32>
   %1 = vector.bitcast %0 : vector<4x8xi32> to vector<4x16xi16>
   return %1 : vector<4x16xi16>
 }
+
+// -----
+
+// CHECK-LABEL: @insert_strided_slice_splat
+//  CHECK-SAME: (%[[ARG:.*]]: f32)
+//  CHECK-NEXT:   %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<8x16xf32>
+//  CHECK-NEXT:   return %[[SPLAT]] : vector<8x16xf32>
+func.func @insert_strided_slice_splat(%x: f32) -> (vector<8x16xf32>) {
+  %splat0 = vector.splat %x : vector<4x4xf32>
+  %splat1 = vector.splat %x : vector<8x16xf32>
+  %0 = vector.insert_strided_slice %splat0, %splat1 {offsets = [2, 2], strides = [1, 1]}
+    : vector<4x4xf32> into vector<8x16xf32>
+  return %0 : vector<8x16xf32>
+}