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

[llvmlistbot at llvm.org]

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

`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

>From 41919e43354b33eb48c238bebba1bd492ca85a36 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Wed, 24 Jun 2026 15:09:36 +0800
Subject: [PATCH] Fix crash in vectorizeOpPrecondition when vector sizes array
 is too short

---
 .gitignore                                    |  1 +
 .../TransformOps/LinalgTransformOps.cpp       |  2 +-
 .../Linalg/Transforms/Vectorization.cpp       | 12 ++++++++++
 .../vectorize-dynamic-mixed-sizes.mlir        | 22 +++++++++++++++++++
 4 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Linalg/transform/vectorize-dynamic-mixed-sizes.mlir

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
+    }
+  }
+}