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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Alessandro Potenza (alepot55)

<details>
<summary>Changes</summary>

`convert-vector-to-gpu` aborts on a `vector.contract` inside an `scf.if`:

```
Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.
  [To = mlir::scf::ForOp, From = mlir::Operation]
```

`supportsMMaMatrixType` accepts every `scf::YieldOp` with no check on its parent, so a yield terminating an `scf.if` enters the conversion set. `convertYieldOp` then opens with `cast<scf::ForOp>(op->getParentOp())` and rewrites the yielded values against the loop's init args. Without assertions the cast is unchecked.

The fix guards the yield arm on the parent being an `scf.for`, keeping a non-convertible yield out of the set rather than bailing out further down.

Verified on two builds of `294c9c8b3`, md5 `dde121e8e553` and `89f261ed8654`: the reduced case aborts on the first and lowers cleanly on the second, leaving the `scf.if` unconverted. All three files under `mlir/test/Conversion/VectorToGPU/` pass with the patch, each run through its own RUN line.

Fixes #<!-- -->205068.

Disclosure per the AI tool policy: AI assistance was used to write this patch. I understand the change and can answer questions about it in review.

---
Full diff: https://github.com/llvm/llvm-project/pull/218226.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp (+5-1) 
- (modified) mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir (+28) 


``````````diff
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.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
+}

``````````

</details>


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