[Mlir-commits] [mlir] [mlir][Vector] add vector.insert canonicalization pattern for vectors created from ub.poison (PR #142944)

[Yang Bai]

================
@@ -3250,6 +3262,128 @@ class InsertSplatToSplat final : public OpRewritePattern<InsertOp> {
     return success();
   }
 };
+
+/// Pattern to optimize a chain of insertions.
+///
+/// This pattern identifies chains of vector.insert operations that:
+/// 1. Only insert values at static positions.
+/// 2. Completely initialize all elements in the resulting vector.
+/// 3. All intermediate insert operations have only one use.
+///
+/// When these conditions are met, the entire chain can be replaced with a
+/// single vector.from_elements operation.
+///
+/// Example transformation:
+///   %poison = ub.poison : vector<2xi32>
+///   %0 = vector.insert %c1, %poison[0] : i32 into vector<2xi32>
+///   %1 = vector.insert %c2, %0[1] : i32 into vector<2xi32>
+/// ->
+///   %result = vector.from_elements %c1, %c2 : vector<2xi32>
+class InsertChainFullyInitialized final : public OpRewritePattern<InsertOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(InsertOp op,
+                                PatternRewriter &rewriter) const override {
+
+    VectorType destTy = op.getDestVectorType();
+    if (destTy.isScalable())
+      return failure();
+    // Check if the result is used as the dest operand of another vector.insert
+    // Only care about the last op in a chain of insertions.
+    for (Operation *user : op.getResult().getUsers())
+      if (auto insertOp = dyn_cast<InsertOp>(user))
+        if (insertOp.getDest() == op.getResult())
+          return failure();
+
+    InsertOp currentOp = op;
+    SmallVector<InsertOp> chainInsertOps;
+    while (currentOp) {
+      // Dynamic position is not supported.
+      if (currentOp.hasDynamicPosition())
+        return failure();
+
+      chainInsertOps.push_back(currentOp);
+      currentOp = currentOp.getDest().getDefiningOp<InsertOp>();
+      // Check that intermediate inserts have only one use to avoid an explosion
+      // of vectors.
+      if (currentOp && !currentOp->hasOneUse())
+        return failure();
----------------
yangtetris wrote:

These are two different checks.
The check on lines L3293-L3296 is to ensure that there are no more insert ops **after** the current op. The purpose of that check is to skip the O(n) time complexity part for intermediate insert ops.
In contrast, this `hasOneUse` check is to handle the following case:
```
%v1 = %vector.insert %c1, %v0[0] : i64 into vector<3xi64>
%v2 = %vector.insert %c2, %v1[1] : i64 into vector<3xi64>
%v3_3 = %vector.insert %c3, %v2[2] : i64 into vector<3xi64>
%v3_4 = %vector.insert %c4, %v2[2] : i64 into vector<3xi64>
%v3_5 = %vector.insert %c5, %v2[2] : i64 into vector<3xi64>
```
The key point is that when `%v1` or `%v2` has multiple users, we should not introduce new from_elements ops because the insert chain cannot be completely eliminated. This IR is also an example of the 'explosion' you asked about later. We want to avoid generating three new form_elements ops for `%v3_3`, `%v3_4` and `%v3_5`.

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