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

[llvmlistbot at llvm.org]

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

>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 1/2] 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
+    }
+  }
+}

>From f97bab1b088d8d2548b326f3e2540bc0c359cbd3 Mon Sep 17 00:00:00 2001
From: LouisLu0602 <yaolu0602 at gmail.com>
Date: Wed, 24 Jun 2026 16:08:30 +0800
Subject: [PATCH 2/2] [MLIR][Vector] Add regression test for bug #204100 (mixed
 static/dynamic vector sizes)

---
 .../Vector/transform-op-vector-to-llvm.mlir   | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/mlir/test/Dialect/Vector/transform-op-vector-to-llvm.mlir b/mlir/test/Dialect/Vector/transform-op-vector-to-llvm.mlir
index 271cdf0e059f4..a74e52dfd4bdb 100644
--- a/mlir/test/Dialect/Vector/transform-op-vector-to-llvm.mlir
+++ b/mlir/test/Dialect/Vector/transform-op-vector-to-llvm.mlir
@@ -19,3 +19,29 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// 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
+    }
+  }
+}
+
+
+
+