[Mlir-commits] [mlir] [MLIR][Vector] Add unroll pattern for vector.shape_cast (PR #164010)

[Diego Caballero]

================
@@ -75,6 +75,45 @@ static SmallVector<Value> sliceLoadStoreIndices(PatternRewriter &rewriter,
   return indices;
 }
 
+/// Creates a result tile by extracting individual elements from the source
+/// and inserting them at the correct positions in the tile.
+static Value createTileFromElements(PatternRewriter &rewriter, Location loc,
+                                    Value source, ArrayRef<int64_t> sourceShape,
+                                    ArrayRef<int64_t> resultShape,
+                                    ArrayRef<int64_t> tileOffsets,
+                                    ArrayRef<int64_t> tileShape,
+                                    VectorType tileType) {
+  // Initialize tile with zeros.
+  Value tile = arith::ConstantOp::create(rewriter, loc, tileType,
+                                         rewriter.getZeroAttr(tileType));
+
+  // Calculate strides for source, result, and tile shapes.
+  SmallVector<int64_t> sourceStrides = computeStrides(sourceShape);
+  SmallVector<int64_t> resultStrides = computeStrides(resultShape);
+  SmallVector<int64_t> tileStrides = computeStrides(tileShape);
+  int64_t numElementsInTile = computeProduct(tileShape);
+
+  // Iterate over all positions in the tile using linear indexing.
+  for (int64_t linearTileIdx = 0; linearTileIdx < numElementsInTile;
+       ++linearTileIdx) {
+    // Convert linear tile index to multi-dimensional tile position.
+    SmallVector<int64_t> tilePosition = delinearize(linearTileIdx, tileStrides);
+
+    // Calculate the global position in the result.
+    SmallVector<int64_t> globalResultPos;
+    globalResultPos.reserve(tileOffsets.size());
+    for (auto [offset, pos] : llvm::zip_equal(tileOffsets, tilePosition)) {
+      globalResultPos.push_back(offset + pos);
+    }
+
+    int64_t linearIndex = linearize(globalResultPos, resultStrides);
+    SmallVector<int64_t> sourcePos = delinearize(linearIndex, sourceStrides);
+    Value element = vector::ExtractOp::create(rewriter, loc, source, sourcePos);
+    tile = vector::InsertOp::create(rewriter, loc, element, tile, tilePosition);
----------------
dcaballe wrote:

could we use `vector.to_elements` and `vector.from_elements` instead of `vectort.extract` and `vector.insert`? That should simplify this loop and reduce code size significantly

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