[spirv] Support 64-bit interlocked intrinsics. (#3321)

* [spirv] Support 64-bit interlocked intrinsics.

* Fix build error.

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -597,6 +597,14 @@ bool CapabilityVisitor::visit(SpirvExtInst *instr) {
   return visitInstruction(instr);
 }
 
+bool CapabilityVisitor::visit(SpirvAtomic *instr) {
+  if (instr->hasValue() && SpirvType::isOrContainsType<IntegerType, 64>(
+                               instr->getValue()->getResultType())) {
+    addCapability(spv::Capability::Int64Atomics, instr->getSourceLocation());
+  }
+  return true;
+}
+
 bool CapabilityVisitor::visit(SpirvDemoteToHelperInvocationEXT *inst) {
   addCapability(spv::Capability::DemoteToHelperInvocationEXT,
                 inst->getSourceLocation());

tools/clang/lib/SPIRV/CapabilityVisitor.h

@@ -37,6 +37,7 @@ public:
   bool visit(SpirvImageSparseTexelsResident *) override;
   bool visit(SpirvExtInstImport *) override;
   bool visit(SpirvExtInst *) override;
+  bool visit(SpirvAtomic *) override;
   bool visit(SpirvDemoteToHelperInvocationEXT *) override;
   bool visit(SpirvReadClock *) override;
 

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -3046,7 +3046,6 @@ SpirvInstruction *SpirvEmitter::processRWByteAddressBufferAtomicMethods(
   // The signature of RWByteAddressBuffer atomic methods are largely:
   // void Interlocked*(in UINT dest, in UINT value);
   // void Interlocked*(in UINT dest, in UINT value, out UINT original_value);
-
   const auto *object = expr->getImplicitObjectArgument();
   auto *objectInfo = loadIfAliasVarRef(object);
 
@@ -7656,7 +7655,8 @@ SpirvEmitter::processIntrinsicInterlockedMethod(const CallExpr *expr,
     const Expr *valueExpr = callExpr->getArg(argIndex);
     if (const auto *castExpr = dyn_cast<ImplicitCastExpr>(valueExpr))
       if (castExpr->getCastKind() == CK_IntegralCast &&
-          castExpr->getSubExpr()->getType() == baseType)
+          castExpr->getSubExpr()->getType()->getCanonicalTypeUnqualified() ==
+              baseType)
         valueExpr = castExpr->getSubExpr();
 
     auto *argInstr = doExpr(valueExpr);

tools/clang/test/CodeGenSPIRV/intrinsics.64bit-interlocked-methods.cs.hlsl

@@ -0,0 +1,80 @@
+// Run: %dxc -T cs_6_0 -E main
+
+groupshared int64_t dest_i;
+groupshared uint64_t dest_u;
+
+RWStructuredBuffer<uint64_t> buff;
+
+RWStructuredBuffer<uint64_t> getDest() {
+  return buff;
+}
+
+[numthreads(1,1,1)]
+void main()
+{
+  uint64_t original_u_val;
+  int64_t original_i_val;
+
+  int64_t   val1_i64;
+  int64_t   val2_i64;
+  uint64_t  val3_u64;
+
+  ////////////////////////////////////////////////////////
+  ///////      Test all Interlocked* functions      //////
+  ////////////////////////////////////////////////////////
+
+// CHECK: OpCapability Int64
+// CHECK: OpCapability Int64Atomics
+
+// CHECK:        [[val1_i64:%\d+]] = OpLoad %long %val1_i64
+// CHECK-NEXT: [[atomic_add:%\d+]] = OpAtomicIAdd %long %dest_i %uint_1 %uint_0 [[val1_i64]]
+// CHECK-NEXT:                       OpStore %original_i_val [[atomic_add]]
+  InterlockedAdd(dest_i, val1_i64, original_i_val);
+
+// CHECK:      [[fn_call_result:%\d+]] = OpFunctionCall %_ptr_Uniform_type_RWStructuredBuffer_uint64 %getDest
+// CHECK-NEXT:            [[ptr:%\d+]] = OpAccessChain %_ptr_Uniform_ulong [[fn_call_result]] %int_0 %uint_0
+// CHECK-NEXT:       [[val3_u64:%\d+]] = OpLoad %ulong %val3_u64
+// CHECK-NEXT:     [[atomic_add:%\d+]] = OpAtomicIAdd %ulong [[ptr]] %uint_1 %uint_0 [[val3_u64]]
+// CHECK-NEXT:                           OpStore %original_u_val [[atomic_add]]
+  InterlockedAdd(getDest()[0], val3_u64, original_u_val);
+
+// CHECK:        [[val3_u64:%\d+]] = OpLoad %ulong %val3_u64
+// CHECK-NEXT: [[atomic_and:%\d+]] = OpAtomicAnd %ulong %dest_u %uint_1 %uint_0 [[val3_u64]]
+// CHECK-NEXT:                       OpStore %original_u_val [[atomic_and]]
+  InterlockedAnd(dest_u, val3_u64,  original_u_val);
+
+// CHECK:        [[val1_i64:%\d+]] = OpLoad %long %val1_i64
+// CHECK-NEXT: [[atomic_max:%\d+]] = OpAtomicSMax %long %dest_i %uint_1 %uint_0 [[val1_i64]]
+// CHECK-NEXT:                       OpStore %original_i_val [[atomic_max]]
+  InterlockedMax(dest_i, val1_i64,  original_i_val);
+
+// CHECK:        [[val3_u64:%\d+]] = OpLoad %ulong %val3_u64
+// CHECK-NEXT: [[atomic_min:%\d+]] = OpAtomicUMin %ulong %dest_u %uint_1 %uint_0 [[val3_u64]]
+// CHECK-NEXT:                       OpStore %original_u_val [[atomic_min]]
+  InterlockedMin(dest_u, val3_u64,  original_u_val);
+
+// CHECK:       [[val2_i64:%\d+]] = OpLoad %long %val2_i64
+// CHECK-NEXT: [[atomic_or:%\d+]] = OpAtomicOr %long %dest_i %uint_1 %uint_0 [[val2_i64:%\d+]]
+// CHECK-NEXT:                      OpStore %original_i_val [[atomic_or]]
+  InterlockedOr (dest_i, val2_i64, original_i_val);
+
+// CHECK:        [[val3_u64:%\d+]] = OpLoad %ulong %val3_u64
+// CHECK-NEXT: [[atomic_xor:%\d+]] = OpAtomicXor %ulong %dest_u %uint_1 %uint_0 [[val3_u64]]
+// CHECK-NEXT:                       OpStore %original_u_val [[atomic_xor]]
+  InterlockedXor(dest_u, val3_u64,  original_u_val);
+
+// CHECK:      [[val1_i64:%\d+]] = OpLoad %long %val1_i64
+// CHECK-NEXT: [[val2_i64:%\d+]] = OpLoad %long %val2_i64
+// CHECK-NEXT:          {{%\d+}} = OpAtomicCompareExchange %long %dest_i %uint_1 %uint_0 %uint_0 [[val2_i64]] [[val1_i64]]
+  InterlockedCompareStore(dest_i, val1_i64, val2_i64);
+
+// CHECK:      [[ace:%\d+]] = OpAtomicCompareExchange %ulong %dest_u %uint_1 %uint_0 %uint_0 %ulong_20 %ulong_15
+// CHECK-NEXT:                OpStore %original_u_val [[ace]]
+  InterlockedCompareExchange(dest_u, 15u, 20u, original_u_val);
+
+// CHECK:      [[val2_i64:%\d+]] = OpLoad %long %val2_i64
+// CHECK-NEXT:       [[ae:%\d+]] = OpAtomicExchange %long %dest_i %uint_1 %uint_0 [[val2_i64]]
+// CHECK-NEXT:                     OpStore %original_i_val [[ae]]
+  InterlockedExchange(dest_i, val2_i64, original_i_val);
+}
+

tools/clang/test/CodeGenSPIRV/intrinsics.64bit-interlocked-methods.ps.hlsl

@@ -0,0 +1,33 @@
+// Run: %dxc -T ps_6_0 -E main
+
+// CHECK: OpCapability Int64
+// CHECK: OpCapability Int64Atomics
+
+// 64-bit atomics were introduced in SM 6.6.
+RWStructuredBuffer<uint64_t> rwb_u64;
+RWStructuredBuffer<int64_t> rwb_i64;
+
+void main() {
+  uint64_t out_u64;
+  uint64_t in_u64;
+
+  int64_t out_i64;
+  int64_t in_i64;
+
+  uint  index;
+
+// CHECK:           [[index:%\d+]] = OpLoad %uint %index
+// CHECK-NEXT:        [[ptr:%\d+]] = OpAccessChain %_ptr_Uniform_ulong %rwb_u64 %int_0 [[index]]
+// CHECK-NEXT:     [[in_u64:%\d+]] = OpLoad %ulong %in_u64
+// CHECK-NEXT: [[atomic_add:%\d+]] = OpAtomicIAdd %ulong [[ptr]] %uint_1 %uint_0 [[in_u64]]
+// CHECK-NEXT:                       OpStore %out_u64 [[atomic_add]]
+  InterlockedAdd(rwb_u64[index], in_u64, out_u64);
+
+// CHECK:           [[index:%\d+]] = OpLoad %uint %index
+// CHECK-NEXT:        [[ptr:%\d+]] = OpAccessChain %_ptr_Uniform_long %rwb_i64 %int_0 [[index]]
+// CHECK-NEXT:     [[in_i64:%\d+]] = OpLoad %long %in_i64
+// CHECK-NEXT: [[atomic_add:%\d+]] = OpAtomicIAdd %long [[ptr]] %uint_1 %uint_0 [[in_i64]]
+// CHECK-NEXT:                       OpStore %out_i64 [[atomic_add]]
+  InterlockedAdd(rwb_i64[index], in_i64, out_i64);
+}
+

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1131,6 +1131,13 @@ TEST_F(FileTest, IntrinsicsIsFinite) {
 TEST_F(FileTest, IntrinsicsInterlockedMethodsPS) {
   runFileTest("intrinsics.interlocked-methods.ps.hlsl");
 }
+TEST_F(FileTest, Intrinsics64BitInterlockedMethodsPS) {
+  runFileTest("intrinsics.64bit-interlocked-methods.ps.hlsl");
+}
+TEST_F(FileTest, Intrinsics64BitInterlockedMethodsCS) {
+  setBeforeHLSLLegalization();
+  runFileTest("intrinsics.64bit-interlocked-methods.cs.hlsl");
+}
 TEST_F(FileTest, IntrinsicsInterlockedMethodsCS) {
   runFileTest("intrinsics.interlocked-methods.cs.hlsl");
 }