[OpenMP] Add optimization to remove the RPC client (#70683)

Summary:
Part of the work done in the `libc` project is to provide host services
for things like `printf` or `malloc`, or generally any syscall-like
behaviour. This scheme works by emitting an externally visible global
called `__llvm_libc_rpc_client` that the host runtime can pick up to get
a handle to the global memory associated with the client. We use the
presence of this symbol to indicate whether or not we need to run an RPC
server. Normally, this symbol is only present if something requiring an
RPC server was linked in, such as `printf`. However, if this call to
`printf` was subsequently optimizated out, the symbol would remain and
cannot be removed (rightfully so) because of its linkage. This patch
adds a special-case optimization to remove this symbol so we can
indicate that an RPC server is no longer needed.

This patch puts this logic in `OpenMPOpt` as the most readily available
place for it. In the future, we should think how to move this somewhere
more generic. Furthermore, we use a hard-coded runtime name (which isn't
uncommon given all the other magic symbol names). But it might be nice
to abstract that part away.

Author: jhuber6

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

@@ -978,6 +978,9 @@ struct OpenMPOpt {
       }
     }
 
+    if (OMPInfoCache.OpenMPPostLink)
+      Changed |= removeRuntimeSymbols();
+
     return Changed;
   }
 
@@ -1507,6 +1510,37 @@ struct OpenMPOpt {
     return Changed;
   }
 
+  /// Tries to remove known runtime symbols that are optional from the module.
+  bool removeRuntimeSymbols() {
+    // The RPC client symbol is defined in `libc` and indicates that something
+    // required an RPC server. If its users were all optimized out then we can
+    // safely remove it.
+    // TODO: This should be somewhere more common in the future.
+    if (GlobalVariable *GV = M.getNamedGlobal("__llvm_libc_rpc_client")) {
+      if (!GV->getType()->isPointerTy())
+        return false;
+
+      Constant *C = GV->getInitializer();
+      if (!C)
+        return false;
+
+      // Check to see if the only user of the RPC client is the external handle.
+      GlobalVariable *Client = dyn_cast<GlobalVariable>(C->stripPointerCasts());
+      if (!Client || Client->getNumUses() > 1 ||
+          Client->user_back() != GV->getInitializer())
+        return false;
+
+      Client->replaceAllUsesWith(PoisonValue::get(Client->getType()));
+      Client->eraseFromParent();
+
+      GV->replaceAllUsesWith(PoisonValue::get(GV->getType()));
+      GV->eraseFromParent();
+
+      return true;
+    }
+    return false;
+  }
+
   /// Tries to hide the latency of runtime calls that involve host to
   /// device memory transfers by splitting them into their "issue" and "wait"
   /// versions. The "issue" is moved upwards as much as possible. The "wait" is

llvm/test/Transforms/OpenMP/keep_rpc_client.ll

@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-globals
+; RUN: opt -S -passes=openmp-opt-postlink < %s | FileCheck %s --check-prefix=POSTLINK
+; RUN: opt -S -passes=openmp-opt < %s | FileCheck %s --check-prefix=PRELINK
+
+@client = internal addrspace(1) global i64 zeroinitializer, align 8
+@__llvm_libc_rpc_client = protected local_unnamed_addr addrspace(1) global ptr addrspacecast (ptr addrspace(1) @client to ptr), align 8
+
+;.
+; POSTLINK: @[[CLIENT:[a-zA-Z0-9_$"\\.-]+]] = internal addrspace(1) global i64 0, align 8
+; POSTLINK: @[[__LLVM_LIBC_RPC_CLIENT:[a-zA-Z0-9_$"\\.-]+]] = protected local_unnamed_addr addrspace(1) global ptr addrspacecast (ptr addrspace(1) @client to ptr), align 8
+;.
+; PRELINK: @[[CLIENT:[a-zA-Z0-9_$"\\.-]+]] = internal addrspace(1) global i64 0, align 8
+; PRELINK: @[[__LLVM_LIBC_RPC_CLIENT:[a-zA-Z0-9_$"\\.-]+]] = protected local_unnamed_addr addrspace(1) global ptr addrspacecast (ptr addrspace(1) @client to ptr), align 8
+;.
+define i64 @a() {
+; POSTLINK-LABEL: define {{[^@]+}}@a
+; POSTLINK-SAME: () #[[ATTR0:[0-9]+]] {
+; POSTLINK-NEXT:    [[RETVAL:%.*]] = load i64, ptr addrspace(1) @client, align 8
+; POSTLINK-NEXT:    ret i64 [[RETVAL]]
+;
+; PRELINK-LABEL: define {{[^@]+}}@a
+; PRELINK-SAME: () #[[ATTR0:[0-9]+]] {
+; PRELINK-NEXT:    [[RETVAL:%.*]] = load i64, ptr addrspace(1) @client, align 8
+; PRELINK-NEXT:    ret i64 [[RETVAL]]
+;
+  %retval = load i64, ptr addrspace(1) @client, align 8
+  ret i64 %retval
+}
+
+!llvm.module.flags = !{!0, !1, !2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"openmp", i32 50}
+!2 = !{i32 7, !"openmp-device", i32 50}

llvm/test/Transforms/OpenMP/remove_rpc_client.ll

@@ -0,0 +1,29 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-globals
+; RUN: opt -S -passes=openmp-opt-postlink < %s | FileCheck %s --check-prefix=POSTLINK
+; RUN: opt -S -passes=openmp-opt < %s | FileCheck %s --check-prefix=PRELINK
+
+@client = internal addrspace(1) global i32 zeroinitializer, align 8
+@__llvm_libc_rpc_client = protected local_unnamed_addr addrspace(1) global ptr addrspacecast (ptr addrspace(1) @client to ptr), align 8
+
+;.
+; POSTLINK-NOT: @[[CLIENT:[a-zA-Z0-9_$"\\.-]+]] = internal addrspace(1) global i32 0, align 8
+; POSTLINK-NOT: @[[__LLVM_LIBC_RPC_CLIENT:[a-zA-Z0-9_$"\\.-]+]] = protected local_unnamed_addr addrspace(1) global ptr addrspacecast (ptr addrspace(1) @client to ptr), align 8
+;.
+; PRELINK: @[[CLIENT:[a-zA-Z0-9_$"\\.-]+]] = internal addrspace(1) global i32 0, align 8
+; PRELINK: @[[__LLVM_LIBC_RPC_CLIENT:[a-zA-Z0-9_$"\\.-]+]] = protected local_unnamed_addr addrspace(1) global ptr addrspacecast (ptr addrspace(1) @client to ptr), align 8
+;.
+define void @a() {
+; POSTLINK-LABEL: define {{[^@]+}}@a() {
+; POSTLINK-NEXT:    ret void
+;
+; PRELINK-LABEL: define {{[^@]+}}@a() {
+; PRELINK-NEXT:    ret void
+;
+  ret void
+}
+
+!llvm.module.flags = !{!0, !1, !2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"openmp", i32 50}
+!2 = !{i32 7, !"openmp-device", i32 50}