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

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-04T14:36:27Z
New Revision: 56e0b6af1d3948fa2ed56149a02707adaa7a4074

URL: https://github.com/llvm/llvm-project/commit/56e0b6af1d3948fa2ed56149a02707adaa7a4074
DIFF: https://github.com/llvm/llvm-project/commit/56e0b6af1d3948fa2ed56149a02707adaa7a4074.diff

LOG: [mlir][affine] Fix crash in vectorizeAffineLoopNest test utility for reduction loops (#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

Added: 
    

Modified: 
    mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir
    mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp

Removed: 
    


################################################################################
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);