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.

Author: aleino-nv

source/slang/slang-emit.cpp

@@ -1322,7 +1322,9 @@ Result linkAndOptimizeIR(
             byteAddressBufferOptions);
     }
 
-    validateAtomicOperations(sink, irModule->getModuleInst());
+    // For SPIR-V, this function is called elsewhere, so that it can happen after address space specialization
+    if (target != CodeGenTarget::SPIRV && target != CodeGenTarget::SPIRVAssembly)
+        validateAtomicOperations(sink, irModule->getModuleInst());
 
     // For CUDA targets only, we will need to turn operations
     // the implicitly reference the "active mask" into ones

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,8 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
         // Specalize address space for all pointers.
         SpirvAddressSpaceAssigner addressSpaceAssigner;
         specializeAddressSpace(m_module, &addressSpaceAssigner);
+
+        validateAtomicOperations(m_sink, m_module->getModuleInst());
     }
 
     void updateFunctionTypes()

source/slang/slang-ir-validate.h

@@ -38,6 +38,13 @@ 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); }
 void validateAtomicOperations(DiagnosticSink* sink, IRInst* inst);
 
 } // namespace Slang