[libc] Install a single LLVM-IR version of the GPU library (#82791)

Summary:
Recent patches have allowed us to treat these libraries as direct
builds. This makes it easier to simply build them to a single LLVM-IR
file. This matches the way these files are presented by the ROCm and
CUDA toolchains and makes it easier to work with.

3 files changed, 80 insertions, 13 deletions

diff --git a/libc/cmake/modules/LLVMLibCLibraryRules.cmake b/libc/cmake/modules/LLVMLibCLibraryRules.cmake
index f15ffd5..9fba51f8 100644
--- a/libc/cmake/modules/LLVMLibCLibraryRules.cmake
+++ b/libc/cmake/modules/LLVMLibCLibraryRules.cmake
@@ -89,7 +89,7 @@ endfunction()
 #     add_gpu_entrypoint_library(
 #       DEPENDS <list of add_entrypoint_object targets>
 #     )
-function(add_gpu_entrypoint_library target_name)
+function(add_gpu_entrypoint_library target_name base_target_name)
   cmake_parse_arguments(
     "ENTRYPOINT_LIBRARY"
     "" # No optional arguments
@@ -127,7 +127,7 @@ function(add_gpu_entrypoint_library target_name)
       COMMAND ${LIBC_CLANG_OFFLOAD_PACKAGER}
               "${prefix},file=$<JOIN:${object},,file=>" -o 
               ${CMAKE_CURRENT_BINARY_DIR}/binary/${name}.gpubin
-      DEPENDS ${dep}
+      DEPENDS ${dep} ${base_target_name}
       COMMENT "Packaging LLVM offloading binary for '${object}'"
     )
     add_custom_target(${dep}.__gpubin__ DEPENDS ${dep}
@@ -140,7 +140,7 @@ function(add_gpu_entrypoint_library target_name)
       OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/stubs/${name}.cpp"
       COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}/stubs
       COMMAND ${CMAKE_COMMAND} -E touch ${CMAKE_CURRENT_BINARY_DIR}/stubs/${name}.cpp
-      DEPENDS ${dep} ${dep}.__gpubin__
+      DEPENDS ${dep} ${dep}.__gpubin__ ${base_target_name}
     )
     add_custom_target(${dep}.__stub__ 
                       DEPENDS ${dep}.__gpubin__ "${CMAKE_CURRENT_BINARY_DIR}/stubs/${name}.cpp")
@@ -156,7 +156,8 @@ function(add_gpu_entrypoint_library target_name)
     target_compile_options(${dep}.__fatbin__ PRIVATE 
       --target=${LLVM_HOST_TRIPLE}
       "SHELL:-Xclang -fembed-offload-object=${CMAKE_CURRENT_BINARY_DIR}/binary/${name}.gpubin")
-    add_dependencies(${dep}.__fatbin__ ${dep} ${dep}.__stub__ ${dep}.__gpubin__)
+    add_dependencies(${dep}.__fatbin__
+                     ${dep} ${dep}.__stub__ ${dep}.__gpubin__ ${base_target_name})
 
     # Set the list of newly create fat binaries containing embedded device code.
     list(APPEND objects $<TARGET_OBJECTS:${dep}.__fatbin__>)
@@ -170,6 +171,46 @@ function(add_gpu_entrypoint_library target_name)
   set_target_properties(${target_name} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${LIBC_LIBRARY_DIR})
 endfunction(add_gpu_entrypoint_library)
 
+# A rule to build a library from a collection of entrypoint objects and bundle
+# it in a single LLVM-IR bitcode file.
+# Usage:
+#     add_gpu_entrypoint_library(
+#       DEPENDS <list of add_entrypoint_object targets>
+#     )
+function(add_bitcode_entrypoint_library target_name base_target_name)
+  cmake_parse_arguments(
+    "ENTRYPOINT_LIBRARY"
+    "" # No optional arguments
+    "" # No single value arguments
+    "DEPENDS" # Multi-value arguments
+    ${ARGN}
+  )
+  if(NOT ENTRYPOINT_LIBRARY_DEPENDS)
+    message(FATAL_ERROR "'add_entrypoint_library' target requires a DEPENDS list "
+                        "of 'add_entrypoint_object' targets.")
+  endif()
+
+  get_fq_deps_list(fq_deps_list ${ENTRYPOINT_LIBRARY_DEPENDS})
+  get_all_object_file_deps(all_deps "${fq_deps_list}")
+
+  set(objects "")
+  foreach(dep IN LISTS all_deps)
+    set(object $<$<STREQUAL:$<TARGET_NAME_IF_EXISTS:${dep}>,${dep}>:$<TARGET_OBJECTS:${dep}>>)
+    list(APPEND objects ${object})
+  endforeach()
+
+  set(output ${CMAKE_CURRENT_BINARY_DIR}/${target_name}.bc)
+  add_custom_command(
+    OUTPUT ${output}
+    COMMAND ${LIBC_LLVM_LINK} ${objects} -o ${output}
+    DEPENDS ${all_deps} ${base_target_name}
+    COMMENT "Linking LLVM-IR bitcode for ${base_target_name}"
+    COMMAND_EXPAND_LISTS
+  )
+  add_custom_target(${target_name} DEPENDS ${output} ${all_deps})
+  set_target_properties(${target_name} PROPERTIES TARGET_OBJECT ${output})
+endfunction(add_bitcode_entrypoint_library)
+
 # A rule to build a library from a collection of entrypoint objects.
 # Usage:
 #     add_entrypoint_library(
diff --git a/libc/cmake/modules/prepare_libc_gpu_build.cmake b/libc/cmake/modules/prepare_libc_gpu_build.cmake
index 548990a..7e9fc74 100644
--- a/libc/cmake/modules/prepare_libc_gpu_build.cmake
+++ b/libc/cmake/modules/prepare_libc_gpu_build.cmake
@@ -26,6 +26,14 @@ if(NOT LIBC_CLANG_OFFLOAD_PACKAGER)
                       "build")
 endif()
 
+# Identify llvm-link program so we can merge the output IR into a single blob.
+find_program(LIBC_LLVM_LINK
+             NAMES llvm-link NO_DEFAULT_PATH
+             PATHS ${LLVM_BINARY_DIR}/bin ${compiler_path})
+if(NOT LIBC_LLVM_LINK)
+  message(FATAL_ERROR "Cannot find 'llvm-link' for the GPU build")
+endif()
+
 # Optionally set up a job pool to limit the number of GPU tests run in parallel.
 # This is sometimes necessary as running too many tests in parallel can cause
 # the GPU or driver to run out of resources.
diff --git a/libc/lib/CMakeLists.txt b/libc/lib/CMakeLists.txt
index 615f427..e5ebd1e1 100644
--- a/libc/lib/CMakeLists.txt
+++ b/libc/lib/CMakeLists.txt
@@ -40,22 +40,33 @@ foreach(archive IN ZIP_LISTS
   # Add the offloading version of the library for offloading languages. These
   # are installed in the standard search path separate from the other libraries.
   if(LIBC_TARGET_OS_IS_GPU)
-    set(libc_gpu_archive_target ${archive_1}gpu)
-    set(libc_gpu_archive_name ${archive_0}gpu-${LIBC_TARGET_ARCHITECTURE})
-
     add_gpu_entrypoint_library(
-      ${libc_gpu_archive_target}
+      ${archive_1}gpu
+      ${archive_1}
       DEPENDS
         ${${archive_2}}
     )
     set_target_properties(
-      ${libc_gpu_archive_target}
+      ${archive_1}gpu
       PROPERTIES
-        ARCHIVE_OUTPUT_NAME ${libc_gpu_archive_name}
+        ARCHIVE_OUTPUT_NAME ${archive_0}gpu-${LIBC_TARGET_ARCHITECTURE}
+        ARCHIVE_OUTPUT_DIRECTORY ${LLVM_LIBRARY_OUTPUT_INTDIR}
+    )
+    list(APPEND added_gpu_archive_targets ${archive_1}gpu)
+
+    add_bitcode_entrypoint_library(
+      ${archive_1}bitcode
+      ${archive_1}
+      DEPENDS
+        ${${archive_2}}
     )
-    set_target_properties(${libc_gpu_archive_target} PROPERTIES 
-                          ARCHIVE_OUTPUT_DIRECTORY ${LLVM_LIBRARY_OUTPUT_INTDIR})
-    list(APPEND added_gpu_archive_targets ${libc_gpu_archive_target})
+    set_target_properties(
+      ${archive_1}bitcode
+      PROPERTIES
+        OUTPUT_NAME ${archive_1}.bc
+    )
+    add_dependencies(${archive_1}gpu ${archive_1}bitcode)
+    list(APPEND added_gpu_bitcode_targets ${archive_1}bitcode)
   endif()
 endforeach()
 
@@ -71,6 +82,13 @@ if(LIBC_TARGET_OS_IS_GPU)
     ARCHIVE DESTINATION lib${LLVM_LIBDIR_SUFFIX}
     COMPONENT libc
   )
+  foreach(file ${added_gpu_bitcode_targets})
+    install(FILES $<TARGET_PROPERTY:${file},TARGET_OBJECT>
+            DESTINATION ${LIBC_INSTALL_LIBRARY_DIR}
+            RENAME $<TARGET_PROPERTY:${file},OUTPUT_NAME>
+            COMPONENT libc
+    )
+  endforeach()
 endif()
 
 if(NOT LIBC_TARGET_OS_IS_BAREMETAL)