[Mlir-commits] [llvm] [mlir] [Linalg] Fix crash in vectorizeScalableVectorPrecondition with undersized vector sizes (PR #205493)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: LouisLu060211

<details>
<summary>Changes</summary>

`vectorizeScalableVectorPrecondition` iterates over the Linalg op's iterator
types and indexes into the user-provided `inputScalableVecDims` and
`inputVectorSizes` arrays without first checking that the arrays are long
enough. When the arrays are shorter than the number of loops. This results in 
the user passes `vector_sizes [8, [16], 4]` to a two-loop `linalg.add`—the
SmallVector `operator[]` assertion `idx < size()` fails and mlir-opt aborts.

Fix:
Add a bounds check that compares the array sizes against `linalgOp.getNumLoops()`
and returns a clean failure instead of crashing. The check is placed after the
LinalgOp interface cast so it only runs when the op is known to be a LinalgOp,
and before any loop that accesses the arrays by index.

`VectorizeOp::apply` emits the failure diagnostic on the target op's location
(`target->getLoc()`), but the regression test expects the error on the
transform op itself via `// expected-error @<!-- -->above`. Change the failure path to
use `getLoc()` so the diagnostic appears on the `transform.structured.vectorize`
line, matching the test directive.

The existing test file
`mlir/test/Dialect/Linalg/transform/vectorize-dynamic-mixed-sizes.mlir`
reproduces the original crash and now passes with `-verify-diagnostics`.

Fixes #<!-- -->204100.
Assisted-by: DeepSeek V4 Pro

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


4 Files Affected:

- (modified) .gitignore (+1) 
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+1-1) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+12) 
- (added) mlir/test/Dialect/Linalg/transform/vectorize-dynamic-mixed-sizes.mlir (+22) 


``````````diff
diff --git a/.gitignore b/.gitignore
index a4382c9ea7390..30e994aefd8bd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -90,3 +90,4 @@ pythonenv*
 /clang/utils/analyzer/projects/*/RefScanBuildResults
 # automodapi puts generated documentation files here.
 /lldb/docs/python_api/
+mlir_venv/
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index f44693096b26b..602f9607ad618 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4171,7 +4171,7 @@ DiagnosedSilenceableFailure transform::VectorizeOp::apply(
   // TODO: Check that the correct number of vectorSizes was provided.
   for (Operation *target : targets) {
     if (!linalg::hasVectorizationImpl(target)) {
-      return mlir::emitSilenceableFailure(target->getLoc())
+      return mlir::emitSilenceableFailure(getLoc())
              << "Unsupported Op, cannot vectorize";
     }
     FailureOr<VectorizationResult> vectorResults =
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index b57e66a1c3580..07c5bef17c851 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2381,6 +2381,18 @@ vectorizeScalableVectorPrecondition(Operation *op,
     return success(isa<linalg::UnPackOp>(op));
   }
 
+  // Ensure that the number of vector sizes and scalable flags provided by the
+  // user does not exceed the number of loops in the target Linalg op.
+  // Accessing iterator types with an out-of-bounds index would cause an
+  // assertion failure (SmallVector::operator[]). This check converts such a
+  // crash into a clean failure, allowing the transform interpreter to report
+  // an error gracefully.
+  // Regression test:
+  // mlir/test/Dialect/Linalg/transform/vectorize-dynamic-mixed-sizes.mlir
+  if (inputScalableVecDims.size() > linalgOp.getNumLoops() ||
+      inputVectorSizes.size() > linalgOp.getNumLoops())
+    return failure();
+
   // Cond 2: There's been no need for more than 2 scalable dims so far
   if (numOfScalableDims > 2)
     return failure();
diff --git a/mlir/test/Dialect/Linalg/transform/vectorize-dynamic-mixed-sizes.mlir b/mlir/test/Dialect/Linalg/transform/vectorize-dynamic-mixed-sizes.mlir
new file mode 100644
index 0000000000000..2af5e26898fb9
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/transform/vectorize-dynamic-mixed-sizes.mlir
@@ -0,0 +1,22 @@
+// RUN: mlir-opt %s -transform-interpreter -verify-diagnostics
+
+// Regression test for bug #204100.
+// Assertion idx < size() in SmallVector.h used to happen here.
+
+// CHECK-LABEL: func @add_dynamic
+module {
+  func.func @add_dynamic(%arg0: memref<?x?xbf16>, %arg1: memref<?x?xbf16>, %arg2: memref<?x?xbf16>) {
+    linalg.add ins(%arg0, %arg1 : memref<?x?xbf16>, memref<?x?xbf16>) outs(%arg2 : memref<?x?xbf16>)
+    return
+  }
+
+  module attributes {transform.with_named_sequence} {
+    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+      %0 = transform.structured.match ops{["linalg.add"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+      // This combination was crashing (static/dynamic mismatch)
+      transform.structured.vectorize %0 vector_sizes [8, [16], 4] : !transform.any_op
+      // expected-error @above {{Attempted to vectorize, but failed}}
+      transform.yield
+    }
+  }
+}

``````````

</details>


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