[mlir][linalg] Fix vectorizer generating invalid vector.gather for 0-D tensor.extract

[edg-l]

Vectorizing a rank-0 linalg.generic whose body contains tensor.extract with data-dependent indices hits the Gather classification in getTensorExtractMemoryAccessPattern because isOutput1DVector returns false for a 0-D result. This produces an invalid vector.gather where operand #2 must be a vector of index values but gets a scalar index instead.

Fix classifies a 0-D result as ScalarBroadcast rather than Gather, and skips mask generation for 0-D in that path.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: Edgar (edg-l)

Changes

When vectorizing a rank-0 linalg.generic whose body contains tensor.extract with data-dependent indices, getTensorExtractMemoryAccessPattern fell through to the Gather classification because isOutput1DVector returns false for a 0-D result. This produced an invalid vector.gather where operand #2 must be a vector of index values but got a scalar index.

Error seen in practice:

'vector.gather' op operand #<!-- -->2 must be vector of integer or index values, but got 'index'

Fix in Vectorization.cpp:

1. Early-return ScalarBroadcast in getTensorExtractMemoryAccessPattern when resType.getRank() == 0, before the isOutput1DVector check.
2. Skip the vector<1xi1> masking step in the ScalarBroadcast handler when dstRank == 0, since 0-D vectors don't support masking.

Reproducer: ONNX Gather ops lowered to linalg.generic + tensor.extract on a rank-0 iteration space, as seen during GPT-2 model compilation via ONNX-MLIR.

---

Full diff: https://github.com/llvm/llvm-project/pull/187085.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+18-11)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 0477815f329bf..d2439ef1f2bf4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1093,6 +1093,11 @@ getTensorExtractMemoryAccessPattern(tensor::ExtractOp extractOp,
   if (inputShape.getShape().empty())
     return VectorMemoryAccessKind::ScalarBroadcast;
 
+  // 0a. Is the result a 0-D vector? If yes, there are no iteration dimensions
+  // so the tensor.extract is a single scalar load regardless of the index.
+  if (resType.getRank() == 0)
+    return VectorMemoryAccessKind::ScalarBroadcast;
+
   // True for vectors that are effectively 1D, e.g. `vector<1x4x1xi32>`, false
   // otherwise.
   bool isOutput1DVector =
@@ -1254,19 +1259,21 @@ vectorizeTensorExtract(RewriterBase &rewriter, VectorizationState &state,
         rewriter, loc, resultType, extractOp.getTensor(), transferReadIdxs,
         /*padding=*/std::nullopt, permutationMap, inBounds);
 
-    // Mask this broadcasting xfer_read here rather than relying on the generic
-    // path (the generic path assumes identity masking map, which wouldn't be
-    // valid here).
-    SmallVector<int64_t> readMaskShape = {1};
-    auto readMaskType = VectorType::get(readMaskShape, rewriter.getI1Type());
-    auto allTrue = vector::ConstantMaskOp::create(
-        rewriter, loc, readMaskType, vector::ConstantMaskKind::AllTrue);
-    auto *maskedReadOp =
-        mlir::vector::maskOperation(rewriter, transferReadOp, allTrue);
+    Operation *resultOp = transferReadOp;
+    if (dstRank > 0) {
+      // Mask this broadcasting xfer_read here rather than relying on the
+      // generic path (the generic path assumes identity masking map, which
+      // wouldn't be valid here).
+      SmallVector<int64_t> readMaskShape = {1};
+      auto readMaskType = VectorType::get(readMaskShape, rewriter.getI1Type());
+      auto allTrue = vector::ConstantMaskOp::create(
+          rewriter, loc, readMaskType, vector::ConstantMaskKind::AllTrue);
+      resultOp =
+          mlir::vector::maskOperation(rewriter, transferReadOp, allTrue);
+    }
 
     LDBG() << "Vectorised as scalar broadcast load: " << extractOp;
-    return VectorizationHookResult{VectorizationHookStatus::NewOp,
-                                   maskedReadOp};
+    return VectorizationHookResult{VectorizationHookStatus::NewOp, resultOp};
   }
 
   // 2b. Handle contiguous access.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[banach-space]

Thanks for the fix! Please add tests :)

[edg-l]

Thanks for the fix! Please add tests :)

added

[banach-space]

LGTM % nits

Thanks!

[banach-space]

@edg-l , do you have commit access to land this?

[edg-l]

@edg-l , do you have commit access to land this?

no i dont have any permission