[LLVM] Maintain element type of @llvm.compiler.used/@llvm.used if it already exists

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -75,6 +75,7 @@
 #include "llvm/Support/SHA1.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Triple.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
@@ -5847,28 +5848,19 @@ static const char *getSectionNameForCommandline(const Triple &T) {
   llvm_unreachable("Unimplemented ObjectFormatType");
 }
 
+static auto globalInUsedHasName(StringRef Name) {
+  return [Name](Constant *C) { return C->getName() == Name; };
+}
+
 void llvm::embedBitcodeInModule(llvm::Module &M, llvm::MemoryBufferRef Buf,
                                 bool EmbedBitcode, bool EmbedCmdline,
                                 const std::vector<uint8_t> &CmdArgs) {
-  // Save llvm.compiler.used and remove it.
-  SmallVector<Constant *, 2> UsedArray;
-  SmallVector<GlobalValue *, 4> UsedGlobals;
-  GlobalVariable *Used = collectUsedGlobalVariables(M, UsedGlobals, true);
-  Type *UsedElementType = Used ? Used->getValueType()->getArrayElementType()
-                               : PointerType::getUnqual(M.getContext());
-  for (auto *GV : UsedGlobals) {
-    if (GV->getName() != "llvm.embedded.module" &&
-        GV->getName() != "llvm.cmdline")
-      UsedArray.push_back(
-          ConstantExpr::getPointerBitCastOrAddrSpaceCast(GV, UsedElementType));
-  }
-  if (Used)
-    Used->eraseFromParent();
 
   // Embed the bitcode for the llvm module.
   std::string Data;
   ArrayRef<uint8_t> ModuleData;
   Triple T(M.getTargetTriple());
+  SmallVector<GlobalValue *, 2> NewGlobals;
 
   if (EmbedBitcode) {
     if (Buf.getBufferSize() == 0 ||
@@ -5887,23 +5879,23 @@ void llvm::embedBitcodeInModule(llvm::Module &M, llvm::MemoryBufferRef Buf,
   }
   llvm::Constant *ModuleConstant =
       llvm::ConstantDataArray::get(M.getContext(), ModuleData);
-  llvm::GlobalVariable *GV = new llvm::GlobalVariable(
+  llvm::GlobalVariable *EmbeddedModule = new llvm::GlobalVariable(
       M, ModuleConstant->getType(), true, llvm::GlobalValue::PrivateLinkage,
       ModuleConstant);
-  GV->setSection(getSectionNameForBitcode(T));
+  EmbeddedModule->setSection(getSectionNameForBitcode(T));
   // Set alignment to 1 to prevent padding between two contributions from input
   // sections after linking.
-  GV->setAlignment(Align(1));
-  UsedArray.push_back(
-      ConstantExpr::getPointerBitCastOrAddrSpaceCast(GV, UsedElementType));
+  EmbeddedModule->setAlignment(Align(1));
+  NewGlobals.push_back(EmbeddedModule);
   if (llvm::GlobalVariable *Old =
           M.getGlobalVariable("llvm.embedded.module", true)) {
+    removeFromUsedLists(M, globalInUsedHasName("llvm.embedded.module"));
     assert(Old->hasZeroLiveUses() &&
            "llvm.embedded.module can only be used once in llvm.compiler.used");
-    GV->takeName(Old);
+    EmbeddedModule->takeName(Old);
     Old->eraseFromParent();
   } else {
-    GV->setName("llvm.embedded.module");
+    EmbeddedModule->setName("llvm.embedded.module");
   }
 
   // Skip if only bitcode needs to be embedded.
@@ -5913,30 +5905,21 @@ void llvm::embedBitcodeInModule(llvm::Module &M, llvm::MemoryBufferRef Buf,
                               CmdArgs.size());
     llvm::Constant *CmdConstant =
         llvm::ConstantDataArray::get(M.getContext(), CmdData);
-    GV = new llvm::GlobalVariable(M, CmdConstant->getType(), true,
-                                  llvm::GlobalValue::PrivateLinkage,
-                                  CmdConstant);
-    GV->setSection(getSectionNameForCommandline(T));
-    GV->setAlignment(Align(1));
-    UsedArray.push_back(
-        ConstantExpr::getPointerBitCastOrAddrSpaceCast(GV, UsedElementType));
+    GlobalVariable *CmdLine = new llvm::GlobalVariable(
+        M, CmdConstant->getType(), true, llvm::GlobalValue::PrivateLinkage,
+        CmdConstant);
+    CmdLine->setSection(getSectionNameForCommandline(T));
+    CmdLine->setAlignment(Align(1));
     if (llvm::GlobalVariable *Old = M.getGlobalVariable("llvm.cmdline", true)) {
+      removeFromUsedLists(M, globalInUsedHasName("llvm.cmdline"));
       assert(Old->hasZeroLiveUses() &&
              "llvm.cmdline can only be used once in llvm.compiler.used");
-      GV->takeName(Old);
+      CmdLine->takeName(Old);
       Old->eraseFromParent();
     } else {
-      GV->setName("llvm.cmdline");
+      CmdLine->setName("llvm.cmdline");
     }
+    NewGlobals.push_back(CmdLine);
+    appendToCompilerUsed(M, NewGlobals);
   }
-
-  if (UsedArray.empty())
-    return;
-
-  // Recreate llvm.compiler.used.
-  ArrayType *ATy = ArrayType::get(UsedElementType, UsedArray.size());
-  auto *NewUsed = new GlobalVariable(
-      M, ATy, false, llvm::GlobalValue::AppendingLinkage,
-      llvm::ConstantArray::get(ATy, UsedArray), "llvm.compiler.used");
-  NewUsed->setSection("llvm.metadata");
 }

llvm/lib/Bitcode/Writer/CMakeLists.txt

@@ -15,4 +15,5 @@ add_llvm_component_library(LLVMBitWriter
   ProfileData
   Support
   TargetParser
+  TransformUtils
   )

llvm/lib/Transforms/Utils/ModuleUtils.cpp

@@ -138,10 +138,11 @@ static void appendToUsedList(Module &M, StringRef Name, ArrayRef<GlobalValue *>
 
   SmallSetVector<Constant *, 16> Init;
   collectUsedGlobals(GV, Init);
+  Type *ArrayEltTy = GV ? GV->getValueType()->getArrayElementType()
+                        : PointerType::getUnqual(M.getContext());
   if (GV)
     GV->eraseFromParent();
 
-  Type *ArrayEltTy = llvm::PointerType::getUnqual(M.getContext());
   for (auto *V : Values)
     Init.insert(ConstantExpr::getPointerBitCastOrAddrSpaceCast(V, ArrayEltTy));
 

llvm/unittests/Transforms/Utils/ModuleUtilsTest.cpp

@@ -69,6 +69,23 @@ TEST(ModuleUtils, AppendToUsedList2) {
   EXPECT_EQ(1, getListSize(*M, "llvm.used"));
 }
 
+TEST(ModuleUtils, AppendToUsedList3) {
+  LLVMContext C;
+
+  std::unique_ptr<Module> M = parseIR(C, R"(
+          @x = addrspace(1) global [2 x i32] zeroinitializer, align 4
+          @y = addrspace(2) global [2 x i32] zeroinitializer, align 4
+          @llvm.compiler.used = appending global [1 x ptr addrspace(3)] [ptr addrspace(3) addrspacecast (ptr addrspace(1) @x to ptr addrspace(3))]
+      )");
+  GlobalVariable *X = M->getNamedGlobal("x");
+  GlobalVariable *Y = M->getNamedGlobal("y");
+  EXPECT_EQ(1, getListSize(*M, "llvm.compiler.used"));
+  appendToCompilerUsed(*M, X);
+  EXPECT_EQ(1, getListSize(*M, "llvm.compiler.used"));
+  appendToCompilerUsed(*M, Y);
+  EXPECT_EQ(2, getListSize(*M, "llvm.compiler.used"));
+}
+
 using AppendFnType = decltype(&appendToGlobalCtors);
 using TransformFnType = decltype(&transformGlobalCtors);
 using ParamType = std::tuple<StringRef, AppendFnType, TransformFnType>;