[CFI][annotation] Leave alone function pointers in function annotations

[yozhu]

Function annotation, as part of llvm.metadata, is for the function itself and doesn't apply to its corresponding jump table entry, so with CFI we shouldn't replace function pointer in function annotation with pointer to its corresponding jump table entry.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-transforms

Author: None (yozhu)

Changes

Function annotation, as part of llvm.metadata, is for the function itself and doesn't apply to its corresponding jump table entry, so with CFI we shouldn't replace function pointer in function annotation with pointer to its corresponding jump table entry.

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/IPO/LowerTypeTests.cpp (+29-1)

diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 733f290b1bc93..f6630019eaec6 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -470,6 +470,9 @@ class LowerTypeTestsModule {
 
   Function *WeakInitializerFn = nullptr;
 
+  GlobalVariable *GlobalAnnotation;
+  std::set<void *> FunctionAnnotations;
+
   bool shouldExportConstantsAsAbsoluteSymbols();
   uint8_t *exportTypeId(StringRef TypeId, const TypeIdLowering &TIL);
   TypeIdLowering importTypeId(StringRef TypeId);
@@ -531,6 +534,10 @@ class LowerTypeTestsModule {
   /// replace each use, which is a direct function call.
   void replaceDirectCalls(Value *Old, Value *New);
 
+  bool isFunctionAnnotation(void *GV) {
+    return FunctionAnnotations.count(GV) != 0;
+  }
+
 public:
   LowerTypeTestsModule(Module &M, ModuleAnalysisManager &AM,
                        ModuleSummaryIndex *ExportSummary,
@@ -1377,8 +1384,11 @@ void LowerTypeTestsModule::replaceWeakDeclarationWithJumpTablePtr(
   // (all?) targets. Switch to a runtime initializer.
   SmallSetVector<GlobalVariable *, 8> GlobalVarUsers;
   findGlobalVariableUsersOf(F, GlobalVarUsers);
-  for (auto *GV : GlobalVarUsers)
+  for (auto *GV : GlobalVarUsers) {
+    if (GV == GlobalAnnotation)
+      continue;
     moveInitializerToModuleConstructor(GV);
+  }
 
   // Can not RAUW F with an expression that uses F. Replace with a temporary
   // placeholder first.
@@ -1392,6 +1402,10 @@ void LowerTypeTestsModule::replaceWeakDeclarationWithJumpTablePtr(
   // Don't use range based loop, because use list will be modified.
   while (!PlaceholderFn->use_empty()) {
     Use &U = *PlaceholderFn->use_begin();
+    if (isFunctionAnnotation(U.getUser())) {
+      U.set(F);
+      continue;
+    }
     auto *InsertPt = dyn_cast<Instruction>(U.getUser());
     assert(InsertPt && "Non-instruction users should have been eliminated");
     auto *PN = dyn_cast<PHINode>(InsertPt);
@@ -1837,6 +1851,16 @@ LowerTypeTestsModule::LowerTypeTestsModule(
   }
   OS = TargetTriple.getOS();
   ObjectFormat = TargetTriple.getObjectFormat();
+
+  // Function annotation describes or applies to function itself, and
+  // shouldn't be associated with jump table thunk generated for CFI.
+  GlobalAnnotation = M.getGlobalVariable("llvm.global.annotations");
+  auto *C = dyn_cast_or_null<Constant>(GlobalAnnotation);
+  if (C && C->getNumOperands() == 1) {
+    C = cast<Constant>(C->getOperand(0));
+    for (auto &Op : C->operands())
+      FunctionAnnotations.insert(Op.get());
+  }
 }
 
 bool LowerTypeTestsModule::runForTesting(Module &M, ModuleAnalysisManager &AM) {
@@ -1900,6 +1924,10 @@ void LowerTypeTestsModule::replaceCfiUses(Function *Old, Value *New,
     if (isDirectCall(U) && (Old->isDSOLocal() || !IsJumpTableCanonical))
       continue;
 
+    // Skip function annotation
+    if (isFunctionAnnotation(U.getUser()))
+      continue;
+
     // Must handle Constants specially, we cannot call replaceUsesOfWith on a
     // constant because they are uniqued.
     if (auto *C = dyn_cast<Constant>(U.getUser())) {

[nikic]

Missing tests.

[nikic]

Implementation looks fine, the test needs more work.

[yozhu]

CC @nikic: thanks for all the comments so far. Any other feedback? If the latest version looks good to you, could you approve it?

[yozhu]

CC: @nikic Unnecessary things have been removed from the test ..

[nikic]

After your last update, the test will also pass without your changes. I believe the issue is that you dropped the !2 = !{i32 4, !"CFI Canonical Jump Tables", i32 0} metadata, which is necessary.

[yozhu]

The issue turns out to be that I dropped attribute and type specifications:

• With the following the problem will reproduce:

define i32 @bar(i32 noundef %0) #0 !type !8 !type !9 {

• And with the following the problem doesn't reproduce:

define i32 @bar(i32 noundef %0) {

#0, !8 and !9 are defined as:

attributes #0 = { noinline nounwind optnone uwtable "frame-pointer"="non-leaf" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+fp-armv8,+neon,+v8a,-fmv" }
 60 !8 = !{i64 0, !"_ZTSFiiE"}
 61 !9 = !{i64 0, !"_ZTSFiiE.generalized"}

Also it appears that the test code can't be made too simple because otherwise indirect call target will be easy to figure out by optimizations and then CFI has nowhere to apply.

[yozhu]

CC @nikic on the updated test .. could we please get this wrapped up this week? Thanks.

[nikic]

LGTM

[yozhu]

LGTM

Thanks for all your feedback and comments! Much appreciated!