[Mlir-commits] [mlir] [mlir][gpu] Lower the remaining subgroup_mma_elementwise ops to NVVM (PR #221288)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 4 10:11:13 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hung-Kuan Tseng (Tim096)
<details>
<summary>Changes</summary>
`convert-gpu-to-nvvm` handles five of the fifteen
`gpu.subgroup_mma_elementwise` ops and hits `llvm_unreachable` on the
other ten, so a `convert-vector-to-gpu` kernel with a `subf`, `negf` or
`truncf` epilogue on the WMMA path aborts in `-gpu-lower-to-nvvm-pipeline`:
```
UNREACHABLE executed at mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp:365!
```
`addf` on an f32 multiplicand fragment builds an `llvm.fadd` on the `i32`
registers that hold the tf32 bit patterns and fails verification instead.
This lowers `subf`, `negatef`, `addi`, `subi`, `muli`, `divs`, `divu` and
`negates` register by register like the existing ops, which the
"Manipulating fragment contents" rules of the PTX ISA allow, and rejects
the rest with a match failure: `extf` and `truncf`, because the PTX ISA
leaves conversions between f16 and f32 accumulator fragments undefined,
and any op on a fragment whose registers pack several elements (s8, u8,
tf32). The SPIR-V lowering handles all of these ops on cooperative
matrices, so `convert-vector-to-gpu` keeps emitting them.
Tests: lit tests for the new lowerings and for the rejected cases, and a
tensor core integration test that runs a matmul with a `subf` + `negf`
epilogue through `convert-vector-to-gpu` and the NVVM pipeline. It aborts
before this change and passes on an RTX 3070 (sm_86).
---
Full diff: https://github.com/llvm/llvm-project/pull/221288.diff
4 Files Affected:
- (modified) mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp (+54-12)
- (added) mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-wmma-elementwise.mlir (+46)
- (modified) mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir (+49)
- (added) mlir/test/Integration/GPU/CUDA/TensorCore/wmma-matmul-f16-elementwise.mlir (+84)
``````````diff
diff --git a/mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp b/mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp
index 7c27238e57b1b..a7934a571d56f 100644
--- a/mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp
@@ -345,25 +345,56 @@ static Value createMinMaxF(OpBuilder &builder, Location loc, Value lhs,
return LLVM::SelectOp::create(builder, loc, isNan, nan, sel);
}
+/// Returns true if `op` can be applied register by register to a WMMA
+/// fragment. The PTX ISA leaves the result of converting between f16 and f32
+/// accumulator fragments undefined, so there is nothing to lower EXTF and
+/// TRUNCF to.
+static bool isSupportedElementwiseOp(gpu::MMAElementwiseOp op) {
+ return op != gpu::MMAElementwiseOp::EXTF &&
+ op != gpu::MMAElementwiseOp::TRUNCF;
+}
+
static Value createScalarOp(OpBuilder &builder, Location loc,
gpu::MMAElementwiseOp op,
ArrayRef<Value> operands) {
+ Type type = operands[0].getType();
switch (op) {
case gpu::MMAElementwiseOp::ADDF:
- return LLVM::FAddOp::create(builder, loc, operands[0].getType(), operands);
+ return LLVM::FAddOp::create(builder, loc, type, operands);
+ case gpu::MMAElementwiseOp::SUBF:
+ return LLVM::FSubOp::create(builder, loc, type, operands);
case gpu::MMAElementwiseOp::MULF:
- return LLVM::FMulOp::create(builder, loc, operands[0].getType(), operands);
+ return LLVM::FMulOp::create(builder, loc, type, operands);
case gpu::MMAElementwiseOp::DIVF:
- return LLVM::FDivOp::create(builder, loc, operands[0].getType(), operands);
+ return LLVM::FDivOp::create(builder, loc, type, operands);
+ case gpu::MMAElementwiseOp::NEGATEF:
+ return LLVM::FNegOp::create(builder, loc, type, operands);
case gpu::MMAElementwiseOp::MAXF:
return createMinMaxF(builder, loc, operands[0], operands[1],
/*isMin=*/false);
case gpu::MMAElementwiseOp::MINF:
return createMinMaxF(builder, loc, operands[0], operands[1],
/*isMin=*/true);
- default:
- llvm_unreachable("unknown op");
+ case gpu::MMAElementwiseOp::ADDI:
+ return LLVM::AddOp::create(builder, loc, type, operands);
+ case gpu::MMAElementwiseOp::SUBI:
+ return LLVM::SubOp::create(builder, loc, type, operands);
+ case gpu::MMAElementwiseOp::MULI:
+ return LLVM::MulOp::create(builder, loc, type, operands);
+ case gpu::MMAElementwiseOp::DIVS:
+ return LLVM::SDivOp::create(builder, loc, type, operands);
+ case gpu::MMAElementwiseOp::DIVU:
+ return LLVM::UDivOp::create(builder, loc, type, operands);
+ case gpu::MMAElementwiseOp::NEGATES: {
+ Value zero =
+ LLVM::ConstantOp::create(builder, loc, type, builder.getZeroAttr(type));
+ return LLVM::SubOp::create(builder, loc, zero, operands[0]);
+ }
+ case gpu::MMAElementwiseOp::EXTF:
+ case gpu::MMAElementwiseOp::TRUNCF:
+ break;
}
+ llvm_unreachable("rejected by isSupportedElementwiseOp");
}
/// Convert GPU MMA elementwise ops to extract + op + insert.
@@ -381,19 +412,32 @@ struct WmmaElementwiseOpToNVVMLowering
return failure();
Location loc = subgroupMmaElementwiseOp.getLoc();
size_t numOperands = adaptor.getOperands().size();
- Type destType = convertMMAToLLVMType(
- cast<gpu::MMAMatrixType>(subgroupMmaElementwiseOp.getType()));
+ gpu::MMAElementwiseOp opType = subgroupMmaElementwiseOp.getOpType();
+ if (!isSupportedElementwiseOp(opType))
+ return rewriter.notifyMatchFailure(
+ subgroupMmaElementwiseOp,
+ "no NVVM lowering for this elementwise op on WMMA fragments");
+ gpu::MMAMatrixType matrixType = subgroupMmaElementwiseOp.getType();
+ Type destType = convertMMAToLLVMType(matrixType);
// If the element is not a struct, it means it's a scalar f64.
LLVM::LLVMStructType structDestTy =
dyn_cast<LLVM::LLVMStructType>(destType);
+
+ // The op is applied to each fragment register, which only computes the
+ // right thing when a register holds elements of the matrix element type.
+ // s8, u8 and tf32 fragments pack their elements into i32 registers.
+ Type registerType = structDestTy ? structDestTy.getBody()[0] : destType;
+ if (getElementTypeOrSelf(registerType) != matrixType.getElementType())
+ return rewriter.notifyMatchFailure(
+ subgroupMmaElementwiseOp,
+ "fragment registers do not hold the matrix element type");
if (!structDestTy) {
SmallVector<Value> operands;
for (auto operand : adaptor.getOperands()) {
operands.push_back(operand);
}
- Value element = createScalarOp(
- rewriter, loc, subgroupMmaElementwiseOp.getOpType(), operands);
+ Value element = createScalarOp(rewriter, loc, opType, operands);
rewriter.replaceOp(subgroupMmaElementwiseOp, element);
return success();
}
@@ -404,9 +448,7 @@ struct WmmaElementwiseOpToNVVMLowering
extractedOperands.push_back(LLVM::ExtractValueOp::create(
rewriter, loc, adaptor.getOperands()[opIdx], i));
}
- Value element =
- createScalarOp(rewriter, loc, subgroupMmaElementwiseOp.getOpType(),
- extractedOperands);
+ Value element = createScalarOp(rewriter, loc, opType, extractedOperands);
matrixStruct =
LLVM::InsertValueOp::create(rewriter, loc, matrixStruct, element, i);
}
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-wmma-elementwise.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-wmma-elementwise.mlir
new file mode 100644
index 0000000000000..e6e3f3206bdab
--- /dev/null
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-wmma-elementwise.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-opt %s -convert-gpu-to-nvvm -split-input-file -verify-diagnostics
+
+// The PTX ISA does not define a conversion between f16 and f32 accumulator
+// fragments.
+
+gpu.module @test_module {
+ func.func @wmma_elementwise_extf(%A : !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp"> {
+ // expected-error @+1 {{failed to legalize operation 'gpu.subgroup_mma_elementwise' that was explicitly marked illegal}}
+ %0 = gpu.subgroup_mma_elementwise extf %A : (!gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">
+ return %0 : !gpu.mma_matrix<16x16xf32, "COp">
+ }
+}
+
+// -----
+
+gpu.module @test_module {
+ func.func @wmma_elementwise_truncf(%A : !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp"> {
+ // expected-error @+1 {{failed to legalize operation 'gpu.subgroup_mma_elementwise' that was explicitly marked illegal}}
+ %0 = gpu.subgroup_mma_elementwise truncf %A : (!gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
+ return %0 : !gpu.mma_matrix<16x16xf16, "COp">
+ }
+}
+
+// -----
+
+// s8 multiplicand fragments pack four elements into each i32 register.
+
+gpu.module @test_module {
+ func.func @wmma_elementwise_packed_s8(%A : !gpu.mma_matrix<16x16xsi8, "AOp">) -> !gpu.mma_matrix<16x16xsi8, "AOp"> {
+ // expected-error @+1 {{failed to legalize operation 'gpu.subgroup_mma_elementwise' that was explicitly marked illegal}}
+ %0 = gpu.subgroup_mma_elementwise addi %A, %A : (!gpu.mma_matrix<16x16xsi8, "AOp">, !gpu.mma_matrix<16x16xsi8, "AOp">) -> !gpu.mma_matrix<16x16xsi8, "AOp">
+ return %0 : !gpu.mma_matrix<16x16xsi8, "AOp">
+ }
+}
+
+// -----
+
+// f32 multiplicand fragments hold tf32 bit patterns in i32 registers.
+
+gpu.module @test_module {
+ func.func @wmma_elementwise_packed_tf32(%A : !gpu.mma_matrix<16x8xf32, "AOp">) -> !gpu.mma_matrix<16x8xf32, "AOp"> {
+ // expected-error @+1 {{failed to legalize operation 'gpu.subgroup_mma_elementwise' that was explicitly marked illegal}}
+ %0 = gpu.subgroup_mma_elementwise addf %A, %A : (!gpu.mma_matrix<16x8xf32, "AOp">, !gpu.mma_matrix<16x8xf32, "AOp">) -> !gpu.mma_matrix<16x8xf32, "AOp">
+ return %0 : !gpu.mma_matrix<16x8xf32, "AOp">
+ }
+}
diff --git a/mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir
index 31bdca01d7ae7..f650f6b4b399b 100644
--- a/mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir
@@ -370,3 +370,52 @@ gpu.module @test_module {
return %D : !gpu.mma_matrix<16x16xf16, "COp">
}
}
+
+// -----
+
+gpu.module @test_module {
+
+// CHECK-LABEL: func @gpu_wmma_elementwise_subf_negatef
+// CHECK: %[[A0:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
+// CHECK: %[[B0:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
+// CHECK: %[[C0:.*]] = llvm.fsub %[[A0]], %[[B0]] : vector<2xf16>
+// CHECK: llvm.insertvalue %[[C0]], %{{.*}}[0] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
+// CHECK-COUNT-3: llvm.fsub %{{.*}}, %{{.*}} : vector<2xf16>
+// CHECK: %[[S0:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
+// CHECK: %[[N0:.*]] = llvm.fneg %[[S0]] : vector<2xf16>
+// CHECK: llvm.insertvalue %[[N0]], %{{.*}}[0] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
+// CHECK-COUNT-3: llvm.fneg %{{.*}} : vector<2xf16>
+ func.func @gpu_wmma_elementwise_subf_negatef(%A : !gpu.mma_matrix<16x16xf16, "COp">, %B : !gpu.mma_matrix<16x16xf16, "COp">) ->(!gpu.mma_matrix<16x16xf16, "COp">) {
+ %C = gpu.subgroup_mma_elementwise subf %A, %B : (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
+ %D = gpu.subgroup_mma_elementwise negatef %C : (!gpu.mma_matrix<16x16xf16, "COp">) -> !gpu.mma_matrix<16x16xf16, "COp">
+ return %D : !gpu.mma_matrix<16x16xf16, "COp">
+ }
+}
+
+// -----
+
+gpu.module @test_module {
+
+// CHECK-LABEL: func @gpu_wmma_elementwise_int
+// CHECK: %[[A0:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(i32, i32, i32, i32, i32, i32, i32, i32)>
+// CHECK: %[[B0:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(i32, i32, i32, i32, i32, i32, i32, i32)>
+// CHECK: %[[C0:.*]] = llvm.add %[[A0]], %[[B0]] : i32
+// CHECK: llvm.insertvalue %[[C0]], %{{.*}}[0] : !llvm.struct<(i32, i32, i32, i32, i32, i32, i32, i32)>
+// CHECK-COUNT-7: llvm.add %{{.*}}, %{{.*}} : i32
+// CHECK-COUNT-8: llvm.sub %{{.*}}, %{{.*}} : i32
+// CHECK-COUNT-8: llvm.mul %{{.*}}, %{{.*}} : i32
+// CHECK-COUNT-8: llvm.sdiv %{{.*}}, %{{.*}} : i32
+// CHECK-COUNT-8: llvm.udiv %{{.*}}, %{{.*}} : i32
+// CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK: llvm.sub %[[ZERO]], %{{.*}} : i32
+// CHECK-COUNT-7: llvm.sub %{{.*}}, %{{.*}} : i32
+ func.func @gpu_wmma_elementwise_int(%A : !gpu.mma_matrix<16x16xi32, "COp">, %B : !gpu.mma_matrix<16x16xi32, "COp">) ->(!gpu.mma_matrix<16x16xi32, "COp">) {
+ %C = gpu.subgroup_mma_elementwise addi %A, %B : (!gpu.mma_matrix<16x16xi32, "COp">, !gpu.mma_matrix<16x16xi32, "COp">) -> !gpu.mma_matrix<16x16xi32, "COp">
+ %D = gpu.subgroup_mma_elementwise subi %C, %B : (!gpu.mma_matrix<16x16xi32, "COp">, !gpu.mma_matrix<16x16xi32, "COp">) -> !gpu.mma_matrix<16x16xi32, "COp">
+ %E = gpu.subgroup_mma_elementwise muli %D, %B : (!gpu.mma_matrix<16x16xi32, "COp">, !gpu.mma_matrix<16x16xi32, "COp">) -> !gpu.mma_matrix<16x16xi32, "COp">
+ %F = gpu.subgroup_mma_elementwise divs %E, %B : (!gpu.mma_matrix<16x16xi32, "COp">, !gpu.mma_matrix<16x16xi32, "COp">) -> !gpu.mma_matrix<16x16xi32, "COp">
+ %G = gpu.subgroup_mma_elementwise divu %F, %B : (!gpu.mma_matrix<16x16xi32, "COp">, !gpu.mma_matrix<16x16xi32, "COp">) -> !gpu.mma_matrix<16x16xi32, "COp">
+ %H = gpu.subgroup_mma_elementwise negates %G : (!gpu.mma_matrix<16x16xi32, "COp">) -> !gpu.mma_matrix<16x16xi32, "COp">
+ return %H : !gpu.mma_matrix<16x16xi32, "COp">
+ }
+}
diff --git a/mlir/test/Integration/GPU/CUDA/TensorCore/wmma-matmul-f16-elementwise.mlir b/mlir/test/Integration/GPU/CUDA/TensorCore/wmma-matmul-f16-elementwise.mlir
new file mode 100644
index 0000000000000..bbd5adc48c986
--- /dev/null
+++ b/mlir/test/Integration/GPU/CUDA/TensorCore/wmma-matmul-f16-elementwise.mlir
@@ -0,0 +1,84 @@
+// RUN: mlir-opt %s -convert-vector-to-gpu \
+// RUN: | mlir-opt -gpu-lower-to-nvvm-pipeline="cubin-format=%gpu_compilation_format" \
+// RUN: | mlir-runner \
+// RUN: --shared-libs=%mlir_cuda_runtime \
+// RUN: --shared-libs=%mlir_runner_utils \
+// RUN: --entry-point-result=void \
+// RUN: | FileCheck %s
+// Check that arith.subf and arith.negf applied to the result of a matmul are
+// lowered through gpu.subgroup_mma_elementwise to NVVM. With A[i][j] = j and
+// C[i][j] = i, the kernel computes -(C - (A * A + C)) = A * A, whose rows are
+// all [-0, 120, 240, ...] (the first column is the negation of 0).
+#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @main() {
+ %0 = memref.alloc() : memref<16x16xf16>
+ %22 = memref.alloc() : memref<16x16xf16>
+ %1 = memref.alloc() : memref<16x16xf32>
+
+ %f0 = arith.constant 0.0e+00 : f16
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+
+ // Initialize the input matrix with the column index and the accumulator
+ // with the row index.
+ scf.for %arg0 = %c0 to %c16 step %c1 {
+ scf.for %arg1 = %c0 to %c16 step %c1 {
+ %2 = arith.index_cast %arg1 : index to i16
+ %3 = arith.sitofp %2 : i16 to f16
+ memref.store %3, %0[%arg0, %arg1] : memref<16x16xf16>
+ %4 = arith.index_cast %arg0 : index to i16
+ %5 = arith.sitofp %4 : i16 to f16
+ memref.store %5, %22[%arg0, %arg1] : memref<16x16xf16>
+ }
+ }
+
+ %3 = memref.cast %1 : memref<16x16xf32> to memref<*xf32>
+
+ // Copy the input and the accumulator to the device.
+ %token = gpu.wait async
+ %d0, %t0 = gpu.alloc async [%token] () : memref<16x16xf16>
+ %d22, %t1 = gpu.alloc async [%token] () : memref<16x16xf16>
+ %x = gpu.memcpy async [%token] %d0, %0 : memref<16x16xf16>, memref<16x16xf16>
+ %y = gpu.memcpy async [%token] %d22, %22 : memref<16x16xf16>, memref<16x16xf16>
+ gpu.wait [%token]
+
+ gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
+ threads(%tx, %ty, %tz) in (%block_x = %c32, %block_y = %c1, %block_z = %c1) {
+ %A = vector.transfer_read %d0[%c0, %c0], %f0 {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
+ %B = vector.transfer_read %d0[%c0, %c0], %f0 {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
+ %C = vector.transfer_read %d22[%c0, %c0], %f0 {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
+ %D = vector.contract {indexing_maps = [#map_a, #map_b, #map_c], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
+ %E = arith.subf %C, %D : vector<16x16xf16>
+ %F = arith.negf %E : vector<16x16xf16>
+ vector.transfer_write %F, %d22[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<16x16xf16>
+ gpu.terminator
+ }
+
+ // Copy the result back to the host.
+ %token2 = gpu.wait async
+ %z = gpu.memcpy async [%token2] %22, %d22 : memref<16x16xf16>, memref<16x16xf16>
+ %w = gpu.dealloc async [%token2] %d0 : memref<16x16xf16>
+ %v = gpu.dealloc async [%token2] %d22 : memref<16x16xf16>
+ gpu.wait [%token2]
+
+ // Convert the results from f16 to f32 for printing.
+ scf.for %arg0 = %c0 to %c16 step %c1 {
+ scf.for %arg1 = %c0 to %c16 step %c1 {
+ %6 = memref.load %22[%arg0, %arg1] : memref<16x16xf16>
+ %7 = arith.extf %6 : f16 to f32
+ memref.store %7, %1[%arg0, %arg1] : memref<16x16xf32>
+ }
+ }
+
+ // Print the memref after computation.
+ call @printMemrefF32(%3) : (memref<*xf32>) -> ()
+ // CHECK-COUNT-16: [-0, 120, 240, 360, 480, 600, 720, 840, 960, 1080, 1200, 1320, 1440, 1560, 1680, 1800]
+ return
+}
+
+func.func private @printMemrefF32(memref<*xf32>)
``````````
</details>
https://github.com/llvm/llvm-project/pull/221288
More information about the Mlir-commits
mailing list