64-bit atomics usage resource flags (#3518)

When a resource comes from a binding, we can't determine if it is on a
heap or not, so we set a flag on it to let the runtime sort it out.

When a resource comes from the heap directly, we know it's on the heap
so we set the shader flag to check for platform support.

include/dxc/DXIL/DxilConstants.h

@@ -1416,8 +1416,9 @@ namespace DXIL {
   const uint64_t ShaderFeatureInfo_AtomicInt64OnTypedResource = 0x400000;
   const uint64_t ShaderFeatureInfo_AtomicInt64OnGroupShared = 0x800000;
   const uint64_t ShaderFeatureInfo_DerivativesInMeshAndAmpShaders = 0x1000000;
+  const uint64_t ShaderFeatureInfo_AtomicInt64OnHeapResource = 0x2000000;
 
-  const unsigned ShaderFeatureInfoCount = 25;
+  const unsigned ShaderFeatureInfoCount = 26;
 
   // DxilSubobjectType must match D3D12_STATE_SUBOBJECT_TYPE, with
   // certain values reserved, since they cannot be used from Dxil.

include/dxc/DXIL/DxilMetadataHelper.h

@@ -199,6 +199,7 @@ public:
   static const unsigned kDxilTypedBufferElementTypeTag            = 0;
   static const unsigned kDxilStructuredBufferElementStrideTag     = 1;
   static const unsigned kDxilSamplerFeedbackKindTag               = 2;
+  static const unsigned kDxilAtomic64UseTag                       = 3;
 
   // Type system.
   static const char kDxilTypeSystemMDName[];

include/dxc/DXIL/DxilResourceProperties.h

@@ -89,8 +89,7 @@ llvm::Constant *getAsConstant(const DxilResourceProperties &, llvm::Type *Ty,
                               const ShaderModel &);
 DxilResourceProperties loadPropsFromConstant(const llvm::Constant &C);
 DxilResourceProperties
-loadPropsFromAnnotateHandle(DxilInst_AnnotateHandle &annotateHandle, llvm::Type *Ty,
-                       const ShaderModel &);
+loadPropsFromAnnotateHandle(DxilInst_AnnotateHandle &annotateHandle, const ShaderModel &);
 DxilResourceProperties loadPropsFromResourceBase(const DxilResourceBase *);
 
 } // namespace resource_helper

include/dxc/DXIL/DxilShaderFlags.h

@@ -120,6 +120,9 @@ namespace hlsl {
     void SetAtomicInt64OnGroupShared(bool flag) { m_bAtomicInt64OnGroupShared = flag; }
     bool GetAtomicInt64OnGroupShared() const { return m_bAtomicInt64OnGroupShared; }
 
+    void SetAtomicInt64OnHeapResource(bool flag) { m_bAtomicInt64OnHeapResource = flag; }
+    bool GetAtomicInt64OnHeapResource() const { return m_bAtomicInt64OnHeapResource; }
+
     void SetDerivativesInMeshAndAmpShaders(bool flag) { m_bDerivativesInMeshAndAmpShaders = flag; }
     bool GetDerivativesInMeshAndAmpShaders() { return m_bDerivativesInMeshAndAmpShaders; }
 
@@ -162,11 +165,13 @@ namespace hlsl {
     unsigned m_bSamplerFeedback : 1; // SHADER_FEATURE_SAMPLER_FEEDBACK
 
     unsigned m_bAtomicInt64OnTypedResource : 1; // SHADER_FEATURE_ATOMIC_INT64_ON_TYPED_RESOURCE
-    unsigned m_bAtomicInt64OnGroupShared : 1;//SHADER_FEATURE_ATOMIC_INT64_ON_GROUP_SHARED
+    unsigned m_bAtomicInt64OnGroupShared : 1; // SHADER_FEATURE_ATOMIC_INT64_ON_GROUP_SHARED
 
     unsigned m_bDerivativesInMeshAndAmpShaders : 1; //SHADER_FEATURE_DERIVATIVES_IN_MESH_AND_AMPLIFICATION_SHADERS
 
-    unsigned m_align0 : 2;        // align to 32 bit.
+    unsigned m_bAtomicInt64OnHeapResource : 1; // SHADER_FEATURE_ATOMIC_INT64_ON_DESCRIPTOR_HEAP_RESOURCE
+
+    unsigned m_align0 : 1;        // align to 32 bit.
     uint32_t m_align1;            // align to 64 bit.
   };
 

lib/DXIL/DxilMetadataHelper.cpp

@@ -2244,6 +2244,11 @@ void DxilExtraPropertyHelper::EmitUAVProperties(const DxilResource &UAV, std::ve
     MDVals.emplace_back(DxilMDHelper::Uint32ToConstMD(DxilMDHelper::kDxilSamplerFeedbackKindTag, m_Ctx));
     MDVals.emplace_back(DxilMDHelper::Uint32ToConstMD((unsigned)UAV.GetSamplerFeedbackType(), m_Ctx));
   }
+  // Whether resource is used for 64-bit atomic op
+  if (DXIL::CompareVersions(m_ValMajor, m_ValMinor, 1, 6) >= 0 && UAV.HasAtomic64Use()) {
+    MDVals.emplace_back(DxilMDHelper::Uint32ToConstMD(DxilMDHelper::kDxilAtomic64UseTag, m_Ctx));
+    MDVals.emplace_back(DxilMDHelper::Uint32ToConstMD((unsigned)true, m_Ctx));
+  }
 }
 
 void DxilExtraPropertyHelper::LoadUAVProperties(const MDOperand &MDO, DxilResource &UAV) {
@@ -2275,6 +2280,9 @@ void DxilExtraPropertyHelper::LoadUAVProperties(const MDOperand &MDO, DxilResour
       DXASSERT_NOMSG(UAV.IsFeedbackTexture());
       UAV.SetSamplerFeedbackType((DXIL::SamplerFeedbackType)DxilMDHelper::ConstMDToUint32(MDO));
       break;
+    case DxilMDHelper::kDxilAtomic64UseTag:
+      UAV.SetHasAtomic64Use(DxilMDHelper::ConstMDToBool(MDO));
+      break;
     default:
       DXASSERT(false, "Unknown resource record tag");
       m_bExtraMetadata = true;

lib/DXIL/DxilResourceProperties.cpp

@@ -130,7 +130,7 @@ DxilResourceProperties loadPropsFromConstant(const Constant &C) {
 
 DxilResourceProperties
 loadPropsFromAnnotateHandle(DxilInst_AnnotateHandle &annotateHandle,
-                            llvm::Type *Ty, const ShaderModel &SM) {
+                            const ShaderModel &SM) {
   Constant *ResProp = cast<Constant>(annotateHandle.get_props());
   return loadPropsFromConstant(*ResProp);
 }

lib/DXIL/DxilShaderFlags.cpp

@@ -11,6 +11,7 @@
 #include "dxc/DXIL/DxilShaderFlags.h"
 #include "dxc/DXIL/DxilOperations.h"
 #include "dxc/DXIL/DxilResource.h"
+#include "dxc/DXIL/DxilResourceBinding.h"
 #include "dxc/Support/Global.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Instructions.h"
@@ -55,6 +56,7 @@ ShaderFlags::ShaderFlags():
 , m_bAtomicInt64OnTypedResource(false)
 , m_bAtomicInt64OnGroupShared(false)
 , m_bDerivativesInMeshAndAmpShaders(false)
+, m_bAtomicInt64OnHeapResource(false)
 , m_align0(0)
 , m_align1(0)
 {
@@ -109,6 +111,7 @@ uint64_t ShaderFlags::GetFeatureInfo() const {
   Flags |= m_bSamplerFeedback ? hlsl::DXIL::ShaderFeatureInfo_SamplerFeedback : 0;
   Flags |= m_bAtomicInt64OnTypedResource ? hlsl::DXIL::ShaderFeatureInfo_AtomicInt64OnTypedResource : 0;
   Flags |= m_bAtomicInt64OnGroupShared ? hlsl::DXIL::ShaderFeatureInfo_AtomicInt64OnGroupShared : 0;
+  Flags |= m_bAtomicInt64OnHeapResource ? hlsl::DXIL::ShaderFeatureInfo_AtomicInt64OnHeapResource : 0;
   Flags |= m_bDerivativesInMeshAndAmpShaders ? hlsl::DXIL::ShaderFeatureInfo_DerivativesInMeshAndAmpShaders : 0;
 
   return Flags;
@@ -166,6 +169,7 @@ uint64_t ShaderFlags::GetShaderFlagsRawForCollection() {
   Flags.SetSamplerFeedback(true);
   Flags.SetAtomicInt64OnTypedResource(true);
   Flags.SetAtomicInt64OnGroupShared(true);
+  Flags.SetAtomicInt64OnHeapResource(true);
   Flags.SetDerivativesInMeshAndAmpShaders(true);
   return Flags.GetShaderFlagsRaw();
 }
@@ -261,14 +265,60 @@ DxilResourceProperties GetResourcePropertyFromHandleCall(const hlsl::DxilModule
     }
   } else if (handleOp == DXIL::OpCode::AnnotateHandle) {
     DxilInst_AnnotateHandle annotateHandle(cast<Instruction>(handleCall));
-    Type *ResPropTy = M->GetOP()->GetResourcePropertiesType();
 
-    RP = resource_helper::loadPropsFromAnnotateHandle(annotateHandle, ResPropTy, *M->GetShaderModel());
+    RP = resource_helper::loadPropsFromAnnotateHandle(annotateHandle, *M->GetShaderModel());
   }
 
   return RP;
 }
 
+struct ResourceKey {
+  uint8_t Class;
+  uint32_t Space;
+  uint32_t LowerBound;
+  uint32_t UpperBound;
+};
+
+struct ResKeyEq : public std::binary_function<ResourceKey, ResourceKey, bool> {
+   bool operator()(const ResourceKey& k1, const ResourceKey& k2) const {
+     return k1.Class == k2.Class && k1.Space == k2.Space &&
+       k1.LowerBound == k2.LowerBound && k1.UpperBound == k2.UpperBound;
+   }
+};
+
+struct ResKeyHash : public std::unary_function<ResourceKey, std::size_t> {
+   std::size_t operator()(const ResourceKey& k) const {
+     return std::hash<uint32_t>()(k.LowerBound) ^ (std::hash<uint32_t>()(k.UpperBound)<<1) ^
+       (std::hash<uint32_t>()(k.Space)<<2) ^ (std::hash<uint8_t>()(k.Class)<<3);
+   }
+};
+
+// 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,
+                                     std::unordered_map<ResourceKey, DxilResource *, ResKeyHash, ResKeyEq> resMap) {
+  DxilResource *resource = nullptr;
+
+  ConstantInt *HandleOpCodeConst = cast<ConstantInt>(
+      handleCall->getArgOperand(DXIL::OperandIndex::kOpcodeIdx));
+  DXIL::OpCode handleOp = static_cast<DXIL::OpCode>(HandleOpCodeConst->getLimitedValue());
+  if (handleOp == DXIL::OpCode::AnnotateHandle) {
+    DxilInst_AnnotateHandle annotateHandle(cast<Instruction>(handleCall));
+    CallInst *createCall = cast<CallInst>(annotateHandle.get_res());
+    ConstantInt *HandleOpCodeConst = cast<ConstantInt>(
+            createCall->getArgOperand(DXIL::OperandIndex::kOpcodeIdx));
+    DXIL::OpCode handleOp = static_cast<DXIL::OpCode>(HandleOpCodeConst->getLimitedValue());
+    if (handleOp == DXIL::OpCode::CreateHandleFromBinding) {
+      DxilInst_CreateHandleFromBinding fromBind(createCall);
+      DxilResourceBinding B = resource_helper::loadBindingFromConstant(*cast<Constant>(fromBind.get_bind()));
+      ResourceKey key = {B.resourceClass, B.spaceID, B.rangeLowerBound, B.rangeUpperBound};
+      resource = resMap[key];
+    }
+  }
+
+  return resource;
+}
+
 
 ShaderFlags ShaderFlags::CollectShaderFlags(const Function *F,
                                            const hlsl::DxilModule *M) {
@@ -298,6 +348,7 @@ ShaderFlags ShaderFlags::CollectShaderFlags(const Function *F,
   bool hasRaytracingTier1_1 = false;
   bool hasAtomicInt64OnTypedResource = false;
   bool hasAtomicInt64OnGroupShared = false;
+  bool hasAtomicInt64OnHeapResource = false;
   bool hasDerivativesInMeshAndAmpShaders = false;
 
   // Try to maintain compatibility with a v1.0 validator if that's what we have.
@@ -309,6 +360,15 @@ ShaderFlags ShaderFlags::CollectShaderFlags(const Function *F,
   Type *int16Ty = Type::getInt16Ty(F->getContext());
   Type *int64Ty = Type::getInt64Ty(F->getContext());
 
+
+  // Set up resource to binding handle map for 64-bit atomics usage
+  std::unordered_map<ResourceKey, DxilResource *, ResKeyHash, ResKeyEq> resMap;
+  for (auto &res : M->GetUAVs()) {
+    ResourceKey key = {(uint8_t)res->GetClass(), res->GetSpaceID(),
+                       res->GetLowerBound(), res->GetUpperBound()};
+    resMap.insert({key, res.get()});
+  }
+
   for (const BasicBlock &BB : F->getBasicBlockList()) {
     for (const Instruction &I : BB.getInstList()) {
       // Skip none dxil function call.
@@ -420,6 +480,13 @@ ShaderFlags ShaderFlags::CollectShaderFlags(const Function *F,
             DxilResourceProperties RP = GetResourcePropertyFromHandleCall(M, handleCall);
             if (DXIL::IsTyped(RP.getResourceKind()))
                 hasAtomicInt64OnTypedResource = true;
+            // set uses 64-bit flag if relevant
+            if (DxilResource *res = GetResourceFromAnnotateHandle(handleCall, resMap)) {
+              res->SetHasAtomic64Use(true);
+            } else {
+              // Assuming CreateHandleFromHeap, which indicates a descriptor
+              hasAtomicInt64OnHeapResource = true;
+            }
           }
           break;
         case DXIL::OpCode::DerivFineX:
@@ -543,6 +610,7 @@ ShaderFlags ShaderFlags::CollectShaderFlags(const Function *F,
   flag.SetRaytracingTier1_1(hasRaytracingTier1_1);
   flag.SetAtomicInt64OnTypedResource(hasAtomicInt64OnTypedResource);
   flag.SetAtomicInt64OnGroupShared(hasAtomicInt64OnGroupShared);
+  flag.SetAtomicInt64OnHeapResource(hasAtomicInt64OnHeapResource);
   flag.SetDerivativesInMeshAndAmpShaders(hasDerivativesInMeshAndAmpShaders);
 
   return flag;

lib/HLSL/DxilCondenseResources.cpp

@@ -2290,7 +2290,7 @@ bool DxilLowerCreateHandleForLib::PatchDynamicTBuffers(DxilModule &DM) {
     CallInst *CI = cast<CallInst>(U);
     DxilInst_AnnotateHandle annot(CI);
     DxilResourceProperties RP = resource_helper::loadPropsFromAnnotateHandle(
-        annot, hlslOP->GetResourcePropertiesType(), *DM.GetShaderModel());
+        annot, *DM.GetShaderModel());
 
     if (RP.getResourceKind() != DXIL::ResourceKind::TBuffer)
       continue;

lib/HLSL/DxilValidation.cpp

@@ -600,7 +600,6 @@ struct ValidationContext {
         }
       }
     }
-    Type *ResPropTy = hlslOP->GetResourcePropertiesType();
     const ShaderModel &SM = *DxilMod.GetShaderModel();
 
     for (auto &it : hlslOP->GetOpFuncList(DXIL::OpCode::AnnotateHandle)) {
@@ -612,7 +611,7 @@ struct ValidationContext {
         CallInst *CI = cast<CallInst>(U);
         DxilInst_AnnotateHandle hdl(CI);
         DxilResourceProperties RP =
-            resource_helper::loadPropsFromAnnotateHandle(hdl, ResPropTy, SM);
+            resource_helper::loadPropsFromAnnotateHandle(hdl, SM);
         if (RP.getResourceKind() == DXIL::ResourceKind::Invalid) {
           EmitInstrError(CI, ValidationRule::InstrOpConstRange);
           continue;

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

@@ -0,0 +1,14 @@
+// RUN: %dxc -T lib_6_6 %s | FileCheck %s
+// Test atomic operations on dynamic resources
+
+// CHECK: RGInt64OnDescriptorHeapIndex
+// CHCK: Note: shader requires additional functionality:
+// CHCK: 64-bit Atomics on Typed Resources
+// CHCK: 64-bit Atomics on Heap Resources
+
+[shader("raygeneration")]
+void RGInt64OnDescriptorHeapIndex()
+{
+    RWTexture2D<int64_t> myTexture = ResourceDescriptorHeap[7];
+    InterlockedAdd(myTexture[int2(0,0)], 1);
+}

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

@@ -0,0 +1,15 @@
+// RUN: %dxc -T ps_6_6 %s | FileCheck %s
+// Test atomic operations on heap resources
+
+// For now just check that it compiles.
+// TODO: improve disassembly output for resources to verify the flag
+
+// CHECK: Note: shader requires additional functionality:
+// CHECK: 64-bit Atomics on Typed Resources
+
+RWTexture2D<int64_t> myTexture : register (u0);
+
+void main() : SV_Target
+{
+    InterlockedAdd(myTexture[int2(0,0)], 1);
+}

tools/clang/tools/dxcompiler/dxcdisassembler.cpp

@@ -342,7 +342,8 @@ PCSTR g_pFeatureInfoNames[] = {
     "Sampler feedback",
     "64-bit Atomics on Typed Resources",
     "64-bit Atomics on Group Shared",
-    "Derivatives in mesh and amplification shaders"
+    "Derivatives in mesh and amplification shaders",
+    "64-bit Atomics on Heap Resources",
 };
 static_assert(_countof(g_pFeatureInfoNames) == ShaderFeatureInfoCount, "g_pFeatureInfoNames needs to be updated");