[SPIRV] Support Peeled Array Layouts for HLSL CBuffers

This commit adds support for 'peeled arrays' in HLSL constant buffers.
HLSL CBuffers may have padding between array elements but not after the
last element. This is represented in LLVM IR as {[N-1 x {T, pad}], T}.

Changes include:
- Recognition of the peeled array pattern.
- Logic to reconstitute these into SPIR-V compatible arrays.
- Support for spirv.Padding type in GlobalRegistry and Builtins.
- Updates to SPIRVCBufferAccess to correctly calculate member offsets
  in these padded structures.

Depends on #169076

Author: s-perron

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

@@ -3373,6 +3373,8 @@ SPIRVType *lowerBuiltinType(const Type *OpaqueType,
     TargetType = getInlineSpirvType(BuiltinType, MIRBuilder, GR);
   } else if (Name == "spirv.VulkanBuffer") {
     TargetType = getVulkanBufferType(BuiltinType, MIRBuilder, GR);
+  } else if (Name == "spirv.Padding") {
+    TargetType = GR->getOrCreatePaddingType(MIRBuilder);
   } else if (Name == "spirv.Layout") {
     TargetType = getLayoutType(BuiltinType, MIRBuilder, GR);
   } else {

llvm/lib/Target/SPIRV/SPIRVCBufferAccess.cpp

@@ -79,15 +79,20 @@ static bool replaceCBufferAccesses(Module &M) {
     // The handle definition should dominate all uses of the cbuffer members.
     // We'll insert our getpointer calls right after it.
     IRBuilder<> Builder(HandleDef->getNextNode());
+    auto *HandleTy = cast<TargetExtType>(Mapping.Handle->getValueType());
+    auto *LayoutTy = cast<StructType>(HandleTy->getTypeParameter(0));
+    const StructLayout *SL = M.getDataLayout().getStructLayout(LayoutTy);
 
-    for (uint32_t Index = 0; Index < Mapping.Members.size(); ++Index) {
-      GlobalVariable *MemberGV = Mapping.Members[Index].GV;
+    for (const hlsl::CBufferMember &Member : Mapping.Members) {
+      GlobalVariable *MemberGV = Member.GV;
       if (MemberGV->use_empty()) {
         continue;
       }
 
+      uint32_t IndexInStruct = SL->getElementContainingOffset(Member.Offset);
+
       // Create the getpointer intrinsic call.
-      Value *IndexVal = Builder.getInt32(Index);
+      Value *IndexVal = Builder.getInt32(IndexInStruct);
       Type *PtrType = MemberGV->getType();
       Value *GetPointerCall = Builder.CreateIntrinsic(
           PtrType, Intrinsic::spv_resource_getpointer, {HandleDef, IndexVal});

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

@@ -841,6 +841,7 @@ Type *SPIRVEmitIntrinsics::deduceElementTypeHelper(
           uint32_t Index = cast<ConstantInt>(II->getOperand(1))->getZExtValue();
           Ty = cast<StructType>(Ty)->getElementType(Index);
         }
+        Ty = reconstitutePeeledArrayType(Ty);
       } else {
         llvm_unreachable("Unknown handle type for spv_resource_getpointer.");
       }

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp

@@ -889,6 +889,17 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeStruct(
     const StructType *Ty, MachineIRBuilder &MIRBuilder,
     SPIRV::AccessQualifier::AccessQualifier AccQual,
     StructOffsetDecorator Decorator, bool EmitIR) {
+  Type *OriginalElementType = nullptr;
+  uint64_t TotalSize = 0;
+  if (matchPeeledArrayPattern(Ty, OriginalElementType, TotalSize)) {
+    SPIRVType *ElementSPIRVType = findSPIRVType(
+        OriginalElementType, MIRBuilder, AccQual,
+        /* ExplicitLayoutRequired= */ Decorator != nullptr, EmitIR);
+    return getOpTypeArray(TotalSize, ElementSPIRVType, MIRBuilder,
+                          /*ExplicitLayoutRequired=*/Decorator != nullptr,
+                          EmitIR);
+  }
+
   const SPIRVSubtarget &ST =
       cast<SPIRVSubtarget>(MIRBuilder.getMF().getSubtarget());
   SmallVector<Register, 4> FieldTypes;
@@ -1405,6 +1416,18 @@ SPIRVType *SPIRVGlobalRegistry::getOrCreateVulkanBufferType(
   return R;
 }
 
+SPIRVType *
+SPIRVGlobalRegistry::getOrCreatePaddingType(MachineIRBuilder &MIRBuilder) {
+  auto Key = SPIRV::irhandle_padding();
+  if (const MachineInstr *MI = findMI(Key, &MIRBuilder.getMF()))
+    return MI;
+  auto *T = Type::getInt8Ty(MIRBuilder.getContext());
+  SPIRVType *R = getOrCreateSPIRVIntegerType(8, MIRBuilder);
+  finishCreatingSPIRVType(T, R);
+  add(Key, R);
+  return R;
+}
+
 SPIRVType *SPIRVGlobalRegistry::getOrCreateLayoutType(
     MachineIRBuilder &MIRBuilder, const TargetExtType *T, bool EmitIr) {
   auto Key = SPIRV::handle(T);

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h

@@ -611,6 +611,8 @@ class SPIRVGlobalRegistry : public SPIRVIRMapping {
                                          SPIRV::StorageClass::StorageClass SC,
                                          bool IsWritable, bool EmitIr = false);
 
+  SPIRVType *getOrCreatePaddingType(MachineIRBuilder &MIRBuilder);
+
   SPIRVType *getOrCreateLayoutType(MachineIRBuilder &MIRBuilder,
                                    const TargetExtType *T, bool EmitIr = false);
 

llvm/lib/Target/SPIRV/SPIRVIRMapping.h

@@ -64,6 +64,7 @@ enum SpecialTypeKind {
   STK_Value,
   STK_MachineInstr,
   STK_VkBuffer,
+  STK_Padding,
   STK_ExplictLayoutType,
   STK_Last = -1
 };
@@ -149,6 +150,10 @@ inline IRHandle irhandle_vkbuffer(const Type *ElementType,
                          SpecialTypeKind::STK_VkBuffer);
 }
 
+inline IRHandle irhandle_padding() {
+  return std::make_tuple(nullptr, 0, SpecialTypeKind::STK_Padding);
+}
+
 inline IRHandle irhandle_explict_layout_type(const Type *Ty) {
   const Type *WrpTy = unifyPtrType(Ty);
   return irhandle_ptr(WrpTy, Ty->getTypeID(), STK_ExplictLayoutType);

llvm/lib/Target/SPIRV/SPIRVUtils.cpp

@@ -1042,6 +1042,71 @@ getFirstValidInstructionInsertPoint(MachineBasicBlock &BB) {
                                                                      : VarPos;
 }
 
+bool matchPeeledArrayPattern(const StructType *Ty, Type *&OriginalElementType,
+                             uint64_t &TotalSize) {
+  // An array of N padded structs is represented as {[N-1 x <{T, pad}>], T}.
+  if (!Ty->isStructTy() || Ty->getStructNumElements() != 2)
+    return false;
+
+  Type *FirstElement = Ty->getStructElementType(0);
+  Type *SecondElement = Ty->getStructElementType(1);
+
+  if (!FirstElement->isArrayTy())
+    return false;
+
+  Type *ArrayElementType = FirstElement->getArrayElementType();
+  if (!ArrayElementType->isStructTy() ||
+      ArrayElementType->getStructNumElements() != 2)
+    return false;
+
+  Type *T_in_struct = ArrayElementType->getStructElementType(0);
+  if (T_in_struct != SecondElement)
+    return false;
+
+  const uint64_t ArraySize = FirstElement->getArrayNumElements();
+  TotalSize = ArraySize + 1;
+  OriginalElementType = ArrayElementType;
+  return true;
+}
+
+Type *reconstitutePeeledArrayType(Type *Ty) {
+  if (!Ty->isStructTy())
+    return Ty;
+
+  auto *STy = cast<StructType>(Ty);
+  Type *OriginalElementType = nullptr;
+  uint64_t TotalSize = 0;
+  if (matchPeeledArrayPattern(STy, OriginalElementType, TotalSize)) {
+    Type *ResultTy = ArrayType::get(
+        reconstitutePeeledArrayType(OriginalElementType), TotalSize);
+    return ResultTy;
+  }
+
+  SmallVector<Type *, 4> NewElementTypes;
+  bool Changed = false;
+  for (Type *ElementTy : STy->elements()) {
+    Type *NewElementTy = reconstitutePeeledArrayType(ElementTy);
+    if (NewElementTy != ElementTy)
+      Changed = true;
+    NewElementTypes.push_back(NewElementTy);
+  }
+
+  if (Changed) {
+    Type *ResultTy;
+    if (STy->isLiteral())
+      ResultTy =
+          StructType::get(STy->getContext(), NewElementTypes, STy->isPacked());
+    else {
+      auto *NewTy = StructType::create(STy->getContext(), STy->getName());
+      NewTy->setBody(NewElementTypes, STy->isPacked());
+      ResultTy = NewTy;
+    }
+    return ResultTy;
+  }
+
+  return Ty;
+}
+
 std::optional<SPIRV::LinkageType::LinkageType>
 getSpirvLinkageTypeFor(const SPIRVSubtarget &ST, const GlobalValue &GV) {
   if (GV.hasLocalLinkage() || GV.hasHiddenVisibility())

llvm/lib/Target/SPIRV/SPIRVUtils.h

@@ -321,6 +321,21 @@ Type *parseBasicTypeName(StringRef &TypeName, LLVMContext &Ctx);
 // Returns true if the function was changed.
 bool sortBlocks(Function &F);
 
+// Check for peeled array structs and recursively reconstitute them. In HLSL
+// CBuffers, arrays may have padding between the elements, but not after the
+// last element. To represent this in LLVM IR an array [N x T] will be
+// represented as {[N-1 x {T, spirv.Padding}], T}. The function
+// matchPeeledArrayPattern recognizes this pattern retrieving the type {T,
+// spirv.Padding}, and the size N.
+bool matchPeeledArrayPattern(const StructType *Ty, Type *&OriginalElementType,
+                             uint64_t &TotalSize);
+
+// This function will turn the type {[N-1 x {T, spirv.Padding}], T} back into
+// [N x {T, spirv.Padding}]. So it can be translated into SPIR-V. The offset
+// decorations will be such that there will be no padding after the array when
+// relevant.
+Type *reconstitutePeeledArrayType(Type *Ty);
+
 inline bool hasInitializer(const GlobalVariable *GV) {
   return GV->hasInitializer() && !isa<UndefValue>(GV->getInitializer());
 }

llvm/test/CodeGen/SPIRV/hlsl-resources/cbuffer-peeled-array-minimal.ll

@@ -0,0 +1,65 @@
+; RUN: llc -O0 -mtriple=spirv-unknown-vulkan-compute %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv-unknown-vulkan-compute %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: %[[FLOAT:[0-9]+]] = OpTypeFloat 32
+; CHECK: %[[STRUCT:[0-9]+]] = OpTypeStruct %[[FLOAT]]
+; CHECK: %[[I8:[0-9]+]] = OpTypeInt 8 0
+; CHECK: %[[STRUCT_PAD:[0-9]+]] = OpTypeStruct %[[STRUCT]] %[[I8]]
+; CHECK: %[[UINT:[0-9]+]] = OpTypeInt 32 0
+; CHECK: %[[CONST_4:[0-9]+]] = OpConstant %[[UINT]] 4
+; CHECK: %[[ARRAY:[0-9]+]] = OpTypeArray %[[STRUCT_PAD]] %[[CONST_4]]
+; CHECK: %[[CBLAYOUT:[0-9]+]] = OpTypeStruct %[[ARRAY]] %[[FLOAT]]
+; CHECK: %[[WRAPPER:[0-9]+]] = OpTypeStruct %[[CBLAYOUT]]
+; CHECK: %[[PTR_WRAPPER:[0-9]+]] = OpTypePointer Uniform %[[WRAPPER]]
+; CHECK: %[[ZERO:[0-9]+]] = OpConstant %[[UINT]] 0
+; CHECK: %[[MYCBUFFER:[0-9]+]] = OpVariable %[[PTR_WRAPPER]] Uniform
+
+%__cblayout_MyCBuffer = type <{ <{ [3 x <{ %OrigType, target("spirv.Padding", 12) }>], %OrigType }>, float }>
+%OrigType = type <{ float }>
+
+@MyCBuffer.cb = local_unnamed_addr global target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) poison
+@myArray = external hidden local_unnamed_addr addrspace(12) global <{ [3 x <{ %OrigType, target("spirv.Padding", 12) }>], %OrigType }>, align 1
+@MyCBuffer.str = private unnamed_addr constant [10 x i8] c"MyCBuffer\00", align 1
+@.str = private unnamed_addr constant [7 x i8] c"output\00", align 1
+
+declare target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) @llvm.spv.resource.handlefromimplicitbinding.tspirv.VulkanBuffer_s___cblayout_MyCBuffers_2_0t(i32, i32, i32, i32, ptr)
+
+define void @main() #1 {
+entry:
+; CHECK: %[[BUFFER_HANDLE:[0-9]+]] = OpCopyObject %[[PTR_WRAPPER]] %[[MYCBUFFER]]
+; CHECK: %[[ACCESS_ARRAY:[0-9]+]] = OpAccessChain {{%[0-9]+}} %[[BUFFER_HANDLE]] %[[ZERO]] %[[ZERO]]
+  %MyCBuffer.cb_h.i.i = tail call target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) @llvm.spv.resource.handlefromimplicitbinding.tspirv.VulkanBuffer_s___cblayout_MyCBuffers_2_0t(i32 0, i32 0, i32 1, i32 0, ptr nonnull @MyCBuffer.str)
+  store target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) %MyCBuffer.cb_h.i.i, ptr @MyCBuffer.cb, align 8
+
+  %0 = tail call target("spirv.Image", float, 5, 2, 0, 0, 2, 1) @llvm.spv.resource.handlefromimplicitbinding.tspirv.Image_f32_5_2_0_0_2_1t(i32 1, i32 0, i32 1, i32 0, ptr nonnull @.str)
+  %1 = tail call i32 @llvm.spv.thread.id.i32(i32 0)
+  %rem.i = and i32 %1, 3
+
+; CHECK: %[[IDX_CONV:[0-9]+]] = OpUConvert {{.*}}
+  %idxprom.i = zext nneg i32 %rem.i to i64
+
+; CHECK: %[[PTR_ELEM:[0-9]+]] = OpAccessChain {{%[0-9]+}} %[[ACCESS_ARRAY]] %[[IDX_CONV]]
+  %cbufferidx.i = getelementptr <{ %OrigType, target("spirv.Padding", 12) }>, ptr addrspace(12) @myArray, i64 %idxprom.i
+
+; CHECK: %[[PTR_FIELD:[0-9]+]] = OpAccessChain {{%[0-9]+}} %[[PTR_ELEM]] %[[ZERO]] %[[ZERO]]
+; CHECK: %[[VAL_FLOAT:[0-9]+]] = OpLoad %[[FLOAT]] %[[PTR_FIELD]] Aligned 4
+  %2 = load float, ptr addrspace(12) %cbufferidx.i, align 4
+
+  %vecinit4.i = insertelement <4 x float> <float poison, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00>, float %2, i64 0
+  %3 = tail call noundef align 16 dereferenceable(16) ptr addrspace(11) @llvm.spv.resource.getpointer.p11.tspirv.Image_f32_5_2_0_0_2_1t(target("spirv.Image", float, 5, 2, 0, 0, 2, 1) %0, i32 0)
+  store <4 x float> %vecinit4.i, ptr addrspace(11) %3, align 16
+; CHECK: OpImageWrite {{%[0-9]+}} {{%[0-9]+}} {{%[0-9]+}}
+  ret void
+}
+
+declare i32 @llvm.spv.thread.id.i32(i32)
+
+declare target("spirv.Image", float, 5, 2, 0, 0, 2, 1) @llvm.spv.resource.handlefromimplicitbinding.tspirv.Image_f32_5_2_0_0_2_1t(i32, i32, i32, i32, ptr)
+
+declare ptr addrspace(11) @llvm.spv.resource.getpointer.p11.tspirv.Image_f32_5_2_0_0_2_1t(target("spirv.Image", float, 5, 2, 0, 0, 2, 1), i32)
+
+attributes #1 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+!hlsl.cbs = !{!0}
+
+!0 = distinct !{ptr @MyCBuffer.cb, ptr addrspace(12) @myArray, null}

llvm/test/CodeGen/SPIRV/hlsl-resources/cbuffer-peeled-array.ll

@@ -0,0 +1,64 @@
+; RUN: llc -O0 -mtriple=spirv-unknown-vulkan-compute %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv-unknown-vulkan-compute %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: %[[FLOAT:[0-9]+]] = OpTypeFloat 32
+; CHECK: %[[VEC3:[0-9]+]] = OpTypeVector %[[FLOAT]] 3
+; CHECK: %[[I8:[0-9]+]] = OpTypeInt 8 0
+; CHECK: %[[STRUCT_PAD:[0-9]+]] = OpTypeStruct %[[VEC3]] %[[I8]]
+; CHECK: %[[UINT:[0-9]+]] = OpTypeInt 32 0
+; CHECK: %[[CONST_3:[0-9]+]] = OpConstant %[[UINT]] 3
+; CHECK: %[[ARRAY:[0-9]+]] = OpTypeArray %[[STRUCT_PAD]] %[[CONST_3]]
+; CHECK: %[[CBLAYOUT:[0-9]+]] = OpTypeStruct %[[ARRAY]]
+; CHECK: %[[WRAPPER:[0-9]+]] = OpTypeStruct %[[CBLAYOUT]]
+; CHECK: %[[PTR_WRAPPER:[0-9]+]] = OpTypePointer Uniform %[[WRAPPER]]
+; CHECK: %[[ZERO:[0-9]+]] = OpConstant %[[UINT]] 0
+; CHECK: %[[MYCBUFFER:[0-9]+]] = OpVariable %[[PTR_WRAPPER]] Uniform
+
+%__cblayout_MyCBuffer = type <{ <{ [2 x <{ <3 x float>, target("spirv.Padding", 4) }>], <3 x float> }> }>
+
+@MyCBuffer.cb = local_unnamed_addr global target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) poison
+@myArray = external hidden local_unnamed_addr addrspace(12) global <{ [2 x <{ <3 x float>, target("spirv.Padding", 4) }>], <3 x float> }>, align 16
+@MyCBuffer.str = private unnamed_addr constant [10 x i8] c"MyCBuffer\00", align 1
+@.str = private unnamed_addr constant [7 x i8] c"output\00", align 1
+
+declare target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_s___cblayout_MyCBuffers_2_0t(i32, i32, i32, i32, ptr)
+
+define void @main() #1 {
+entry:
+; CHECK: %[[BUFFER_HANDLE:[0-9]+]] = OpCopyObject %[[PTR_WRAPPER]] %[[MYCBUFFER]]
+; CHECK: %[[ACCESS_ARRAY:[0-9]+]] = OpAccessChain {{%[0-9]+}} %[[BUFFER_HANDLE]] %[[ZERO]] %[[ZERO]]
+  %MyCBuffer.cb_h.i.i = tail call target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_s___cblayout_MyCBuffers_2_0t(i32 0, i32 0, i32 1, i32 0, ptr nonnull @MyCBuffer.str)
+  store target("spirv.VulkanBuffer", %__cblayout_MyCBuffer, 2, 0) %MyCBuffer.cb_h.i.i, ptr @MyCBuffer.cb, align 8
+
+  %0 = tail call target("spirv.VulkanBuffer", [0 x <3 x float>], 12, 1) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_a0v3f32_12_1t(i32 0, i32 0, i32 1, i32 0, ptr nonnull @.str)
+  %1 = tail call i32 @llvm.spv.thread.id.i32(i32 0)
+
+; CHECK: %[[IDX:[0-9]+]] = OpUMod %[[UINT]] {{%[0-9]+}} %[[CONST_3]]
+  %rem.i = urem i32 %1, 3
+
+; CHECK: %[[IDX_CONV:[0-9]+]] = OpUConvert {{.*}} %[[IDX]]
+  %idxprom.i = zext nneg i32 %rem.i to i64
+
+; CHECK: %[[PTR_ELEM:[0-9]+]] = OpAccessChain {{%[0-9]+}} %[[ACCESS_ARRAY]] %[[IDX_CONV]]
+  %cbufferidx.i = getelementptr <{ <3 x float>, target("spirv.Padding", 4) }>, ptr addrspace(12) @myArray, i64 %idxprom.i
+
+; CHECK: %[[PTR_FIELD:[0-9]+]] = OpAccessChain {{%[0-9]+}} %[[PTR_ELEM]] {{.*}}
+; CHECK: %[[VAL_VEC3:[0-9]+]] = OpLoad %[[VEC3]] %[[PTR_FIELD]] Aligned 16
+  %2 = load <3 x float>, ptr addrspace(12) %cbufferidx.i, align 16
+
+  %3 = tail call noundef align 16 dereferenceable(16) ptr addrspace(11) @llvm.spv.resource.getpointer.p11.tspirv.VulkanBuffer_a0v3f32_12_1t(target("spirv.VulkanBuffer", [0 x <3 x float>], 12, 1) %0, i32 %1)
+  store <3 x float> %2, ptr addrspace(11) %3, align 16
+  ret void
+}
+
+declare i32 @llvm.spv.thread.id.i32(i32)
+
+declare target("spirv.VulkanBuffer", [0 x <3 x float>], 12, 1) @llvm.spv.resource.handlefrombinding.tspirv.VulkanBuffer_a0v3f32_12_1t(i32, i32, i32, i32, ptr)
+
+declare ptr addrspace(11) @llvm.spv.resource.getpointer.p11.tspirv.VulkanBuffer_a0v3f32_12_1t(target("spirv.VulkanBuffer", [0 x <3 x float>], 12, 1), i32)
+
+attributes #1 = { "hlsl.numthreads"="8,1,1" "hlsl.shader"="compute" }
+
+!hlsl.cbs = !{!0}
+
+!0 = !{ptr @MyCBuffer.cb, ptr addrspace(12) @myArray}