[HLSL] Implement explicit layout for default constant buffer ($Globals) (#128991)

Processes `HLSLResourceBindingAttr` attributes that represent
`register(c#)` annotations on default constant buffer declarations and
applies its value to the buffer layout. Any default buffer declarations
without an explicit `register(c#)` annotation are placed after the
elements with explicit layout.

This PR also adds a test case for a `cbuffer` that does not have
`packoffset` on all declarations. Same layout rules apply here as well.

Fixes #126791

Author: hekota

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);
 

clang/lib/CodeGen/CGHLSLRuntime.h

@@ -146,7 +146,7 @@ class CGHLSLRuntime {
 
   llvm::Type *
   convertHLSLSpecificType(const Type *T,
-                          SmallVector<unsigned> *Packoffsets = nullptr);
+                          SmallVector<int32_t> *Packoffsets = nullptr);
 
   void annotateHLSLResource(const VarDecl *D, llvm::GlobalVariable *GV);
   void generateGlobalCtorDtorCalls();

clang/lib/CodeGen/HLSLBufferLayoutBuilder.cpp

@@ -10,6 +10,7 @@
 #include "CGHLSLRuntime.h"
 #include "CodeGenModule.h"
 #include "clang/AST/Type.h"
+#include <climits>
 
 //===----------------------------------------------------------------------===//
 // Implementation of constant buffer layout common between DirectX and
@@ -18,6 +19,7 @@
 
 using namespace clang;
 using namespace clang::CodeGen;
+using llvm::hlsl::CBufferRowSizeInBytes;
 
 namespace {
 
@@ -51,16 +53,22 @@ namespace clang {
 namespace CodeGen {
 
 // Creates a layout type for given struct with HLSL constant buffer layout
-// taking into account Packoffsets, if provided.
+// taking into account PackOffsets, if provided.
 // Previously created layout types are cached by CGHLSLRuntime.
 //
 // The function iterates over all fields of the StructType (including base
 // classes) and calls layoutField to converts each field to its corresponding
 // LLVM type and to calculate its HLSL constant buffer layout. Any embedded
 // structs (or arrays of structs) are converted to target layout types as well.
+//
+// When PackOffsets are specified the elements will be placed based on the
+// user-specified offsets. Not all elements must have a packoffset/register(c#)
+// annotation though. For those that don't, the PackOffsets array will contain
+// -1 value instead. These elements must be placed at the end of the layout
+// after all of the elements with specific offset.
 llvm::TargetExtType *HLSLBufferLayoutBuilder::createLayoutType(
     const RecordType *StructType,
-    const llvm::SmallVector<unsigned> *Packoffsets) {
+    const llvm::SmallVector<int32_t> *PackOffsets) {
 
   // check if we already have the layout type for this struct
   if (llvm::TargetExtType *Ty =
@@ -72,6 +80,8 @@ llvm::TargetExtType *HLSLBufferLayoutBuilder::createLayoutType(
   unsigned Index = 0; // packoffset index
   unsigned EndOffset = 0;
 
+  SmallVector<std::pair<const FieldDecl *, unsigned>> DelayLayoutFields;
+
   // reserve first spot in the layout vector for buffer size
   Layout.push_back(0);
 
@@ -84,22 +94,55 @@ llvm::TargetExtType *HLSLBufferLayoutBuilder::createLayoutType(
            "HLSL doesn't support multiple inheritance");
     RecordTypes.push_back(D->bases_begin()->getType()->getAs<RecordType>());
   }
+
+  unsigned FieldOffset;
+  llvm::Type *FieldType;
+
   while (!RecordTypes.empty()) {
     const RecordType *RT = RecordTypes.back();
     RecordTypes.pop_back();
 
     for (const auto *FD : RT->getDecl()->fields()) {
-      assert((!Packoffsets || Index < Packoffsets->size()) &&
-             "number of elements in layout struct does not "
-             "match number of packoffset annotations");
-
-      if (!layoutField(FD, EndOffset, Layout, LayoutElements,
-                       Packoffsets ? (*Packoffsets)[Index] : -1))
-        return nullptr;
-      Index++;
+      assert((!PackOffsets || Index < PackOffsets->size()) &&
+             "number of elements in layout struct does not match number of "
+             "packoffset annotations");
+
+      // No PackOffset info at all, or have a valid packoffset/register(c#)
+      // annotations value -> layout the field.
+      const int PO = PackOffsets ? (*PackOffsets)[Index++] : -1;
+      if (!PackOffsets || PO != -1) {
+        if (!layoutField(FD, EndOffset, FieldOffset, FieldType, PO))
+          return nullptr;
+        Layout.push_back(FieldOffset);
+        LayoutElements.push_back(FieldType);
+        continue;
+      }
+      // Have PackOffset info, but there is no packoffset/register(cX)
+      // annotation on this field. Delay the layout until after all of the
+      // other elements with packoffsets/register(cX) are processed.
+      DelayLayoutFields.emplace_back(FD, LayoutElements.size());
+      // reserve space for this field in the layout vector and elements list
+      Layout.push_back(UINT_MAX);
+      LayoutElements.push_back(nullptr);
     }
   }
 
+  // process delayed layouts
+  for (auto I : DelayLayoutFields) {
+    const FieldDecl *FD = I.first;
+    const unsigned IndexInLayoutElements = I.second;
+    // the first item in layout vector is size, so we need to offset the index
+    // by 1
+    const unsigned IndexInLayout = IndexInLayoutElements + 1;
+    assert(Layout[IndexInLayout] == UINT_MAX &&
+           LayoutElements[IndexInLayoutElements] == nullptr);
+
+    if (!layoutField(FD, EndOffset, FieldOffset, FieldType))
+      return nullptr;
+    Layout[IndexInLayout] = FieldOffset;
+    LayoutElements[IndexInLayoutElements] = FieldType;
+  }
+
   // set the size of the buffer
   Layout[0] = EndOffset;
 
@@ -122,16 +165,19 @@ llvm::TargetExtType *HLSLBufferLayoutBuilder::createLayoutType(
 // The function converts a single field of HLSL Buffer to its corresponding
 // LLVM type and calculates it's layout. Any embedded structs (or
 // arrays of structs) are converted to target layout types as well.
-// The converted type is appended to the LayoutElements list, the element
-// offset is added to the Layout list and the EndOffset updated to the offset
-// just after the lay-ed out element (which is basically the size of the
-// buffer).
+// The converted type is set to the FieldType parameter, the element
+// offset is set to the FieldOffset parameter. The EndOffset (=size of the
+// buffer) is also updated accordingly to the offset just after the placed
+// element, unless the incoming EndOffset already larger (may happen in case
+// of unsorted packoffset annotations).
 // Returns true if the conversion was successful.
 // The packoffset parameter contains the field's layout offset provided by the
 // user or -1 if there was no packoffset (or register(cX)) annotation.
-bool HLSLBufferLayoutBuilder::layoutField(
-    const FieldDecl *FD, unsigned &EndOffset, SmallVector<unsigned> &Layout,
-    SmallVector<llvm::Type *> &LayoutElements, int Packoffset) {
+bool HLSLBufferLayoutBuilder::layoutField(const FieldDecl *FD,
+                                          unsigned &EndOffset,
+                                          unsigned &FieldOffset,
+                                          llvm::Type *&FieldType,
+                                          int Packoffset) {
 
   // Size of element; for arrays this is a size of a single element in the
   // array. Total array size of calculated as (ArrayCount-1) * ArrayStride +
@@ -141,8 +187,7 @@ bool HLSLBufferLayoutBuilder::layoutField(
   unsigned ArrayCount = 1;
   unsigned ArrayStride = 0;
 
-  const unsigned BufferRowAlign = 16U;
-  unsigned NextRowOffset = llvm::alignTo(EndOffset, BufferRowAlign);
+  unsigned NextRowOffset = llvm::alignTo(EndOffset, CBufferRowSizeInBytes);
 
   llvm::Type *ElemLayoutTy = nullptr;
   QualType FieldTy = FD->getType();
@@ -172,7 +217,7 @@ bool HLSLBufferLayoutBuilder::layoutField(
           getScalarOrVectorSizeInBytes(CGM.getTypes().ConvertTypeForMem(Ty));
       ElemLayoutTy = CGM.getTypes().ConvertTypeForMem(FieldTy);
     }
-    ArrayStride = llvm::alignTo(ElemSize, BufferRowAlign);
+    ArrayStride = llvm::alignTo(ElemSize, CBufferRowSizeInBytes);
     ElemOffset = (Packoffset != -1) ? Packoffset : NextRowOffset;
 
   } else if (FieldTy->isStructureType()) {
@@ -220,8 +265,8 @@ bool HLSLBufferLayoutBuilder::layoutField(
   EndOffset = std::max<unsigned>(EndOffset, NewEndOffset);
 
   // add the layout element and offset to the lists
-  Layout.push_back(ElemOffset);
-  LayoutElements.push_back(ElemLayoutTy);
+  FieldOffset = ElemOffset;
+  FieldType = ElemLayoutTy;
   return true;
 }
 

clang/lib/CodeGen/HLSLBufferLayoutBuilder.h

@@ -35,13 +35,12 @@ class HLSLBufferLayoutBuilder {
   // the Layout is the size followed by offsets for each struct element.
   llvm::TargetExtType *
   createLayoutType(const RecordType *StructType,
-                   const llvm::SmallVector<unsigned> *Packoffsets = nullptr);
+                   const llvm::SmallVector<int32_t> *Packoffsets = nullptr);
 
 private:
   bool layoutField(const clang::FieldDecl *FD, unsigned &EndOffset,
-                   llvm::SmallVector<unsigned> &Layout,
-                   llvm::SmallVector<llvm::Type *> &LayoutElements,
-                   int Packoffset);
+                   unsigned &FieldOffset, llvm::Type *&FieldType,
+                   int Packoffset = -1);
 };
 
 } // namespace CodeGen

clang/lib/CodeGen/TargetInfo.h

@@ -441,7 +441,7 @@ class TargetCodeGenInfo {
   /// Return an LLVM type that corresponds to a HLSL type
   virtual llvm::Type *
   getHLSLType(CodeGenModule &CGM, const Type *T,
-              const SmallVector<unsigned> *Packoffsets = nullptr) const {
+              const SmallVector<int32_t> *Packoffsets = nullptr) const {
     return nullptr;
   }
 

clang/lib/CodeGen/Targets/DirectX.cpp

@@ -29,14 +29,14 @@ class DirectXTargetCodeGenInfo : public TargetCodeGenInfo {
   DirectXTargetCodeGenInfo(CodeGen::CodeGenTypes &CGT)
       : TargetCodeGenInfo(std::make_unique<DefaultABIInfo>(CGT)) {}
 
-  llvm::Type *getHLSLType(
-      CodeGenModule &CGM, const Type *T,
-      const SmallVector<unsigned> *Packoffsets = nullptr) const override;
+  llvm::Type *
+  getHLSLType(CodeGenModule &CGM, const Type *T,
+              const SmallVector<int32_t> *Packoffsets = nullptr) const override;
 };
 
 llvm::Type *DirectXTargetCodeGenInfo::getHLSLType(
     CodeGenModule &CGM, const Type *Ty,
-    const SmallVector<unsigned> *Packoffsets) const {
+    const SmallVector<int32_t> *Packoffsets) const {
   auto *ResType = dyn_cast<HLSLAttributedResourceType>(Ty);
   if (!ResType)
     return nullptr;

clang/lib/CodeGen/Targets/SPIR.cpp

@@ -52,9 +52,9 @@ class CommonSPIRTargetCodeGenInfo : public TargetCodeGenInfo {
 
   unsigned getOpenCLKernelCallingConv() const override;
   llvm::Type *getOpenCLType(CodeGenModule &CGM, const Type *T) const override;
-  llvm::Type *getHLSLType(
-      CodeGenModule &CGM, const Type *Ty,
-      const SmallVector<unsigned> *Packoffsets = nullptr) const override;
+  llvm::Type *
+  getHLSLType(CodeGenModule &CGM, const Type *Ty,
+              const SmallVector<int32_t> *Packoffsets = nullptr) const override;
   llvm::Type *getSPIRVImageTypeFromHLSLResource(
       const HLSLAttributedResourceType::Attributes &attributes,
       llvm::Type *ElementType, llvm::LLVMContext &Ctx) const;
@@ -371,7 +371,7 @@ llvm::Type *CommonSPIRTargetCodeGenInfo::getOpenCLType(CodeGenModule &CGM,
 
 llvm::Type *CommonSPIRTargetCodeGenInfo::getHLSLType(
     CodeGenModule &CGM, const Type *Ty,
-    const SmallVector<unsigned> *Packoffsets) const {
+    const SmallVector<int32_t> *Packoffsets) const {
   auto *ResType = dyn_cast<HLSLAttributedResourceType>(Ty);
   if (!ResType)
     return nullptr;

clang/lib/Sema/SemaHLSL.cpp

@@ -1961,6 +1961,17 @@ void SemaHLSL::ActOnEndOfTranslationUnit(TranslationUnitDecl *TU) {
     SemaRef.getCurLexicalContext()->addDecl(DefaultCBuffer);
     createHostLayoutStructForBuffer(SemaRef, DefaultCBuffer);
 
+    // Set HasValidPackoffset if any of the decls has a register(c#) annotation;
+    for (const Decl *VD : DefaultCBufferDecls) {
+      const HLSLResourceBindingAttr *RBA =
+          VD->getAttr<HLSLResourceBindingAttr>();
+      if (RBA &&
+          RBA->getRegisterType() == HLSLResourceBindingAttr::RegisterType::C) {
+        DefaultCBuffer->setHasValidPackoffset(true);
+        break;
+      }
+    }
+
     DeclGroupRef DG(DefaultCBuffer);
     SemaRef.Consumer.HandleTopLevelDecl(DG);
   }

clang/test/CodeGenHLSL/cbuffer_with_packoffset.hlsl

@@ -3,6 +3,7 @@
 // RUN:   -emit-llvm -disable-llvm-passes -o - | FileCheck %s
 
 // CHECK: %__cblayout_CB = type <{ float, double, <2 x i32> }>
+// CHECK: %__cblayout_CB_1 = type <{ float, <2 x float> }>
 
 // CHECK: @CB.cb = external constant target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 176, 16, 168, 88))
 // CHECK: @a = external addrspace(2) global float, align 4
@@ -15,6 +16,17 @@ cbuffer CB : register(b1, space3) {
   int2 c : packoffset(c5.z);
 }
 
+// CHECK: @CB.cb.1 = external constant target("dx.CBuffer", target("dx.Layout", %__cblayout_CB_1, 92, 88, 80))
+// CHECK: @x = external addrspace(2) global float, align 4
+// CHECK: @y = external addrspace(2) global <2 x float>, align 8
+
+// Missing packoffset annotation will produce a warning.
+// Element x will be placed after the element y that has an explicit packoffset.
+cbuffer CB : register(b0) {
+  float x;
+  float2 y : packoffset(c5);
+}
+
 // CHECK: define internal void @_init_resource_CB.cb()
 // CHECK-NEXT: entry:
 // CHECK-NEXT: %CB.cb_h = call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 176, 16, 168, 88))
@@ -34,5 +46,6 @@ void main() {
   foo();
 }
 
-// CHECK: !hlsl.cbs = !{![[CB:[0-9]+]]}
-// CHECK: ![[CB]] = !{ptr @CB.cb, ptr addrspace(2) @a, ptr addrspace(2) @b, ptr addrspace(2) @c}
+// CHECK: !hlsl.cbs = !{![[CB1:[0-9]+]], ![[CB2:[0-9]+]]}
+// CHECK: ![[CB1]] = !{ptr @CB.cb, ptr addrspace(2) @a, ptr addrspace(2) @b, ptr addrspace(2) @c}
+// CHECK: ![[CB2]] = !{ptr @CB.cb.1, ptr addrspace(2) @x, ptr addrspace(2) @y}