[mlir][vector] Add foldInsertUseChain folder function to insert op

[linuxlonelyeagle]

When the result of an insert op is used by an insert op, and the subsequent insert op is inserted at the same location as the previous insert op, replaces the dest of the subsequent insert op with the dest of the previous insert op.This is because the previous insert op does not affect subsequent insert ops.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

Changes

---

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 --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>  
+}

[linuxlonelyeagle]

Subsequently my emails seem to be Pending, resubmit them. @joker-eph

[banach-space]

Thanks!

[linuxlonelyeagle]

I've read everyone's advice, I'm just not sure I've accurately understood everyone's realisation.I have made the changes to the test that I understand should be made.

Record the main elements of the three test tests via email.

• fold_insert_use_chain
  The test is a complete insert use chain with all the same insert position.
• no_fold_insert_use_chain
  The test is a complete insert use chain, but the insert position is different.
• fold_insert_use_chain_add_float
  This is a complete insert use chain, but the user of the first insert contains other ops (in this case the arith.addf op).

[joker-eph]

LGTM, but I see no use for the third test right now.

[dcaballe]

Thanks! LGTM % ongoing comments

[banach-space]

LGTM, thanks, this is a very nice improvement!

[linuxlonelyeagle]

Thank you all.