[Mlir-commits] [mlir] [mlir][gpu] Lower the remaining subgroup_mma_elementwise ops to NVVM (PR #221288)

Hung-Kuan Tseng llvmlistbot at llvm.org
Fri Sep 4 10:10:33 PDT 2026


https://github.com/Tim096 created https://github.com/llvm/llvm-project/pull/221288

`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).


>From 704d5b9c54b2cd5db76d6859fbfa354e6239fa71 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sat, 5 Sep 2026 00:57:38 +0800
Subject: [PATCH] [mlir][gpu] Lower the remaining subgroup_mma_elementwise ops
 to NVVM

`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`. `addf` on an f32 multiplicand fragment
builds an `llvm.fadd` on the `i32` registers that hold the tf32 bit
patterns and fails verification instead.

Lower `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 reject
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.
---
 .../Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp    | 66 ++++++++++++---
 .../gpu-to-nvvm-invalid-wmma-elementwise.mlir | 46 ++++++++++
 .../GPUToNVVM/wmma-ops-to-nvvm.mlir           | 49 +++++++++++
 .../wmma-matmul-f16-elementwise.mlir          | 84 +++++++++++++++++++
 4 files changed, 233 insertions(+), 12 deletions(-)
 create mode 100644 mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-wmma-elementwise.mlir
 create mode 100644 mlir/test/Integration/GPU/CUDA/TensorCore/wmma-matmul-f16-elementwise.mlir

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>)



More information about the Mlir-commits mailing list