[mlir][AMDGPU] Make sure gfx942 and previous archs emulate bf16 atomics

Author: dhernandez0

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

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
+}