From 04353fb7602b7eb6a8b86193510ebe0c670b7724 Mon Sep 17 00:00:00 2001
From: Anders Leino <aleino@nvidia.com>
Date: Wed, 22 Jan 2025 19:10:35 +0200
Subject: 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>
---
 source/slang/slang-ir-validate.cpp | 100 +++++++++++++++++++++++++++++++++++++
 1 file changed, 100 insertions(+)

(limited to 'source/slang/slang-ir-validate.cpp')

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
-- 
cgit v1.2.3