[libc] Install a single LLVM-IR version of the GPU library

[jhuber6]

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.

[llvmbot]

@llvm/pr-subscribers-libc

Author: Joseph Huber (jhuber6)

Changes

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.

---

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

3 Files Affected:

• (modified) libc/cmake/modules/LLVMLibCLibraryRules.cmake (+40)
• (modified) libc/cmake/modules/prepare_libc_gpu_build.cmake (+8)
• (modified) libc/lib/CMakeLists.txt (+27-16)

diff --git a/libc/cmake/modules/LLVMLibCLibraryRules.cmake b/libc/cmake/modules/LLVMLibCLibraryRules.cmake
index f15ffd5f9c2187..9a8082b447a589 100644
--- a/libc/cmake/modules/LLVMLibCLibraryRules.cmake
+++ b/libc/cmake/modules/LLVMLibCLibraryRules.cmake
@@ -170,6 +170,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)
+  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}
+    COMMENT "Linking LLVM-IR bitcode for ${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 752182f67cc019..a37c86be55c9d0 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 615f4270646fb5..d8478eda3c3513 100644
--- a/libc/lib/CMakeLists.txt
+++ b/libc/lib/CMakeLists.txt
@@ -37,25 +37,41 @@ foreach(archive IN ZIP_LISTS
   endif()
   list(APPEND added_archive_targets ${archive_1})
 
-  # Add the offloading version of the library for offloading languages. These
-  # are installed in the standard search path separate from the other libraries.
+  # The GPU build additionally exports the libraries as both a fat binary
+  # archive and an LLVM-IR bitcode blob.
+  # FIXME: These don't properly re-run if the source gets modified. We willy
+  #        likely need source file dependencies for them.
   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}.__fatbin__
       DEPENDS
         ${${archive_2}}
     )
     set_target_properties(
-      ${libc_gpu_archive_target}
+      ${archive_1}.__fatbin__
       PROPERTIES
-        ARCHIVE_OUTPUT_NAME ${libc_gpu_archive_name}
+        ARCHIVE_OUTPUT_NAME ${archive_0}gpu-${LIBC_TARGET_ARCHITECTURE}
+        ARCHIVE_OUTPUT_DIRECTORY ${LLVM_LIBRARY_OUTPUT_INTDIR}
+    )
+
+    add_bitcode_entrypoint_library(
+      ${archive_1}.__bitcode__
+      DEPENDS
+        ${${archive_2}}
+    )
+    add_dependencies(${archive_1} ${archive_1}.__bitcode__)
+
+    install(
+      TARGETS ${added_gpu_archive_targets}
+      ARCHIVE DESTINATION lib${LLVM_LIBDIR_SUFFIX}
+      COMPONENT libc
+    )
+
+    install(FILES $<TARGET_PROPERTY:${archive_1}.__bitcode__,TARGET_OBJECT>
+            DESTINATION ${LIBC_INSTALL_LIBRARY_DIR}
+            RENAME ${archive_1}.bc
+            COMPONENT libc
     )
-    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})
   endif()
 endforeach()
 
@@ -66,11 +82,6 @@ install(
 )
 
 if(LIBC_TARGET_OS_IS_GPU)
-  install(
-    TARGETS ${added_gpu_archive_targets}
-    ARCHIVE DESTINATION lib${LLVM_LIBDIR_SUFFIX}
-    COMPONENT libc
-  )
 endif()
 
 if(NOT LIBC_TARGET_OS_IS_BAREMETAL)

[jhuber6]

I fixed the dependency issue. These tasks are always a pain to automate with CMake, but I found the easiest way to do it is just to make these tools depend on the main libc.a so that they will rebuild whenver that changes using the more correct CMake dependency handling.

[jplehr]

LG from my side