[llvm] [SPIR-V] Allow bfloat vector atomics lowering without scalarization (PR #202083)

[Dmitry Sidorov via llvm-commits]

https://github.com/MrSidims updated https://github.com/llvm/llvm-project/pull/202083

>From 615ff33bda619a65eea8fb8b5ca7c3cf54714dba Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Sat, 6 Jun 2026 19:36:03 +0200
Subject: [PATCH 1/2] [SPIR-V] Allow bfloat vector atomics lowering without
 scalarization

This is a workaround only for AMD triple - to use
SPV_NV_shader_atomic_fp16_vector for this lowering.
---
 llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp |  5 +-
 .../atomicrmw_vector_bfloat16_amd.ll          | 55 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/SPIRV/extensions/SPV_NV_shader_atomic_fp16_vector/atomicrmw_vector_bfloat16_amd.ll

diff --git a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
index 217e7619a9a3c..cf95d3a57c8a0 100644
--- a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
@@ -1159,7 +1159,10 @@ static void AddAtomicVectorFloatRequirements(const MachineInstr &MI,
         "The element type for the result type of an atomic vector float "
         "instruction must be a 16-bit floating-point scalar");
 
-  if (isBFloat16Type(EltTypeDef))
+  // The extension is defined for fp16, but the AMD target lets a bf16 vector use
+  // the same instruction so it can lower to a packed bf16 atomic.
+  if (isBFloat16Type(EltTypeDef) &&
+      ST.getTargetTriple().getVendor() != Triple::AMD)
     reportFatalUsageError(
         "The element type for the result type of an atomic vector float "
         "instruction cannot be a bfloat16 scalar");
diff --git a/llvm/test/CodeGen/SPIRV/extensions/SPV_NV_shader_atomic_fp16_vector/atomicrmw_vector_bfloat16_amd.ll b/llvm/test/CodeGen/SPIRV/extensions/SPV_NV_shader_atomic_fp16_vector/atomicrmw_vector_bfloat16_amd.ll
new file mode 100644
index 0000000000000..d25335eb1f5d7
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/extensions/SPV_NV_shader_atomic_fp16_vector/atomicrmw_vector_bfloat16_amd.ll
@@ -0,0 +1,55 @@
+; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv64-amd-amdhsa %s -o - | FileCheck %s
+; RUN: not llc -verify-machineinstrs -O0 -mtriple=spirv64-unknown-unknown --spirv-ext=+SPV_KHR_bfloat16,+SPV_NV_shader_atomic_fp16_vector,+SPV_INTEL_bfloat16_arithmetic %s -o - 2>&1 | FileCheck %s --check-prefix=NOAMD
+
+; SPV_NV_shader_atomic_fp16_vector extension is meant for fp16,
+; but the bf16 type uses the same vector atomic so the AMD backend can lower it
+; to a packed bf16 atomic.
+
+; Anywhere else the bf16 element type is not allowed on the vector atomic, so the
+; same module is rejected.
+
+; NOAMD: cannot be a bfloat16 scalar
+
+; CHECK-DAG: Capability BFloat16TypeKHR
+; CHECK-DAG: Capability AtomicFloat16VectorNV
+; CHECK-DAG: Extension "SPV_KHR_bfloat16"
+; CHECK-DAG: Extension "SPV_NV_shader_atomic_fp16_vector"
+; CHECK-DAG: %[[#BF16:]] = OpTypeFloat 16 0
+; CHECK-DAG: %[[#V2:]] = OpTypeVector %[[#BF16]] 2
+; CHECK-DAG: %[[#V4:]] = OpTypeVector %[[#BF16]] 4
+
+; CHECK: OpFunction %[[#V2]] None %[[#]]
+; CHECK: OpAtomicFAddEXT %[[#V2]]
+define <2 x bfloat> @test_fadd_v2(ptr addrspace(1) %ptr, <2 x bfloat> %val) {
+  %r = atomicrmw fadd ptr addrspace(1) %ptr, <2 x bfloat> %val seq_cst
+  ret <2 x bfloat> %r
+}
+
+; CHECK: OpFunction %[[#V2]] None %[[#]]
+; CHECK: %[[#Neg:]] = OpFNegate %[[#V2]]
+; CHECK: OpAtomicFAddEXT %[[#V2]] %[[#]] %[[#]] %[[#]] %[[#Neg]]
+define <2 x bfloat> @test_fsub_v2(ptr addrspace(1) %ptr, <2 x bfloat> %val) {
+  %r = atomicrmw fsub ptr addrspace(1) %ptr, <2 x bfloat> %val seq_cst
+  ret <2 x bfloat> %r
+}
+
+; CHECK: OpFunction %[[#V2]] None %[[#]]
+; CHECK: OpAtomicFMinEXT %[[#V2]]
+define <2 x bfloat> @test_fmin_v2(ptr addrspace(1) %ptr, <2 x bfloat> %val) {
+  %r = atomicrmw fmin ptr addrspace(1) %ptr, <2 x bfloat> %val seq_cst
+  ret <2 x bfloat> %r
+}
+
+; CHECK: OpFunction %[[#V2]] None %[[#]]
+; CHECK: OpAtomicFMaxEXT %[[#V2]]
+define <2 x bfloat> @test_fmax_v2(ptr addrspace(1) %ptr, <2 x bfloat> %val) {
+  %r = atomicrmw fmax ptr addrspace(1) %ptr, <2 x bfloat> %val seq_cst
+  ret <2 x bfloat> %r
+}
+
+; CHECK: OpFunction %[[#V4]] None %[[#]]
+; CHECK: OpAtomicFAddEXT %[[#V4]]
+define <4 x bfloat> @test_fadd_v4(ptr addrspace(1) %ptr, <4 x bfloat> %val) {
+  %r = atomicrmw fadd ptr addrspace(1) %ptr, <4 x bfloat> %val seq_cst
+  ret <4 x bfloat> %r
+}

>From 9cc5fd787f7bb39e50a72a45d868a1926e907abb Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Sat, 6 Jun 2026 22:12:04 +0200
Subject: [PATCH 2/2] fix

---
 llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
index cf95d3a57c8a0..6f641074ecaa7 100644
--- a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
@@ -1159,8 +1159,8 @@ static void AddAtomicVectorFloatRequirements(const MachineInstr &MI,
         "The element type for the result type of an atomic vector float "
         "instruction must be a 16-bit floating-point scalar");
 
-  // The extension is defined for fp16, but the AMD target lets a bf16 vector use
-  // the same instruction so it can lower to a packed bf16 atomic.
+  // The extension is defined for fp16, but the AMD target lets a bf16 vector
+  // use the same instruction so it can lower to a packed bf16 atomic.
   if (isBFloat16Type(EltTypeDef) &&
       ST.getTargetTriple().getVendor() != Triple::AMD)
     reportFatalUsageError(