[MergeFunctions] Preserve symbols used llvm.used/llvm.compiler.used

[Amanieu]

Summary

llvm.used and llvm.compiler.used are often used with inline assembly that refers to a specific symbol so that the symbol is kept through to the linker even though there are no references to it from LLVM IR. This fixes the MergeFunctions pass to preserve references to these symbols in llvm.used/llvm.compiler.used so they are not deleted from the IR. This doesn't prevent these functions from being merged, but guarantees that an alias or thunk with the expected symbol name is kept in the IR.

[nikic]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

862

So does this just leave the old function behind, effectively not doing a merge? I think what you actually want is to go into the alias/thunk code path in this case.

llvm/test/Transforms/MergeFunc/merge-used.ll

6

Please include full check lines (update_test_checks.py with --check-globals flag), it's unclear what the result here is.

[Amanieu]

Updated test.

llvm/lib/Transforms/IPO/MergeFunctions.cpp

862

The writeThunkOrAlias call below will replace the old function with a thunk or alias. Adding the old function back to llvm.used will cause the G->use_empty() condition below to fail which prevents the function from being removed from the IR.

[nikic]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

862

I see. In that case I'd suggest to move a condition (on a single combined set) next to G->hasGlobalUnnamedAddr(), so we only do replaceDirectCallers() replacement. This might replace the function in slightly less cases, but I'd prefer that over the explicit manipulation of the used globals list here.

[nikic]

This looks good to me. Maybe @MaskRay or @efriedma want to take a look: I'd mainly like confirmation that the current transform (merging an unnamed addr, internal linkage function that is part of llvm.used by replacing all uses) is indeed illegal, because I'm a bit hazy on the precise guarantees in this area.

llvm/test/Transforms/MergeFunc/merge-used.ll

15

Might make sense to including another function that calls @a, just to show that that call does get replaced with @b and doesn't call into the thunk.

[efriedma]

We can't rename or erase a function used by @llvm.used; the symbol has to exist for use by inline asm etc. So this is correct in that respect.

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

I'm a bit suspicious of the hasGlobalUnnamedAddr() check here. hasGlobalUnnamedAddr() generally implies it's legal to merge two constants/functions, but we don't actually figure out if we're going to perform an actual merge at this point; that doesn't happen until we try to discard it, or call writeThunkOrAlias. (unnamed_addr only allows rewriting all the uses of a variable/function at once, not just some of them. Otherwise you end up with weird inconsistencies where a value isn't equal to itself.)

[Amanieu]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

I think the code is valid as it is:

• If hasGlobalUnnamedAddr() == true then we replace all uses of the function with replaceAllUsesWith.
• Otherwise we only replace direct caller, which is equivalent to inlining the thunk/alias.

[efriedma]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

replaceAllUsesWith() only replaces uses in the current module, not uses in other modules. So despite the name, it's not replacing "all" the uses in the sense we need here.

If the function has internal linkage, we're fine because we're actually rewriting all the uses. If we replace the function with an alias, we're fine because every use refers to the same address. Otherwise, there's a problem: we're replacing the uses in the current module, but not the uses in other modules.

(Making sure unnamed_addr works correctly in this respect is more likely to be a practical issue with constant variables, as opposed to functions, but I'd rather keep this consistent if possible.)

[MaskRay]

llvm/test/Transforms/MergeFunc/merge-used.ll

34

Add a comment that this is preserved due to llvm.compiler.used

[MaskRay]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

hasGlobalUnnamedAddr was from D34805. So what's the recommendation here? Further restrict the condition to local linkage functions?

llvm/test/Transforms/MergeFunc/merge-used.ll

2

Use -passes=mergefunc

[nikic]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

Trying to understand what the concern here is based on an example. Would it be something like this https://llvm.godbolt.org/z/e3WYWr69P? Where in this module we observe that @a == @b, and then we have another module with linkonce_odr a,b (and lets say those definitions get picked by the linker) and we might observe @a != @b in that module?

[efriedma]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

Trying to understand what the concern here is based on an example.

The fundamental invariant we want to ensure here is basically @a == @a. Or more generally, "%a == %a" for any pointer that isn't poison/undef. This should remain true even if one side of the comparison comes from a different module.

That also implies @a == @b is the same in all modules.

For a real-world example of how things can explode if unnamed_addr is mishandled in this respect, see https://github.com/llvm/llvm-project/issues/32127 .

hasGlobalUnnamedAddr was from D34805. So what's the recommendation here? Further restrict the condition to local linkage functions?

Local linkage should be safe to optimize like this. Situations where canCreateAliasFor() is true should also be safe. (Assuming the global isn't referred to by llvm.used.)

[Amanieu]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

I'm not sure what solution you are suggesting here exactly, but it seems to me that this issue was already present beforehand even without this change. Should this be handled separately so this bugfix can land?

[nikic]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

Yes, I think this is fine to land, these problems aren't really related.

I don't really see a way to resolve @efriedma's issue short of restricting to local linkage. Even if canCreateAliasFor() return true that alias will simply get dropped later, if the function has discardable linkage.

[efriedma]

llvm/lib/Transforms/IPO/MergeFunctions.cpp

846

Sure, I won't object to landing this as-is.

llvm/lib/Transforms/IPO/MergeFunctions.cpp

 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/FunctionComparator.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <algorithm>
 #include <cassert>
 #include <iterator>
 #include <set>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
private:
   using FnTreeType = std::set<FunctionNode, FunctionNodeCmp>;
 
   GlobalNumberState GlobalNumbers;
 
   /// A work queue of functions that may have been modified and should be
   /// analyzed again.
   std::vector<WeakTrackingVH> Deferred;
 
+  /// Set of values marked as used in llvm.used and llvm.compiler.used.
+  SmallPtrSet<GlobalValue *, 4> Used;
+  SmallPtrSet<GlobalValue *, 4> CompilerUsed;
+
 #ifndef NDEBUG
   /// Checks the rules of order relation introduced among functions set.
   /// Returns true, if check has been passed, and false if failed.
   bool doFunctionalCheck(std::vector<WeakTrackingVH> &Worklist);
 #endif
 
   /// Insert a ComparableFunction into the FnTree, or merge it away if it's
   /// equal to one that's already present.

 /// Check whether \p F is eligible for function merging.
 static bool isEligibleForMerging(Function &F) {
   return !F.isDeclaration() && !F.hasAvailableExternallyLinkage();
 }
 
 bool MergeFunctions::runOnModule(Module &M) {
   bool Changed = false;
 
+  // Keep track of functions referred to by llvm.used/llvm.compiler.used.
+  SmallVector<GlobalValue *, 4> UsedV;
+  collectUsedGlobalVariables(M, UsedV, /*CompilerUsed=*/false);
+  Used.insert(UsedV.begin(), UsedV.end());
+  SmallVector<GlobalValue *, 4> CompilerUsedV;
+  collectUsedGlobalVariables(M, CompilerUsedV, /*CompilerUsed=*/true);
+  CompilerUsed.insert(CompilerUsedV.begin(), CompilerUsedV.end());
+
   // All functions in the module, ordered by hash. Functions with a unique
   // hash value are easily eliminated.
   std::vector<std::pair<FunctionComparator::FunctionHash, Function *>>
     HashedFuncs;
   for (Function &Func : M) {
     if (isEligibleForMerging(Func)) {
       HashedFuncs.push_back({FunctionComparator::functionHash(Func), &Func});
     }
for (WeakTrackingVH &I : Worklist) {
       }
     }
     LLVM_DEBUG(dbgs() << "size of FnTree: " << FnTree.size() << '\n');
   } while (!Deferred.empty());
 
   FnTree.clear();
   FNodesInTree.clear();
   GlobalNumbers.clear();
+  Used.clear();
+  CompilerUsed.clear();
 
   return Changed;
 }
 
 // Replace direct callers of Old with New.
 void MergeFunctions::replaceDirectCallers(Function *Old, Function *New) {
   Constant *BitcastNew = ConstantExpr::getBitCast(New, Old->getType());
   for (Use &U : llvm::make_early_inc_range(Old->uses())) {
if (F->isInterposable()) {
     ++NumFunctionsMerged;
   } else {
     // For better debugability, under MergeFunctionsPDI, we do not modify G's
     // call sites to point to F even when within the same translation unit.
     if (!G->isInterposable() && !MergeFunctionsPDI) {
       if (G->hasGlobalUnnamedAddr()) {
         // G might have been a key in our GlobalNumberState, and it's illegal
         // to replace a key in ValueMap<GlobalValue *> with a non-global.
         GlobalNumbers.erase(G);
         // If G's address is not significant, replace it entirely.
         Constant *BitcastF = ConstantExpr::getBitCast(F, G->getType());
         removeUsers(G);
         G->replaceAllUsesWith(BitcastF);
+
+        // Functions referred to by llvm.used/llvm.compiler.used are special:
+        // there are uses of the symbol name that are not visible to LLVM,
+        // usually from inline asm. We need to handle these by keeping the old
+        // symbol in the used list as an alias or thunk.
+        //
+        // RAUW will have removed references to G from the existing list, make
+        // sure to add it back.
+        if (Used.contains(G))
+          appendToUsed(*G->getParent(), {G});
+        if (CompilerUsed.contains(G))
+          appendToCompilerUsed(*G->getParent(), {G});
       } else {
         // Redirect direct callers of G to F. (See note on MergeFunctionsPDI
         // above).
         replaceDirectCallers(G, F);
       }
     }
 
     // If G was internal then we may have replaced all uses of G with F. If so,

llvm/test/Transforms/MergeFunc/merge-used.ll

+; RUN: opt -S -mergefunc < %s | FileCheck %s
+
[e-mail] = appending global [1 x i8*] [i8* bitcast (i32 (i32)* @a to i8*)], section "llvm.metadata"
+
+; CHECK: @a
+; CHECK: @b
+
+define internal i32 @a(i32 %a) unnamed_addr {
+  %b = xor i32 %a, 0
+  %c = xor i32 %b, 0
+  ret i32 %c
+}
+
+define i32 @b(i32 %a) unnamed_addr {
+  %b = xor i32 %a, 0
+  %c = xor i32 %b, 0
+  ret i32 %c
+}