[SPIRV] Allow globallycoherent on DescriptorHeap accesses (#7740)

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -32,6 +32,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/SaveAndRestore.h"
 
 #ifdef SUPPORT_QUERY_GIT_COMMIT_INFO
 #include "clang/Basic/Version.h"
@@ -1559,6 +1560,9 @@ bool SpirvEmitter::handleNodePayloadArrayType(const ParmVarDecl *decl,
 }
 
 void SpirvEmitter::doFunctionDecl(const FunctionDecl *decl) {
+  llvm::SaveAndRestore<const FunctionDecl *> savedCurFunctionDecl(
+      curFunctionDecl, decl);
+
   // Forward declaration of a function inside another.
   if (!decl->isThisDeclarationADefinition()) {
     addFunctionToWorkQueue(spvContext.getCurrentShaderModelKind(), decl,
@@ -6658,6 +6662,62 @@ SpirvEmitter::doCXXOperatorCallExpr(const CXXOperatorCallExpr *expr,
       auto *decl = cast<VarDecl>(declRefExpr->getDecl());
       auto *var = declIdMapper.createResourceHeap(decl, resourceType);
 
+      // Decide whether the destination of this heap access is
+      // globallycoherent. Three sources:
+      //   (a) An enclosing DeclStmt VarDecl initializer
+      //   (b) A file-scope VarDecl initializer
+      //   (c) A ReturnStmt inside a globallycoherent returning helper
+      auto isCoherentDest = [](const VarDecl *vd) {
+        return vd && vd->hasAttr<HLSLGloballyCoherentAttr>();
+      };
+
+      bool destIsCoherent = false;
+
+      // (a) and (c): walk Stmt parents through transparent wrappers.
+      {
+        const Stmt *cur = parentExpr;
+        while (cur) {
+          const Stmt *p = parentMap->getParent(cur);
+          if (!p)
+            break;
+          if (isa<CastExpr>(p) || isa<ParenExpr>(p)) {
+            cur = p;
+            continue;
+          }
+          if (const auto *ds = dyn_cast<DeclStmt>(p)) {
+            for (const Decl *d : ds->decls())
+              if (const auto *vd = dyn_cast<VarDecl>(d))
+                if (isCoherentDest(vd)) {
+                  destIsCoherent = true;
+                  break;
+                }
+          } else if (isa<ReturnStmt>(p) && curFunctionDecl &&
+                     curFunctionDecl->hasAttr<HLSLGloballyCoherentAttr>()) {
+            destIsCoherent = true;
+          }
+          break;
+        }
+      }
+      // (b) File-scope init: VarDecl is the AST parent of the top cast.
+      if (!destIsCoherent) {
+        const Expr *top = parentExpr;
+        while (true) {
+          auto parents = astContext.getParents(*top);
+          if (parents.empty())
+            break;
+          if (const auto *castParent = parents[0].get<CastExpr>()) {
+            top = castParent;
+            continue;
+          }
+          if (const auto *vd = parents[0].get<VarDecl>()) {
+            if (isCoherentDest(vd))
+              destIsCoherent = true;
+          }
+          break;
+        }
+      }
+      if (destIsCoherent)
+        spvBuilder.decorateCoherent(var, baseExpr->getExprLoc());
       auto *index = doExpr(indexExpr);
 
       if (spirvOptions.useDescriptorHeap) {

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -1648,6 +1648,10 @@ class SpirvEmitter : public ASTConsumer {
 
   /// ParentMap of the current function.
   std::unique_ptr<ParentMap> parentMap = nullptr;
+
+  /// AST FunctionDecl currently being lowered. Used by the descriptor-heap
+  /// access path to detect globallycoherent returning helpers.
+  const FunctionDecl *curFunctionDecl = nullptr;
 };
 
 void SpirvEmitter::doDeclStmt(const DeclStmt *declStmt) {

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.globallycoherent.hlsl

@@ -0,0 +1,54 @@
+// RUN: %dxc -T cs_6_6 -E CSMain -spirv -DTYPE=0 %s | FileCheck %s --check-prefix=TYPE0
+// RUN: %dxc -T cs_6_6 -E CSMain -spirv -DTYPE=1 %s | FileCheck %s --check-prefix=TYPE1
+// RUN: %dxc -T cs_6_6 -E CSMain -spirv -DTYPE=2 %s | FileCheck %s --check-prefix=TYPE2
+// RUN: %dxc -T cs_6_6 -E CSMain -spirv -DTYPE=3 %s | FileCheck %s --check-prefix=TYPE3
+
+// TYPE=0: direct ResourceDescriptorHeap[] init of a `globallycoherent` static.
+// TYPE0-DAG: OpDecorate %ResourceDescriptorHeap{{(_[0-9]+)?}} DescriptorSet 0
+// TYPE0-DAG: OpDecorate %ResourceDescriptorHeap{{(_[0-9]+)?}} Binding 0
+// TYPE0-DAG: OpDecorate %ResourceDescriptorHeap{{(_[0-9]+)?}} Coherent
+
+// TYPE=1: bindless array carries the qualifier.
+// TYPE1-DAG: OpDecorate %FakeHeapOfA DescriptorSet 0
+// TYPE1-DAG: OpDecorate %FakeHeapOfA Binding 0
+// TYPE1-DAG: OpDecorate %FakeHeapOfA Coherent
+
+// TYPE=2: stand-alone `globallycoherent` UAV.
+// TYPE2-DAG: OpDecorate %A DescriptorSet 0
+// TYPE2-DAG: OpDecorate %A Binding 0
+// TYPE2-DAG: OpDecorate %A Coherent
+
+// TYPE=3: heap access lives inside a `globallycoherent`-returning helper, and 
+// a `globallycoherent static` captures the result.
+// TYPE3-DAG: OpDecorate %ResourceDescriptorHeap{{(_[0-9]+)?}} DescriptorSet 0
+// TYPE3-DAG: OpDecorate %ResourceDescriptorHeap{{(_[0-9]+)?}} Binding {{[0-9]+}}
+// TYPE3-DAG: OpDecorate %ResourceDescriptorHeap{{(_[0-9]+)?}} Coherent
+// TYPE3-NOT: OpDecorate %Buf Coherent
+
+#if TYPE == 0
+globallycoherent static RWStructuredBuffer<uint> A = ResourceDescriptorHeap[0];
+#elif TYPE == 1
+globallycoherent RWStructuredBuffer<uint> FakeHeapOfA[];
+globallycoherent static RWStructuredBuffer<uint> A = FakeHeapOfA[0];
+#elif TYPE == 2
+globallycoherent RWStructuredBuffer<uint> A;
+#elif TYPE == 3
+uint BindlessUAV_Buf;
+typedef RWByteAddressBuffer SafeTypeBuf;
+globallycoherent SafeTypeBuf GetBuf() { return ResourceDescriptorHeap[BindlessUAV_Buf]; }
+static const globallycoherent SafeTypeBuf A = GetBuf();
+#endif
+
+[numthreads(64, 1, 1)]
+void CSMain(uint3 ThreadId : SV_DispatchThreadId)
+{
+#if TYPE == 3
+    A.InterlockedAdd(0, 1);
+    AllMemoryBarrierWithGroupSync();
+    A.Store(0, 42);
+#else
+    InterlockedAdd(A[0], 1);
+    AllMemoryBarrierWithGroupSync();
+    InterlockedAdd(A[1], A[0]);
+#endif
+}