[llvm][CFI] Ensure COFF comdat renaming applies for imported functions

[PiJoules]

I ran into the same issue as
#139962 regarding the comdat corresponding to a renamed key function but for thinlto. My last patch had not considered the thinlto case, so this applies the same fix for imported functions.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: None (PiJoules)

Changes

I ran into the same issue as
#139962 regarding the comdat corresponding to a renamed key function but for thinlto. My last patch had not considered the thinlto case, so this applies the same fix for imported functions.

---

Full diff: https://github.com/llvm/llvm-project/pull/143421.diff

3 Files Affected:

• (modified) llvm/lib/Transforms/IPO/LowerTypeTests.cpp (+21-13)
• (added) llvm/test/Transforms/LowerTypeTests/Inputs/import-thinlto-funcs.yaml (+5)
• (modified) llvm/test/Transforms/LowerTypeTests/cfi-coff-comdat-rename.ll (+2)

diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 8d8a73729e74c..e5f9401fda276 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -561,6 +561,8 @@ class LowerTypeTestsModule {
     return FunctionAnnotations.contains(V);
   }
 
+  void maybeReplaceComdat(Function *F, StringRef OriginalName);
+
 public:
   LowerTypeTestsModule(Module &M, ModuleAnalysisManager &AM,
                        ModuleSummaryIndex *ExportSummary,
@@ -1082,6 +1084,23 @@ void LowerTypeTestsModule::importTypeTest(CallInst *CI) {
   }
 }
 
+void LowerTypeTestsModule::maybeReplaceComdat(Function *F,
+                                              StringRef OriginalName) {
+  // For COFF we should also rename the comdat if this function also
+  // happens to be the key function. Even if the comdat name changes, this
+  // should still be fine since comdat and symbol resolution happens
+  // before LTO, so all symbols which would prevail have been selected.
+  if (F->hasComdat() && ObjectFormat == Triple::COFF &&
+      F->getComdat()->getName() == OriginalName) {
+    Comdat *OldComdat = F->getComdat();
+    Comdat *NewComdat = M.getOrInsertComdat(F->getName());
+    for (GlobalObject &GO : M.global_objects()) {
+      if (GO.getComdat() == OldComdat)
+        GO.setComdat(NewComdat);
+    }
+  }
+}
+
 // ThinLTO backend: the function F has a jump table entry; update this module
 // accordingly. isJumpTableCanonical describes the type of the jump table entry.
 void LowerTypeTestsModule::importFunction(
@@ -1115,6 +1134,7 @@ void LowerTypeTestsModule::importFunction(
     FDecl->setVisibility(GlobalValue::HiddenVisibility);
   } else {
     F->setName(Name + ".cfi");
+    maybeReplaceComdat(F, Name);
     F->setLinkage(GlobalValue::ExternalLinkage);
     FDecl = Function::Create(F->getFunctionType(), GlobalValue::ExternalLinkage,
                              F->getAddressSpace(), Name, &M);
@@ -1734,19 +1754,7 @@ void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
       FAlias->takeName(F);
       if (FAlias->hasName()) {
         F->setName(FAlias->getName() + ".cfi");
-        // For COFF we should also rename the comdat if this function also
-        // happens to be the key function. Even if the comdat name changes, this
-        // should still be fine since comdat and symbol resolution happens
-        // before LTO, so all symbols which would prevail have been selected.
-        if (F->hasComdat() && ObjectFormat == Triple::COFF &&
-            F->getComdat()->getName() == FAlias->getName()) {
-          Comdat *OldComdat = F->getComdat();
-          Comdat *NewComdat = M.getOrInsertComdat(F->getName());
-          for (GlobalObject &GO : M.global_objects()) {
-            if (GO.getComdat() == OldComdat)
-              GO.setComdat(NewComdat);
-          }
-        }
+        maybeReplaceComdat(F, FAlias->getName());
       }
       replaceCfiUses(F, FAlias, IsJumpTableCanonical);
       if (!F->hasLocalLinkage())
diff --git a/llvm/test/Transforms/LowerTypeTests/Inputs/import-thinlto-funcs.yaml b/llvm/test/Transforms/LowerTypeTests/Inputs/import-thinlto-funcs.yaml
new file mode 100644
index 0000000000000..459d45032b0c4
--- /dev/null
+++ b/llvm/test/Transforms/LowerTypeTests/Inputs/import-thinlto-funcs.yaml
@@ -0,0 +1,5 @@
+---
+CfiFunctionDefs:
+  - f1
+  - f2
+...
diff --git a/llvm/test/Transforms/LowerTypeTests/cfi-coff-comdat-rename.ll b/llvm/test/Transforms/LowerTypeTests/cfi-coff-comdat-rename.ll
index 7dda7f6df10c3..7eede8b7322f8 100644
--- a/llvm/test/Transforms/LowerTypeTests/cfi-coff-comdat-rename.ll
+++ b/llvm/test/Transforms/LowerTypeTests/cfi-coff-comdat-rename.ll
@@ -1,5 +1,6 @@
 ; REQUIRES: x86-registered-target
 ; RUN: opt -S -passes=lowertypetests %s | FileCheck %s
+; RUN: opt -S -passes=lowertypetests -lowertypetests-summary-action=import -lowertypetests-read-summary=%p/Inputs/import-thinlto-funcs.yaml %s | FileCheck %s
 
 ;; This is a check to assert we don't crash with:
 ;;
@@ -7,6 +8,7 @@
 ;;
 ;; So this just needs to exit normally.
 ; RUN: opt -S -passes=lowertypetests %s | llc -asm-verbose=false
+; RUN: opt -S -passes=lowertypetests -lowertypetests-summary-action=import -lowertypetests-read-summary=%p/Inputs/import-thinlto-funcs.yaml %s | llc -asm-verbose=false
 
 target datalayout = "e-p:64:64"
 target triple = "x86_64-pc-windows-msvc"

[ilovepi]

How will this affect FatLTO? We run a version of LowerTypeTests as part of the non-LTO compilation that basically drops everything related to tests after we apply some cleanups (

llvm-project/llvm/lib/Passes/PassBuilderPipelines.cpp

Line 1694 in 6cbd91e

LowerTypeTestsPass(nullptr, nullptr, lowertypetests::DropTestKind::All));

). I don't think the renaming will happen in this case, but I'm not 100% certain.

[PiJoules]

How will this affect FatLTO? We run a version of LowerTypeTests as part of the non-LTO compilation that basically drops everything related to tests after we apply some cleanups (

llvm-project/llvm/lib/Passes/PassBuilderPipelines.cpp

Line 1694 in 6cbd91e

LowerTypeTestsPass(nullptr, nullptr, lowertypetests::DropTestKind::All));

). I don't think the renaming will happen in this case, but I'm not 100% certain.

I think the renaming shouldn't happen since LowerTypeTestsModule::lower should immediately return if DropTypeTests != DropTestKind::None and all the renaming and actual lowering happens after that check.

[ilovepi]

you may want to give @pcc a chance to chime in before landing.

[PiJoules]

@pcc does this lgty?

[pcc]

LGTM