add initial f16 type support for atomicrmw in llvm-spirv translator (#2210)

VyacheslavLevytskyy · sys-ce-bb · commit 84d92de62ed7 · 2023-11-16T06:36:13.000-08:00
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 Original commit: KhronosGroup/SPIRV-LLVM-Translator@1aae8db
diff --git a/llvm-spirv/include/LLVMSPIRVExtensions.inc b/llvm-spirv/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_EXT_image_raw10_raw12)
diff --git a/llvm-spirv/lib/SPIRV/libSPIRV/SPIRVInstruction.h b/llvm-spirv/lib/SPIRV/libSPIRV/SPIRVInstruction.h
@@ -2746,16 +2746,22 @@ class SPIRVAtomicStoreInst : public SPIRVAtomicInstBase {
 class SPIRVAtomicFAddEXTInst : public SPIRVAtomicInstBase {
 public:
   std::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");
   }
 };
 
diff --git a/llvm-spirv/lib/SPIRV/libSPIRV/SPIRVModule.cpp b/llvm-spirv/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) {
@@ -690,7 +701,8 @@ SPIRVConstant *SPIRVModuleImpl::getLiteralAsConstant(unsigned Literal) {
   if (Loc != LiteralMap.end())
     return Loc->second;
   auto *Ty = addIntegerType(32);
-  auto *V = new SPIRVConstant(this, Ty, getId(), static_cast<uint64_t>(Literal));
+  auto *V =
+      new SPIRVConstant(this, Ty, getId(), static_cast<uint64_t>(Literal));
   LiteralMap[Literal] = V;
   addConstant(V);
   return V;
@@ -1227,8 +1239,9 @@ SPIRVValue *SPIRVModuleImpl::addSpecConstantComposite(
     End = ((Elements.end() - End) > MaxNumElements) ? End + MaxNumElements
                                                     : Elements.end();
     Slice.assign(Start, End);
-    auto *Continued = static_cast<SPIRVSpecConstantComposite::ContinuedInstType>(
-        addSpecConstantCompositeContinuedINTEL(Slice));
+    auto *Continued =
+        static_cast<SPIRVSpecConstantComposite::ContinuedInstType>(
+            addSpecConstantCompositeContinuedINTEL(Slice));
     Res->addContinuedInstruction(Continued);
   }
   return Res;
@@ -1440,7 +1453,7 @@ SPIRVValue *SPIRVModuleImpl::addAsmINTEL(SPIRVTypeFunction *TheType,
                                          const std::string &TheInstructions,
                                          const std::string &TheConstraints) {
   auto *Asm = new SPIRVAsmINTEL(this, TheType, getId(), TheTarget,
-                               TheInstructions, TheConstraints);
+                                TheInstructions, TheConstraints);
   return add(Asm);
 }
 
@@ -1725,8 +1738,9 @@ SPIRVInstruction *SPIRVModuleImpl::addExpectKHRInst(SPIRVType *ResultTy,
 // Create AliasDomainDeclINTEL/AliasScopeDeclINTEL/AliasScopeListDeclINTEL
 // instructions
 template <typename AliasingInstType>
-SPIRVEntry *SPIRVModuleImpl::getOrAddMemAliasingINTELInst(
-    std::vector<SPIRVId> Args, llvm::MDNode *MD) {
+SPIRVEntry *
+SPIRVModuleImpl::getOrAddMemAliasingINTELInst(std::vector<SPIRVId> Args,
+                                              llvm::MDNode *MD) {
   assert(MD && "noalias/alias.scope metadata can't be null");
   // Don't duplicate aliasing instruction. For that use a map with a MDNode key
   if (AliasInstMDMap.find(MD) != AliasInstMDMap.end())
@@ -1737,20 +1751,23 @@ SPIRVEntry *SPIRVModuleImpl::getOrAddMemAliasingINTELInst(
 }
 
 // Create AliasDomainDeclINTEL instruction
-SPIRVEntry *SPIRVModuleImpl::getOrAddAliasDomainDeclINTELInst(
-    std::vector<SPIRVId> Args, llvm::MDNode *MD) {
+SPIRVEntry *
+SPIRVModuleImpl::getOrAddAliasDomainDeclINTELInst(std::vector<SPIRVId> Args,
+                                                  llvm::MDNode *MD) {
   return getOrAddMemAliasingINTELInst<SPIRVAliasDomainDeclINTEL>(Args, MD);
 }
 
 // Create AliasScopeDeclINTEL instruction
-SPIRVEntry *SPIRVModuleImpl::getOrAddAliasScopeDeclINTELInst(
-    std::vector<SPIRVId> Args, llvm::MDNode *MD) {
+SPIRVEntry *
+SPIRVModuleImpl::getOrAddAliasScopeDeclINTELInst(std::vector<SPIRVId> Args,
+                                                 llvm::MDNode *MD) {
   return getOrAddMemAliasingINTELInst<SPIRVAliasScopeDeclINTEL>(Args, MD);
 }
 
 // Create AliasScopeListDeclINTEL instruction
-SPIRVEntry *SPIRVModuleImpl::getOrAddAliasScopeListDeclINTELInst(
-    std::vector<SPIRVId> Args, llvm::MDNode *MD) {
+SPIRVEntry *
+SPIRVModuleImpl::getOrAddAliasScopeListDeclINTELInst(std::vector<SPIRVId> Args,
+                                                     llvm::MDNode *MD) {
   return getOrAddMemAliasingINTELInst<SPIRVAliasScopeListDeclINTEL>(Args, MD);
 }
 
diff --git a/llvm-spirv/test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fadd_half.ll b/llvm-spirv/test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fadd_half.ll
@@ -0,0 +1,34 @@
+; RUN: llvm-as < %s -o %t.bc
+; RUN: llvm-spirv --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 %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}
diff --git a/llvm-spirv/test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fsub_half.ll b/llvm-spirv/test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fsub_half.ll
@@ -0,0 +1,39 @@
+; RUN: llvm-as < %s -o %t.bc
+; RUN: llvm-spirv --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 %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 ptr 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{{.*}}(ptr 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}
