[Backport to 16] Add SPV_INTEL_16bit_atomics extension (KhronosGroup#3424)

This continues KhronosGroup#3343 and reflects specification update, including extension renaming.
Specification: intel/llvm#20009

Author: vmaksimo

include/LLVMSPIRVExtensions.inc

@@ -79,7 +79,7 @@ EXT(SPV_INTEL_bfloat16_arithmetic)
 EXT(SPV_INTEL_ternary_bitwise_function)
 EXT(SPV_INTEL_int4)
 EXT(SPV_INTEL_function_variants)
-EXT(SPV_INTEL_shader_atomic_bfloat16)
+EXT(SPV_INTEL_16bit_atomics)
 EXT(SPV_EXT_float8)
 EXT(SPV_INTEL_predicated_io)
 EXT(SPV_INTEL_sigmoid)

lib/SPIRV/libSPIRV/SPIRVEnum.h

@@ -218,6 +218,16 @@ template <> inline void SPIRVMap<SPIRVCapabilityKind, SPIRVCapVec>::init() {
                {CapabilityInt4TypeINTEL, CapabilityCooperativeMatrixKHR});
   ADD_VEC_INIT(internal::CapabilityBFloat16ArithmeticINTEL,
                {CapabilityBFloat16TypeKHR});
+  ADD_VEC_INIT(internal::CapabilityAtomicInt16CompareExchangeINTEL,
+               {CapabilityInt16});
+  ADD_VEC_INIT(internal::CapabilityInt16AtomicsINTEL,
+               {internal::CapabilityAtomicInt16CompareExchangeINTEL});
+  ADD_VEC_INIT(internal::CapabilityAtomicBFloat16LoadStoreINTEL,
+               {CapabilityBFloat16TypeKHR});
+  ADD_VEC_INIT(internal::CapabilityAtomicBFloat16AddINTEL,
+               {CapabilityBFloat16TypeKHR});
+  ADD_VEC_INIT(internal::CapabilityAtomicBFloat16MinMaxINTEL,
+               {CapabilityBFloat16TypeKHR});
 }
 
 template <> inline void SPIRVMap<SPIRVExecutionModelKind, SPIRVCapVec>::init() {

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -2833,8 +2833,16 @@ class SPIRVAtomicInstBase : public SPIRVInstTemplateBase {
     // Besides, OpAtomicCompareExchangeWeak, OpAtomicFlagTestAndSet and
     // OpAtomicFlagClear instructions require the "kernel" capability. But this
     // capability should be added by setting the OpenCL memory model.
-    if (hasType() && getType()->isTypeInt(64))
-      return {CapabilityInt64Atomics};
+    if (hasType()) {
+      if (getType()->isTypeInt(64))
+        return {CapabilityInt64Atomics};
+      if (getType()->isTypeInt(16) &&
+          Module->isAllowedToUseExtension(
+              ExtensionID::SPV_INTEL_16bit_atomics)) {
+        Module->addExtension(ExtensionID::SPV_INTEL_16bit_atomics);
+        return {internal::CapabilityInt16AtomicsINTEL};
+      }
+    }
     return {};
   }
 
@@ -2848,7 +2856,24 @@ class SPIRVAtomicInstBase : public SPIRVInstTemplateBase {
   }
 };
 
-class SPIRVAtomicStoreInst : public SPIRVAtomicInstBase {
+// This specialization will handle smaller set of compare-and-swap instructions
+// that require only one capability. The instructions are: OpAtomicLoad,
+// OpAtomicStore, OpAtomicExchange, OpAtomicCompareExchange and
+// OpAtomicCompareExchangeWeak.
+class SPIRVAtomicCompareExchangeInstructions : public SPIRVAtomicInstBase {
+public:
+  SPIRVCapVec getRequiredCapability() const override {
+    if (hasType() && getType()->isTypeInt(16) &&
+        this->getModule()->isAllowedToUseExtension(
+            ExtensionID::SPV_INTEL_16bit_atomics)) {
+      Module->addExtension(ExtensionID::SPV_INTEL_16bit_atomics);
+      return {internal::CapabilityAtomicInt16CompareExchangeINTEL};
+    }
+    return SPIRVAtomicInstBase::getRequiredCapability();
+  }
+};
+
+class SPIRVAtomicStoreInst : public SPIRVAtomicCompareExchangeInstructions {
 public:
   // Overriding the following method because of 'const'-related
   // issues with overriding getRequiredCapability(). TODO: Resolve.
@@ -2865,7 +2890,7 @@ class SPIRVAtomicFAddEXTInst : public SPIRVAtomicInstBase {
   std::optional<ExtensionID> getRequiredExtension() const override {
     assert(hasType());
     if (getType()->isTypeFloat(16, FPEncodingBFloat16KHR))
-      return ExtensionID::SPV_INTEL_shader_atomic_bfloat16;
+      Module->addExtension(ExtensionID::SPV_INTEL_16bit_atomics);
     if (getType()->isTypeFloat(16))
       return ExtensionID::SPV_EXT_shader_atomic_float16_add;
     return ExtensionID::SPV_EXT_shader_atomic_float_add;
@@ -2890,7 +2915,7 @@ class SPIRVAtomicFMinMaxEXTBase : public SPIRVAtomicInstBase {
 public:
   std::optional<ExtensionID> getRequiredExtension() const override {
     if (getType()->isTypeFloat(16, FPEncodingBFloat16KHR))
-      return ExtensionID::SPV_INTEL_shader_atomic_bfloat16;
+      Module->addExtension(ExtensionID::SPV_INTEL_16bit_atomics);
     return ExtensionID::SPV_EXT_shader_atomic_float_min_max;
   }
 
@@ -2914,10 +2939,6 @@ class SPIRVAtomicFMinMaxEXTBase : public SPIRVAtomicInstBase {
 // Atomic builtins
 _SPIRV_OP(AtomicFlagTestAndSet, true, 6)
 _SPIRV_OP(AtomicFlagClear, false, 4)
-_SPIRV_OP(AtomicLoad, true, 6)
-_SPIRV_OP(AtomicExchange, true, 7)
-_SPIRV_OP(AtomicCompareExchange, true, 9)
-_SPIRV_OP(AtomicCompareExchangeWeak, true, 9)
 _SPIRV_OP(AtomicIIncrement, true, 6)
 _SPIRV_OP(AtomicIDecrement, true, 6)
 _SPIRV_OP(AtomicIAdd, true, 7)
@@ -2934,7 +2955,11 @@ _SPIRV_OP(MemoryBarrier, false, 3)
 #define _SPIRV_OP(x, BaseClass, ...)                                           \
   typedef SPIRVInstTemplate<SPIRV##BaseClass, Op##x, __VA_ARGS__> SPIRV##x;
 // Specialized atomic builtins
+_SPIRV_OP(AtomicLoad, AtomicCompareExchangeInstructions, true, 6)
 _SPIRV_OP(AtomicStore, AtomicStoreInst, false, 5)
+_SPIRV_OP(AtomicExchange, AtomicCompareExchangeInstructions, true, 7)
+_SPIRV_OP(AtomicCompareExchange, AtomicCompareExchangeInstructions, true, 9)
+_SPIRV_OP(AtomicCompareExchangeWeak, AtomicCompareExchangeInstructions, true, 9)
 _SPIRV_OP(AtomicFAddEXT, AtomicFAddEXTInst, true, 7)
 _SPIRV_OP(AtomicFMinEXT, AtomicFMinMaxEXTBase, true, 7)
 _SPIRV_OP(AtomicFMaxEXT, AtomicFMinMaxEXTBase, true, 7)

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -616,9 +616,6 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(CapabilityLongCompositesINTEL, "LongCompositesINTEL");
   add(CapabilityOptNoneEXT, "OptNoneEXT");
   add(CapabilityAtomicFloat16AddEXT, "AtomicFloat16AddEXT");
-  add(internal::CapabilityAtomicBFloat16AddINTEL, "AtomicBFloat16AddINTEL");
-  add(internal::CapabilityAtomicBFloat16MinMaxINTEL,
-      "AtomicBFloat16MinMaxINTEL");
   add(CapabilityDebugInfoModuleINTEL, "DebugInfoModuleINTEL");
   add(CapabilitySplitBarrierINTEL, "SplitBarrierINTEL");
   add(CapabilityGlobalVariableFPGADecorationsINTEL,
@@ -676,6 +673,14 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(internal::CapabilityFloat4E2M1CooperativeMatrixINTEL,
       "Float4E2M1CooperativeMatrixINTEL");
   add(internal::CapabilityFloatConversionsINTEL, "FloatConversionsINTEL");
+  add(internal::CapabilityAtomicBFloat16AddINTEL, "AtomicBFloat16AddINTEL");
+  add(internal::CapabilityAtomicBFloat16MinMaxINTEL,
+      "AtomicBFloat16MinMaxINTEL");
+  add(internal::CapabilityAtomicInt16CompareExchangeINTEL,
+      "AtomicInt16CompareExchangeINTEL");
+  add(internal::CapabilityInt16AtomicsINTEL, "Int16AtomicsINTEL");
+  add(internal::CapabilityAtomicBFloat16LoadStoreINTEL,
+      "AtomicBFloat16LoadStoreINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -118,6 +118,9 @@ enum InternalCapability {
   ICapabilityAtomicBFloat16AddINTEL = 6255,
   ICapabilityAtomicBFloat16MinMaxINTEL = 6256,
   ICapabilityPredicatedIOINTEL = 6257,
+  ICapabilityAtomicInt16CompareExchangeINTEL = 6260,
+  ICapabilityInt16AtomicsINTEL = 6261,
+  ICapabilityAtomicBFloat16LoadStoreINTEL = 6262,
   ICapabilityCooperativeMatrixPrefetchINTEL = 6411,
   ICapabilityComplexFloatMulDivINTEL = 6414,
   ICapabilityTensorFloat32RoundingINTEL = 6425,
@@ -277,6 +280,13 @@ constexpr Capability CapabilityGlobalVariableDecorationsINTEL =
 constexpr Capability CapabilityBFloat16ArithmeticINTEL =
     static_cast<Capability>(ICapabilityBFloat16ArithmeticINTEL);
 
+constexpr Capability CapabilityAtomicInt16CompareExchangeINTEL =
+    static_cast<Capability>(ICapabilityAtomicInt16CompareExchangeINTEL);
+constexpr Capability CapabilityInt16AtomicsINTEL =
+    static_cast<Capability>(ICapabilityInt16AtomicsINTEL);
+constexpr Capability CapabilityAtomicBFloat16LoadStoreINTEL =
+    static_cast<Capability>(ICapabilityAtomicBFloat16LoadStoreINTEL);
+
 } // namespace internal
 } // namespace spv
 

…_shader_atomic_bfloat16/AtomicFAddEXT.ll …SPV_INTEL_16bit_atomics/AtomicFAddEXT.lltest/extensions/INTEL/SPV_INTEL_shader_atomic_bfloat16/AtomicFAddEXT.ll renamed to test/extensions/INTEL/SPV_INTEL_16bit_atomics/AtomicFAddEXT.ll test/extensions/INTEL/SPV_INTEL_shader_atomic_bfloat16/AtomicFAddEXT.ll renamed to test/extensions/INTEL/SPV_INTEL_16bit_atomics/AtomicFAddEXT.ll

@@ -1,5 +1,5 @@
 ; RUN: llvm-as %s -o %t.bc
-; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_shader_atomic_bfloat16,+SPV_KHR_bfloat16 -o %t.spv
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_16bit_atomics,+SPV_KHR_bfloat16,+SPV_EXT_shader_atomic_float_add -o %t.spv
 ; RUN: llvm-spirv -to-text %t.spv -o %t.spt
 ; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
 
@@ -11,7 +11,7 @@ target triple = "spir64-unknown-unknown"
 
 ; CHECK-SPIRV-DAG: Capability AtomicBFloat16AddINTEL
 ; CHECK-SPIRV-DAG: Capability BFloat16TypeKHR
-; CHECK-SPIRV-DAG: Extension "SPV_INTEL_shader_atomic_bfloat16"
+; CHECK-SPIRV-DAG: Extension "SPV_INTEL_16bit_atomics"
 ; CHECK-SPIRV-DAG: Extension "SPV_KHR_bfloat16"
 
 ; CHECK-SPIRV: TypeFloat [[BFLOAT:[0-9]+]] 16 0

…_shader_atomic_bfloat16/AtomicFMaxEXT.ll …SPV_INTEL_16bit_atomics/AtomicFMaxEXT.lltest/extensions/INTEL/SPV_INTEL_shader_atomic_bfloat16/AtomicFMaxEXT.ll renamed to test/extensions/INTEL/SPV_INTEL_16bit_atomics/AtomicFMaxEXT.ll test/extensions/INTEL/SPV_INTEL_shader_atomic_bfloat16/AtomicFMaxEXT.ll renamed to test/extensions/INTEL/SPV_INTEL_16bit_atomics/AtomicFMaxEXT.ll

@@ -1,5 +1,5 @@
 ; RUN: llvm-as %s -o %t.bc
-; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_shader_atomic_bfloat16,+SPV_KHR_bfloat16 -o %t.spv
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_16bit_atomics,+SPV_KHR_bfloat16,+SPV_EXT_shader_atomic_float_min_max -o %t.spv
 ; RUN: llvm-spirv -to-text %t.spv -o %t.spt
 ; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
 
@@ -11,7 +11,7 @@ target triple = "spir64-unknown-unknown"
 
 ; CHECK-SPIRV-DAG: Capability AtomicBFloat16MinMaxINTEL
 ; CHECK-SPIRV-DAG: Capability BFloat16TypeKHR
-; CHECK-SPIRV-DAG: Extension "SPV_INTEL_shader_atomic_bfloat16"
+; CHECK-SPIRV-DAG: Extension "SPV_INTEL_16bit_atomics"
 ; CHECK-SPIRV-DAG: Extension "SPV_KHR_bfloat16"
 
 ; CHECK-SPIRV: TypeFloat [[BFLOAT:[0-9]+]] 16 0

…_shader_atomic_bfloat16/AtomicFMinEXT.ll …SPV_INTEL_16bit_atomics/AtomicFMinEXT.lltest/extensions/INTEL/SPV_INTEL_shader_atomic_bfloat16/AtomicFMinEXT.ll renamed to test/extensions/INTEL/SPV_INTEL_16bit_atomics/AtomicFMinEXT.ll test/extensions/INTEL/SPV_INTEL_shader_atomic_bfloat16/AtomicFMinEXT.ll renamed to test/extensions/INTEL/SPV_INTEL_16bit_atomics/AtomicFMinEXT.ll

@@ -1,5 +1,5 @@
 ; RUN: llvm-as %s -o %t.bc
-; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_shader_atomic_bfloat16,+SPV_KHR_bfloat16 -o %t.spv
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_16bit_atomics,+SPV_KHR_bfloat16,+SPV_EXT_shader_atomic_float_min_max -o %t.spv
 ; RUN: llvm-spirv -to-text %t.spv -o %t.spt
 ; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
 
@@ -11,7 +11,7 @@ target triple = "spir64-unknown-unknown"
 
 ; CHECK-SPIRV-DAG: Capability AtomicBFloat16MinMaxINTEL
 ; CHECK-SPIRV-DAG: Capability BFloat16TypeKHR
-; CHECK-SPIRV-DAG: Extension "SPV_INTEL_shader_atomic_bfloat16"
+; CHECK-SPIRV-DAG: Extension "SPV_INTEL_16bit_atomics"
 ; CHECK-SPIRV-DAG: Extension "SPV_KHR_bfloat16"
 
 ; CHECK-SPIRV: TypeFloat [[BFLOAT:[0-9]+]] 16 0

test/extensions/INTEL/SPV_INTEL_16bit_atomics/AtomicInt16CompareExchangeCapability.ll

@@ -0,0 +1,55 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv --spirv-ext=+SPV_INTEL_16bit_atomics
+; RUN: llvm-spirv -to-text %t.spv -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv -r --spirv-target-env=CL2.0 %t.spv -o %t.rev.bc
+; RUN: llvm-dis %t.rev.bc 
+; RUN: FileCheck < %t.rev.ll %s --check-prefixes=CHECK-LLVM
+
+; Check that without extension we don't use its capabilities - there is no
+; limitation on using i16 with atomic instruction in the core specification.
+; RUN: llvm-spirv %t.bc -o %t.noext.spv
+; RUN: spirv-val %t.noext.spv
+; RUN: llvm-spirv -to-text %t.noext.spv -o %t.noext.spt
+; RUN: FileCheck < %t.noext.spt %s --check-prefix=CHECK-SPIRV-NOEXT
+
+; CHECK-SPIRV: Capability Int16
+; CHECK-SPIRV: Capability AtomicInt16CompareExchangeINTEL
+; CHECK-SPIRV-NOT: Capability Int16AtomicsINTEL
+; CHECK-SPIRV: Extension "SPV_INTEL_16bit_atomics"
+
+; CHECK-SPIRV-NOEXT: Capability Int16
+; CHECK-SPIRV-NOEXT-NOT: Capability AtomicInt16CompareExchangeINTEL
+; CHECK-SPIRV-NOEXT-NOT: Capability Int16AtomicsINTEL
+; CHECK-SPIRV-NOEXT-NOT: Extension "SPV_INTEL_16bit_atomics"
+
+; CHECK-SPIRV-DAG: Constant [[#]] [[#DeviceScope:]] 1
+; CHECK-SPIRV-DAG: Constant [[#]] [[#Release:]] 4
+; CHECK-SPIRV-DAG: Constant [[#]] [[#SequentiallyConsistent:]] 16
+; CHECK-SPIRV-DAG: Constant [[#]] [[#Acquire:]] 2
+
+; CHECK-LLVM: call spir_func void @_Z21atomic_store_explicitPU3AS4VU7_Atomicss12memory_order12memory_scope
+; CHECK-LLVM: call spir_func i16 @_Z20atomic_load_explicitPU3AS4VU7_Atomics12memory_order12memory_scope
+; CHECK-LLVM: call spir_func i16 @_Z24atomic_exchange_explicitPU3AS4VU7_Atomicss12memory_order12memory_scope
+; CHECK-LLVM: call spir_func i1 @_Z39atomic_compare_exchange_strong_explicitPU3AS4VU7_AtomicsPU3AS4ss12memory_orderS4_12memory_scope
+
+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"
+
+@ui = common dso_local addrspace(1) global i16 0, align 4
+; Function Attrs: nounwind
+define dso_local spir_func void @test() {
+entry:
+; CHECK-SPIRV: Variable [[#]] [[#PTR:]] 7
+  %0 = alloca i16
+; CHECK-SPIRV: AtomicStore [[#PTR]] [[#DeviceScope]] [[#Release]] [[#]]
+  store atomic i16 0, ptr %0 release, align 4
+; CHECK-SPIRV: AtomicLoad [[#]] [[#]] [[#PTR]] [[#DeviceScope]] [[#Acquire]]
+  %1 = load atomic i16, ptr %0 acquire, align 4
+; CHECK-SPIRV: AtomicExchange [[#]] [[#]] [[#]] [[#DeviceScope]] [[#]] {{.+}}
+  %2 = atomicrmw xchg ptr addrspace(1) @ui, i16 42 acq_rel
+; CHECK-SPIRV: AtomicCompareExchange [[#]] [[#]] [[#]] [[#DeviceScope]] [[#SequentiallyConsistent]] [[#Acquire]] {{.+}}
+  %3 = cmpxchg ptr %0, i16 128, i16 456 seq_cst acquire
+  ret void
+}

test/extensions/INTEL/SPV_INTEL_16bit_atomics/Int16AtomicsINTEL.ll

@@ -0,0 +1,44 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv --spirv-ext=+SPV_INTEL_16bit_atomics
+; RUN: llvm-spirv -to-text %t.spv -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv -r --spirv-target-env=CL2.0 %t.spv -o %t.rev.bc
+; RUN: llvm-dis %t.rev.bc 
+; RUN: FileCheck < %t.rev.ll %s --check-prefixes=CHECK-LLVM
+
+; RUN: llvm-spirv -r --spirv-target-env="SPV-IR" %t.spv -o %t.rev.bc
+; RUN: llvm-dis %t.rev.bc 
+; RUN: FileCheck < %t.rev.ll %s --check-prefixes=CHECK-LLVM-SPV-IR
+
+; Check that without extension we don't use its capabilities - there is no
+; limitation on using i16 with atomic instruction in the core specification.
+; RUN: llvm-spirv %t.bc -o %t.noext.spv
+; RUN: spirv-val %t.noext.spv
+; RUN: llvm-spirv -to-text %t.noext.spv -o %t.noext.spt
+; RUN: FileCheck < %t.noext.spt %s --check-prefix=CHECK-SPIRV-NOEXT
+
+; CHECK-SPIRV: Capability Int16
+; CHECK-SPIRV: Capability AtomicInt16CompareExchangeINTEL
+; CHECK-SPIRV: Capability Int16AtomicsINTEL
+; CHECK-SPIRV: Extension "SPV_INTEL_16bit_atomics"
+; CHECK-SPIRV: AtomicOr
+
+; CHECK-SPIRV-NOEXT: Capability Int16
+; CHECK-SPIRV-NOEXT-NOT: Capability AtomicInt16CompareExchangeINTEL
+; CHECK-SPIRV-NOEXT-NOT: Capability Int16AtomicsINTEL
+; CHECK-SPIRV-NOEXT-NOT: Extension "SPV_INTEL_16bit_atomics"
+
+; CHECK-LLVM: call spir_func i16 @_Z24atomic_fetch_or_explicitPU3AS4VU7_Atomicss12memory_order12memory_scope
+; CHECK-LLVM-SPV-IR: call spir_func i16 @_Z16__spirv_AtomicOrPU3AS1siis
+
+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"
+
+@ui = common dso_local addrspace(1) global i16 0, align 4
+
+define dso_local spir_func void @test() {
+entry:
+  %0 = atomicrmw or ptr addrspace(1) @ui, i16 42 release
+  ret void
+}