[Mlir-commits] [mlir] [mlir][vector-to-gpu] Lower transposed transfer_write to a transposed MMA store (PR #221248)
Hung-Kuan Tseng
llvmlistbot at llvm.org
Fri Sep 4 08:11:12 PDT 2026
https://github.com/Tim096 created https://github.com/llvm/llvm-project/pull/221248
`convert-vector-to-gpu` lowers a `vector.transfer_read` whose permutation
map puts the most minor dimension first to a `transpose` MMA load since
18ecdbfe6c74, but still rejects the same shape of `vector.transfer_write`:
```mlir
// Lowered to gpu.subgroup_mma_load_matrix ... leadDimension 3 transpose
%0 = vector.transfer_read %a[%c0, %c0], %f {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : memref<5x3xf16>, vector<3x5xf16>
// Left alone
vector.transfer_write %1, %b[%c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : vector<3x5xf16>, memref<5x3xf16>
```
The TODO in `transferWriteSupportsMMAMatrixType` waits for a transpose
attribute on the GPU dialect op. `gpu.subgroup_mma_store_matrix` has had
one since 3d35546cd168, and both the NVVM and the SPIR-V lowerings honour
it, so this classifies the write the same way as the read and sets the
attribute. `leadDimension` already comes from the position of the map
dimensions, so the stride helper is unchanged.
Tests: the existing `no_convert_write_transpose` becomes a positive test,
plus the 3-D and 4-D strided cases and a not-last-dim negative case,
mirroring the read side.
>From defdf4e2b54b44b9a09920b36b52db00fd4cd031 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Fri, 4 Sep 2026 23:00:11 +0800
Subject: [PATCH] [mlir][vector-to-gpu] Lower transposed transfer_write to a
transposed MMA store
convert-vector-to-gpu lowers a transfer_read whose permutation map puts
the most minor dimension first to a transpose MMA load since
18ecdbfe6c74, but still rejects the same shape of transfer_write. The
TODO waits for a transpose attribute on the GPU dialect op, which
gpu.subgroup_mma_store_matrix has had since 3d35546cd168 and which both
the NVVM and the SPIR-V lowerings honour. Classify the write the same
way as the read and set the attribute.
---
.../Conversion/VectorToGPU/VectorToGPU.cpp | 14 +++--
.../VectorToGPU/vector-to-mma-ops.mlir | 51 +++++++++++++++++--
2 files changed, 57 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index 975fe28399609..29c9fb880405a 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -210,8 +210,7 @@ transferWriteSupportsMMAMatrixType(vector::TransferWriteOp writeOp) {
MLIRContext *ctx = writeOp.getContext();
AffineExpr innerDim = getAffineDimExpr(permutationMap.getNumDims() - 1, ctx);
- // TODO: Support transpose once it is added to GPU dialect ops.
- return permutationMap.getResult(1) == innerDim;
+ return llvm::is_contained(permutationMap.getResults(), innerDim);
}
/// Return true if the constant is a splat to a 2D vector so that it can be
@@ -623,13 +622,19 @@ convertTransferWriteOp(RewriterBase &rewriter, vector::TransferWriteOp op,
rewriter.setInsertionPoint(op);
assert(transferWriteSupportsMMAMatrixType(op));
+ AffineMap permutationMap = op.getPermutationMap();
std::optional<int64_t> stride =
- getStaticallyKnownRowStride(op.getShapedType(), op.getPermutationMap());
+ getStaticallyKnownRowStride(op.getShapedType(), permutationMap);
if (!stride.has_value()) {
LDBG() << "no stride";
return rewriter.notifyMatchFailure(op, "no stride");
}
+ // As for transfer_read, transferWriteSupportsMMAMatrixType ensures that
+ // either of the map results is the most minor dimension, so the first result
+ // being that dimension means a transposed store.
+ const bool isTranspose = isFirstResultLastMapDimension(permutationMap);
+
auto it = valueMapping.find(op.getVector());
if (it == valueMapping.end()) {
LDBG() << "no mapping";
@@ -639,7 +644,8 @@ convertTransferWriteOp(RewriterBase &rewriter, vector::TransferWriteOp op,
Value matrix = it->second;
auto store = gpu::SubgroupMmaStoreMatrixOp::create(
rewriter, op.getLoc(), matrix, op.getBase(), op.getIndices(),
- rewriter.getIndexAttr(*stride), /*transpose=*/UnitAttr());
+ rewriter.getIndexAttr(*stride),
+ isTranspose ? rewriter.getUnitAttr() : UnitAttr());
(void)store;
LDBG() << "transfer write to: " << store;
diff --git a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
index 00e5966eb2a34..3af86a6dedff7 100644
--- a/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
+++ b/mlir/test/Conversion/VectorToGPU/vector-to-mma-ops.mlir
@@ -660,15 +660,58 @@ func.func @no_convert_read_transpose_not_last_dim(%arg0: memref<2x2x2xf16>, %arg
// -----
-// Transpose write is not supported.
-// CHECK-LABEL: func @no_convert_write_transpose
+// CHECK-LABEL: func @write_transpose
+func.func @write_transpose(%arg0: memref<3x5xf16>, %arg1: memref<5x3xf16>) {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<3x5xf16>, vector<3x5xf16>
+ %B = arith.addf %A, %A : vector<3x5xf16>
+ // CHECK: gpu.subgroup_mma_store_matrix %{{.*}} leadDimension 3 transpose : !gpu.mma_matrix<3x5xf16, "COp">, memref<5x3xf16>
+ vector.transfer_write %B, %arg1[%c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : vector<3x5xf16>, memref<5x3xf16>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @write_transpose_with_strides_3d
+func.func @write_transpose_with_strides_3d(%arg0: memref<3x5xf16>, %arg1: memref<5x7x3xf16>, %arg2: memref<2x5x3xf16>) {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<3x5xf16>, vector<3x5xf16>
+ %B = arith.addf %A, %A : vector<3x5xf16>
+ // CHECK: gpu.subgroup_mma_store_matrix %{{.*}} leadDimension 21 transpose : !gpu.mma_matrix<3x5xf16, "COp">, memref<5x7x3xf16>
+ vector.transfer_write %B, %arg1[%c0, %c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2) -> (d2, d0)>} : vector<3x5xf16>, memref<5x7x3xf16>
+ // CHECK: gpu.subgroup_mma_store_matrix %{{.*}} leadDimension 3 transpose : !gpu.mma_matrix<3x5xf16, "COp">, memref<2x5x3xf16>
+ vector.transfer_write %B, %arg2[%c0, %c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2) -> (d2, d1)>} : vector<3x5xf16>, memref<2x5x3xf16>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @write_transpose_with_strides_4d
+func.func @write_transpose_with_strides_4d(%arg0: memref<3x5xf16>, %arg1: memref<5x7x11x3xf16>, %arg2: memref<2x5x11x3xf16>) {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<3x5xf16>, vector<3x5xf16>
+ %B = arith.addf %A, %A : vector<3x5xf16>
+ // CHECK: gpu.subgroup_mma_store_matrix %{{.*}} leadDimension 231 transpose : !gpu.mma_matrix<3x5xf16, "COp">, memref<5x7x11x3xf16>
+ vector.transfer_write %B, %arg1[%c0, %c0, %c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d0)>} : vector<3x5xf16>, memref<5x7x11x3xf16>
+ // CHECK: gpu.subgroup_mma_store_matrix %{{.*}} leadDimension 33 transpose : !gpu.mma_matrix<3x5xf16, "COp">, memref<2x5x11x3xf16>
+ vector.transfer_write %B, %arg2[%c0, %c0, %c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d1)>} : vector<3x5xf16>, memref<2x5x11x3xf16>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @no_convert_write_transpose_not_last_dim
// CHECK-NOT: gpu
-func.func @no_convert_write_transpose(%arg0: memref<2x2xf16>) {
+func.func @no_convert_write_transpose_not_last_dim(%arg0: memref<2x2xf16>, %arg1: memref<2x2x2xf16>) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.0 : f16
%A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<2x2xf16>, vector<2x2xf16>
%B = arith.addf %A, %A : vector<2x2xf16>
- vector.transfer_write %B, %arg0[%c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : vector<2x2xf16>, memref<2x2xf16>
+ // Legal map, but does not map the last memref dim so should not be lowered to an MMA store.
+ vector.transfer_write %B, %arg1[%c0, %c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2) -> (d1, d0)>} : vector<2x2xf16>, memref<2x2x2xf16>
return
}
More information about the Mlir-commits
mailing list