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 shader-slang#5989.

Author: aleino-nv

source/slang/slang-diagnostic-defs.h

@@ -2251,6 +2251,12 @@ 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.
 //

source/slang/slang-emit.cpp

@@ -1322,6 +1322,8 @@ Result linkAndOptimizeIR(
             byteAddressBufferOptions);
     }
 
+    validateAtomicOperations(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.

source/slang/slang-ir-insts.h

@@ -2415,6 +2415,7 @@ struct IRLoad : IRInst
     IRUse ptr;
     IR_LEAF_ISA(Load)
 
+    IRInst* getAddress() { return getOperand(0); }
     IRInst* getPtr() { return ptr.get(); }
 };
 
@@ -2491,6 +2492,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);

source/slang/slang-ir-validate.cpp

@@ -410,4 +410,96 @@ 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(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))
+        if (auto outType = as<IROutTypeBase>(param->getDataType()))
+            if (outType->getAddressSpace() == AddressSpace::GroupShared)
+                return true;
+    if (auto getElementPtr = as<IRGetElementPtr>(dst))
+        return isValidAtomicDest(getElementPtr->getBase());
+    if (auto getOffsetPtr = as<IRGetOffsetPtr>(dst))
+        return isValidAtomicDest(getOffsetPtr->getBase());
+    if (auto fieldAddress = as<IRFieldAddress>(dst))
+        return isValidAtomicDest(fieldAddress->getBase());
+
+    return false;
+}
+
+void validateAtomicOperations(DiagnosticSink* sink, IRInst* inst)
+{
+    // There may be unused functions containing violations after address space specialization.
+    if (auto func = as<IRFunc>(inst))
+        if(!(func->hasUses() || func->findDecoration<IREntryPointDecoration>()))
+            return;
+
+    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(destinationPtr))
+                sink->diagnose(inst->sourceLoc, Diagnostics::invalidAtomicDestinationPointer);
+        }
+        break;
+
+    default:
+        break;
+    }
+
+    for (auto child : inst->getModifiableChildren())
+    {
+        validateAtomicOperations(sink, child);
+    }
+}
+
 } // namespace Slang

source/slang/slang-ir-validate.h

@@ -38,4 +38,6 @@ void validateIRModuleIfEnabled(CodeGenContext* codeGenContext, IRModule* module)
 void disableIRValidationAtInsert();
 void enableIRValidationAtInsert();
 
+void validateAtomicOperations(DiagnosticSink* sink, IRInst* inst);
+
 } // namespace Slang