Add SPV_NV_shader_atomic_fp16_vector extension (#3772)

It allowed fp16 atomics translation.

Spec:

https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/NV/SPV_NV_shader_atomic_fp16_vector.asciidoc

Author: MrSidims

include/LLVMSPIRVExtensions.inc

@@ -96,3 +96,4 @@ EXT(SPV_INTEL_float4)
 EXT(SPV_INTEL_fp_conversions)
 EXT(SPV_KHR_float_controls2)
 EXT(SPV_AMD_weak_linkage)
+EXT(SPV_NV_shader_atomic_fp16_vector)

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -3135,6 +3135,9 @@ class SPIRVAtomicFAddEXTInst : public SPIRVAtomicInstBase {
 public:
   std::optional<ExtensionID> getRequiredExtension() const override {
     assert(hasType());
+    if (getType()->isTypeVector() &&
+        getType()->getVectorComponentType()->isTypeFloat(16))
+      return ExtensionID::SPV_NV_shader_atomic_fp16_vector;
     if (getType()->isTypeFloat(16, FPEncodingBFloat16KHR))
       Module->addExtension(ExtensionID::SPV_INTEL_16bit_atomics);
     if (getType()->isTypeFloat(16))
@@ -3144,6 +3147,9 @@ class SPIRVAtomicFAddEXTInst : public SPIRVAtomicInstBase {
 
   SPIRVCapVec getRequiredCapability() const override {
     assert(hasType());
+    if (getType()->isTypeVector() &&
+        getType()->getVectorComponentType()->isTypeFloat(16))
+      return {CapabilityAtomicFloat16VectorNV};
     if (getType()->isTypeFloat(16, FPEncodingBFloat16KHR))
       return {internal::CapabilityAtomicBFloat16AddINTEL};
     if (getType()->isTypeFloat(16))
@@ -3160,13 +3166,19 @@ class SPIRVAtomicFAddEXTInst : public SPIRVAtomicInstBase {
 class SPIRVAtomicFMinMaxEXTBase : public SPIRVAtomicInstBase {
 public:
   std::optional<ExtensionID> getRequiredExtension() const override {
+    if (getType()->isTypeVector() &&
+        getType()->getVectorComponentType()->isTypeFloat(16))
+      return ExtensionID::SPV_NV_shader_atomic_fp16_vector;
     if (getType()->isTypeFloat(16, FPEncodingBFloat16KHR))
       Module->addExtension(ExtensionID::SPV_INTEL_16bit_atomics);
     return ExtensionID::SPV_EXT_shader_atomic_float_min_max;
   }
 
   SPIRVCapVec getRequiredCapability() const override {
     assert(hasType());
+    if (getType()->isTypeVector() &&
+        getType()->getVectorComponentType()->isTypeFloat(16))
+      return {CapabilityAtomicFloat16VectorNV};
     if (getType()->isTypeFloat(16, FPEncodingBFloat16KHR))
       return {internal::CapabilityAtomicBFloat16MinMaxINTEL};
     if (getType()->isTypeFloat(16))

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -630,6 +630,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(CapabilityLongCompositesINTEL, "LongCompositesINTEL");
   add(CapabilityOptNoneEXT, "OptNoneEXT");
   add(CapabilityAtomicFloat16AddEXT, "AtomicFloat16AddEXT");
+  add(CapabilityAtomicFloat16VectorNV, "AtomicFloat16VectorNV");
   add(internal::CapabilityAtomicBFloat16AddINTEL, "AtomicBFloat16AddINTEL");
   add(internal::CapabilityAtomicBFloat16MinMaxINTEL,
       "AtomicBFloat16MinMaxINTEL");

test/extensions/NV/SPV_NV_shader_atomic_fp16_vector/atomicrmw_fp16_vector.ll

@@ -0,0 +1,41 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_NV_shader_atomic_fp16_vector -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis %t.rev.bc -o - | FileCheck %s --check-prefix=CHECK-LLVM
+
+; SPV_NV_shader_atomic_fp16_vector lets OpAtomicFAddEXT, OpAtomicFMinEXT and
+; OpAtomicFMaxEXT operate on a 2- or 4-component vector of fp16. fsub has no
+; atomic opcode, so it is emitted as OpFNegate followed by OpAtomicFAddEXT.
+
+; CHECK-SPIRV-DAG: Capability AtomicFloat16VectorNV
+; CHECK-SPIRV-DAG: Extension "SPV_NV_shader_atomic_fp16_vector"
+; CHECK-SPIRV: TypeFloat [[#Half:]] 16
+; CHECK-SPIRV: TypeVector [[#V2:]] [[#Half]] 2
+; CHECK-SPIRV: TypeVector [[#V4:]] [[#Half]] 4
+; CHECK-SPIRV: AtomicFAddEXT [[#V2]]
+; CHECK-SPIRV: FNegate [[#V2]] [[#Neg:]]
+; CHECK-SPIRV: AtomicFAddEXT [[#V2]] [[#]] [[#]] [[#]] [[#]] [[#Neg]]
+; CHECK-SPIRV: AtomicFMinEXT [[#V2]]
+; CHECK-SPIRV: AtomicFMaxEXT [[#V2]]
+; CHECK-SPIRV: AtomicFAddEXT [[#V4]]
+
+; CHECK-LLVM: call spir_func <2 x half> @{{.*}}atomic_add{{.*}}(ptr addrspace(1) %p, <2 x half> %v2)
+; CHECK-LLVM: fneg <2 x half> %v2
+; CHECK-LLVM: call spir_func <2 x half> @{{.*}}atomic_add
+; CHECK-LLVM: call spir_func <2 x half> @{{.*}}atomic_min
+; CHECK-LLVM: call spir_func <2 x half> @{{.*}}atomic_max
+; CHECK-LLVM: call spir_func <4 x half> @{{.*}}atomic_add
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir64"
+
+define spir_func void @test(ptr addrspace(1) %p, <2 x half> %v2, <4 x half> %v4) {
+entry:
+  %a = atomicrmw fadd ptr addrspace(1) %p, <2 x half> %v2 seq_cst
+  %s = atomicrmw fsub ptr addrspace(1) %p, <2 x half> %v2 seq_cst
+  %mn = atomicrmw fmin ptr addrspace(1) %p, <2 x half> %v2 seq_cst
+  %mx = atomicrmw fmax ptr addrspace(1) %p, <2 x half> %v2 seq_cst
+  %a4 = atomicrmw fadd ptr addrspace(1) %p, <4 x half> %v4 seq_cst
+  ret void
+}

test/extensions/NV/SPV_NV_shader_atomic_fp16_vector/negative.ll

@@ -0,0 +1,17 @@
+; RUN: llvm-as < %s -o %t.bc
+; RUN: not llvm-spirv %t.bc 2>&1 | FileCheck %s
+
+; Without SPV_NV_shader_atomic_fp16_vector an fp16 vector atomic cannot be
+; translated.
+
+; CHECK: RequiresExtension: Feature requires the following SPIR-V extension:
+; CHECK-NEXT: SPV_NV_shader_atomic_fp16_vector
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir64"
+
+define spir_func void @test(ptr addrspace(1) %p, <2 x half> %v2) {
+entry:
+  %a = atomicrmw fadd ptr addrspace(1) %p, <2 x half> %v2 seq_cst
+  ret void
+}