ThinLTOBitcodeWriter: Emit __cfi_check to full LTO part of bitcode file.

[pcc]

The CrossDSOCFI pass runs on the full LTO module and fills in the
body of __cfi_check. This function must have the correct attributes in
order to be compatible with the rest of the program. For example, when
building with -mbranch-protection=standard, the function must have the
branch-target-enforcement attribute, which is normally added by Clang.
When __cfi_check is missing, CrossDSOCFI will give it the default set
of attributes, which are likely incorrect. Therefore, emit __cfi_check
to the full LTO part, where CrossDSOCFI will see it.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Peter Collingbourne (pcc)

Changes

The CrossDSOCFI pass runs on the full LTO module and fills in the
body of __cfi_check. This function must have the correct attributes in
order to be compatible with the rest of the program. For example, when
building with -mbranch-protection=standard, the function must have the
branch-target-enforcement attribute, which is normally added by Clang.
When __cfi_check is missing, CrossDSOCFI will give it the default set
of attributes, which are likely incorrect. Therefore, emit __cfi_check
to the full LTO part, where CrossDSOCFI will see it.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp (+11-1)
• (added) llvm/test/Transforms/ThinLTOBitcodeWriter/cfi-check.ll (+19)

diff --git a/llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp b/llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp
index 4387c3827c015..838f97c8f49af 100644
--- a/llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp
+++ b/llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp
@@ -350,12 +350,20 @@ void splitAndWriteThinLTOBitcode(
       });
     }
 
+  auto MustEmitToMergedModule = [](const GlobalValue *GV) {
+    // The __cfi_check definition is filled in by the CrossDSOCFI pass which
+    // runs only in the merged module.
+    return GV->getName() == "__cfi_check";
+  };
+
   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 (MustEmitToMergedModule(GV))
+          return true;
         if (auto *F = dyn_cast<Function>(GV))
           return EligibleVirtualFns.count(F);
         if (auto *GVar =
@@ -372,7 +380,7 @@ void splitAndWriteThinLTOBitcode(
   cloneUsedGlobalVariables(M, *MergedM, /*CompilerUsed*/ true);
 
   for (Function &F : *MergedM)
-    if (!F.isDeclaration()) {
+    if (!F.isDeclaration() && !MustEmitToMergedModule(&F)) {
       // 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.
@@ -398,6 +406,8 @@ void splitAndWriteThinLTOBitcode(
     if (const auto *C = GV->getComdat())
       if (MergedMComdats.count(C))
         return false;
+    if (MustEmitToMergedModule(GV))
+      return false;
     return true;
   });
 
diff --git a/llvm/test/Transforms/ThinLTOBitcodeWriter/cfi-check.ll b/llvm/test/Transforms/ThinLTOBitcodeWriter/cfi-check.ll
new file mode 100644
index 0000000000000..b927af6b92f7c
--- /dev/null
+++ b/llvm/test/Transforms/ThinLTOBitcodeWriter/cfi-check.ll
@@ -0,0 +1,19 @@
+; RUN: opt -thinlto-bc -thinlto-split-lto-unit -o %t %s
+; RUN: llvm-modextract -b -n 0 -o - %t | llvm-dis | FileCheck --check-prefix=M0 %s
+; RUN: llvm-modextract -b -n 1 -o - %t | llvm-dis | FileCheck --check-prefix=M1 %s
+
+; Check that __cfi_check is emitted on the full LTO side with
+; attributes preserved.
+
+; M0: define void @f()
+define void @f() !type !{!"f1", i32 0} {
+  ret void 
+}
+
+; M1: define void @__cfi_check() #0
+define void @__cfi_check() #0 {
+  ret void
+}
+
+; M1: attributes #0 = { "branch-target-enforcement" }
+attributes #0 = { "branch-target-enforcement" }

[efriedma-quic]

As discussed in #98673, I think we should come up with a better way of specifying appropriate target features for compiler/linker-generated functions, as opposed to making every translation unit define a function named __cfi_check with the right target features attached, but fixing the bug in the existing approach seems fine for now.

I'm not an expert on ThinLTO linking, but the code looks right, and the fix seems to work for my testcases as far as I can tell.

LGTM

[efriedma-quic]

/cherry-pick ff85dbd

[llvmbot]

/pull-request #154859