Support register binding on resource in cbuffer. (#2582)

include/dxc/HLSL/HLModule.h

@@ -254,6 +254,9 @@ public:
   DxilSubobjects *ReleaseSubobjects();
   void ResetSubobjects(DxilSubobjects *subobjects);
 
+  // Reg binding for resource in cb.
+  void AddRegBinding(unsigned CbID, unsigned ConstantIdx, unsigned Srv, unsigned Uav, unsigned Sampler);
+
 private:
   // Signatures.
   std::vector<uint8_t> m_SerializedRootSignature;
@@ -273,6 +276,11 @@ private:
 
   // Resource type annotation.
   std::unordered_map<llvm::Type *, std::pair<DXIL::ResourceClass, DXIL::ResourceKind>> m_ResTypeAnnotation;
+  // Resource bindings for res in cb.
+  // Key = CbID << 32 | ConstantIdx. Val is reg binding.
+  std::unordered_map<uint64_t, unsigned> m_SrvBindingInCB;
+  std::unordered_map<uint64_t, unsigned> m_UavBindingInCB;
+  std::unordered_map<uint64_t, unsigned> m_SamplerBindingInCB;
 
 private:
   llvm::LLVMContext &m_Ctx;

lib/HLSL/HLModule.cpp

@@ -14,6 +14,7 @@
 #include "dxc/DXIL/DxilCBuffer.h"
 #include "dxc/HLSL/HLModule.h"
 #include "dxc/DXIL/DxilTypeSystem.h"
+#include "dxc/DXIL/DxilUtil.h"
 #include "dxc/Support/WinAdapter.h"
 
 #include "llvm/ADT/STLExtras.h"
@@ -805,16 +806,140 @@ DxilResourceBase *HLModule::AddResourceWithGlobalVariableAndMDNode(llvm::Constan
   return R;
 }
 
+static uint64_t getRegBindingKey(unsigned CbID, unsigned ConstantIdx) {
+  return (uint64_t)(CbID) << 32 | ConstantIdx;
+}
+
+void HLModule::AddRegBinding(unsigned CbID, unsigned ConstantIdx, unsigned Srv, unsigned Uav,
+                             unsigned Sampler) {
+  uint64_t Key = getRegBindingKey(CbID, ConstantIdx);
+  m_SrvBindingInCB[Key] = Srv;
+  m_UavBindingInCB[Key] = Uav;
+  m_SamplerBindingInCB[Key] = Sampler;
+}
+
+static DXIL::ResourceClass GetRCFromType(Type *ResTy, Module &M) {
+  MDNode *MD = HLModule::GetDxilResourceAttrib(ResTy, M);
+  if (!MD)
+    return DXIL::ResourceClass::Invalid;
+  DxilResource::Class RC =
+      static_cast<DxilResource::Class>(DxilMDHelper::ConstMDToUint32(
+          MD->getOperand(DxilMDHelper::kHLDxilResourceAttributeClass)));
+  return RC;
+}
+
+static unsigned CountResNum(Module &M, Type *Ty, DXIL::ResourceClass RC) {
+  // Count num of RCs.
+  unsigned ArraySize = 1;
+  while (ArrayType *AT = dyn_cast<ArrayType>(Ty)) {
+    ArraySize *= AT->getNumElements();
+    Ty = AT->getElementType();
+  }
+
+  if (!Ty->isAggregateType())
+    return 0;
+
+  StructType *ST = dyn_cast<StructType>(Ty);
+  DXIL::ResourceClass TmpRC = GetRCFromType(ST, M);
+  if (TmpRC == RC)
+    return ArraySize;
+
+  unsigned Size = 0;
+  for (Type *EltTy : ST->elements()) {
+    Size += CountResNum(M, EltTy, RC);
+  }
+
+  return Size * ArraySize;
+}
+// Note: the rule for register binding on struct array is like this:
+// struct X {
+//   Texture2D x;
+//   SamplerState s ;
+//   Texture2D y;
+// };
+// X x[2] : register(t3) : register(s3);
+// x[0].x t3
+// x[0].s s3
+// x[0].y t4
+// x[1].x t5
+// x[1].s s4
+// x[1].y t6
+// So x[0].x and x[1].x not in an array.
+static unsigned CalcRegBinding(gep_type_iterator GEPIt, gep_type_iterator E,
+                               Module &M, DXIL::ResourceClass RC) {
+  unsigned NumRC = 0;
+  // Count GEP offset when only count RC size.
+  for (; GEPIt != E; GEPIt++) {
+    Type *Ty = *GEPIt;
+    Value *idx = GEPIt.getOperand();
+    Constant *constIdx = dyn_cast<Constant>(idx);
+    unsigned immIdx = constIdx->getUniqueInteger().getLimitedValue();
+    // Not support dynamic indexing.
+    // Array should be just 1d res array as global res.
+    if (ArrayType *AT = dyn_cast<ArrayType>(Ty)) {
+      NumRC += immIdx * CountResNum(M, AT->getElementType(), RC);
+    } else if (StructType *ST = dyn_cast<StructType>(Ty)) {
+      for (unsigned i=0;i<immIdx;i++) {
+        NumRC += CountResNum(M, ST->getElementType(i), RC);
+      }
+    }
+  }
+  return NumRC;
+}
+
 unsigned HLModule::GetBindingForResourceInCB(GetElementPtrInst *CbPtr,
                                              GlobalVariable *CbGV) {
-  DXIL::ResourceClass RC = GetResourceClass(CbPtr->getResultElementType());
+  if (!CbPtr->hasAllConstantIndices()) {
+    // Not support dynmaic indexing resource array inside cb.
+    string ErrorMsg("Index for resource array inside cbuffer must be a literal expression");
+    dxilutil::EmitErrorOnInstruction(
+        CbPtr,
+        ErrorMsg);
+    return UINT_MAX;
+  }
+  Module &M = *m_pModule;
+  Type *ResTy = CbPtr->getResultElementType();
+  DxilResource::Class RC = GetRCFromType(ResTy, M);
+
+  unsigned RegBinding = UINT_MAX;
   for (auto &CB : m_CBuffers) {
     if (CbGV != CB->GetGlobalSymbol())
       continue;
-    RC = DXIL::ResourceClass::Invalid;
+
+    gep_type_iterator GEPIt = gep_type_begin(CbPtr), E = gep_type_end(CbPtr);
+    // The pointer index.
+    GEPIt++;
+    unsigned ID = CB->GetID();
+    unsigned idx = cast<ConstantInt>(GEPIt.getOperand())->getLimitedValue();
+    // The first level index to get current constant.
+    GEPIt++;
+
+    uint64_t Key = getRegBindingKey(ID, idx);
+    switch (RC) {
+    default:
+      break;
+    case DXIL::ResourceClass::SRV:
+      if (m_SrvBindingInCB.count(Key))
+        RegBinding = m_SrvBindingInCB[Key];
+      break;
+    case DXIL::ResourceClass::UAV:
+      if (m_UavBindingInCB.count(Key))
+        RegBinding = m_UavBindingInCB[Key];
+      break;
+    case DXIL::ResourceClass::Sampler:
+      if (m_SamplerBindingInCB.count(Key))
+        RegBinding = m_SamplerBindingInCB[Key];
+      break;
+    }
+    if (RegBinding == UINT_MAX)
+      break;
+
+    // Calc RegBinding.
+    RegBinding += CalcRegBinding(GEPIt, E, M, RC);
+
     break;
   }
-  return UINT_MAX;
+  return RegBinding;
 }
 
 // TODO: Don't check names.

lib/HLSL/HLOperationLower.cpp

@@ -317,8 +317,9 @@ private:
     IRBuilder<> Builder(HLM.GetCtx());
     Value *zero = Builder.getInt32(0);
     for (; GEPIt != E; ++GEPIt, ++i) {
-      if (GEPIt->isArrayTy()) {
-        // Change array idx to 0 to make sure all array ptr share same key.
+      ConstantInt *ImmIdx = dyn_cast<ConstantInt>(GEPIt.getOperand());
+      if (!ImmIdx) {
+        // Remove dynamic indexing to avoid crash.
         idxList[i] = zero;
       }
     }
@@ -342,6 +343,12 @@ private:
     for (; GEPIt != E; ++GEPIt, ++i) {
       if (GEPIt->isArrayTy()) {
         arraySize *= GEPIt->getArrayNumElements();
+        if (!Name.empty())
+          Name += ".";
+        if (ConstantInt *ImmIdx = dyn_cast<ConstantInt>(GEPIt.getOperand())) {
+          unsigned idx = ImmIdx->getLimitedValue();
+          Name += std::to_string(idx);
+        }
       } else if (GEPIt->isStructTy()) {
         DxilStructAnnotation *typeAnnot =
             typeSys.GetStructAnnotation(cast<StructType>(*GEPIt));
@@ -356,9 +363,7 @@ private:
     }
 
     Type *Ty = CbPtr->getResultElementType();
-    if (arraySize > 1) {
-      Ty = ArrayType::get(Ty, arraySize);
-    }
+    // Not support resource array in cbuffer.
     unsigned ResBinding = HLM.GetBindingForResourceInCB(CbPtr, CbGV);
     return CreateResourceGV(Ty, Name, MD, ResBinding);
   }

tools/clang/lib/CodeGen/CGHLSLMS.cpp

@@ -102,6 +102,10 @@ private:
   llvm::DenseMap<HLSLBufferDecl *, uint32_t> constantBufMap;
   // Map for resource type to resource metadata value.
   std::unordered_map<llvm::Type *, MDNode*> resMetadataMap;
+  // Map from Constant to register bindings.
+  llvm::DenseMap<llvm::Constant *,
+                 llvm::SmallVector<std::pair<DXIL::ResourceClass, unsigned>, 1>>
+      constantRegBindingMap;
 
   bool  m_bDebugInfo;
   bool  m_bIsLib;
@@ -3104,6 +3108,7 @@ void CGMSHLSLRuntime::AddConstant(VarDecl *constDecl, HLCBuffer &CB) {
     return;
   }
   llvm::Constant *constVal = CGM.GetAddrOfGlobalVar(constDecl);
+  auto &regBindings = constantRegBindingMap[constVal];
 
   bool isGlobalCB = CB.GetID() == globalCBIndex;
   uint32_t offset = 0;
@@ -3130,8 +3135,8 @@ void CGMSHLSLRuntime::AddConstant(VarDecl *constDecl, HLCBuffer &CB) {
       break;
     }
     case hlsl::UnusualAnnotation::UA_RegisterAssignment: {
+      RegisterAssignment *ra = cast<RegisterAssignment>(it);
       if (isGlobalCB) {
-        RegisterAssignment *ra = cast<RegisterAssignment>(it);
         if (ra->RegisterSpace.hasValue()) {
           DiagnosticsEngine& Diags = CGM.getDiags();
           unsigned DiagID = Diags.getCustomDiagID(
@@ -3144,6 +3149,22 @@ void CGMSHLSLRuntime::AddConstant(VarDecl *constDecl, HLCBuffer &CB) {
         offset <<= 2;
         userOffset = true;
       }
+      switch (ra->RegisterType) {
+      default:
+        break;
+      case 't':
+        regBindings.emplace_back(
+            std::make_pair(DXIL::ResourceClass::SRV, ra->RegisterNumber));
+        break;
+      case 'u':
+        regBindings.emplace_back(
+            std::make_pair(DXIL::ResourceClass::UAV, ra->RegisterNumber));
+        break;
+      case 's':
+        regBindings.emplace_back(
+            std::make_pair(DXIL::ResourceClass::Sampler, ra->RegisterNumber));
+        break;
+      }
       break;
     }
     case hlsl::UnusualAnnotation::UA_SemanticDecl:
@@ -3342,6 +3363,45 @@ static unsigned AllocateDxilConstantBuffer(HLCBuffer &CB,
   return offset;
 }
 
+static void AddRegBindingsForResourceInConstantBuffer(
+    HLModule *pHLModule,
+    llvm::DenseMap<llvm::Constant *,
+                   llvm::SmallVector<std::pair<DXIL::ResourceClass, unsigned>,
+                                     1>> &constantRegBindingMap) {
+  for (unsigned i = 0; i < pHLModule->GetCBuffers().size(); i++) {
+    HLCBuffer &CB = *static_cast<HLCBuffer *>(&(pHLModule->GetCBuffer(i)));
+    auto &Constants = CB.GetConstants();
+    for (unsigned j = 0; j < Constants.size(); j++) {
+      const std::unique_ptr<DxilResourceBase> &C = Constants[j];
+      Constant *CGV = C->GetGlobalSymbol();
+      auto &regBindings = constantRegBindingMap[CGV];
+      if (regBindings.empty())
+        continue;
+      unsigned Srv = UINT_MAX;
+      unsigned Uav = UINT_MAX;
+      unsigned Sampler = UINT_MAX;
+      for (auto it : regBindings) {
+        unsigned RegNum = it.second;
+        switch (it.first) {
+        case DXIL::ResourceClass::SRV:
+          Srv = RegNum;
+          break;
+        case DXIL::ResourceClass::UAV:
+          Uav = RegNum;
+          break;
+        case DXIL::ResourceClass::Sampler:
+          Sampler = RegNum;
+          break;
+        default:
+          DXASSERT(0, "invalid resource class");
+          break;
+        }
+      }
+      pHLModule->AddRegBinding(CB.GetID(), j, Srv, Uav, Sampler);
+    }
+  }
+}
+
 static void AllocateDxilConstantBuffers(HLModule *pHLModule,
   std::unordered_map<Constant*, DxilFieldAnnotation> &constVarAnnotationMap) {
   for (unsigned i = 0; i < pHLModule->GetCBuffers().size(); i++) {
@@ -5203,6 +5263,9 @@ void CGMSHLSLRuntime::FinishCodeGen() {
     }
   }
 
+  // Add Reg bindings for resource in cb.
+  AddRegBindingsForResourceInConstantBuffer(m_pHLModule, constantRegBindingMap);
+
   // Allocate constant buffers.
   AllocateDxilConstantBuffers(m_pHLModule, m_ConstVarAnnotationMap);
   // TODO: create temp variable for constant which has store use.

tools/clang/test/HLSLFileCheck/hlsl/objects/CbufferLegacy/resource-in-cb3.hlsl

@@ -1,6 +1,6 @@
 // RUN: %dxc -E main -T ps_6_0 %s  | FileCheck %s
 
-// CHECK: Tex1                              texture     f32          2d      T0             t0     2
+// CHECK: Tex1.0                              texture     f32          2d      T0             t0     1
 
 SamplerState Samp;
 cbuffer CB

tools/clang/test/HLSLFileCheck/hlsl/objects/CbufferLegacy/resource-in-cb4.hlsl

@@ -1,6 +1,6 @@
 // RUN: %dxc -E main -T ps_6_0 %s  | FileCheck %s
 
-// CHECK: Tex1                              texture     f32          2d      T0             t0     4
+// CHECK:Index for resource array inside cbuffer must be a literal expression
 
 
 SamplerState Samp;

tools/clang/test/HLSLFileCheck/hlsl/resource_binding/res_in_cb1.hlsl

@@ -0,0 +1,36 @@
+// RUN: %dxc -E main -T ps_6_0 %s  | FileCheck %s
+
+// Make sure register binding on struct works.
+//CHECK: tx0.s                             sampler      NA          NA      S0             s0     1
+//CHECK: tx1.s                             sampler      NA          NA      S1             s1     1
+//CHECK: s                                 sampler      NA          NA      S2             s3     1
+//CHECK: tx0.t2                            texture     f32          2d      T0             t1     1
+//CHECK: tx0.t                             texture     f32          2d      T1             t0     1
+//CHECK: tx1.t2                            texture     f32          2d      T2             t6     1
+//CHECK: tx1.t                             texture     f32          2d      T3             t5     1
+//CHECK: x                                 texture     f32          2d      T4             t3     1
+
+struct LegacyTex
+{
+	Texture2D t;
+        Texture2D t2;
+	SamplerState  s;
+};
+
+LegacyTex tx0 : register(t0) : register(s0);
+LegacyTex tx1 : register(t5) : register(s1);
+
+float4 tex2D(LegacyTex tx, float2 uv)
+{
+	return tx.t.Sample(tx.s,uv) + tx.t2.Sample(tx.s, uv);
+}
+
+cbuffer n {
+   Texture2D x: register(t3);
+   SamplerState s : register(s3);
+}
+
+float4 main(float2 uv:UV) : SV_Target
+{
+	return tex2D(tx0,uv) + tex2D(tx1,uv) + x.Sample(s,uv);
+}

tools/clang/test/HLSLFileCheck/hlsl/resource_binding/res_in_cb2.hlsl

@@ -0,0 +1,26 @@
+// RUN: %dxc -E main -T ps_6_0 %s  | FileCheck %s
+
+// Make sure array of struct with resource works.
+
+//CHECK: x.1.s                             sampler      NA          NA      S0             s4     1
+//CHECK: x.0.s                             sampler      NA          NA      S1             s3     1
+//CHECK: x.1.x                             texture     f32          2d      T0             t5     1
+//CHECK: x.0.y                             texture     f32          2d      T1             t4     1
+//CHECK: m                                 texture     f32          2d      T2             t9     1
+
+struct X {
+   Texture2D x;
+   SamplerState s ;
+   Texture2D y;
+};
+
+X x[2] : register(t3) : register(s3);
+
+cbuffer A {
+  Texture2D m : register(t9);
+}
+
+float4 main(float2 uv:UV, uint i:I) : SV_Target
+{
+	return x[1].x.Sample(x[1].s,uv) + x[0].y.Sample(x[0].s, uv) + m.Sample(x[0].s, uv);
+}

tools/clang/test/HLSLFileCheck/hlsl/resource_binding/res_in_cb3.hlsl

@@ -0,0 +1,17 @@
+// RUN: %dxc -E main -T ps_6_0 %s  | FileCheck %s
+
+// Make sure report error when dynamic indexing on resource array inside cbuffer.
+//CHECK:Index for resource array inside cbuffer must be a literal expression
+
+struct X {
+   Texture2D x;
+   SamplerState s ;
+   Texture2D y;
+};
+
+X x[2] : register(t3) : register(s3);
+
+float4 main(float2 uv:UV, uint i:I) : SV_Target
+{
+	return x[1].x.Sample(x[1].s,uv) + x[0].y.Sample(x[i].s, uv);
+}

tools/clang/test/HLSLFileCheck/hlsl/resource_binding/resource-in-struct3.hlsl

@@ -1,6 +1,6 @@
 // RUN: %dxc -E main -T ps_6_0 %s  | FileCheck %s
 
-// CHECK: var.res.Tex1                      texture     f32          2d      T0             t0     2
+// CHECK: var.res.Tex1.0                      texture     f32          2d      T0             t0     1
 
 SamplerState Samp;
 struct Resource