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

[Jakub Kuderski]

================
@@ -75,6 +75,54 @@ 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 = rewriter.create<arith::ConstantOp>(
+      loc, tileType, rewriter.getZeroAttr(tileType));
+
+  // Calculate strides for both source and result shapes.
+  SmallVector<int64_t> sourceStrides = computeStrides(sourceShape);
+  SmallVector<int64_t> resultStrides = computeStrides(resultShape);
+
+  // Iterate over all positions in the tile using linear indexing.
+  for (int64_t linearTileIdx = 0; linearTileIdx < computeProduct(tileShape);
+       ++linearTileIdx) {
+    // Convert linear tile index to multi-dimensional tile position.
+    SmallVector<int64_t> tilePosition =
+        delinearize(linearTileIdx, computeStrides(tileShape));
+
+    // Calculate the global position in the result.
+    SmallVector<int64_t> globalResultPos;
+    globalResultPos.reserve(tileOffsets.size());
+    for (auto [offset, pos] : llvm::zip(tileOffsets, tilePosition)) {
+      globalResultPos.push_back(offset + pos);
+    }
+
+    // Convert result position to linear index.
+    int64_t linearIndex = linearize(globalResultPos, resultStrides);
+
+    // Convert linear index to source position.
+    SmallVector<int64_t> sourcePos =
+        delinearize(linearIndex, computeStrides(sourceShape));
+
+    // Extract element from source.
+    Value element = vector::ExtractOp::create(rewriter, loc, source, sourcePos);
+
+    // Insert element into tile.
+    tile = vector::InsertOp::create(rewriter, loc, element, tile, tilePosition);
----------------
kuhar wrote:

```suggestion
    // Convert result position to linear index.
    int64_t linearIndex = linearize(globalResultPos, resultStrides);
    // Convert linear index to source position.
    SmallVector<int64_t> sourcePos =
        delinearize(linearIndex, computeStrides(sourceShape));
    // Extract element from source.
    Value element = vector::ExtractOp::create(rewriter, loc, source, sourcePos);
    // Insert element into tile.
    tile = vector::InsertOp::create(rewriter, loc, element, tile, tilePosition);
```

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