[Mlir-commits] [mlir] [mlir][affine] Fix SuperVectorizer crash on uniforms from multi-result ops (PR #234531)

[Mattéo Rizza Murgier]

https://github.com/matteo-rm updated https://github.com/llvm/llvm-project/pull/234531

>From 7e809b06e9afd1f12bfff345db2952a3173c8f1c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matt=C3=A9o=20Rizza=20Murgier?=
 <matteo.rizza-murgier at sipearl.com>
Date: Wed, 21 Oct 2026 20:56:10 +0200
Subject: [PATCH 1/2] [mlir][affine] Fix SuperVectorizer crash on uniforms from
 multi-result ops

---
 .../Affine/Transforms/SuperVectorize.cpp        |  8 +++-----
 .../SuperVectorize/uniform_divergent.mlir       | 17 +++++++++++++++++
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 0bb344b0484f06..8fd78fe711fc2d 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -808,16 +808,14 @@ void VectorizationState::registerOpVectorReplacement(Operation *replaced,
 /// scalar version still exists after vectorization), like uniforms.
 ///
 /// Example:
-///   * 'replaced': block argument or operation outside of the vectorized loop.
+///   * 'replaced': block argument or operation result defined outside of the
+//       vectorized loop.
 ///   * 'replacement': %0 = vector.broadcast %1 : f32 to vector<128xf32>
 void VectorizationState::registerValueVectorReplacement(
     Value replaced, Operation *replacement) {
   assert(replacement->getNumResults() == 1 &&
          "Expected single-result replacement");
-  if (Operation *defOp = replaced.getDefiningOp())
-    registerOpVectorReplacement(defOp, replacement);
-  else
-    registerValueVectorReplacementImpl(replaced, replacement->getResult(0));
+  registerValueVectorReplacementImpl(replaced, replacement->getResult(0));
 }
 
 /// Registers the vector replacement of a block argument (e.g., iter_args).
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir b/mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir
index 545f2d08134846..d9f7d20372a116 100644
--- a/mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir
+++ b/mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir
@@ -58,3 +58,20 @@ func.func @uniform_load(%A : memref<?x?xf32>, %C : memref<?x?xf32>) {
 // CHECK-NEXT:   %[[bcast:.*]] = vector.broadcast %[[uniform_ld]] : f32 to vector<128xf32>
 // CHECK-NEXT:   affine.for
 // CHECK:          arith.addf %[[bcast]], %{{.*}} : vector<128xf32>
+
+// -----
+
+// CHECK-LABEL: @uniform_multi_result
+func.func @uniform_multi_result(%in : memref<512xi32>, %a : i32, %b : i32) {
+  %sum, %overflow = arith.addui_extended %a, %b : i32, i1
+  affine.for %i = 0 to 512 {
+    %ld = affine.load %in[%i] : memref<512xi32>
+    %add = arith.addi %ld, %sum : i32
+  }
+  return
+}
+
+// CHECK:      %[[sum:.*]], %{{.*}} = arith.addui_extended
+// CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[sum]] : i32 to vector<128xi32>
+// CHECK-NEXT: affine.for
+// CHECK:        arith.addi %{{.*}}, %[[bcast]] : vector<128xi32>

>From 48bd1cced85267ec6caecafe258a51ce05807ddc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matt=C3=A9o=20Rizza=20Murgier?=
 <matteo.rizza-murgier at sipearl.com>
Date: Thu, 22 Oct 2026 11:44:27 +0200
Subject: [PATCH 2/2] Fix doc comment

---
 mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 8fd78fe711fc2d..e2c0a3c6ab4fa9 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -809,7 +809,7 @@ void VectorizationState::registerOpVectorReplacement(Operation *replaced,
 ///
 /// Example:
 ///   * 'replaced': block argument or operation result defined outside of the
-//       vectorized loop.
+///     vectorized loop.
 ///   * 'replacement': %0 = vector.broadcast %1 : f32 to vector<128xf32>
 void VectorizationState::registerValueVectorReplacement(
     Value replaced, Operation *replacement) {