[Mlir-commits] [mlir] [mlir][linalg] Collapse ND tensor.extract sources before gathering (PR #237308)

[Federico Bruzzone]

================
@@ -1205,10 +1205,25 @@ vectorizeTensorExtract(RewriterBase &rewriter, VectorizationState &state,
   if (memAccessKind == VectorMemoryAccessKind::Gather) {
     Value offset = calculateGatherOffset(rewriter, state, extractOp, bvm);
 
+    // vector.gather adds each element of its index vector to the innermost
+    // base coordinate. Flatten the source to match the linearized indices.
+    Value source = extractOp.getTensor();
+    auto sourceType = cast<RankedTensorType>(source.getType());
+    if (sourceType.getRank() > 1) {
+      SmallVector<ReassociationIndices> reassociation{
----------------
FedericoBruzzone wrote:

The `SmallVector` should not be needed 🤔

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