[Mlir-commits] [mlir] [mlir][vector] Add InsertInsertToInsert to insert op canonicalize patterns (PR #147045)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/147045.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+24-1) 
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+14) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 1fb8c7a928e06..2d090bcce45ab 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -3335,6 +3335,28 @@ class InsertSplatToSplat final : public OpRewritePattern<InsertOp> {
   }
 };
 
+/// Pattern to rewrite a InsertOp(InsertOp) to InsertOp.
+class InsertInsertToInsert final : public OpRewritePattern<InsertOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(InsertOp op,
+                                PatternRewriter &rewriter) const override {
+    auto destInsert = op.getDest().getDefiningOp<InsertOp>();
+    if (!destInsert)
+      return failure();
+
+    if (!destInsert->hasOneUse())
+      return failure();
+
+    if (op.getMixedPosition() != destInsert.getMixedPosition())
+      return failure();
+
+    rewriter.replaceOpWithNewOp<InsertOp>(
+        op, op.getValueToStore(), destInsert.getDest(), op.getMixedPosition());
+    return success();
+  }
+};
+
 } // namespace
 
 static Attribute
@@ -3389,7 +3411,8 @@ foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,
 
 void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
-  results.add<InsertToBroadcast, BroadcastFolder, InsertSplatToSplat>(context);
+  results.add<InsertToBroadcast, BroadcastFolder, InsertSplatToSplat,
+              InsertInsertToInsert>(context);
 }
 
 OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) {
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 0282e9cac5e02..73bced6149ff6 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3446,3 +3446,17 @@ func.func @fold_insert_constant_indices(%arg : vector<4x1xi32>) -> vector<4x1xi3
   %res = vector.insert %1, %arg[%0, %0] : i32 into vector<4x1xi32>
   return %res : vector<4x1xi32>
 }
+
+// -----
+
+// CHECK-LABEL: @insert_insert_to_insert(
+//  CHECK-SAME:   %[[ARG:.*]]: vector<4xf32>,
+//  CHECK-SAME:   %[[VAL:.*]]: f32) -> vector<4xf32> {
+//       CHECK:   %[[RES:.*]] = vector.insert %[[VAL]], %[[ARG]] [0] : f32 into vector<4xf32>
+//       CHECK:    return %[[RES]] : vector<4xf32>
+func.func @insert_insert_to_insert(%v : vector<4xf32>, %value : f32) -> vector<4xf32> {
+  %v_0 = vector.insert %value, %v[0] : f32 into vector<4xf32>
+  %v_1 = vector.insert %value, %v_0[0] : f32 into vector<4xf32>
+  %v_2 = vector.insert %value, %v_1[0] : f32 into vector<4xf32>
+  return %v_2 : vector<4xf32>  
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/147045