ThinLTOBitcodeWriter: keep comdats together, rename if leader is renamed

[inglorion]

Summary

COFF requires that every comdat contain a symbol with the same name as
the comdat. ThinLTOBitcodeWriter renames symbols, which may cause this
requirement to be violated. This change avoids such violations by
renaming comdats if their leaders are renamed. It also keeps comdats
together when splitting modules.

[inglorion]

This is an alternative to D31735. Instead of deleting comdats that are invalidated, it keeps comdats valid by keeping their members in the same module and renaming the comdat if its leader is renamed. I think this is a cleaner solution.

[pcc]

lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp

319 ↗

(On Diff #94908)

I don't think you need to extract this lambda. The appropriate test is just whether the GlobalValue is a GlobalVariable with type metadata. The canonical definitions of any virtual functions identified by this lambda live in the thin LTO module, and the copies in the merged module are used for optimization only and are not comdat members (see line 347 below). So your code for identifying comdats could be added to the loop above.

366 ↗

(On Diff #94908)

The code above should guarantee that if a module contains one member of a comdat it will also contain all other members. So I think you can now apply comdat renaming in promoteInternals.

test/Transforms/ThinLTOBitcodeWriter/comdat.ll

5–26 ↗

(On Diff #94908)

Please move these checks next to the corresponding input IR (where possible).

[inglorion]

lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp

319 ↗

(On Diff #94908)

Yes. The way I wrote this actually introduces a bug where we can end up with functions available_externally in MergedM, and no definition in M. I'm working on a fix.

[inglorion]

test/Transforms/ThinLTOBitcodeWriter/comdat.ll

5–26 ↗

(On Diff #94908)

What is a good way to do that? When I wrote the checks next to the IR, I got errors, because the IR was emitted in a different order from how I wrote it. Specifically, all the comdat declarations come first, then all the globals, then all the functions. So I could reorder my code to be in that same order, but that would mix things from different comdats together, which, as you pointed out elsewhere, made the test harder to read. I thought this was a reasonable compromise - putting all the same comdats together in the IR, and putting the same modules together in the checks. But if you have a better suggestion, I'm all ears.

[inglorion]

fixed bug that caused function definitions to disappear

[pcc]

test/Transforms/ThinLTOBitcodeWriter/comdat.ll

5–26 ↗

(On Diff #94908)

The main thing that made your previous test hard to read was the fact that the global variables from different comdats were not grouped, combined with the choice of names. It's fine if you separate comdats from global variables from aliases from functions, as long as the global variables are grouped and all names are prefixed.

41 ↗

(On Diff #94908)

I think you need to add a comdat($lwt) here if you intend this to be a member of lwt. It also doesn't look like you have any tests for this function.

59 ↗

(On Diff #94908)

Same here.

74 ↗

(On Diff #94908)

Same here.

[inglorion]

perform comdat renaming in promoteInternals

[pcc]

lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp

308 ↗

(On Diff #94914)

Nit: I think you can collect Comdat pointers rather than StringRefs here.

310–312 ↗

(On Diff #94914)

As I mentioned you could move lines 310 and 312 into the loop starting on line 290 above.

322 ↗

(On Diff #94914)

Remove else

[inglorion]

reordered test and put checks next to the IR they apply to

[pcc]

test/Transforms/ThinLTOBitcodeWriter/comdat.ll

24 ↗

(On Diff #94917)

If this isn't being renamed (as I'd expect), can you please replace this regex with a straightforward match against the original name?

33 ↗

(On Diff #94917)

Same here.

56 ↗

(On Diff #94917)

Is this meant to be @nlwt_aliasee?

[inglorion]

Implemented more of @pcc's suggestions:

• Use Comdat* instead of StringRef in MergedMComdats.
• Populate MergedMComdats in existing loop instead of creating a new one.
• Remove regexes when matching unpromoted names.

[inglorion]

removed unnecessary else

[inglorion]

I think that's all your comments, @pcc. The functions are just there to make sure the globals are used. I've added a comment to that effect. The functions are purposely not part of the comdats, so that they can be externally visible without causing the comdats to have externally visible members.

[pcc]

LGTM

llvm/trunk/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp

std::string getModuleId(Module *M) {
   SmallString<32> Str;
   MD5::stringifyResult(R, Str);
   return ("$" + Str).str();
 }
 
 // Promote each local-linkage entity defined by ExportM and used by ImportM by
 // changing visibility and appending the given ModuleId.
 void promoteInternals(Module &ExportM, Module &ImportM, StringRef ModuleId) {
+  DenseMap<const Comdat *, Comdat *> RenamedComdats;
   for (auto &ExportGV : ExportM.global_values()) {
     if (!ExportGV.hasLocalLinkage())
       continue;
 
-    GlobalValue *ImportGV = ImportM.getNamedValue(ExportGV.getName());
+    auto Name = ExportGV.getName();
+    GlobalValue *ImportGV = ImportM.getNamedValue(Name);
     if (!ImportGV || ImportGV->use_empty())
       continue;
 
-    std::string NewName = (ExportGV.getName() + ModuleId).str();
+    std::string NewName = (Name + ModuleId).str();
+
+    if (const auto *C = ExportGV.getComdat())
+      if (C->getName() == Name)
+        RenamedComdats.try_emplace(C, ExportM.getOrInsertComdat(NewName));
 
     ExportGV.setName(NewName);
     ExportGV.setLinkage(GlobalValue::ExternalLinkage);
     ExportGV.setVisibility(GlobalValue::HiddenVisibility);
 
     ImportGV->setName(NewName);
     ImportGV->setVisibility(GlobalValue::HiddenVisibility);
   }
+
+  if (!RenamedComdats.empty())
+    for (auto &GO : ExportM.global_objects())
+      if (auto *C = GO.getComdat()) {
+        auto Replacement = RenamedComdats.find(C);
+        if (Replacement != RenamedComdats.end())
+          GO.setComdat(Replacement->second);
+      }
 }
 
 // Promote all internal (i.e. distinct) type ids used by the module by replacing
 // them with external type ids formed using the module id.
 //
 // Note that this needs to be done before we clone the module because each clone
 // will receive its own set of distinct metadata nodes.
 void promoteTypeIds(Module &M, StringRef ModuleId) {
void splitAndWriteThinLTOBitcode(
   //
   // Note that we test whether this copy of the function is readnone, rather
   // than testing function attributes, which must hold for any copy of the
   // function, even a less optimized version substituted at link time. This is
   // sound because the virtual constant propagation optimizations effectively
   // inline all implementations of the virtual function into each call site,
   // rather than using function attributes to perform local optimization.
   std::set<const Function *> EligibleVirtualFns;
+  // If any member of a comdat lives in MergedM, put all members of that
+  // comdat in MergedM to keep the comdat together.
+  DenseSet<const Comdat *> MergedMComdats;
   for (GlobalVariable &GV : M.globals())
-    if (HasTypeMetadata(&GV))
+    if (HasTypeMetadata(&GV)) {
+      if (const auto *C = GV.getComdat())
+        MergedMComdats.insert(C);
       forEachVirtualFunction(GV.getInitializer(), [&](Function *F) {
         auto *RT = dyn_cast<IntegerType>(F->getReturnType());
         if (!RT || RT->getBitWidth() > 64 || F->arg_empty() ||
             !F->arg_begin()->use_empty())
           return;
         for (auto &Arg : make_range(std::next(F->arg_begin()), F->arg_end())) {
           auto *ArgT = dyn_cast<IntegerType>(Arg.getType());
           if (!ArgT || ArgT->getBitWidth() > 64)
             return;
         }
         if (computeFunctionBodyMemoryAccess(*F, AARGetter(*F)) == MAK_ReadNone)
           EligibleVirtualFns.insert(F);
       });
+    }
 
   ValueToValueMapTy VMap;
   std::unique_ptr<Module> MergedM(
       CloneModule(&M, VMap, [&](const GlobalValue *GV) -> bool {
+        if (const auto *C = GV->getComdat())
+          if (MergedMComdats.count(C))
+            return true;
         if (auto *F = dyn_cast<Function>(GV))
           return EligibleVirtualFns.count(F);
         if (auto *GVar = dyn_cast_or_null<GlobalVariable>(GV->getBaseObject()))
           return HasTypeMetadata(GVar);
         return false;
       }));
   StripDebugInfo(*MergedM);
 
   for (Function &F : *MergedM)
     if (!F.isDeclaration()) {
       // Reset the linkage of all functions eligible for virtual constant
       // propagation. The canonical definitions live in the thin LTO module so
       // that they can be imported.
       F.setLinkage(GlobalValue::AvailableExternallyLinkage);
       F.setComdat(nullptr);
     }
 
-  // Remove all globals with type metadata, as well as aliases pointing to them,
-  // from the thin LTO module.
+  // Remove all globals with type metadata, globals with comdats that live in
+  // MergedM, and aliases pointing to such globals from the thin LTO module.
   filterModule(&M, [&](const GlobalValue *GV) {
     if (auto *GVar = dyn_cast_or_null<GlobalVariable>(GV->getBaseObject()))
-      return !HasTypeMetadata(GVar);
+      if (HasTypeMetadata(GVar))
+        return false;
+    if (const auto *C = GV->getComdat())
+      if (MergedMComdats.count(C))
+        return false;
     return true;
   });
 
   promoteInternals(*MergedM, M, ModuleId);
   promoteInternals(M, *MergedM, ModuleId);
 
   simplifyExternals(*MergedM);
 

llvm/trunk/test/Transforms/ThinLTOBitcodeWriter/comdat.ll

+; RUN: opt -thinlto-bc -o %t %s
+; RUN: llvm-modextract -n 0 -o - %t | llvm-dis | FileCheck --check-prefix=THIN %s
+; RUN: llvm-modextract -n 1 -o - %t | llvm-dis | FileCheck --check-prefix=MERGED %s
+
+target datalayout = "e-m:w-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.0.24215"
+
+; Internal comdat leader with type metadata. All comdat members need to live
+; in the merged module, and the comdat needs to be renamed.
+; MERGED: ${{"?lwt[^ ]+}} = comdat any
+$lwt = comdat any
+
+; External comdat leader, type metadata on non-leader. All comdat
+; members need to live in the merged module, internal members need to
+; be renamed.
+; MERGED: $nlwt = comdat any
+$nlwt = comdat any
+
+; Comdat with two members without type metadata. All comdat members live in
+; the ThinLTO module and no renaming needs to take place.
+; THIN: $nt = comdat any
+$nt = comdat any
+
+; MERGED: @lwt_aliasee = private unnamed_addr global
+; MERGED-SAME: comdat(${{"?lwt[^ ]+}})
+@lwt_aliasee = private unnamed_addr global [1 x i8*] [i8* null], comdat($lwt), !type !0
+
+; MERGED: {{@"?lwt_nl[^ ]+}} = hidden unnamed_addr global
+; MERGED-SAME: comdat(${{"?lwt[^ ]+}})
+; THIN: {{@"?lwt_nl[^ ]+}} = external hidden
+@lwt_nl = internal unnamed_addr global i32 0, comdat($lwt)
+
+; MERGED: @nlwt_aliasee = private unnamed_addr global
+; MERGED-SAME: comdat($nlwt)
+@nlwt_aliasee = private unnamed_addr global [1 x i8*] [i8* null], comdat($nlwt), !type !0
+
+; MERGED: @nlwt = unnamed_addr global
+; MERGED-SAME: comdat
+; THIN: @nlwt = external
+@nlwt = unnamed_addr global i32 0, comdat
+
+; THIN: @nt = internal
+; THIN-SAME: comdat
+@nt = internal unnamed_addr global [1 x i8*] [i8* null], comdat
+
+; THIN: @nt_nl = internal
+; THIN-SAME: comdat($nt)
+@nt_nl = internal unnamed_addr global i32 0, comdat($nt)
+
+; MERGED: {{@"?lwt[^ ]+}} = hidden unnamed_addr alias
+; THIN: {{@"?lwt[^ ]+}} = external hidden
+@lwt = internal unnamed_addr alias [1 x i8*], [1 x i8*]* @lwt_aliasee
+
+; MERGED: {{@"?nlwt_nl[^ ]+}} = hidden unnamed_addr alias
+; THIN: {{@"?nlwt_nl[^ ]+}} = external hidden
+@nlwt_nl = internal unnamed_addr alias [1 x i8*], [1 x i8*]* @nlwt_aliasee
+
+; The functions below exist just to make sure the globals are used.
+define i8* @lwt_fun() {
+  %1 = load i32, i32* @lwt_nl
+  %2 = getelementptr inbounds [1 x i8*], [1 x i8*]* @lwt, i32 0, i32 %1
+  %3 = load i8*, i8** %2
+  ret i8* %3
+}
+
+define i8* @nlwt_fun() {
+  %1 = load i32, i32* @nlwt
+  %2 = getelementptr inbounds [1 x i8*], [1 x i8*]* @nlwt_nl, i32 0, i32 %1
+  %3 = load i8*, i8** %2
+  ret i8* %3
+}
+
+define i8* @nt_fun() {
+  %1 = load i32, i32* @nt_nl
+  %2 = getelementptr inbounds [1 x i8*], [1 x i8*]* @nt, i32 0, i32 %1
+  %3 = load i8*, i8** %2
+  ret i8* %3
+}
+
+!0 = !{i64 8, !"?AVA@@"}