[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{
+          llvm::to_vector<2>(llvm::seq<int64_t>(0, sourceType.getRank()))};
+      // Infer the type without an encoding that may depend on the source rank.
+      auto collapsedType = tensor::CollapseShapeOp::inferCollapsedType(
+          sourceType, reassociation);
+      source = tensor::CollapseShapeOp::create(rewriter, loc, collapsedType,
+                                               source, reassociation);
+      baseIndices.resize(1);
----------------
FedericoBruzzone wrote:

Is this the correct position for this line?

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