Add validation for destination of atomic operations (#6093)

* Reimplement the GLSL atomic* functions in terms of __intrinsic_op

Many of these functions map directly to atomic IR instructions.
The functions taking atomic_uint are left as they are.

This helps to address #5989, since the destination pointer type validation can then be
written only for the atomic IR instructions.

* Add validation for atomic operations

Diagnose an error if the destination of the atomic operation is not appropriate, where
appropriate means it's either:
- 'groupshared'
- from a device buffer

This closes #5989.

* Add tests for GLSL atomics destination validation

Attempting to use the GLSL atomic functions on destinations that are neither groupshared
nor from a device buffer should fail with the following error:

error 41403: cannot perform atomic operation because destination is neither groupshared
             nor from a device buffer.

* Validate atomic operations after address space specialization

Address space specialization for SPIR-V is not done as part of `linkAndOptimizeIR`, as it
is for e.g. Metal, so opt out and add a separate call for SPIR-V.

* Allow unchecked in/inout parameters for non-SPIRV targets

* Clean up callees left without uses during address space specialziation

* format code

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>
Co-authored-by: Yong He <yonghe@outlook.com>

9 files changed, 225 insertions, 75 deletions

diff --git a/source/slang/glsl.meta.slang b/source/slang/glsl.meta.slang
index 1b44c7825..dd1c5a907 100644
--- a/source/slang/glsl.meta.slang
+++ b/source/slang/glsl.meta.slang
@@ -8470,21 +8470,24 @@ for (const auto& item : atomics)
 __glsl_version(430)
 [ForceInline]
 [require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicAdd))
+$(item.name) atomicAddWithOrder(inout $(item.name) mem, $(item.name) data, MemoryOrder order);
+
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
 public $(item.name) atomicAdd(inout $(item.name) mem, $(item.name) data)
 {
     typeRequireChecks_atomic_using_float1_tier<$(item.name)>();
     typeRequireChecks_atomic_using_add<$(item.name)>();
-    __target_switch
-    {
-    case glsl: __intrinsic_asm "atomicAdd($0, $1)";
-    case spirv:
-        return spirv_asm
-        {
-            OpAtomic$(item.classType)Add$(item.suffix) $$$(item.name) result &mem Device UniformMemory $data
-        };
-    }
+    return atomicAddWithOrder(mem, data, MemoryOrder::Relaxed);
 }
 
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicMin))
+$(item.name) atomicMinWithOrder(inout $(item.name) mem, $(item.name) data, MemoryOrder order);
 
 __glsl_version(430)
 [ForceInline]
@@ -8493,17 +8496,14 @@ public $(item.name) atomicMin(inout $(item.name) mem, $(item.name) data)
 {
     typeRequireChecks_atomic_using_float2_tier<$(item.name)>();
     typeRequireChecks_atomic_using_MinMax<$(item.name)>();
-    __target_switch
-    {
-    case glsl: __intrinsic_asm "atomicMin($0, $1)";
-    case spirv:
-        return spirv_asm
-        {
-            OpAtomic$(item.subclassType)Min$(item.suffix) $$$(item.name) result &mem Device UniformMemory $data
-        };
-    }
+    return atomicMinWithOrder(mem, data, MemoryOrder::Relaxed);
 }
 
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicMax))
+$(item.name) atomicMaxWithOrder(inout $(item.name) mem, $(item.name) data, MemoryOrder order);
 
 __glsl_version(430)
 [ForceInline]
@@ -8512,17 +8512,14 @@ public $(item.name) atomicMax(inout $(item.name) mem, $(item.name) data)
 {
     typeRequireChecks_atomic_using_float2_tier<$(item.name)>();
     typeRequireChecks_atomic_using_MinMax<$(item.name)>();
-    __target_switch
-    {
-    case glsl: __intrinsic_asm "atomicMax($0, $1)";
-    case spirv:
-        return spirv_asm
-        {
-            OpAtomic$(item.subclassType)Max$(item.suffix) $$$(item.name) result &mem Device UniformMemory $data
-        };
-    }
+    return atomicMaxWithOrder(mem, data, MemoryOrder::Relaxed);
 }
 
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicExchange))
+$(item.name) atomicExchangeWithOrder(inout $(item.name) mem, $(item.name) data, MemoryOrder order);
 
 __glsl_version(430)
 [ForceInline]
@@ -8530,15 +8527,7 @@ __glsl_version(430)
 public $(item.name) atomicExchange(inout $(item.name) mem, $(item.name) data)
 {
     typeRequireChecks_atomic_using_float1_tier<$(item.name)>();
-    __target_switch
-    {
-    case glsl: __intrinsic_asm "atomicExchange($0, $1)";
-    case spirv:
-        return spirv_asm
-        {
-            OpAtomicExchange $$$(item.name) result &mem Device UniformMemory $data
-        };
-    }
+    return atomicExchangeWithOrder(mem, data, MemoryOrder::Relaxed);
 }
 
 ${{{{
@@ -8550,24 +8539,24 @@ if(item.isFloat)
 __glsl_version(430)
 [ForceInline]
 [require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicAnd))
+$(item.name) atomicAndWithOrder(inout $(item.name) mem, $(item.name) data, MemoryOrder order);
+
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
 public $(item.name) atomicAnd(inout $(item.name) mem, $(item.name) data)
 {
     typeRequireChecks_atomic_using_float0_tier<$(item.name)>();
     typeRequireChecks_atomic_using_Logical_CAS<$(item.name)>();
-    __target_switch
-    {
-    case glsl:
-    {
-        __intrinsic_asm "atomicAnd($0, $1)";
-    }
-    case spirv:
-        return spirv_asm
-        {
-            OpAtomicAnd $$$(item.name) result &mem Device UniformMemory $data
-        };
-    }
+    return atomicAndWithOrder(mem, data, MemoryOrder::Relaxed);
 }
 
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicOr))
+$(item.name) atomicOrWithOrder(inout $(item.name) mem, $(item.name) data, MemoryOrder order);
 
 __glsl_version(430)
 [ForceInline]
@@ -8576,36 +8565,31 @@ public $(item.name) atomicOr(inout $(item.name) mem, $(item.name) data)
 {
     typeRequireChecks_atomic_using_float0_tier<$(item.name)>();
     typeRequireChecks_atomic_using_Logical_CAS<$(item.name)>();
-    __target_switch
-    {
-    case glsl: __intrinsic_asm "atomicOr($0, $1)";
-    case spirv:
-        return spirv_asm
-        {
-            OpAtomicOr $$$(item.name) result &mem Device UniformMemory $data
-        };
-    }
+    return atomicOrWithOrder(mem, data, MemoryOrder::Relaxed);
 }
 
 
 __glsl_version(430)
 [ForceInline]
 [require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicXor))
+$(item.name) atomicXorWithOrder(inout $(item.name) mem, $(item.name) data, MemoryOrder order);
+
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
 public $(item.name) atomicXor(inout $(item.name) mem, $(item.name) data)
 {
     typeRequireChecks_atomic_using_float0_tier<$(item.name)>();
     typeRequireChecks_atomic_using_Logical_CAS<$(item.name)>();
-    __target_switch
-    {
-    case glsl: __intrinsic_asm "atomicXor($0, $1)";
-    case spirv:
-        return spirv_asm
-        {
-            OpAtomicXor $$$(item.name) result &mem Device UniformMemory $data
-        };
-    }
+    return atomicXorWithOrder(mem, data, MemoryOrder::Relaxed);
 }
 
+__glsl_version(430)
+[ForceInline]
+[require(glsl_spirv, atomic_glsl)]
+__intrinsic_op($(kIROp_AtomicCompareExchange))
+$(item.name) atomicCompSwapWithOrder(inout $(item.name) mem, $(item.name) compare, $(item.name) data, MemoryOrder successOrder, MemoryOrder failOrder);
 
 __glsl_version(430)
 [ForceInline]
@@ -8614,15 +8598,7 @@ public $(item.name) atomicCompSwap(inout $(item.name) mem, $(item.name) compare,
 {
     typeRequireChecks_atomic_using_float0_tier<$(item.name)>();
     typeRequireChecks_atomic_using_Logical_CAS<$(item.name)>();
-    __target_switch
-    {
-    case glsl: __intrinsic_asm "atomicCompSwap($0, $1, $2)";
-    case spirv:
-        return spirv_asm
-        {
-            result:$$$(item.name) = OpAtomicCompareExchange &mem Device None None $data $compare
-        };
-    }
+    return atomicCompSwapWithOrder(mem, compare, data, MemoryOrder::Relaxed, MemoryOrder::Relaxed);
 }
 
 ${{{{
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 915ed9b42..fa96b89f9 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -2255,6 +2255,13 @@ DIAGNOSTIC(
     multiSampledTextureDoesNotAllowWrites,
     "cannot write to a multisampled texture with target '$0'.")
 
+DIAGNOSTIC(
+    41403,
+    Error,
+    invalidAtomicDestinationPointer,
+    "cannot perform atomic operation because destination is neither groupshared nor from a device "
+    "buffer.")
+
 //
 // 5xxxx - Target code generation.
 //
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index ea209b598..cf4070df4 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -1320,6 +1320,14 @@ Result linkAndOptimizeIR(
             byteAddressBufferOptions);
     }
 
+    // For SPIR-V, this function is called elsewhere, so that it can happen after address space
+    // specialization
+    if (target != CodeGenTarget::SPIRV && target != CodeGenTarget::SPIRVAssembly)
+    {
+        bool skipFuncParamValidation = true;
+        validateAtomicOperations(skipFuncParamValidation, sink, irModule->getModuleInst());
+    }
+
     // For CUDA targets only, we will need to turn operations
     // the implicitly reference the "active mask" into ones
     // that use (and pass around) an explicit mask instead.
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index a311e024c..caee45ba8 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -2502,6 +2502,13 @@ struct IRGetElementPtr : IRInst
     IRInst* getIndex() { return getOperand(1); }
 };
 
+struct IRGetOffsetPtr : IRInst
+{
+    IR_LEAF_ISA(GetOffsetPtr);
+    IRInst* getBase() { return getOperand(0); }
+    IRInst* getOffset() { return getOperand(1); }
+};
+
 struct IRRWStructuredBufferGetElementPtr : IRInst
 {
     IR_LEAF_ISA(RWStructuredBufferGetElementPtr);
diff --git a/source/slang/slang-ir-specialize-address-space.cpp b/source/slang/slang-ir-specialize-address-space.cpp
index a3c45617d..ae0542734 100644
--- a/source/slang/slang-ir-specialize-address-space.cpp
+++ b/source/slang/slang-ir-specialize-address-space.cpp
@@ -13,6 +13,7 @@ struct AddressSpaceContext : public AddressSpaceSpecializationContext
 
     Dictionary<IRInst*, AddressSpace> mapInstToAddrSpace;
     InitialAddressSpaceAssigner* addrSpaceAssigner;
+    HashSet<IRFunc*> functionsToConsiderRemoving;
 
     AddressSpaceContext(IRModule* inModule, InitialAddressSpaceAssigner* inAddrSpaceAssigner)
         : module(inModule), addrSpaceAssigner(inAddrSpaceAssigner)
@@ -279,6 +280,8 @@ struct AddressSpaceContext : public AddressSpaceSpecializationContext
                                     callInst = as<IRCall>(builder.replaceOperand(
                                         callInst->getOperands(),
                                         specializedCallee));
+                                    // At this point, the original callee may be left without uses.
+                                    functionsToConsiderRemoving.add(callee);
                                 }
                                 auto callResultAddrSpace =
                                     getFuncResultAddrSpace(specializedCallee);
@@ -394,6 +397,13 @@ struct AddressSpaceContext : public AddressSpaceSpecializationContext
         }
 
         applyAddressSpaceToInstType();
+
+        for (IRFunc* func : functionsToConsiderRemoving)
+        {
+            SLANG_ASSERT(!func->findDecoration<IREntryPointDecoration>());
+            if (!func->hasUses())
+                func->removeAndDeallocate();
+        }
     }
 };
 
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 737209207..63b16080f 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -22,6 +22,7 @@
 #include "slang-ir-simplify-cfg.h"
 #include "slang-ir-specialize-address-space.h"
 #include "slang-ir-util.h"
+#include "slang-ir-validate.h"
 #include "slang-ir.h"
 #include "slang-legalize-types.h"
 
@@ -1931,6 +1932,11 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
         // Specalize address space for all pointers.
         SpirvAddressSpaceAssigner addressSpaceAssigner;
         specializeAddressSpace(m_module, &addressSpaceAssigner);
+
+        // For SPIR-V, we don't skip this validation, because we might then be generating invalid
+        // SPIR-V.
+        bool skipFuncParamValidation = false;
+        validateAtomicOperations(skipFuncParamValidation, m_sink, m_module->getModuleInst());
     }
 
     void updateFunctionTypes()
diff --git a/source/slang/slang-ir-validate.cpp b/source/slang/slang-ir-validate.cpp
index 4652b011f..11587d600 100644
--- a/source/slang/slang-ir-validate.cpp
+++ b/source/slang/slang-ir-validate.cpp
@@ -410,4 +410,104 @@ void validateIRModuleIfEnabled(CodeGenContext* codeGenContext, IRModule* module)
     validateIRModule(module, sink);
 }
 
+// Returns whether 'dst' is a valid destination for atomic operations, meaning
+// it leads either to 'groupshared' or 'device buffer' memory.
+static bool isValidAtomicDest(bool skipFuncParamValidation, IRInst* dst)
+{
+    bool isGroupShared = as<IRGroupSharedRate>(dst->getRate());
+    if (isGroupShared)
+        return true;
+
+    if (as<IRRWStructuredBufferGetElementPtr>(dst))
+        return true;
+    if (as<IRImageSubscript>(dst))
+        return true;
+
+    if (auto ptrType = as<IRPtrType>(dst->getDataType()))
+    {
+        switch (ptrType->getAddressSpace())
+        {
+        case AddressSpace::Global:
+        case AddressSpace::GroupShared:
+        case AddressSpace::StorageBuffer:
+        case AddressSpace::UserPointer:
+            return true;
+        default:
+            break;
+        }
+    }
+
+    if (as<IRGlobalParam>(dst))
+    {
+        switch (dst->getDataType()->getOp())
+        {
+        case kIROp_GLSLShaderStorageBufferType:
+        case kIROp_TextureType:
+            return true;
+        default:
+            return false;
+        }
+    }
+
+    if (auto param = as<IRParam>(dst))
+    {
+        auto paramType = param->getDataType();
+        if (auto outType = as<IROutTypeBase>(paramType))
+        {
+            if (outType->getAddressSpace() == AddressSpace::GroupShared)
+            {
+                return true;
+            }
+            else if (skipFuncParamValidation)
+            {
+                // We haven't actually verified that this is a valid atomic operation destination,
+                // but the callee wants to skip this specific validation.
+                return true;
+            }
+        }
+    }
+    if (auto getElementPtr = as<IRGetElementPtr>(dst))
+        return isValidAtomicDest(skipFuncParamValidation, getElementPtr->getBase());
+    if (auto getOffsetPtr = as<IRGetOffsetPtr>(dst))
+        return isValidAtomicDest(skipFuncParamValidation, getOffsetPtr->getBase());
+    if (auto fieldAddress = as<IRFieldAddress>(dst))
+        return isValidAtomicDest(skipFuncParamValidation, fieldAddress->getBase());
+
+    return false;
+}
+
+void validateAtomicOperations(bool skipFuncParamValidation, DiagnosticSink* sink, IRInst* inst)
+{
+    switch (inst->getOp())
+    {
+    case kIROp_AtomicLoad:
+    case kIROp_AtomicStore:
+    case kIROp_AtomicExchange:
+    case kIROp_AtomicCompareExchange:
+    case kIROp_AtomicAdd:
+    case kIROp_AtomicSub:
+    case kIROp_AtomicAnd:
+    case kIROp_AtomicOr:
+    case kIROp_AtomicXor:
+    case kIROp_AtomicMin:
+    case kIROp_AtomicMax:
+    case kIROp_AtomicInc:
+    case kIROp_AtomicDec:
+        {
+            IRInst* destinationPtr = inst->getOperand(0);
+            if (!isValidAtomicDest(skipFuncParamValidation, destinationPtr))
+                sink->diagnose(inst->sourceLoc, Diagnostics::invalidAtomicDestinationPointer);
+        }
+        break;
+
+    default:
+        break;
+    }
+
+    for (auto child : inst->getModifiableChildren())
+    {
+        validateAtomicOperations(skipFuncParamValidation, sink, child);
+    }
+}
+
 } // namespace Slang
diff --git a/source/slang/slang-ir-validate.h b/source/slang/slang-ir-validate.h
index d10aff7a0..dc7a52cee 100644
--- a/source/slang/slang-ir-validate.h
+++ b/source/slang/slang-ir-validate.h
@@ -38,4 +38,15 @@ void validateIRModuleIfEnabled(CodeGenContext* codeGenContext, IRModule* module)
 void disableIRValidationAtInsert();
 void enableIRValidationAtInsert();
 
+// Validate that the destination of an atomic operation is appropriate, meaning it's
+// either 'groupshared' or in a device buffer.
+// Note that validation of atomic operations should be done after address space
+// specialization for targets (e.g. SPIR-V and Metal) which support this kind of use-case:
+//   void atomicOp(inout int array){ InterlockedAdd(array, 1);}
+//   groupshared int gArray;
+//   [numthreads(1, 1, 1)] void main() { atomicOp(gArray); }
+// If 'skipFuncParamValidation' is true, then the validation allows destinations that
+// lead back to in/inout parameters that we can't validate.
+void validateAtomicOperations(bool skipFuncParamValidation, DiagnosticSink* sink, IRInst* inst);
+
 } // namespace Slang
diff --git a/tests/glsl-intrinsic/atomic/invalidDest.slang b/tests/glsl-intrinsic/atomic/invalidDest.slang
new file mode 100644
index 000000000..315321f88
--- /dev/null
+++ b/tests/glsl-intrinsic/atomic/invalidDest.slang
@@ -0,0 +1,25 @@
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK):  -allow-glsl -stage compute -entry computeMain -target glsl -DTARGET_GLSL
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK):  -allow-glsl -stage compute -entry computeMain -target spirv -skip-spirv-validation -emit-spirv-directly -DTARGET_SPIRV
+#version 430
+
+uint notShared;
+
+void computeMain()
+{
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicAdd(notShared, 1u);
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicAnd(notShared, 1u);
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicCompSwap(notShared, 1u, 2u);
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicExchange(notShared, 1u);
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicMax(notShared, 1u);
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicMin(notShared, 1u);
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicOr(notShared, 1u);
+    // CHECK: ([[# @LINE+1]]): error 41403
+    atomicXor(notShared, 1u);
+}