[MLIR][Vector] Fix crash in operatesOnSuperVectorsOf on rank-mismatched shaped

[joker-eph]

The operatesOnSuperVectorsOf function in VectorUtils.cpp contained an assertion that fired when a vector.transfer operation's vector type had a different rank (or non-divisible shape) from the sub-vector type supplied by the caller:

assert((ratio || !mustDivide) &&
"vector.transfer operation in which super-vector size is not an"
" integer multiple of sub-vector size");

This assertion was incorrect because the function's callers (e.g., the affine super-vectorizer) legitimately pass transfer ops whose vector type doesn't match the requested sub-vector shape. In those cases the correct answer is simply that the op does not operate on a super-vector of that sub-vector type, so operatesOnSuperVectorsOf should return false.

Remove the assert return false when computeShapeRatio produces no result, and remove the now-unused mustDivide variable.

Fixes #149327
Fixes #131096

[llvmbot]

@llvm/pr-subscribers-mlir-vector
@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

The operatesOnSuperVectorsOf function in VectorUtils.cpp contained an assertion that fired when a vector.transfer operation's vector type had a different rank (or non-divisible shape) from the sub-vector type supplied by the caller:

assert((ratio || !mustDivide) &&
"vector.transfer operation in which super-vector size is not an"
" integer multiple of sub-vector size");

This assertion was incorrect because the function's callers (e.g., the affine super-vectorizer) legitimately pass transfer ops whose vector type doesn't match the requested sub-vector shape. In those cases the correct answer is simply that the op does not operate on a super-vector of that sub-vector type, so operatesOnSuperVectorsOf should return false.

Remove the assert return false when computeShapeRatio produces no result, and remove the now-unused mustDivide variable.

Fixes #149327

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp (+3-15)
• (modified) mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir (+17)

diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index e123f9e21bbeb..d1ce0fad2fb56 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -209,12 +209,9 @@ bool matcher::operatesOnSuperVectorsOf(Operation &op,
   // explicitly checked for this property.
   /// TODO: there should be a single function for all ops to do this so we
   /// do not have to special case. Maybe a trait, or just a method, unclear atm.
-  bool mustDivide = false;
-  (void)mustDivide;
   VectorType superVectorType;
   if (auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
     superVectorType = transfer.getVectorType();
-    mustDivide = true;
   } else if (op.getNumResults() == 0) {
     if (!isa<func::ReturnOp>(op)) {
       op.emitError("NYI: assuming only return operations can have 0 "
@@ -235,20 +232,11 @@ bool matcher::operatesOnSuperVectorsOf(Operation &op,
     return false;
   }
 
-  // Get the ratio.
+  // Get the ratio. If the shapes are incompatible (e.g., different ranks or
+  // non-integer divisibility), the operation does not operate on a super-vector
+  // of the given sub-vector type.
   auto ratio =
       computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());
-
-  // Sanity check.
-  assert((ratio || !mustDivide) &&
-         "vector.transfer operation in which super-vector size is not an"
-         " integer multiple of sub-vector size");
-
-  // This catches cases that are not strictly necessary to have multiplicity but
-  // still aren't divisible by the sub-vector shape.
-  // This could be useful information if we wanted to reshape at the level of
-  // the vector type (but we would have to look at the compute and distinguish
-  // between parallel, reduction and possibly other cases.
   return ratio.has_value();
 }
 
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir
index bd71164244c00..fcf31daa987b4 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir
@@ -52,6 +52,23 @@ func.func @double_loop_nest(%a: memref<20x30xf32>, %b: memref<20xf32>) {
   return
 }
 
+// Regression test for https://github.com/llvm/llvm-project/issues/149327
+// Verifies no crash when a vector.transfer_read has a 1D vector type but the
+// shape ratio is 2D, causing a rank mismatch in operatesOnSuperVectorsOf.
+// The transfer op should simply not be matched (not crash).
+// CHECK-LABEL: func @transfer_rank_mismatch_no_crash
+func.func @transfer_rank_mismatch_no_crash(%arg0: memref<82x97xf32>) {
+  %0 = ub.poison : f32
+  affine.for %arg1 = 0 to 82 {
+    affine.for %arg2 = 0 to 97 step 128 {
+      // The vector type is 1D but the shape ratio is 2D — no crash.
+      // CHECK-NOT: matched: {{.*}} = vector.transfer_read
+      %1 = vector.transfer_read %arg0[%arg1, %arg2], %0 : memref<82x97xf32>, vector<128xf32>
+    }
+  }
+  return
+}
+
 // VECNEST:       affine.for %{{.*}} = 0 to 20 step 4 {
 // VECNEST:         vector.transfer_read
 // VECNEST-NEXT:    affine.for %{{.*}} = 0 to 30 {

[github-actions]

🐧 Linux x64 Test Results

• 7543 tests passed
• 601 tests skipped

✅ The build succeeded and all tests passed.

[github-actions]

🪟 Windows x64 Test Results

• 3487 tests passed
• 414 tests skipped

✅ The build succeeded and all tests passed.

[banach-space]

Thanks!

I've not really touched this code myself, but the fix makes sense and is quite straightforward.

LGTM

[llvm-ci]

LLVM Buildbot has detected a new failure on builder ppc64le-flang-rhel-clang running on ppc64le-flang-rhel-test while building mlir at step 7 "test-build-unified-tree-check-flang-rt".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/157/builds/45020

Here is the relevant piece of the build log for the reference

Step 7 (test-build-unified-tree-check-flang-rt) failure: 1200 seconds without output running [b'ninja', b'check-flang-rt'], attempting to kill
...
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/229/271 (267 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/238/271 (268 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/239/271 (269 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/228/271 (270 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/236/271 (271 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/90/271 (272 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/244/271 (273 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/243/271 (274 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/234/271 (275 of 276)
PASS: flang-rt-Unit :: Runtime/./RuntimeTests/235/271 (276 of 276)
command timed out: 1200 seconds without output running [b'ninja', b'check-flang-rt'], attempting to kill
process killed by signal 9
program finished with exit code -1
elapsedTime=1588.735265