From caa2a829cdf905a5e8664d96a464d414b2adb42e Mon Sep 17 00:00:00 2001
From: Amanieu d'Antras <amanieu@gmail.com>
Date: Thu, 16 Jun 2022 21:31:33 +0100
Subject: [PATCH] [MergeFunctions] Preserve symbols used
 llvm.used/llvm.compiler.used

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.

Differential Revision: https://reviews.llvm.org/D127751
---
 llvm/lib/Transforms/IPO/MergeFunctions.cpp   | 15 +++++++++++-
 llvm/test/Transforms/MergeFunc/merge-used.ll | 35 ++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/MergeFunc/merge-used.ll

diff --git a/llvm/lib/Transforms/IPO/MergeFunctions.cpp b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
index 2c35dd6..bb85dde 100644
--- a/llvm/lib/Transforms/IPO/MergeFunctions.cpp
+++ b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
@@ -120,6 +120,7 @@
 #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>
@@ -225,6 +226,9 @@ private:
   /// analyzed again.
   std::vector<WeakTrackingVH> Deferred;
 
+  /// Set of values marked as used in llvm.used and llvm.compiler.used.
+  SmallPtrSet<GlobalValue *, 4> Used;
+
 #ifndef NDEBUG
   /// Checks the rules of order relation introduced among functions set.
   /// Returns true, if check has been passed, and false if failed.
@@ -407,6 +411,11 @@ static bool isEligibleForMerging(Function &F) {
 bool MergeFunctions::runOnModule(Module &M) {
   bool Changed = false;
 
+  SmallVector<GlobalValue *, 4> UsedV;
+  collectUsedGlobalVariables(M, UsedV, /*CompilerUsed=*/false);
+  collectUsedGlobalVariables(M, UsedV, /*CompilerUsed=*/true);
+  Used.insert(UsedV.begin(), UsedV.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 *>>
@@ -453,6 +462,7 @@ bool MergeFunctions::runOnModule(Module &M) {
   FnTree.clear();
   FNodesInTree.clear();
   GlobalNumbers.clear();
+  Used.clear();
 
   return Changed;
 }
@@ -825,7 +835,10 @@ void MergeFunctions::mergeTwoFunctions(Function *F, Function *G) {
     // 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()) {
+      // 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.
+      if (G->hasGlobalUnnamedAddr() && !Used.contains(G)) {
         // 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);
diff --git a/llvm/test/Transforms/MergeFunc/merge-used.ll b/llvm/test/Transforms/MergeFunc/merge-used.ll
new file mode 100644
index 0000000..a86e66e
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/merge-used.ll
@@ -0,0 +1,35 @@
+; RUN: opt -S -mergefunc < %s | FileCheck %s
+
+@llvm.compiler.used = appending global [1 x i8*] [i8* bitcast (i32 (i32)* @a to i8*)], section "llvm.metadata"
+
+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
+}
+
+define i32 @c(i32 %a) unnamed_addr {
+  %b = tail call i32 @a(i32 %a)
+  ret i32 %b
+}
+
+; CHECK-LABEL: @llvm.compiler.used = appending global [1 x i8*] [i8* bitcast (i32 (i32)* @a to i8*)], section "llvm.metadata"
+
+; CHECK-LABEL: define i32 @b(i32 %a) unnamed_addr
+; CHECK-NEXT:    xor
+; CHECK-NEXT:    xor
+; CHECK-NEXT:    ret
+
+; CHECK-LABEL: define i32 @c(i32 %a) unnamed_addr
+; CHECK-NEXT:    tail call i32 @b(i32 %a)
+; CHECK-NEXT:    ret
+
+; CHECK-LABEL: define internal i32 @a(i32 %0) unnamed_addr
+; CHECK-NEXT:    tail call i32 @b(i32 %0)
+; CHECK-NEXT:    ret
-- 
2.7.4