[HLSL] Fix global resource initialization (#123394)

Create separate resource initialization function for each resource and
add them to CodeGenModule's `CXXGlobalInits` list.
Fixes #120636 and addresses this [comment
](https://github.com/llvm/llvm-project/pull/119755/files#r1894093603).

4 files changed, 65 insertions, 80 deletions

diff --git a/clang/lib/CodeGen/CGDeclCXX.cpp b/clang/lib/CodeGen/CGDeclCXX.cpp
index 9651751..1c2fece 100644
--- a/clang/lib/CodeGen/CGDeclCXX.cpp
+++ b/clang/lib/CodeGen/CGDeclCXX.cpp
@@ -1131,14 +1131,6 @@ CodeGenFunction::GenerateCXXGlobalInitFunc(llvm::Function *Fn,
       if (Decls[i])
         EmitRuntimeCall(Decls[i]);
 
-    if (getLangOpts().HLSL) {
-      CGHLSLRuntime &CGHLSL = CGM.getHLSLRuntime();
-      if (CGHLSL.needsResourceBindingInitFn()) {
-        llvm::Function *ResInitFn = CGHLSL.createResourceBindingInitFn();
-        Builder.CreateCall(llvm::FunctionCallee(ResInitFn), {});
-      }
-    }
-
     Scope.ForceCleanup();
 
     if (ExitBlock) {
diff --git a/clang/lib/CodeGen/CGHLSLRuntime.cpp b/clang/lib/CodeGen/CGHLSLRuntime.cpp
index 5679bd7..345e218 100644
--- a/clang/lib/CodeGen/CGHLSLRuntime.cpp
+++ b/clang/lib/CodeGen/CGHLSLRuntime.cpp
@@ -536,89 +536,85 @@ void CGHLSLRuntime::generateGlobalCtorDtorCalls() {
   }
 }
 
-void CGHLSLRuntime::handleGlobalVarDefinition(const VarDecl *VD,
-                                              llvm::GlobalVariable *GV) {
-  // If the global variable has resource binding, add it to the list of globals
-  // that need resource binding initialization.
-  const HLSLResourceBindingAttr *RBA = VD->getAttr<HLSLResourceBindingAttr>();
-  if (!RBA)
-    return;
-
-  if (!HLSLAttributedResourceType::findHandleTypeOnResource(
-          VD->getType().getTypePtr()))
-    // FIXME: Only simple declarations of resources are supported for now.
-    // Arrays of resources or resources in user defined classes are
-    // not implemented yet.
-    return;
-
-  ResourcesToBind.emplace_back(VD, GV);
-}
-
-bool CGHLSLRuntime::needsResourceBindingInitFn() {
-  return !ResourcesToBind.empty();
+// Returns true if the type is an HLSL resource class
+static bool isResourceRecordType(const clang::Type *Ty) {
+  return HLSLAttributedResourceType::findHandleTypeOnResource(Ty) != nullptr;
 }
 
-llvm::Function *CGHLSLRuntime::createResourceBindingInitFn() {
-  // No resources to bind
-  assert(needsResourceBindingInitFn() && "no resources to bind");
-
+static void createResourceInitFn(CodeGenModule &CGM, const VarDecl *VD,
+                                 llvm::GlobalVariable *GV, unsigned Slot,
+                                 unsigned Space) {
   LLVMContext &Ctx = CGM.getLLVMContext();
   llvm::Type *Int1Ty = llvm::Type::getInt1Ty(Ctx);
 
-  llvm::Function *InitResBindingsFunc =
-      llvm::Function::Create(llvm::FunctionType::get(CGM.VoidTy, false),
-                             llvm::GlobalValue::InternalLinkage,
-                             "_init_resource_bindings", CGM.getModule());
+  llvm::Function *InitResFunc = llvm::Function::Create(
+      llvm::FunctionType::get(CGM.VoidTy, false),
+      llvm::GlobalValue::InternalLinkage,
+      ("_init_resource_" + VD->getName()).str(), CGM.getModule());
+  InitResFunc->addFnAttr(llvm::Attribute::AlwaysInline);
 
   llvm::BasicBlock *EntryBB =
-      llvm::BasicBlock::Create(Ctx, "entry", InitResBindingsFunc);
+      llvm::BasicBlock::Create(Ctx, "entry", InitResFunc);
   CGBuilderTy Builder(CGM, Ctx);
   const DataLayout &DL = CGM.getModule().getDataLayout();
   Builder.SetInsertPoint(EntryBB);
 
-  for (const auto &[VD, GV] : ResourcesToBind) {
-    for (Attr *A : VD->getAttrs()) {
-      HLSLResourceBindingAttr *RBA = dyn_cast<HLSLResourceBindingAttr>(A);
-      if (!RBA)
-        continue;
-
-      const HLSLAttributedResourceType *AttrResType =
-          HLSLAttributedResourceType::findHandleTypeOnResource(
-              VD->getType().getTypePtr());
-
-      // FIXME: Only simple declarations of resources are supported for now.
-      // Arrays of resources or resources in user defined classes are
-      // not implemented yet.
-      assert(AttrResType != nullptr &&
-             "Resource class must have a handle of HLSLAttributedResourceType");
-
-      llvm::Type *TargetTy =
-          CGM.getTargetCodeGenInfo().getHLSLType(CGM, AttrResType);
-      assert(TargetTy != nullptr &&
-             "Failed to convert resource handle to target type");
-
-      auto *Space = llvm::ConstantInt::get(CGM.IntTy, RBA->getSpaceNumber());
-      auto *Slot = llvm::ConstantInt::get(CGM.IntTy, RBA->getSlotNumber());
+  const HLSLAttributedResourceType *AttrResType =
+      HLSLAttributedResourceType::findHandleTypeOnResource(
+          VD->getType().getTypePtr());
+
+  // FIXME: Only simple declarations of resources are supported for now.
+  // Arrays of resources or resources in user defined classes are
+  // not implemented yet.
+  assert(AttrResType != nullptr &&
+         "Resource class must have a handle of HLSLAttributedResourceType");
+
+  llvm::Type *TargetTy =
+      CGM.getTargetCodeGenInfo().getHLSLType(CGM, AttrResType);
+  assert(TargetTy != nullptr &&
+         "Failed to convert resource handle to target type");
+
+  llvm::Value *Args[] = {
+      llvm::ConstantInt::get(CGM.IntTy, Space), /* reg_space */
+      llvm::ConstantInt::get(CGM.IntTy, Slot),  /* lower_bound */
       // FIXME: resource arrays are not yet implemented
-      auto *Range = llvm::ConstantInt::get(CGM.IntTy, 1);
-      auto *Index = llvm::ConstantInt::get(CGM.IntTy, 0);
+      llvm::ConstantInt::get(CGM.IntTy, 1), /* range_size */
+      llvm::ConstantInt::get(CGM.IntTy, 0), /* index */
       // FIXME: NonUniformResourceIndex bit is not yet implemented
-      auto *NonUniform = llvm::ConstantInt::get(Int1Ty, false);
-      llvm::Value *Args[] = {Space, Slot, Range, Index, NonUniform};
+      llvm::ConstantInt::get(Int1Ty, false) /* non-uniform */
+  };
+  llvm::Value *CreateHandle = Builder.CreateIntrinsic(
+      /*ReturnType=*/TargetTy,
+      CGM.getHLSLRuntime().getCreateHandleFromBindingIntrinsic(), Args, nullptr,
+      Twine(VD->getName()).concat("_h"));
+
+  llvm::Value *HandleRef = Builder.CreateStructGEP(GV->getValueType(), GV, 0);
+  Builder.CreateAlignedStore(CreateHandle, HandleRef,
+                             HandleRef->getPointerAlignment(DL));
+  Builder.CreateRetVoid();
 
-      llvm::Value *CreateHandle = Builder.CreateIntrinsic(
-          /*ReturnType=*/TargetTy, getCreateHandleFromBindingIntrinsic(), Args,
-          nullptr, Twine(VD->getName()).concat("_h"));
+  CGM.AddCXXGlobalInit(InitResFunc);
+}
 
-      llvm::Value *HandleRef =
-          Builder.CreateStructGEP(GV->getValueType(), GV, 0);
-      Builder.CreateAlignedStore(CreateHandle, HandleRef,
-                                 HandleRef->getPointerAlignment(DL));
-    }
-  }
+void CGHLSLRuntime::handleGlobalVarDefinition(const VarDecl *VD,
+                                              llvm::GlobalVariable *GV) {
 
-  Builder.CreateRetVoid();
-  return InitResBindingsFunc;
+  // If the global variable has resource binding, create an init function
+  // for the resource
+  const HLSLResourceBindingAttr *RBA = VD->getAttr<HLSLResourceBindingAttr>();
+  if (!RBA)
+    // FIXME: collect unbound resources for implicit binding resolution later
+    // on?
+    return;
+
+  if (!isResourceRecordType(VD->getType().getTypePtr()))
+    // FIXME: Only simple declarations of resources are supported for now.
+    // Arrays of resources or resources in user defined classes are
+    // not implemented yet.
+    return;
+
+  createResourceInitFn(CGM, VD, GV, RBA->getSlotNumber(),
+                       RBA->getSpaceNumber());
 }
 
 llvm::Instruction *CGHLSLRuntime::getConvergenceToken(BasicBlock &BB) {
diff --git a/clang/lib/CodeGen/CGHLSLRuntime.h b/clang/lib/CodeGen/CGHLSLRuntime.h
index f9dc7b8..032b2de 100644
--- a/clang/lib/CodeGen/CGHLSLRuntime.h
+++ b/clang/lib/CodeGen/CGHLSLRuntime.h
@@ -159,8 +159,6 @@ public:
   void setHLSLFunctionAttributes(const FunctionDecl *FD, llvm::Function *Fn);
   void handleGlobalVarDefinition(const VarDecl *VD, llvm::GlobalVariable *Var);
 
-  bool needsResourceBindingInitFn();
-  llvm::Function *createResourceBindingInitFn();
   llvm::Instruction *getConvergenceToken(llvm::BasicBlock &BB);
 
 private:
@@ -173,9 +171,6 @@ private:
   void addBufferDecls(const DeclContext *DC, Buffer &CB);
   llvm::Triple::ArchType getArch();
   llvm::SmallVector<Buffer> Buffers;
-
-  llvm::SmallVector<std::pair<const VarDecl *, llvm::GlobalVariable *>>
-      ResourcesToBind;
 };
 
 } // namespace CodeGen
diff --git a/clang/lib/CodeGen/CodeGenModule.h b/clang/lib/CodeGen/CodeGenModule.h
index d5ef1a7..1aa5d48 100644
--- a/clang/lib/CodeGen/CodeGenModule.h
+++ b/clang/lib/CodeGen/CodeGenModule.h
@@ -1226,6 +1226,8 @@ public:
 
   llvm::Function *getIntrinsic(unsigned IID, ArrayRef<llvm::Type *> Tys = {});
 
+  void AddCXXGlobalInit(llvm::Function *F) { CXXGlobalInits.push_back(F); }
+
   /// Emit code for a single top level declaration.
   void EmitTopLevelDecl(Decl *D);