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

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -70,7 +70,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.
@@ -179,21 +179,45 @@ 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 quaranteed 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();
+    size_t Offset = -1;
+    if (VD->hasAttrs()) {
+      for (auto *Attr : VD->getAttrs()) {
+        if (auto *POA = dyn_cast<HLSLPackOffsetAttr>(Attr)) {
+          Offset = POA->getOffsetInBytes();
+        } else if (auto *RBA = dyn_cast<HLSLResourceBindingAttr>(Attr)) {
+          if (RBA->getRegisterType() ==
+              HLSLResourceBindingAttr::RegisterType::C) {
+            // size of constant buffer row is 16 bytes
+            Offset = RBA->getSlotNumber() * 16U;
+          }
+        }
+      }
+    }
     Layout.push_back(Offset);
   }
 }
@@ -212,7 +236,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
@@ -58,9 +59,15 @@ namespace CodeGen {
 // 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 constain
+// -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 +79,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);
 
@@ -89,14 +98,57 @@ llvm::TargetExtType *HLSLBufferLayoutBuilder::createLayoutType(
     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");
+      unsigned FieldOffset = UINT_MAX;
+      llvm::Type *FieldType = nullptr;
+
+      if (Packoffsets) {
+        // have packoffset/register(c#) annotations
+        assert(Index < Packoffsets->size() &&
+               "number of elements in layout struct does not match number of "
+               "packoffset annotations");
+        int PO = (*Packoffsets)[Index++];
+        if (PO != -1) {
+          if (!layoutField(FD, EndOffset, FieldOffset, FieldType, PO))
+            return nullptr;
+        } else {
+          // No packoffset/register(cX) annotation on this field;
+          // Delay the layout until after all of the other elements
+          // annotated 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);
+          continue;
+        }
+      } else {
+        if (!layoutField(FD, EndOffset, FieldOffset, FieldType))
+          return nullptr;
+      }
+
+      assert(FieldOffset != UINT_MAX && FieldType != nullptr);
+      Layout.push_back((unsigned)FieldOffset);
+      LayoutElements.push_back(FieldType);
+    }
+  }
 
-      if (!layoutField(FD, EndOffset, Layout, LayoutElements,
-                       Packoffsets ? (*Packoffsets)[Index] : -1))
+  // process delayed layouts
+  if (!DelayLayoutFields.empty()) {
+    for (auto I : DelayLayoutFields) {
+      const FieldDecl *FD = I.first;
+      unsigned IndexInLayoutElements = I.second;
+      // the first item in layout vector is size, so we need to offset the index
+      // by 1
+      unsigned IndexInLayout = IndexInLayoutElements + 1;
+      assert(Layout[IndexInLayout] == UINT_MAX &&
+             LayoutElements[IndexInLayoutElements] == nullptr);
+
+      unsigned FieldOffset = UINT_MAX;
+      llvm::Type *FieldType = nullptr;
+      if (!layoutField(FD, EndOffset, FieldOffset, FieldType))
         return nullptr;
-      Index++;
+
+      Layout[IndexInLayout] = (unsigned)FieldOffset;
+      LayoutElements[IndexInLayoutElements] = FieldType;
     }
   }
 
@@ -122,16 +174,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 +
@@ -220,8 +275,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,18 @@ 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) {
+      if (const HLSLResourceBindingAttr *RBA =
+              VD->getAttr<HLSLResourceBindingAttr>()) {
+        if (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}

clang/test/CodeGenHLSL/default_cbuffer_with_layout.hlsl

@@ -0,0 +1,37 @@
+// RUN: %clang_cc1 -finclude-default-header -triple dxil-pc-shadermodel6.3-compute \
+// RUN:            -emit-llvm -disable-llvm-passes -o - %s | FileCheck %s
+
+// CHECK: %"__cblayout_$Globals" = type <{ i32, float, [4 x double], <4 x i32>, <4 x float>,
+// CHECK-SAME: target("dx.Layout", %S, 8, 0) }>
+// CHECK: %S = type <{ <2 x float> }>
+
+// CHECK-DAG: @b = external addrspace(2) global float, align 4
+// CHECK-DAG: @d = external addrspace(2) global <4 x i32>, align 16
+// CHECK-DAG: @"$Globals.cb" = external constant target("dx.CBuffer",
+// CHECK-DAG-SAME: target("dx.Layout", %"__cblayout_$Globals", 144, 120, 16, 32, 64, 128, 112))
+// CHECK-DAG: @a = external addrspace(2) global i32, align 4
+// CHECK-DAG: @c = external addrspace(2) global [4 x double], align 8
+// CHECK-DAG: @e = external addrspace(2) global <4 x float>, align 16
+// CHECK-DAG: @s = external addrspace(2) global target("dx.Layout", %S, 8, 0), align 8
+
+struct S {
+  float2 v;
+}; 
+
+int a;
+float b : register(c1);
+double c[4] : register(c2);
+int4 d : register(c4);
+float4 e;
+S s : register(c7);
+
+RWBuffer<float> Buf;
+
+[numthreads(4,1,1)]
+void main() {
+  Buf[0] = a;
+}
+
+// CHECK: !hlsl.cbs = !{![[CB:.*]]}
+// CHECK: ![[CB]] = !{ptr @"$Globals.cb", ptr addrspace(2) @a, ptr addrspace(2) @b, ptr addrspace(2) @c,
+// CHECK-SAME: ptr addrspace(2) @d, ptr addrspace(2) @e, ptr addrspace(2) @s}