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

[James Newling]

================
@@ -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);
----------------
newling wrote:

My understanding of how this would work is as follows. Suppose the vector being reshaped has N elements, and the small tile (2x2) has n elements. The N extracts could be reduced to a single from_elements op, before anything else happens.  And n inserts can be replaced by 1 to_elements. Something like 

X = shape_cast Y

becomes

```

elements = from_elements X 
for i = 1 : N / n:
  indices = [...] compute the source elements as currently done
  tile = from_elements {elements at indices} 
  insert_strided tile into big_tile
```

It'd be quite a complex canonicalization to get to this from the current IR, I think. 

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