[Mlir-commits] [mlir] [mlir][AMDGPU] Enable emulating vector buffer_atomic_fadd for bf16 on gfx942 (PR #129029)

[Daniel Hernandez-Juarez]

https://github.com/dhernandez0 updated https://github.com/llvm/llvm-project/pull/129029

>From f7f87a40817502389f225acdf81a29cc37e7cc92 Mon Sep 17 00:00:00 2001
From: Daniel Hernandez-Juarez <dhernandez0 at gmail.com>
Date: Thu, 27 Feb 2025 09:32:05 +0000
Subject: [PATCH] [mlir][AMDGPU] Make sure gfx942 and previous archs emulate
 bf16 atomics

---
 .../AMDGPU/Transforms/EmulateAtomics.cpp      |  8 +++
 .../AMDGPU/amdgpu-emulate-atomics.mlir        | 49 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/mlir/lib/Dialect/AMDGPU/Transforms/EmulateAtomics.cpp b/mlir/lib/Dialect/AMDGPU/Transforms/EmulateAtomics.cpp
index 7459a6503cddf..04dc871941b98 100644
--- a/mlir/lib/Dialect/AMDGPU/Transforms/EmulateAtomics.cpp
+++ b/mlir/lib/Dialect/AMDGPU/Transforms/EmulateAtomics.cpp
@@ -189,6 +189,14 @@ void mlir::amdgpu::populateAmdgpuEmulateAtomicsPatterns(
     } else {
       target.addIllegalOp<RawBufferAtomicFmaxOp>();
     }
+    // gfx950 has bf16 atomics
+    if (chipset < Chipset(9, 5, 0)) {
+      target.addDynamicallyLegalOp<RawBufferAtomicFaddOp>(
+          [](RawBufferAtomicFaddOp op) -> bool {
+            Type elemType = getElementTypeOrSelf(op.getValue().getType());
+            return !isa<BFloat16Type>(elemType);
+          });
+    }
   }
   patterns.add<
       RawBufferAtomicByCasPattern<RawBufferAtomicFaddOp, arith::AddFOp>,
diff --git a/mlir/test/Dialect/AMDGPU/amdgpu-emulate-atomics.mlir b/mlir/test/Dialect/AMDGPU/amdgpu-emulate-atomics.mlir
index 104af58a0d277..7368f853ca293 100644
--- a/mlir/test/Dialect/AMDGPU/amdgpu-emulate-atomics.mlir
+++ b/mlir/test/Dialect/AMDGPU/amdgpu-emulate-atomics.mlir
@@ -1,6 +1,9 @@
 // RUN: mlir-opt -split-input-file -amdgpu-emulate-atomics=chipset=gfx90a %s | FileCheck %s --check-prefixes=CHECK,GFX9
 // RUN: mlir-opt -split-input-file -amdgpu-emulate-atomics=chipset=gfx1030 %s | FileCheck %s --check-prefixes=CHECK,GFX10
 // RUN: mlir-opt -split-input-file -amdgpu-emulate-atomics=chipset=gfx1100 %s | FileCheck %s --check-prefixes=CHECK,GFX11
+// RUN: mlir-opt -split-input-file -amdgpu-emulate-atomics=chipset=gfx1200 %s | FileCheck %s --check-prefixes=CHECK,GFX12
+// RUN: mlir-opt -split-input-file -amdgpu-emulate-atomics=chipset=gfx942 %s | FileCheck %s --check-prefixes=CHECK,GFX942
+// RUN: mlir-opt -split-input-file -amdgpu-emulate-atomics=chipset=gfx950 %s | FileCheck %s --check-prefixes=CHECK,GFX950
 
 // -----
 
@@ -10,6 +13,7 @@ func.func @atomic_fmax(%val: f32, %buffer: memref<?xf32>, %idx: i32) {
 // CHECK: gpu.printf "Begin\0A"
 // GFX10: amdgpu.raw_buffer_atomic_fmax {foo, indexOffset = 4 : i32} [[val]] -> [[buffer]][[[idx]]]
 // GFX11: amdgpu.raw_buffer_atomic_fmax {foo, indexOffset = 4 : i32} [[val]] -> [[buffer]][[[idx]]]
+// GFX12: amdgpu.raw_buffer_atomic_fmax {foo, indexOffset = 4 : i32} [[val]] -> [[buffer]][[[idx]]]
 // GFX9:  [[ld:%.+]] = amdgpu.raw_buffer_load {foo, indexOffset = 4 : i32} [[buffer]][[[idx]]]
 // GFX9:  cf.br [[loop:\^.+]]([[ld]] : f32)
 // GFX9:  [[loop]]([[arg:%.+]]: f32):
@@ -20,6 +24,26 @@ func.func @atomic_fmax(%val: f32, %buffer: memref<?xf32>, %idx: i32) {
 // GFX9:  [[test:%.+]] = arith.cmpi eq, [[resCast]], [[argCast]]
 // GFX9:  cf.cond_br [[test]], [[post:\^.+]], [[loop]]([[atomicRes]] : f32)
 // GFX9:  [[post]]:
+// GFX942:  [[ld:%.+]] = amdgpu.raw_buffer_load {foo, indexOffset = 4 : i32} [[buffer]][[[idx]]]
+// GFX942:  cf.br [[loop:\^.+]]([[ld]] : f32)
+// GFX942:  [[loop]]([[arg:%.+]]: f32):
+// GFX942:  [[operated:%.+]] = arith.maximumf [[val]], [[arg]]
+// GFX942:  [[atomicRes:%.+]] = amdgpu.raw_buffer_atomic_cmpswap {foo, indexOffset = 4 : i32} [[operated]], [[arg]] -> [[buffer]][[[idx]]]
+// GFX942:  [[argCast:%.+]] = arith.bitcast [[arg]] : f32 to i32
+// GFX942:  [[resCast:%.+]] = arith.bitcast [[atomicRes]] : f32 to i32
+// GFX942:  [[test:%.+]] = arith.cmpi eq, [[resCast]], [[argCast]]
+// GFX942:  cf.cond_br [[test]], [[post:\^.+]], [[loop]]([[atomicRes]] : f32)
+// GFX942:  [[post]]:
+// GFX950:  [[ld:%.+]] = amdgpu.raw_buffer_load {foo, indexOffset = 4 : i32} [[buffer]][[[idx]]]
+// GFX950:  cf.br [[loop:\^.+]]([[ld]] : f32)
+// GFX950:  [[loop]]([[arg:%.+]]: f32):
+// GFX950:  [[operated:%.+]] = arith.maximumf [[val]], [[arg]]
+// GFX950:  [[atomicRes:%.+]] = amdgpu.raw_buffer_atomic_cmpswap {foo, indexOffset = 4 : i32} [[operated]], [[arg]] -> [[buffer]][[[idx]]]
+// GFX950:  [[argCast:%.+]] = arith.bitcast [[arg]] : f32 to i32
+// GFX950:  [[resCast:%.+]] = arith.bitcast [[atomicRes]] : f32 to i32
+// GFX950:  [[test:%.+]] = arith.cmpi eq, [[resCast]], [[argCast]]
+// GFX950:  cf.cond_br [[test]], [[post:\^.+]], [[loop]]([[atomicRes]] : f32)
+// GFX950:  [[post]]:
 // CHECK-NEXT: gpu.printf "End\0A"
   gpu.printf "Begin\n"
   amdgpu.raw_buffer_atomic_fmax {foo, indexOffset = 4 : i32} %val -> %buffer[%idx] : f32 -> memref<?xf32>, i32
@@ -36,6 +60,9 @@ func.func @atomic_fmax_f64(%val: f64, %buffer: memref<?xf64>, %idx: i32) {
 // GFX9:  amdgpu.raw_buffer_atomic_fmax [[val]] -> [[buffer]][[[idx]]]
 // GFX10: amdgpu.raw_buffer_atomic_fmax [[val]] -> [[buffer]][[[idx]]]
 // GFX11: amdgpu.raw_buffer_atomic_fmax [[val]] -> [[buffer]][[[idx]]]
+// GFX12: amdgpu.raw_buffer_atomic_fmax [[val]] -> [[buffer]][[[idx]]]
+// GFX942: amdgpu.raw_buffer_atomic_fmax [[val]] -> [[buffer]][[[idx]]]
+// GFX950: amdgpu.raw_buffer_atomic_fmax [[val]] -> [[buffer]][[[idx]]]
 // CHECK-NEXT: gpu.printf "End\0A"
   gpu.printf "Begin\n"
   amdgpu.raw_buffer_atomic_fmax %val -> %buffer[%idx] : f64 -> memref<?xf64>, i32
@@ -51,6 +78,9 @@ func.func @atomic_fadd(%val: f32, %buffer: memref<?xf32>, %idx: i32) {
 // GFX10: amdgpu.raw_buffer_load
 // GFX10: amdgpu.raw_buffer_atomic_cmpswap
 // GFX11: amdgpu.raw_buffer_atomic_fadd
+// GFX12: amdgpu.raw_buffer_atomic_fadd
+// GFX942: amdgpu.raw_buffer_atomic_fadd
+// GFX950: amdgpu.raw_buffer_atomic_fadd
   amdgpu.raw_buffer_atomic_fadd %val -> %buffer[%idx] : f32 -> memref<?xf32>, i32
   func.return
 }
@@ -69,6 +99,25 @@ func.func @atomic_fadd_v2f16(%val: vector<2xf16>, %buffer: memref<?xf16>, %idx:
 // GFX11: %[[vecCastOld:.+]] = vector.bitcast %[[old]] : vector<2xf16> to vector<1xi32>
 // GFX11: %[[scalarOld:.+]] = vector.extract %[[vecCastOld]][0]
 // GFX11: arith.cmpi eq, %[[scalarOld]], %[[scalarExpected]]
+// GFX942: amdgpu.raw_buffer_atomic_fadd
+// GFX12:  amdgpu.raw_buffer_atomic_fadd
+// GFX950:  amdgpu.raw_buffer_atomic_fadd
   amdgpu.raw_buffer_atomic_fadd %val -> %buffer[%idx] : vector<2xf16> -> memref<?xf16>, i32
   func.return
 }
+
+// CHECK: func @atomic_fadd_v2bf16
+func.func @atomic_fadd_v2bf16(%val: vector<2xbf16>, %buffer: memref<?xbf16>, %idx: i32) {
+// GFX9: amdgpu.raw_buffer_load
+// GFX9: amdgpu.raw_buffer_atomic_cmpswap
+// GFX10: amdgpu.raw_buffer_load
+// GFX10: amdgpu.raw_buffer_atomic_cmpswap
+// GFX11: amdgpu.raw_buffer_load
+// GFX11: amdgpu.raw_buffer_atomic_cmpswap
+// GFX942: amdgpu.raw_buffer_load
+// GFX942: amdgpu.raw_buffer_atomic_cmpswap
+// GFX12:  amdgpu.raw_buffer_atomic_fadd
+// GFX950:  amdgpu.raw_buffer_atomic_fadd
+  amdgpu.raw_buffer_atomic_fadd %val -> %buffer[%idx] : vector<2xbf16> -> memref<?xbf16>, i32
+  func.return
+}