[OpenMP] Add optimization to remove the RPC client

[jhuber6]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Joseph Huber (jhuber6)

Changes

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.

---

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

3 Files Affected:

• (modified) llvm/lib/Transforms/IPO/OpenMPOpt.cpp (+37)
• (added) llvm/test/Transforms/OpenMP/keep_rpc_client.ll (+34)
• (added) llvm/test/Transforms/OpenMP/remove_rpc_client.ll (+29)

diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index e9dbcd334321f34..103a782dba2c54f 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -978,6 +978,9 @@ struct OpenMPOpt {
       }
     }
 
+    if (OMPInfoCache.OpenMPPostLink)
+      Changed |= removeRuntimeSymbols();
+
     return Changed;
   }
 
@@ -1507,6 +1510,40 @@ struct OpenMPOpt {
     return Changed;
   }
 
+  /// Tries to remove known runtime symbols that are optional.
+  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;
+
+      GlobalVariable *Client = dyn_cast<GlobalVariable>(C->stripPointerCasts());
+      if (!Client)
+        return false;
+
+      // Check to see if the only user of the RPC client is the external handle.
+      bool Unused = llvm::all_of(Client->uses(), [GV](Use &U) {
+        if (Use *Sym = U.getUser()->getSingleUndroppableUse())
+          return Sym->getUser() == GV;
+        return false;
+      });
+
+      if (!Unused)
+        return false;
+
+      // Making the symbol private prevents it from being exported.
+      GV->setLinkage(GlobalValue::LinkageTypes::PrivateLinkage);
+    }
+    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
diff --git a/llvm/test/Transforms/OpenMP/keep_rpc_client.ll b/llvm/test/Transforms/OpenMP/keep_rpc_client.ll
new file mode 100644
index 000000000000000..8ee7bddd4dc0bd9
--- /dev/null
+++ b/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}
diff --git a/llvm/test/Transforms/OpenMP/remove_rpc_client.ll b/llvm/test/Transforms/OpenMP/remove_rpc_client.ll
new file mode 100644
index 000000000000000..5451fefce6664f1
--- /dev/null
+++ b/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: @[[CLIENT:[a-zA-Z0-9_$"\\.-]+]] = internal addrspace(1) global i32 0, align 8
+; POSTLINK: @[[__LLVM_LIBC_RPC_CLIENT:[a-zA-Z0-9_$"\\.-]+]] = private 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}

[jplehr]

LGTM

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[jdoerfert]

LG