[Backport to 14] add initial f16 type support for atomicrmw in llvm-s…

…pirv translator (KhronosGroup#2210)

This PR aims to add f16 type support for atomicrmw in llvm-spirv translator, with the reference to the extension documented in [1].
There are two concerns related to the subject:

SPIRVAtomicFAddEXTInst::getRequiredExtension() should return a list of required extension to support the requirement to list both SPV_EXT_shader_atomic_float16_add and SPV_EXT_shader_atomic_float_add extensions in the module (see "Extension Name" section of the ref [1]). However, the return type is std::optional<ExtensionID> and returning a vector would need a bigger rework.
Including SPV_EXT_shader_atomic_float16_add into --spirv-ext argument of llvm-spirv doesn't result in producing the correspondent capability (AtomicFloat16AddEXT) and extension in a SPIRV output.
$ llvm-spirv AtomicFAddEXT.ll.tmp.bc --spirv-ext=+SPV_EXT_shader_atomic_float_add,+SPV_EXT_shader_atomic_float16_add -o AtomicFAddEXT.ll.tmp.spv
$ llvm-spirv -to-text AtomicFAddEXT.ll.tmp.spv -o /dev/stdout
...
2 Capability AtomicFloat32AddEXT
2 Capability AtomicFloat64AddEXT
9 Extension "SPV_EXT_shader_atomic_float_add"
...
This prevents extending the test case of AtomicFAddEXT.ll in EXT/SPV_EXT_shader_atomic_float.

References:
[1] https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/EXT/SPV_EXT_shader_atomic_float16_add.asciidoc

include/LLVMSPIRVExtensions.inc

@@ -2,6 +2,7 @@
 #define EXT(X)
 #endif
 
+EXT(SPV_EXT_shader_atomic_float16_add)
 EXT(SPV_EXT_shader_atomic_float_add)
 EXT(SPV_EXT_shader_atomic_float_min_max)
 EXT(SPV_KHR_no_integer_wrap_decoration)

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -2743,16 +2743,22 @@ class SPIRVAtomicStoreInst : public SPIRVAtomicInstBase {
 class SPIRVAtomicFAddEXTInst : public SPIRVAtomicInstBase {
 public:
   llvm::Optional<ExtensionID> getRequiredExtension() const override {
+    assert(hasType());
+    if (getType()->isTypeFloat(16))
+      return ExtensionID::SPV_EXT_shader_atomic_float16_add;
     return ExtensionID::SPV_EXT_shader_atomic_float_add;
   }
 
   SPIRVCapVec getRequiredCapability() const override {
     assert(hasType());
+    if (getType()->isTypeFloat(16))
+      return {CapabilityAtomicFloat16AddEXT};
     if (getType()->isTypeFloat(32))
       return {CapabilityAtomicFloat32AddEXT};
-    assert(getType()->isTypeFloat(64) &&
-           "AtomicFAddEXT can only be generated for f32 or f64 types");
-    return {CapabilityAtomicFloat64AddEXT};
+    if (getType()->isTypeFloat(64))
+      return {CapabilityAtomicFloat64AddEXT};
+    llvm_unreachable(
+        "AtomicFAddEXT can only be generated for f16, f32, f64 types");
   }
 };
 

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -655,6 +655,17 @@ void SPIRVModuleImpl::addExtension(ExtensionID Ext) {
     return;
   }
   SPIRVExt.insert(ExtName);
+
+  // SPV_EXT_shader_atomic_float16_add extends the
+  // SPV_EXT_shader_atomic_float_add extension.
+  // The specification requires both extensions to be added to use
+  // AtomicFloat16AddEXT capability whereas getRequiredExtension()
+  // is able to return a single extensionID.
+  if (Ext == ExtensionID::SPV_EXT_shader_atomic_float16_add) {
+    SPIRVMap<ExtensionID, std::string>::find(
+        ExtensionID::SPV_EXT_shader_atomic_float_add, &ExtName);
+    SPIRVExt.insert(ExtName);
+  }
 }
 
 void SPIRVModuleImpl::addCapability(SPIRVCapabilityKind Cap) {

test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fadd_half.ll

@@ -0,0 +1,34 @@
+; RUN: llvm-as -opaque-pointers < %s -o %t.bc
+; RUN: llvm-spirv -opaque-pointers --spirv-ext=+SPV_EXT_shader_atomic_float16_add %t.bc -o %t.spv
+; RUN: spirv-val %t.spv
+; RUN: llvm-spirv -opaque-pointers -to-text %t.spv -o - | FileCheck %s
+
+; CHECK-DAG: Extension "SPV_EXT_shader_atomic_float16_add"
+; CHECK-DAG: Extension "SPV_EXT_shader_atomic_float_add"
+; CHECK-DAG: Capability AtomicFloat16AddEXT
+; CHECK: TypeInt [[TypeIntID:[0-9]+]] 32 0
+; CHECK-DAG: Constant [[TypeIntID]] [[ScopeDevice:[0-9]+]] 1 {{$}}
+; CHECK-DAG: Constant [[TypeIntID]] [[MemSem_SequentiallyConsistent:[0-9]+]] 16
+; CHECK: TypeFloat [[TypeFloatHalfID:[0-9]+]] 16
+; CHECK: Variable {{[0-9]+}} [[HalfPointer:[0-9]+]]
+; CHECK: Constant [[TypeFloatHalfID]] [[HalfValue:[0-9]+]] 20800
+
+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"
+
+@f = common dso_local local_unnamed_addr addrspace(1) global half 0.000000e+00, align 4
+
+; Function Attrs: nounwind
+define dso_local spir_func void @test_atomicrmw_fadd() local_unnamed_addr #0 {
+entry:
+ %0 = atomicrmw fadd ptr addrspace(1) @f, half 42.000000e+00 seq_cst
+; CHECK: AtomicFAddEXT [[TypeFloatHalfID]] {{[0-9]+}} [[HalfPointer]] [[ScopeDevice]] [[MemSem_SequentiallyConsistent]] [[HalfValue]]
+
+  ret void
+}
+
+attributes #0 = { nounwind "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"wchar_size", i32 4}

test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fsub_half.ll

@@ -0,0 +1,39 @@
+; RUN: llvm-as -opaque-pointers=0 < %s -o %t.bc
+; RUN: llvm-spirv -opaque-pointers=0 --spirv-ext=+SPV_EXT_shader_atomic_float16_add %t.bc -o %t.spv
+; RUN: spirv-val %t.spv
+; RUN: llvm-spirv -to-text %t.spv -o - | FileCheck --check-prefix=CHECK-SPIRV %s
+
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis -opaque-pointers=0 %t.rev.bc -o - | FileCheck %s --check-prefixes=CHECK-LLVM
+
+; CHECK-SPIRV-DAG: Extension "SPV_EXT_shader_atomic_float16_add"
+; CHECK-SPIRV-DAG: Extension "SPV_EXT_shader_atomic_float_add"
+; CHECK-SPIRV-DAG: Capability AtomicFloat16AddEXT
+; CHECK-SPIRV: TypeInt [[Int:[0-9]+]] 32 0
+; CHECK-SPIRV-DAG: Constant [[Int]] [[ScopeDevice:[0-9]+]] 1 {{$}}
+; CHECK-SPIRV-DAG: Constant [[Int]] [[MemSem_SequentiallyConsistent:[0-9]+]] 16
+; CHECK-SPIRV: TypeFloat [[Half:[0-9]+]] 16
+; CHECK-SPIRV: Variable {{[0-9]+}} [[HalfPointer:[0-9]+]]
+; CHECK-SPIRV: Constant [[Half]] [[HalfValue:[0-9]+]] 15360
+
+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"
+
+@f = common dso_local local_unnamed_addr addrspace(1) global half 0.000000e+00, align 4
+
+; Function Attrs: nounwind
+define dso_local spir_func void @test_atomicrmw_fadd() local_unnamed_addr #0 {
+entry:
+ %0 = atomicrmw fsub half addrspace(1)* @f, half 1.0e+00 seq_cst
+; CHECK-SPIRV: FNegate [[Half]] [[NegateValue:[0-9]+]] [[HalfValue]]
+; CHECK-SPIRV: AtomicFAddEXT [[Half]] {{[0-9]+}} [[HalfPointer]] [[ScopeDevice]] [[MemSem_SequentiallyConsistent]] [[NegateValue]]
+; CHECK-LLVM: [[FNegateLLVM:%[0-9]+]] = fneg half 0xH3C00
+; CHECK-LLVM: call spir_func half {{.*}}atomic_add{{.*}}(half addrspace(1)* @f, half [[FNegateLLVM]])
+  ret void
+}
+
+attributes #0 = { nounwind "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"wchar_size", i32 4}