[Mlir-commits] [mlir] [mlir][VectorToGPU] Only convert an scf.yield that terminates an scf.for (PR #218226)

[Alessandro Potenza]

https://github.com/alepot55 updated https://github.com/llvm/llvm-project/pull/218226

>From 6b7b52cae2ce872fe016dfd9baf70aa2263a5f70 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Sun, 23 Aug 2026 12:29:03 +0200
Subject: [PATCH] [mlir][VectorToGPU] Only convert an scf.yield that terminates
 an scf.for

supportsMMaMatrixType() accepts every scf::YieldOp without looking at its
parent, so a yield terminating an scf.if enters the conversion set and reaches
convertYieldOp(), which opens with cast<scf::ForOp>(op->getParentOp()) and
rewrites the yielded values against the loop's init args. With assertions that
aborts; without them the cast is unchecked.

Guard the yield arm on the parent being an scf.for, which keeps a
non-convertible yield out of the set rather than bailing out later.

Fixes #205068.

Assisted-by: Claude (Anthropic)
---
 .../Conversion/VectorToGPU/VectorToGPU.cpp    |  6 +++-
 .../vector-to-mma-ops-mma-sync.mlir           | 28 +++++++++++++++++++
 .../VectorToGPU/vector-to-mma-ops.mlir        | 28 +++++++++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index 975fe28399609..93c6c024b9681 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -297,8 +297,12 @@ extractStridedSliceSupportsMMAMatrixType(vector::ExtractStridedSliceOp op) {
 }
 
 static bool supportsMMaMatrixType(Operation *op, bool useNvGpu) {
-  if (isa<scf::ForOp, scf::YieldOp>(op))
+  if (isa<scf::ForOp>(op))
     return true;
+  // A yield is only convertible as the terminator of an scf.for body:
+  // convertYieldOp rewrites the yielded values against the loop's init args.
+  if (isa<scf::YieldOp>(op))
+    return isa<scf::ForOp>(op->getParentOp());
   if (auto transferRead = dyn_cast<vector::TransferReadOp>(op))
     return useNvGpu ? nvgpu::canLowerToWarpMatrixOperation(transferRead)
                     : transferReadSupportsMMAMatrixType(transferRead);
diff --git a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir
index 988122ae0fd45..b6dfcb5a6780d 100644
--- a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir
+++ b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir
@@ -835,3 +835,31 @@ func.func @unsupported_transposed_store(%arg0: !smem_type,
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1)->(d1, d0)>} : vector<16x8xf16>, !smem_type
   return
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// Same guard on the nvgpu path: getOpToConvert feeds both entry points, and a
+// yield that does not terminate an scf.for body is not convertible in either.
+// CHECK-LABEL: func @no_convert_yield_of_scf_if
+// CHECK: scf.if
+// CHECK-NOT: nvgpu.
+func.func @no_convert_yield_of_scf_if(%arg0: memref<128x128xf16>, %arg1: memref<128x128xf16>, %arg2: memref<128x128xf16>, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f16
+  %C = vector.transfer_read %arg2[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
+  %r = scf.if %cond -> (vector<16x16xf16>) {
+    %a = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
+    %b = vector.transfer_read %arg1[%c0, %c0], %cst {permutation_map = #map0, in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
+    %d = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %a, %b, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
+    scf.yield %d : vector<16x16xf16>
+  } else {
+    scf.yield %C : vector<16x16xf16>
+  }
+  vector.transfer_write %r, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<128x128xf16>
+  return
+}
diff --git a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
index 00e5966eb2a34..fa2359d8a696e 100644
--- a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
+++ b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
@@ -684,3 +684,31 @@ func.func @read_transpose_with_broadcast_3d(%arg0: memref<2x2x2xf16>, %arg1: mem
   vector.transfer_write %B, %arg1[%c0, %c0] {in_bounds = [true, true]} : vector<2x2xf16>, memref<2x2xf16>
   return
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// A yield that terminates something other than an scf.for body is not
+// convertible: convertYieldOp rewrites against the loop's init args.
+// CHECK-LABEL: func @no_convert_yield_of_scf_if
+// CHECK: scf.if
+// CHECK-NOT: gpu.subgroup_mma
+func.func @no_convert_yield_of_scf_if(%arg0: memref<128x128xf16>, %arg1: memref<128x128xf16>, %arg2: memref<128x128xf16>, %cond: i1) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f16
+  %C = vector.transfer_read %arg2[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
+  %r = scf.if %cond -> (vector<16x16xf16>) {
+    %a = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
+    %b = vector.transfer_read %arg1[%c0, %c0], %cst {permutation_map = #map0, in_bounds = [true, true]} : memref<128x128xf16>, vector<16x16xf16>
+    %d = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %a, %b, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
+    scf.yield %d : vector<16x16xf16>
+  } else {
+    scf.yield %C : vector<16x16xf16>
+  }
+  vector.transfer_write %r, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<128x128xf16>
+  return
+}