Fix legalization of `kIROp_GetLegalizedSPIRVGlobalParamAddr`. (#4141)

4 files changed, 81 insertions, 4 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 6b3c5db59..b3e323bfc 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -7432,11 +7432,22 @@ void InterlockedAdd(__ref uint dest, uint value, out uint original_value)
 }
 
 [ForceInline]
+[require(cuda_glsl_hlsl_spirv, atomic_glsl_hlsl_cuda)]
 void InterlockedAdd(__ref  int64_t dest,  int64_t value)
 {
     __target_switch
     {
     case hlsl: __intrinsic_asm "InterlockedAdd";
+    case cuda: __intrinsic_asm "atomicAdd((uint64_t*)$0, $1)";
+    case glsl:
+        __requireGLSLExtension("GL_EXT_shader_atomic_int64");
+        __intrinsic_asm "$atomicAdd($A, $1)";
+    case spirv:
+        spirv_asm
+        {
+            OpCapability Int64Atomics;
+            result:$$int64_t = OpAtomicIAdd &dest Device None $value;
+        };
     }
 }
 
@@ -7446,27 +7457,61 @@ void InterlockedAdd(__ref  int64_t dest,  int64_t value, out  int64_t original_v
     __target_switch
     {
     case hlsl: __intrinsic_asm "InterlockedAdd";
+    case cuda: __intrinsic_asm "atomicAdd((uint64_t*)$0, $1)";
+    case glsl:
+        __requireGLSLExtension("GL_EXT_shader_atomic_int64");
+        __intrinsic_asm "$atomicAdd($A, $1)";
+    case spirv:
+        spirv_asm
+        {
+            OpCapability Int64Atomics;
+            %origin:$$int64_t = OpAtomicIAdd &dest Device None $value;
+            OpStore &original_value %origin
+        };
     }
 }
 
 [ForceInline]
-void InterlockedAdd(__ref  uint64_t dest,  uint64_t value)
+[require(cuda_glsl_hlsl_spirv, atomic_glsl_hlsl_cuda)]
+void InterlockedAdd(__ref uint64_t dest, uint64_t value)
 {
     __target_switch
     {
     case hlsl: __intrinsic_asm "InterlockedAdd";
+    case cuda: __intrinsic_asm "atomicAdd($0, $1)";
+    case glsl:
+        __requireGLSLExtension("GL_EXT_shader_atomic_int64");
+        __intrinsic_asm "$atomicAdd($A, $1)";
+    case spirv:
+        spirv_asm
+        {
+            OpCapability Int64Atomics;
+            result:$$uint64_t = OpAtomicIAdd &dest Device None $value;
+        };
     }
 }
 
 [ForceInline]
-void InterlockedAdd(__ref  uint64_t dest,  uint64_t value, out  uint64_t original_value)
+void InterlockedAdd(__ref uint64_t dest, uint64_t value, out uint64_t original_value)
 {
     __target_switch
     {
     case hlsl: __intrinsic_asm "InterlockedAdd";
+    case cuda: __intrinsic_asm "atomicAdd($0, $1)";
+    case glsl:
+        __requireGLSLExtension("GL_EXT_shader_atomic_int64");
+        __intrinsic_asm "$atomicAdd($A, $1)";
+    case spirv:
+        spirv_asm
+        {
+            OpCapability Int64Atomics;
+            %origin:$$uint64_t = OpAtomicIAdd &dest Device None $value;
+            OpStore &original_value %origin
+        };
     }
 }
 
+
 __glsl_version(430)
 [require(cuda_glsl_hlsl_spirv, atomic_glsl_hlsl_cuda)]
 void InterlockedAnd(__ref  int dest,  int value)
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 48a46f962..e863f279e 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -221,7 +221,7 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
 
             if(user->getOp() == kIROp_GetLegalizedSPIRVGlobalParamAddr)
             {
-                user->replaceUsesWith(use->get());
+                user->replaceUsesWith(addr);
                 user->removeAndDeallocate();
             }
             else if((as<IRGetElement>(user) || as<IRFieldExtract>(user)) &&
@@ -251,10 +251,13 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
             {
                 // Skip load's for referenced `Input` variables since a ref implies
                 // passing as is, which needs to be a pointer (pass as is).
-                if (user->getDataType() 
+                if (user->getDataType()
                     && user->getDataType()->getOp() == kIROp_RefType
                     && storageClass == SpvStorageClassInput)
+                {
+                    builder.replaceOperand(use, addr);
                     continue;
+                }
                 // If this is being used in an asm block, insert the load to
                 // just prior to the block.
                 const auto asmBlock = spirvAsmOperand->getAsmBlock();
@@ -276,7 +279,9 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
         }
 
         for (auto i : instsToRemove)
+        {
             i->removeAndDeallocate();
+        }
     }
 
     // Returns true if the given type that should be decorated as in `UniformConstant` address space.
diff --git a/tests/spirv/image-atomic-array-2.slang b/tests/spirv/image-atomic-array-2.slang
new file mode 100644
index 000000000..f5dc2b101
--- /dev/null
+++ b/tests/spirv/image-atomic-array-2.slang
@@ -0,0 +1,13 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+
+// CHECK: %{{.*}} = OpImageTexelPointer %{{.*}} %{{.*}} %{{.*}} %int_0
+// CHECK: %{{.*}} = OpAtomicIAdd %ulong %{{.*}} %uint_1 %uint_0 %ulong_0
+
+[[vk::binding(0,0)]] RWTexture2D<uint64_t> tex;
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void main()
+{
+    InterlockedAdd(tex[uint2(0,0)], uint64_t(0));
+}
\ No newline at end of file
diff --git a/tests/spirv/image-atomic-array.slang b/tests/spirv/image-atomic-array.slang
new file mode 100644
index 000000000..6cff48f6d
--- /dev/null
+++ b/tests/spirv/image-atomic-array.slang
@@ -0,0 +1,14 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+
+import glsl;
+[[vk::binding(0,0)]] RWTexture2D<uint64_t> tex[];
+
+// CHECK: OpEntryPoint
+// CHECK: %{{.*}} = OpImageTexelPointer %{{.*}} %{{.*}} %{{.*}} %int_0
+// CHECK: %{{.*}} = OpAtomicIAdd %ulong %{{.*}} %uint_1 %uint_0 %ulong_0
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void main()
+{
+    imageAtomicAdd(tex[0], int2(0,0), uint64_t(0));
+}
\ No newline at end of file