1 files changed, 38 insertions, 7 deletions

diff --git a/clang/lib/CodeGen/CGHLSLRuntime.cpp b/clang/lib/CodeGen/CGHLSLRuntime.cpp
index dc34653..5916fa6 100644
--- a/clang/lib/CodeGen/CGHLSLRuntime.cpp
+++ b/clang/lib/CodeGen/CGHLSLRuntime.cpp
@@ -40,6 +40,8 @@ using namespace CodeGen;
 using namespace clang::hlsl;
 using namespace llvm;
 
+using llvm::hlsl::CBufferRowSizeInBytes;
+
 static void createResourceInitFn(CodeGenModule &CGM, llvm::GlobalVariable *GV,
                                  unsigned Slot, unsigned Space);
 
@@ -70,7 +72,7 @@ void addDxilValVersion(StringRef ValVersionStr, llvm::Module &M) {
 
 llvm::Type *
 CGHLSLRuntime::convertHLSLSpecificType(const Type *T,
-                                       SmallVector<unsigned> *Packoffsets) {
+                                       SmallVector<int32_t> *Packoffsets) {
   assert(T->isHLSLSpecificType() && "Not an HLSL specific type!");
 
   // Check if the target has a specific translation for this type first.
@@ -174,21 +176,50 @@ createBufferHandleType(const HLSLBufferDecl *BufDecl) {
   return cast<HLSLAttributedResourceType>(QT.getTypePtr());
 }
 
+// Iterates over all declarations in the HLSL buffer and based on the
+// packoffset or register(c#) annotations it fills outs the Layout
+// vector with the user-specified layout offsets.
+// The buffer offsets can be specified 2 ways:
+// 1. declarations in cbuffer {} block can have a packoffset annotation
+//    (translates to HLSLPackOffsetAttr)
+// 2. default constant buffer declarations at global scope can have
+//    register(c#) annotations (translates to HLSLResourceBindingAttr with
+//    RegisterType::C)
+// It is not guaranteed that all declarations in a buffer have an annotation.
+// For those where it is not specified a -1 value is added to the Layout
+// vector. In the final layout these declarations will be placed at the end
+// of the HLSL buffer after all of the elements with specified offset.
 static void fillPackoffsetLayout(const HLSLBufferDecl *BufDecl,
-                                 SmallVector<unsigned> &Layout) {
+                                 SmallVector<int32_t> &Layout) {
   assert(Layout.empty() && "expected empty vector for layout");
   assert(BufDecl->hasValidPackoffset());
 
-  for (Decl *D : BufDecl->decls()) {
+  for (Decl *D : BufDecl->buffer_decls()) {
     if (isa<CXXRecordDecl, EmptyDecl>(D) || isa<FunctionDecl>(D)) {
       continue;
     }
     VarDecl *VD = dyn_cast<VarDecl>(D);
     if (!VD || VD->getType().getAddressSpace() != LangAS::hlsl_constant)
       continue;
-    assert(VD->hasAttr<HLSLPackOffsetAttr>() &&
-           "expected packoffset attribute on every declaration");
-    size_t Offset = VD->getAttr<HLSLPackOffsetAttr>()->getOffsetInBytes();
+
+    if (!VD->hasAttrs()) {
+      Layout.push_back(-1);
+      continue;
+    }
+
+    int32_t Offset = -1;
+    for (auto *Attr : VD->getAttrs()) {
+      if (auto *POA = dyn_cast<HLSLPackOffsetAttr>(Attr)) {
+        Offset = POA->getOffsetInBytes();
+        break;
+      }
+      auto *RBA = dyn_cast<HLSLResourceBindingAttr>(Attr);
+      if (RBA &&
+          RBA->getRegisterType() == HLSLResourceBindingAttr::RegisterType::C) {
+        Offset = RBA->getSlotNumber() * CBufferRowSizeInBytes;
+        break;
+      }
+    }
     Layout.push_back(Offset);
   }
 }
@@ -207,7 +238,7 @@ void CGHLSLRuntime::addBuffer(const HLSLBufferDecl *BufDecl) {
     return;
 
   // create global variable for the constant buffer
-  SmallVector<unsigned> Layout;
+  SmallVector<int32_t> Layout;
   if (BufDecl->hasValidPackoffset())
     fillPackoffsetLayout(BufDecl, Layout);