[Mlir-commits] [mlir] [mlir][vector] Support multi-dimensional vectors in VectorFromElementsLowering (PR #151175)

[Andrzej Warzyński]

================
@@ -735,10 +735,52 @@ struct LowerVectorToFromElementsToShuffleTreePass
   }
 };
 
+/// Unrolls 2 or more dimensional `vector.from_elements` ops by unrolling the
+/// outermost dimension. For example:
+/// ```
+/// %v = vector.from_elements %e0, %e1, %e2, %e3, %e4, %e5 : vector<2x3xf32>
+///
+/// ==>
+///
+/// %0   = ub.poison : vector<2x3xf32>
+/// %v0  = vector.from_elements %e0, %e1, %e2 : vector<3xf32>
+/// %1   = vector.insert %v0, %0 [0] : vector<3xf32> into vector<2x3xf32>
+/// %v1  = vector.from_elements %e3, %e4, %e5 : vector<3xf32>
+/// %v   = vector.insert %v1, %1 [1] : vector<3xf32> into vector<2x3xf32>
+/// ```
+///
+/// When applied exhaustively, this will produce a sequence of 1-d from_elements
+/// ops.
+struct UnrollFromElements : OpRewritePattern<vector::FromElementsOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::FromElementsOp op,
+                                PatternRewriter &rewriter) const override {
+    ValueRange allElements = op.getElements();
+
+    auto unrollFromElementsFn = [&](PatternRewriter &rewriter, Location loc,
+                                    VectorType subTy, int64_t index) {
+      size_t subTyNumElements = subTy.getNumElements();
+      assert((index + 1) * subTyNumElements <= allElements.size() &&
+             "out of bounds");
+      ValueRange subElements =
+          allElements.slice(index * subTyNumElements, subTyNumElements);
+      return vector::FromElementsOp::create(rewriter, loc, subTy, subElements);
+    };
+
+    return unrollVectorOp(op, rewriter, unrollFromElementsFn);
+  }
+};
+
----------------
banach-space wrote:

The name of this file, [LowerVectorToFromElementsToShuffleTree.cpp](https://github.com/llvm/llvm-project/pull/151175/files#diff-ddfef9af0cf24f4069d46746801503e3c721a6f0958c43d6ff555af70f969a3a), as well as the comment at the top:
```cpp
// This file implements pattern rewrites to lower sequences of
// `vector.to_elements` and `vector.from_elements` operations into a tree of
// `vector.shuffle` operations.
```
imply that the only transformation implemented in this file is "to-shuffle". So either we need to update the filename or create a new file. Given how large and complex "to-shuffle" is, I would just create a new file. Thanks!

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