[libc] Refactor _build_gpu_objects cmake function. (#80631)

2 files changed, 121 insertions, 81 deletions

diff --git a/libc/cmake/modules/LLVMLibCObjectRules.cmake b/libc/cmake/modules/LLVMLibCObjectRules.cmake
index 4213fe1..70e64a6 100644
--- a/libc/cmake/modules/LLVMLibCObjectRules.cmake
+++ b/libc/cmake/modules/LLVMLibCObjectRules.cmake
@@ -145,21 +145,78 @@ function(get_nvptx_compile_options output_var gpu_arch)
   set(${output_var} ${nvptx_options} PARENT_SCOPE)
 endfunction()
 
-# Builds the object target for the GPU.
+# Build the object target for a single GPU arch.
+# Usage:
+#     _build_gpu_object_for_single_arch(
+#       <target_name>
+#       <gpu_arch>
+#       SRCS <list of .cpp files>
+#       HDRS <list of .h files>
+#       DEPENDS <list of dependencies>
+#       COMPILE_OPTIONS <optional list of special compile options for this target>
+#       FLAGS <optional list of flags>
+#     )
+function(_build_gpu_object_for_single_arch fq_target_name gpu_arch)
+  cmake_parse_arguments(
+    "ADD_GPU_OBJ"
+    "" # No optional arguments
+    "NAME;CXX_STANDARD" # Single value arguments
+    "SRCS;HDRS;DEPENDS;COMPILE_OPTIONS;FLAGS"  # Multi value arguments
+    ${ARGN}
+  )
+
+  if(NOT ADD_GPU_OBJ_CXX_STANDARD)
+    set(ADD_GPU_OBJ_CXX_STANDARD ${CMAKE_CXX_STANDARD})
+  endif()
+
+  set(compile_options ${ADD_GPU_OBJ_COMPILE_OPTIONS})
+  # Derive the triple from the specified architecture.
+  if("${gpu_arch}" IN_LIST all_amdgpu_architectures)
+    set(gpu_target_triple ${AMDGPU_TARGET_TRIPLE})
+    list(APPEND compile_options "-mcpu=${gpu_arch}")
+    list(APPEND compile_options "SHELL:-Xclang -mcode-object-version=none")
+    list(APPEND compile_options "-emit-llvm")
+  elseif("${gpu_arch}" IN_LIST all_nvptx_architectures)
+    set(gpu_target_triple ${NVPTX_TARGET_TRIPLE})
+    get_nvptx_compile_options(nvptx_options ${gpu_arch})
+    list(APPEND compile_options "${nvptx_options}")
+  else()
+    message(FATAL_ERROR "Unknown GPU architecture '${gpu_arch}'")
+  endif()
+  list(APPEND compile_options "--target=${gpu_target_triple}")
+
+  # Build the library for this target architecture. We always emit LLVM-IR for
+  # packaged GPU binaries.
+  add_library(${fq_target_name}
+    EXCLUDE_FROM_ALL
+    OBJECT
+    ${ADD_GPU_OBJ_SRCS}
+    ${ADD_GPU_OBJ_HDRS}
+  )
+
+  target_compile_options(${fq_target_name} PRIVATE ${compile_options})
+  target_include_directories(${fq_target_name} SYSTEM PRIVATE ${LIBC_INCLUDE_DIR})
+  target_include_directories(${fq_target_name} PRIVATE ${LIBC_SOURCE_DIR})
+  set_target_properties(${fq_target_name} PROPERTIES CXX_STANDARD ${ADD_GPU_OBJ_CXX_STANDARD})
+  if(ADD_GPU_OBJ_DEPENDS)
+    add_dependencies(${fq_target_name} ${ADD_GPU_OBJ_DEPENDS})
+    set_target_properties(${fq_target_name} PROPERTIES DEPS "${ADD_GPU_OBJ_DEPENDS}")
+  endif()
+endfunction(_build_gpu_object_for_single_arch)
+
+# Build the object target for the GPU.
 # This compiles the target for all supported architectures and embeds it into
-# host binary for installing. The internal target contains the GPU code directly
-# compiled for a single architecture used internally.
+# host binary for installing.
 # Usage:
-#     _build_gpu_objects(
+#     _build_gpu_object_bundle(
 #       <target_name>
-#       <internal_target_name>
 #       SRCS <list of .cpp files>
 #       HDRS <list of .h files>
 #       DEPENDS <list of dependencies>
 #       COMPILE_OPTIONS <optional list of special compile options for this target>
 #       FLAGS <optional list of flags>
 #     )
-function(_build_gpu_objects fq_target_name internal_target_name)
+function(_build_gpu_object_bundle fq_target_name)
   cmake_parse_arguments(
     "ADD_GPU_OBJ"
     "" # No optional arguments
@@ -168,7 +225,6 @@ function(_build_gpu_objects fq_target_name internal_target_name)
     ${ARGN}
   )
 
-  set(common_compile_options ${ADD_GPU_OBJ_COMPILE_OPTIONS})
   if(NOT ADD_GPU_OBJ_CXX_STANDARD)
     set(ADD_GPU_OBJ_CXX_STANDARD ${CMAKE_CXX_STANDARD})
   endif()
@@ -179,49 +235,28 @@ function(_build_gpu_objects fq_target_name internal_target_name)
     foreach(gpu_arch ${LIBC_GPU_ARCHITECTURES})
       get_filename_component(src_name ${add_gpu_obj_src} NAME)
       set(gpu_target_name ${fq_target_name}.${src_name}.${gpu_arch})
-      set(compile_options ${ADD_GPU_OBJ_COMPILE_OPTIONS})
-      # Derive the triple from the specified architecture.
-      if("${gpu_arch}" IN_LIST all_amdgpu_architectures)
-        set(gpu_target_triple "amdgcn-amd-amdhsa")
-        list(APPEND compile_options "-mcpu=${gpu_arch}")
-        list(APPEND compile_options "SHELL:-Xclang -mcode-object-version=none")
-      elseif("${gpu_arch}" IN_LIST all_nvptx_architectures)
-        set(gpu_target_triple "nvptx64-nvidia-cuda")
-        get_nvptx_compile_options(nvptx_options ${gpu_arch})
-        list(APPEND compile_options "${nvptx_options}")
-      else()
-        message(FATAL_ERROR "Unknown GPU architecture '${gpu_arch}'")
-      endif()
-      list(APPEND compile_options "--target=${gpu_target_triple}")
-      list(APPEND compile_options "-emit-llvm")
-
-      # Build the library for this target architecture. We always emit LLVM-IR for
-      # packaged GPU binaries.
-      add_library(${gpu_target_name}
-        EXCLUDE_FROM_ALL
-        OBJECT
-        ${add_gpu_obj_src}
-        ${ADD_GPU_OBJ_HDRS}
-      )
-
-      target_compile_options(${gpu_target_name} PRIVATE ${compile_options})
-      target_include_directories(${gpu_target_name} SYSTEM PRIVATE ${LIBC_INCLUDE_DIR})
-      target_include_directories(${gpu_target_name} PRIVATE ${LIBC_SOURCE_DIR})
-      target_compile_definitions(${gpu_target_name} PRIVATE LIBC_COPT_PUBLIC_PACKAGING)
-      set_target_properties(${gpu_target_name} PROPERTIES CXX_STANDARD ${ADD_GPU_OBJ_CXX_STANDARD})
-      if(ADD_GPU_OBJ_DEPENDS)
-        add_dependencies(${gpu_target_name} ${ADD_GPU_OBJ_DEPENDS})
-      endif()
 
+      _build_gpu_object_for_single_arch(
+        ${gpu_target_name}
+        ${gpu_arch}
+        CXX_STANDARD ${ADD_GPU_OBJ_CXX_STANDARD}
+        HDRS ${ADD_GPU_OBJ_HDRS}
+        SRCS ${add_gpu_obj_src}
+        COMPILE_OPTIONS
+          ${ADD_GPU_OBJ_COMPILE_OPTIONS}
+          "-emit-llvm"
+        DEPENDS ${ADD_GPU_OBJ_DEPENDS}
+      )
       # Append this target to a list of images to package into a single binary.
       set(input_file $<TARGET_OBJECTS:${gpu_target_name}>)
       if("${gpu_arch}" IN_LIST all_nvptx_architectures)
+        get_nvptx_compile_options(nvptx_options ${gpu_arch})
         string(REGEX MATCH "\\+ptx[0-9]+" nvptx_ptx_feature ${nvptx_options})
         list(APPEND packager_images
-             --image=file=${input_file},arch=${gpu_arch},triple=${gpu_target_triple},feature=${nvptx_ptx_feature})
+             --image=file=${input_file},arch=${gpu_arch},triple=${NVPTX_TARGET_TRIPLE},feature=${nvptx_ptx_feature})
       else()
         list(APPEND packager_images
-             --image=file=${input_file},arch=${gpu_arch},triple=${gpu_target_triple})
+             --image=file=${input_file},arch=${gpu_arch},triple=${AMDGPU_TARGET_TRIPLE})
        endif()
       list(APPEND gpu_target_objects ${input_file})
     endforeach()
@@ -269,7 +304,7 @@ function(_build_gpu_objects fq_target_name internal_target_name)
     ${CMAKE_CURRENT_BINARY_DIR}/stubs/${stub_filename}
   )
   target_compile_options(${fq_target_name} BEFORE PRIVATE
-                         ${common_compile_options} -nostdlib)
+                         ${ADD_GPU_OBJ_COMPILE_OPTIONS} -nostdlib)
   foreach(packaged_gpu_binary ${packaged_gpu_binaries})
     target_compile_options(${fq_target_name} PRIVATE
                            "SHELL:-Xclang -fembed-offload-object=${packaged_gpu_binary}")
@@ -278,33 +313,6 @@ function(_build_gpu_objects fq_target_name internal_target_name)
   target_include_directories(${fq_target_name} PRIVATE ${LIBC_SOURCE_DIR})
   add_dependencies(${fq_target_name}
                    ${full_deps_list} ${packaged_gpu_names} ${stub_target_name})
-
-  # We only build the internal target for a single supported architecture.
-  if(LIBC_GPU_TARGET_ARCHITECTURE_IS_AMDGPU OR
-     LIBC_GPU_TARGET_ARCHITECTURE_IS_NVPTX)
-    add_library(
-      ${internal_target_name}
-      EXCLUDE_FROM_ALL
-      OBJECT
-      ${ADD_GPU_OBJ_SRCS}
-      ${ADD_GPU_OBJ_HDRS}
-    )
-    target_compile_options(${internal_target_name} BEFORE PRIVATE
-                           ${common_compile_options} --target=${LIBC_GPU_TARGET_TRIPLE})
-    if(LIBC_GPU_TARGET_ARCHITECTURE_IS_AMDGPU)
-      target_compile_options(${internal_target_name} PRIVATE
-                             "SHELL:-Xclang -mcode-object-version=none"
-                             -mcpu=${LIBC_GPU_TARGET_ARCHITECTURE} -flto)
-    elseif(LIBC_GPU_TARGET_ARCHITECTURE_IS_NVPTX)
-      get_nvptx_compile_options(nvptx_options ${LIBC_GPU_TARGET_ARCHITECTURE})
-      target_compile_options(${internal_target_name} PRIVATE ${nvptx_options})
-    endif()
-    target_include_directories(${internal_target_name} SYSTEM PRIVATE ${LIBC_INCLUDE_DIR})
-    target_include_directories(${internal_target_name} PRIVATE ${LIBC_SOURCE_DIR})
-    if(full_deps_list)
-      add_dependencies(${internal_target_name} ${full_deps_list})
-    endif()
-  endif()
 endfunction()
 
 # Rule which is essentially a wrapper over add_library to compile a set of
@@ -354,8 +362,10 @@ function(create_object_library fq_target_name)
   # The GPU build uses a separate internal file.
   if(LIBC_TARGET_ARCHITECTURE_IS_GPU AND NOT ${ADD_OBJECT_NO_GPU_BUNDLE})
     set(internal_target_name ${fq_target_name}.__internal__)
+    set(public_packaging_for_internal "")
   else()
     set(internal_target_name ${fq_target_name})
+    set(public_packaging_for_internal "-DLIBC_COPT_PUBLIC_PACKAGING")
   endif()
 
   _get_common_compile_options(compile_options "${ADD_OBJECT_FLAGS}")
@@ -363,15 +373,27 @@ function(create_object_library fq_target_name)
 
   # GPU builds require special handling for the objects because we want to
   # export several different targets at once, e.g. for both Nvidia and AMD.
-  if(LIBC_TARGET_ARCHITECTURE_IS_GPU AND NOT ${ADD_OBJECT_NO_GPU_BUNDLE})
-    _build_gpu_objects(
-      ${fq_target_name}
+  if(LIBC_TARGET_ARCHITECTURE_IS_GPU)
+    if(NOT ${ADD_OBJECT_NO_GPU_BUNDLE})
+      _build_gpu_object_bundle(
+        ${fq_target_name}
+        SRCS ${ADD_OBJECT_SRCS}
+        HDRS ${ADD_OBJECT_HDRS}
+        CXX_STANDARD ${ADD_OBJECT_CXX_STANDARD}
+        COMPILE_OPTIONS ${compile_options} "-DLIBC_COPT_PUBLIC_PACKAGING"
+        DEPENDS ${fq_deps_list}
+      )
+    endif()
+    # When the target for GPU is not bundled, internal_target_name is the same
+    # as fq_targetname
+    _build_gpu_object_for_single_arch(
       ${internal_target_name}
+      ${LIBC_GPU_TARGET_ARCHITECTURE}
       SRCS ${ADD_OBJECT_SRCS}
       HDRS ${ADD_OBJECT_HDRS}
-      DEPENDS ${fq_deps_list}
       CXX_STANDARD ${ADD_OBJECT_CXX_STANDARD}
-      COMPILE_OPTIONS ${compile_options}
+      COMPILE_OPTIONS ${compile_options} ${public_packaging_for_internal}
+      DEPENDS ${fq_deps_list}
     )
   else()
     add_library(
@@ -567,9 +589,18 @@ function(create_entrypoint_object fq_target_name)
   # GPU builds require special handling for the objects because we want to
   # export several different targets at once, e.g. for both Nvidia and AMD.
   if(LIBC_TARGET_ARCHITECTURE_IS_GPU)
-    _build_gpu_objects(
+    _build_gpu_object_bundle(
       ${fq_target_name}
+      SRCS ${ADD_ENTRYPOINT_OBJ_SRCS}
+      HDRS ${ADD_ENTRYPOINT_OBJ_HDRS}
+      COMPILE_OPTIONS ${common_compile_options} "-DLIBC_COPT_PUBLIC_PACKAGING"
+      CXX_STANDARD ${ADD_ENTRYPOINT_OBJ_CXX_STANDARD}
+      DEPENDS ${full_deps_list}
+      FLAGS "${ADD_ENTRYPOINT_OBJ_FLAGS}"
+    )
+    _build_gpu_object_for_single_arch(
       ${internal_target_name}
+      ${LIBC_GPU_TARGET_ARCHITECTURE}
       SRCS ${ADD_ENTRYPOINT_OBJ_SRCS}
       HDRS ${ADD_ENTRYPOINT_OBJ_HDRS}
       COMPILE_OPTIONS ${common_compile_options}
diff --git a/libc/cmake/modules/prepare_libc_gpu_build.cmake b/libc/cmake/modules/prepare_libc_gpu_build.cmake
index 05c46a6..2086175 100644
--- a/libc/cmake/modules/prepare_libc_gpu_build.cmake
+++ b/libc/cmake/modules/prepare_libc_gpu_build.cmake
@@ -15,6 +15,8 @@ set(all_gpu_architectures
     "${all_amdgpu_architectures};${all_nvptx_architectures}")
 set(LIBC_GPU_ARCHITECTURES "all" CACHE STRING
     "List of GPU architectures to build the libc for.")
+set(AMDGPU_TARGET_TRIPLE "amdgcn-amd-amdhsa")
+set(NVPTX_TARGET_TRIPLE "nvptx64-nvidia-cuda")
 
 # Ensure the compiler is a valid clang when building the GPU target.
 set(req_ver "${LLVM_VERSION_MAJOR}.${LLVM_VERSION_MINOR}.${LLVM_VERSION_PATCH}")
@@ -96,18 +98,25 @@ elseif(detected_gpu_architectures)
   message(STATUS "Using GPU architecture detected on the system for testing: "
                  "'${gpu_test_architecture}'")
 else()
-  message(STATUS "No GPU architecture set for testing. GPU tests will not be "
-                 "availibe. Set 'LIBC_GPU_TEST_ARCHITECTURE' to override.")
-  return()
+  list(LENGTH LIBC_GPU_ARCHITECTURES n_gpu_archs)
+  if (${n_gpu_archs} EQUAL 1)
+    set(gpu_test_architecture ${LIBC_GPU_ARCHITECTURES})
+    message(STATUS "Using user-specified GPU architecture for testing: "
+                  "'${gpu_test_architecture}'")
+  else()
+    message(STATUS "No GPU architecture set for testing. GPU tests will not be "
+                  "availibe. Set 'LIBC_GPU_TEST_ARCHITECTURE' to override.")
+    return()
+  endif()
 endif()
 
 if("${gpu_test_architecture}" IN_LIST all_amdgpu_architectures)
   set(LIBC_GPU_TARGET_ARCHITECTURE_IS_AMDGPU TRUE)
-  set(LIBC_GPU_TARGET_TRIPLE "amdgcn-amd-amdhsa")
+  set(LIBC_GPU_TARGET_TRIPLE ${AMDGPU_TARGET_TRIPLE})
   set(LIBC_GPU_TARGET_ARCHITECTURE "${gpu_test_architecture}")
 elseif("${gpu_test_architecture}" IN_LIST all_nvptx_architectures)
   set(LIBC_GPU_TARGET_ARCHITECTURE_IS_NVPTX TRUE)
-  set(LIBC_GPU_TARGET_TRIPLE "nvptx64-nvidia-cuda")
+  set(LIBC_GPU_TARGET_TRIPLE ${NVPTX_TARGET_TRIPLE})
   set(LIBC_GPU_TARGET_ARCHITECTURE "${gpu_test_architecture}")
 else()
   message(FATAL_ERROR "Unknown GPU architecture '${gpu_test_architecture}'")