[Mlir-commits] [mlir] [mlir][affine] Fix crash in vectorizeAffineLoopNest test utility for reduction loops (PR #184617)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/184617

The test utility function `testVecAffineLoopNest` called `isLoopParallel` with a `reductions` output parameter, which populates reduction descriptors when the loop performs a reduction. However, these descriptors were never added to `strategy.reductionLoops` before calling `vectorizeAffineLoopNest`. When the vectorizer then processed a loop with `iter_args`, it found no reduction descriptors in the strategy and hit an assertion failure.

Fix by registering the reduction loop descriptors in the strategy before vectorization, matching what the production vectorizer code already does correctly.

Fixes #128334

>From fe74ee3d780ad53531f186a19810a60bb7b14e56 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 05:44:14 -0800
Subject: [PATCH] [mlir][affine] Fix crash in vectorizeAffineLoopNest test
 utility for reduction loops

The test utility function `testVecAffineLoopNest` called `isLoopParallel` with
a `reductions` output parameter, which populates reduction descriptors when the
loop performs a reduction. However, these descriptors were never added to
`strategy.reductionLoops` before calling `vectorizeAffineLoopNest`. When the
vectorizer then processed a loop with `iter_args`, it found no reduction
descriptors in the strategy and hit an assertion failure.

Fix by registering the reduction loop descriptors in the strategy before
vectorization, matching what the production vectorizer code already does
correctly.

Fixes #128334
---
 .../SuperVectorize/vectorize_unsupported.mlir | 25 +++++++++++++++++++
 .../Dialect/Affine/TestVectorizationUtils.cpp |  2 ++
 2 files changed, 27 insertions(+)

diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
index 4f59b7812a668..8ed152a5b0e20 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
@@ -40,3 +40,28 @@ func.func @iv_mapped_to_multiple_indices_unsupported(%arg0: index) -> memref<2x2
 // CHECK:             %[[VAL_6:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_4]]){{\[}}%[[VAL_1]]]
 // CHECK:           }
 // CHECK:         }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/128334
+// The vectorizer test utility used to crash when a reduction loop with a
+// dynamic upper bound was vectorized via 'vectorizeAffineLoopNest', because
+// the reduction descriptors were not added to the vectorization strategy.
+
+// CHECK-LABEL: func.func @reduction_loop_dynamic_bound_vectorized
+// CHECK:         affine.for %{{.*}} iter_args(%{{.*}} = {{.*}}) -> (vector<4xf32>) {
+// CHECK:           vector.transfer_read
+// CHECK:           arith.addf
+// CHECK:           affine.yield
+// CHECK:         }
+// CHECK:         vector.reduction <add>
+func.func @reduction_loop_dynamic_bound_vectorized(%buffer: memref<1024xf32>) {
+  %c10 = arith.constant 10 : index
+  %sum_0 = arith.constant 0.0 : f32
+  affine.for %i = 0 to %c10 iter_args(%sum_iter = %sum_0) -> (f32) {
+    %t = affine.load %buffer[%i] : memref<1024xf32>
+    %sum_next = arith.addf %sum_iter, %t : f32
+    affine.yield %sum_next : f32
+  }
+  return
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp b/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp
index 2c93991386675..88d26d233c5cd 100644
--- a/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp
@@ -247,6 +247,8 @@ void VectorizerTestPass::testVecAffineLoopNest(llvm::raw_ostream &outs) {
     outs << "Outermost loop cannot be parallel\n";
     return;
   }
+  if (!reductions.empty())
+    strategy.reductionLoops[outermostLoop] = reductions;
   std::vector<SmallVector<AffineForOp, 2>> loopsToVectorize;
   loopsToVectorize.push_back({outermostLoop});
   (void)vectorizeAffineLoopNest(loopsToVectorize, strategy);