wip: Update DXILCBufferAccess to padding approach

Author: bogner

clang/lib/CodeGen/TargetInfo.h

@@ -448,6 +448,20 @@ class TargetCodeGenInfo {
     return nullptr;
   }
 
+  virtual llvm::Type *
+  getHLSLPadding(CodeGenModule &CGM, CharUnits NumBytes) const {
+    return llvm::ArrayType::get(llvm::Type::getInt8Ty(CGM.getLLVMContext()),
+                                NumBytes.getQuantity());
+  }
+
+  virtual bool isHLSLPadding(llvm::Type *Ty) const {
+    // TODO: Do we actually want to default these functions like this?
+    if (auto *AT = dyn_cast<llvm::ArrayType>(Ty))
+      if (AT->getElementType() == llvm::Type::getInt8Ty(Ty->getContext()))
+        return true;
+    return false;
+  }
+
   // Set the Branch Protection Attributes of the Function accordingly to the
   // BPI. Remove attributes that contradict with current BPI.
   static void

clang/lib/CodeGen/Targets/DirectX.cpp

@@ -32,6 +32,19 @@ class DirectXTargetCodeGenInfo : public TargetCodeGenInfo {
   llvm::Type *
   getHLSLType(CodeGenModule &CGM, const Type *T,
               const SmallVector<int32_t> *Packoffsets = nullptr) const override;
+
+  llvm::Type *getHLSLPadding(CodeGenModule &CGM,
+                             CharUnits NumBytes) const override {
+    unsigned Size = NumBytes.getQuantity();
+    return llvm::TargetExtType::get(CGM.getLLVMContext(), "dx.Padding", {},
+                                    {Size});
+  }
+
+  bool isHLSLPadding(llvm::Type *Ty) const override {
+    if (auto *TET = dyn_cast<llvm::TargetExtType>(Ty))
+      return TET->getName() == "dx.Padding";
+    return false;
+  }
 };
 
 llvm::Type *DirectXTargetCodeGenInfo::getHLSLType(

clang/lib/CodeGen/Targets/SPIR.cpp

@@ -56,6 +56,20 @@ class CommonSPIRTargetCodeGenInfo : public TargetCodeGenInfo {
   llvm::Type *
   getHLSLType(CodeGenModule &CGM, const Type *Ty,
               const SmallVector<int32_t> *Packoffsets = nullptr) const override;
+
+  llvm::Type *
+  getHLSLPadding(CodeGenModule &CGM, CharUnits NumBytes) const override {
+    unsigned Size = NumBytes.getQuantity();
+    return llvm::TargetExtType::get(CGM.getLLVMContext(), "spirv.Padding", {},
+                                    {Size});
+  }
+
+  bool isHLSLPadding(llvm::Type *Ty) const override {
+    if (auto *TET = dyn_cast<llvm::TargetExtType>(Ty))
+      return TET->getName() == "spirv.Padding";
+    return false;
+  }
+
   llvm::Type *getSPIRVImageTypeFromHLSLResource(
       const HLSLAttributedResourceType::Attributes &attributes,
       QualType SampledType, CodeGenModule &CGM) const;

llvm/include/llvm/Frontend/HLSL/CBuffer.h

@@ -46,7 +46,8 @@ class CBufferMetadata {
   CBufferMetadata(NamedMDNode *MD) : MD(MD) {}
 
 public:
-  static std::optional<CBufferMetadata> get(Module &M);
+  static std::optional<CBufferMetadata>
+  get(Module &M, llvm::function_ref<bool(Type *)> IsPadding);
 
   using iterator = SmallVector<CBufferMapping>::iterator;
   iterator begin() { return Mappings.begin(); }
@@ -55,9 +56,6 @@ class CBufferMetadata {
   void eraseFromModule();
 };
 
-APInt translateCBufArrayOffset(const DataLayout &DL, APInt Offset,
-                               ArrayType *Ty);
-
 } // namespace hlsl
 } // namespace llvm
 

llvm/lib/Frontend/HLSL/CBuffer.cpp

@@ -15,25 +15,28 @@
 using namespace llvm;
 using namespace llvm::hlsl;
 
-static size_t getMemberOffset(GlobalVariable *Handle, size_t Index) {
+static SmallVector<size_t>
+getMemberOffsets(const DataLayout &DL, GlobalVariable *Handle,
+                 llvm::function_ref<bool(Type *)> IsPadding) {
+  SmallVector<size_t> Offsets;
+
   auto *HandleTy = cast<TargetExtType>(Handle->getValueType());
   assert((HandleTy->getName().ends_with(".CBuffer") ||
           HandleTy->getName() == "spirv.VulkanBuffer") &&
          "Not a cbuffer type");
   assert(HandleTy->getNumTypeParameters() == 1 && "Expected layout type");
+  auto *LayoutTy = cast<StructType>(HandleTy->getTypeParameter(0));
 
-  auto *LayoutTy = cast<TargetExtType>(HandleTy->getTypeParameter(0));
-  assert(LayoutTy->getName().ends_with(".Layout") && "Not a layout type");
-
-  // Skip the "size" parameter.
-  size_t ParamIndex = Index + 1;
-  assert(LayoutTy->getNumIntParameters() > ParamIndex &&
-         "Not enough parameters");
+  const StructLayout *SL = DL.getStructLayout(LayoutTy);
+  for (int I = 0, E = LayoutTy->getNumElements(); I < E; ++I)
+    if (!IsPadding(LayoutTy->getElementType(I)))
+      Offsets.push_back(SL->getElementOffset(I));
 
-  return LayoutTy->getIntParameter(ParamIndex);
+  return Offsets;
 }
 
-std::optional<CBufferMetadata> CBufferMetadata::get(Module &M) {
+std::optional<CBufferMetadata>
+CBufferMetadata::get(Module &M, llvm::function_ref<bool(Type *)> IsPadding) {
   NamedMDNode *CBufMD = M.getNamedMetadata("hlsl.cbs");
   if (!CBufMD)
     return std::nullopt;
@@ -47,13 +50,16 @@ std::optional<CBufferMetadata> CBufferMetadata::get(Module &M) {
         cast<ValueAsMetadata>(MD->getOperand(0))->getValue());
     CBufferMapping &Mapping = Result->Mappings.emplace_back(Handle);
 
+    SmallVector<size_t> MemberOffsets =
+        getMemberOffsets(M.getDataLayout(), Handle, IsPadding);
+
     for (int I = 1, E = MD->getNumOperands(); I < E; ++I) {
       Metadata *OpMD = MD->getOperand(I);
       // Some members may be null if they've been optimized out.
       if (!OpMD)
         continue;
       auto *V = cast<GlobalVariable>(cast<ValueAsMetadata>(OpMD)->getValue());
-      Mapping.Members.emplace_back(V, getMemberOffset(Handle, I - 1));
+      Mapping.Members.emplace_back(V, MemberOffsets[I - 1]);
     }
   }
 
@@ -64,10 +70,3 @@ void CBufferMetadata::eraseFromModule() {
   // Remove the cbs named metadata
   MD->eraseFromParent();
 }
-
-APInt hlsl::translateCBufArrayOffset(const DataLayout &DL, APInt Offset,
-                                     ArrayType *Ty) {
-  int64_t TypeSize = DL.getTypeSizeInBits(Ty->getElementType()) / 8;
-  int64_t RoundUp = alignTo(TypeSize, Align(CBufferRowSizeInBytes));
-  return Offset.udiv(TypeSize) * RoundUp;
-}

llvm/lib/IR/Type.cpp

@@ -1008,6 +1008,10 @@ static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
   }
   if (Name == "spirv.IntegralConstant" || Name == "spirv.Literal")
     return TargetTypeInfo(Type::getVoidTy(C));
+  if (Name == "spirv.Padding")
+    return TargetTypeInfo(
+        ArrayType::get(Type::getInt8Ty(C), Ty->getIntParameter(0)),
+        TargetExtType::CanBeGlobal);
   if (Name.starts_with("spirv."))
     return TargetTypeInfo(PointerType::get(C, 0), TargetExtType::HasZeroInit,
                           TargetExtType::CanBeGlobal,

llvm/lib/Target/DirectX/DXILCBufferAccess.cpp

@@ -8,6 +8,7 @@
 
 #include "DXILCBufferAccess.h"
 #include "DirectX.h"
+#include "llvm/Analysis/DXILResource.h"
 #include "llvm/Frontend/HLSL/CBuffer.h"
 #include "llvm/Frontend/HLSL/HLSLResource.h"
 #include "llvm/IR/IRBuilder.h"
@@ -97,10 +98,6 @@ struct CBufferResource {
       (void)Success;
       assert(Success && "Offsets into cbuffer globals must be constant");
 
-      if (auto *ATy = dyn_cast<ArrayType>(Member->getValueType()))
-        ConstantOffset =
-            hlsl::translateCBufArrayOffset(DL, ConstantOffset, ATy);
-
       return ConstantOffset.getZExtValue();
     }
 
@@ -194,102 +191,21 @@ static void replaceLoad(LoadInst *LI, CBufferResource &CBR,
   DeadInsts.push_back(LI);
 }
 
-/// This function recursively copies N array elements from the cbuffer resource
-/// CBR to the MemCpy Destination. Recursion is used to unravel multidimensional
-/// arrays into a sequence of scalar/vector extracts and stores.
-static void copyArrayElemsForMemCpy(IRBuilder<> &Builder, MemCpyInst *MCI,
-                                    CBufferResource &CBR, ArrayType *ArrTy,
-                                    size_t ArrOffset, size_t N,
-                                    const Twine &Name = "") {
-  const DataLayout &DL = MCI->getDataLayout();
-  Type *ElemTy = ArrTy->getElementType();
-  size_t ElemTySize = DL.getTypeAllocSize(ElemTy);
-  for (unsigned I = 0; I < N; ++I) {
-    size_t Offset = ArrOffset + I * ElemTySize;
-
-    // Recursively copy nested arrays
-    if (ArrayType *ElemArrTy = dyn_cast<ArrayType>(ElemTy)) {
-      copyArrayElemsForMemCpy(Builder, MCI, CBR, ElemArrTy, Offset,
-                              ElemArrTy->getNumElements(), Name);
-      continue;
-    }
-
-    // Load CBuffer value and store it in Dest
-    APInt CBufArrayOffset(
-        DL.getIndexTypeSizeInBits(MCI->getSource()->getType()), Offset);
-    CBufArrayOffset =
-        hlsl::translateCBufArrayOffset(DL, CBufArrayOffset, ArrTy);
-    Value *CBufferVal =
-        CBR.loadValue(Builder, ElemTy, CBufArrayOffset.getZExtValue(), Name);
-    Value *GEP =
-        Builder.CreateInBoundsGEP(Builder.getInt8Ty(), MCI->getDest(),
-                                  {Builder.getInt32(Offset)}, Name + ".dest");
-    Builder.CreateStore(CBufferVal, GEP, MCI->isVolatile());
-  }
-}
-
-/// Replace memcpy from a cbuffer global with a memcpy from the cbuffer handle
-/// itself. Assumes the cbuffer global is an array, and the length of bytes to
-/// copy is divisible by array element allocation size.
-/// The memcpy source must also be a direct cbuffer global reference, not a GEP.
-static void replaceMemCpy(MemCpyInst *MCI, CBufferResource &CBR) {
-
-  ArrayType *ArrTy = dyn_cast<ArrayType>(CBR.getValueType());
-  assert(ArrTy && "MemCpy lowering is only supported for array types");
-
-  // This assumption vastly simplifies the implementation
-  if (MCI->getSource() != CBR.Member)
-    reportFatalUsageError(
-        "Expected MemCpy source to be a cbuffer global variable");
-
-  ConstantInt *Length = dyn_cast<ConstantInt>(MCI->getLength());
-  uint64_t ByteLength = Length->getZExtValue();
-
-  // If length to copy is zero, no memcpy is needed
-  if (ByteLength == 0) {
-    MCI->eraseFromParent();
-    return;
-  }
-
-  const DataLayout &DL = CBR.getDataLayout();
-
-  Type *ElemTy = ArrTy->getElementType();
-  size_t ElemSize = DL.getTypeAllocSize(ElemTy);
-  assert(ByteLength % ElemSize == 0 &&
-         "Length of bytes to MemCpy must be divisible by allocation size of "
-         "source/destination array elements");
-  size_t ElemsToCpy = ByteLength / ElemSize;
-
-  IRBuilder<> Builder(MCI);
-  CBR.createAndSetCurrentHandle(Builder);
-
-  copyArrayElemsForMemCpy(Builder, MCI, CBR, ArrTy, 0, ElemsToCpy,
-                          "memcpy." + MCI->getDest()->getName() + "." +
-                              MCI->getSource()->getName());
-
-  MCI->eraseFromParent();
-}
-
 static void replaceAccessesWithHandle(CBufferResource &CBR) {
   SmallVector<WeakTrackingVH> DeadInsts;
 
   SmallVector<User *> ToProcess{CBR.users()};
   while (!ToProcess.empty()) {
     User *Cur = ToProcess.pop_back_val();
+    assert(!isa<MemCpyInst>(Cur) &&
+           "memcpy should have been removed in an earlier pass");
 
     // If we have a load instruction, replace the access.
     if (auto *LI = dyn_cast<LoadInst>(Cur)) {
       replaceLoad(LI, CBR, DeadInsts);
       continue;
     }
 
-    // If we have a memcpy instruction, replace it with multiple accesses and
-    // subsequent stores to the destination
-    if (auto *MCI = dyn_cast<MemCpyInst>(Cur)) {
-      replaceMemCpy(MCI, CBR);
-      continue;
-    }
-
     // Otherwise, walk users looking for a load...
     if (isa<GetElementPtrInst>(Cur) || isa<GEPOperator>(Cur)) {
       ToProcess.append(Cur->user_begin(), Cur->user_end());
@@ -302,7 +218,8 @@ static void replaceAccessesWithHandle(CBufferResource &CBR) {
 }
 
 static bool replaceCBufferAccesses(Module &M) {
-  std::optional<hlsl::CBufferMetadata> CBufMD = hlsl::CBufferMetadata::get(M);
+  std::optional<hlsl::CBufferMetadata> CBufMD = hlsl::CBufferMetadata::get(
+      M, [](Type *Ty) { return isa<llvm::dxil::PaddingExtType>(Ty); });
   if (!CBufMD)
     return false;
 

llvm/test/CodeGen/DirectX/CBufferAccess/array-typedgep.ll

@@ -3,24 +3,24 @@
 ; cbuffer CB : register(b0) {
 ;   float a1[3];
 ; }
-%__cblayout_CB = type <{ [3 x float] }>
+%__cblayout_CB = type <{ [2 x <{ float, [12 x i8] }>], float }>
 
-@CB.cb = global target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 36, 0)) poison
+@CB.cb = global target("dx.CBuffer", %__cblayout_CB) poison
 ; CHECK: @CB.cb =
 ; CHECK-NOT: external {{.*}} addrspace(2) global
-@a1 = external addrspace(2) global [3 x float], align 4
+@a1 = external addrspace(2) global <{ [2 x <{ float, [12 x i8] }>], float }>, align 4
 
 ; CHECK: define void @f
 define void @f(ptr %dst) {
 entry:
-  %CB.cb_h = call target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 36, 0)) @llvm.dx.resource.handlefrombinding.tdx.CBuffer_tdx.Layout_s___cblayout_CBs_36_0tt(i32 0, i32 0, i32 1, i32 0, ptr null)
-  store target("dx.CBuffer", target("dx.Layout", %__cblayout_CB, 36, 0)) %CB.cb_h, ptr @CB.cb, align 4
+  %CB.cb_h = call target("dx.CBuffer", %__cblayout_CB) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr null)
+  store target("dx.CBuffer", %__cblayout_CB) %CB.cb_h, ptr @CB.cb, align 4
 
-  ; CHECK: [[CB:%.*]] = load target("dx.CBuffer", {{.*}})), ptr @CB.cb
-  ; CHECK: [[LOAD:%.*]] = call { float, float, float, float } @llvm.dx.resource.load.cbufferrow.4.{{.*}}(target("dx.CBuffer", {{.*}})) [[CB]], i32 1)
+  ; CHECK: [[CB:%.*]] = load target("dx.CBuffer", %__cblayout_CB), ptr @CB.cb
+  ; CHECK: [[LOAD:%.*]] = call { float, float, float, float } @llvm.dx.resource.load.cbufferrow.4.{{.*}}(target("dx.CBuffer", %__cblayout_CB) [[CB]], i32 1)
   ; CHECK: [[X:%.*]] = extractvalue { float, float, float, float } [[LOAD]], 0
   ; CHECK: store float [[X]], ptr %dst
-  %a1 = load float, ptr addrspace(2) getelementptr inbounds ([3 x float], ptr addrspace(2) @a1, i32 0, i32 1), align 4
+  %a1 = load float, ptr addrspace(2) getelementptr inbounds (<{ [2 x <{ float, [12 x i8] }>], float }>, ptr addrspace(2) @a1, i32 0, i32 0, i32 1), align 4
   store float %a1, ptr %dst, align 32
 
   ret void