Correct atomics on heap usage for libs (#3525)

The previous shader flag code failed to account for the library-specific
handle creation function

Author: Greg Roth

include/dxc/DxilContainer/DxilRuntimeReflection.h

@@ -148,6 +148,7 @@ enum class DxilResourceFlag : uint32_t {
   UAVCounter                = 1 << 1,
   UAVRasterizerOrderedView  = 1 << 2,
   DynamicIndexing           = 1 << 3,
+  Atomics64Use              = 1 << 4,
 };
 
 struct RuntimeDataResourceInfo {

lib/DXIL/DxilShaderFlags.cpp

@@ -297,9 +297,9 @@ struct ResKeyHash {
    }
 };
 
-// Limited to retrieving handles created by CreateHandleFromBinding. returns null otherwise
-// map should contain resources indexed by class, lower, and upper bounds
-DxilResource *GetResourceFromAnnotateHandle(CallInst *handleCall,
+// Limited to retrieving handles created by CreateHandleFromBinding and CreateHandleForLib. returns null otherwise
+// map should contain resources indexed by space, class, lower, and upper bounds
+DxilResource *GetResourceFromAnnotateHandle(const hlsl::DxilModule *M, CallInst *handleCall,
                                      std::unordered_map<ResourceKey, DxilResource *, ResKeyHash, ResKeyEq> resMap) {
   DxilResource *resource = nullptr;
 
@@ -317,6 +317,17 @@ DxilResource *GetResourceFromAnnotateHandle(CallInst *handleCall,
       DxilResourceBinding B = resource_helper::loadBindingFromConstant(*cast<Constant>(fromBind.get_bind()));
       ResourceKey key = {B.resourceClass, B.spaceID, B.rangeLowerBound, B.rangeUpperBound};
       resource = resMap[key];
+    } else if (handleOp == DXIL::OpCode::CreateHandleForLib) {
+      // If library handle, find DxilResource by checking the name
+      if (LoadInst *LI = dyn_cast<LoadInst>(createCall->getArgOperand(
+                                              DXIL::OperandIndex::kCreateHandleForLibResOpIdx))) {
+        Value *resType = LI->getOperand(0);
+        for (auto &&res : M->GetUAVs()) {
+          if (res->GetGlobalSymbol() == resType) {
+            return resource = res.get();
+          }
+        }
+      }
     }
   }
 
@@ -487,7 +498,7 @@ ShaderFlags ShaderFlags::CollectShaderFlags(const Function *F,
             if (DXIL::IsTyped(RP.getResourceKind()))
                 hasAtomicInt64OnTypedResource = true;
             // set uses 64-bit flag if relevant
-            if (DxilResource *res = GetResourceFromAnnotateHandle(handleCall, resMap)) {
+            if (DxilResource *res = GetResourceFromAnnotateHandle(M, handleCall, resMap)) {
               res->SetHasAtomic64Use(true);
             } else {
               // Assuming CreateHandleFromHeap, which indicates a descriptor

lib/DxilContainer/DxilContainerAssembler.cpp

@@ -1179,6 +1179,8 @@ class DxilRDATWriter : public DxilPartWriter {
         info.Flags |= static_cast<uint32_t>(DxilResourceFlag::UAVGloballyCoherent);
       if (pRes->IsROV())
         info.Flags |= static_cast<uint32_t>(DxilResourceFlag::UAVRasterizerOrderedView);
+      if (pRes->HasAtomic64Use())
+        info.Flags |= static_cast<uint32_t>(DxilResourceFlag::Atomics64Use);
       // TODO: add dynamic index flag
     }
     m_pResourceTable->Insert(info);

tools/clang/test/HLSLFileCheck/hlsl/intrinsics/atomic/atomic_i64_root_resource.hlsl

@@ -0,0 +1,15 @@
+// RUN: %dxc -T lib_6_6 %s | FileCheck %s
+// Make sure library targets correctly identify non-dynamic resources
+
+// CHECK: Note: shader requires additional functionality:
+// CHECK-NOT: 64-bit Atomics on Heap Resources
+
+
+RWStructuredBuffer<int64_t> myBuf : register(u0);
+
+[shader("raygeneration")]
+[RootSignature("UAV(u0)")]
+void RGInt64OnDescriptorHeapIndex()
+{
+    InterlockedAdd(myBuf[0], 1);
+}