[Mlir-commits] [mlir] [MLIR][Vector] Fix crash in operatesOnSuperVectorsOf on rank-mismatched shaped (PR #183967)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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
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 {

``````````

</details>


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