[DeviceRTL] Allow IsSPMDMode to be optimized out in LTO mode

[jhuber6]

Summary

A previous patch merged the static and bitcode versions of the
deviceRTL. We previously used the static library's separate compilation
to set a special flag that prevented IsSPMDMode from being put in the
used list and preventing it from being optimized out. When they were
merged we could no longer do this separate compilation that allowed
users of LTO to get more optimal code.

This patch rearranges the code. The IsSPMDMode global is now
transitively used by its inclusion in the changed __keep_alive
function. This allows us to then manually delete the __keep_alive
function from the module when building the static library via
llvm-extract. The result is that the bitcode library correctly will
maintain the needed shared state, while the static library will be able
to internalize it and optimize it out.

[jdoerfert]

LG

openmp/libomptarget/DeviceRTL/CMakeLists.txt

 
 if (LLVM_DIR)
   # Builds that use pre-installed LLVM have LLVM_DIR set.
   # A standalone or LLVM_ENABLE_RUNTIMES=openmp build takes this route
   find_program(CLANG_TOOL clang PATHS ${LLVM_TOOLS_BINARY_DIR} NO_DEFAULT_PATH)
   find_program(PACKAGER_TOOL clang-offload-packager PATHS ${LLVM_TOOLS_BINARY_DIR} NO_DEFAULT_PATH)
   find_program(LINK_TOOL llvm-link PATHS ${LLVM_TOOLS_BINARY_DIR} NO_DEFAULT_PATH)
   find_program(OPT_TOOL opt PATHS ${LLVM_TOOLS_BINARY_DIR} NO_DEFAULT_PATH)
-  if ((NOT CLANG_TOOL) OR (NOT LINK_TOOL) OR (NOT OPT_TOOL))
+  find_program(EXTRACT_TOOL llvm-extract PATHS ${LLVM_TOOLS_BINARY_DIR} NO_DEFAULT_PATH)
+  if ((NOT CLANG_TOOL) OR (NOT LINK_TOOL) OR (NOT OPT_TOOL) OR (NOT EXTRACT_TOOL))
     libomptarget_say("Not building DeviceRTL. Missing clang: ${CLANG_TOOL}, llvm-link: ${LINK_TOOL} or opt: ${OPT_TOOL}")
     return()
   else()
     libomptarget_say("Building DeviceRTL. Using clang: ${CLANG_TOOL}, llvm-link: ${LINK_TOOL} and opt: ${OPT_TOOL}")
   endif()
 elseif (LLVM_TOOL_CLANG_BUILD AND NOT CMAKE_CROSSCOMPILING AND NOT OPENMP_STANDALONE_BUILD)
   # LLVM in-tree builds may use CMake target names to discover the tools.
   # A LLVM_ENABLE_PROJECTS=openmp build takes this route
   set(CLANG_TOOL $<TARGET_FILE:clang>)
   set(PACKAGER_TOOL $<TARGET_FILE:clang-offload-packager>)
   set(LINK_TOOL $<TARGET_FILE:llvm-link>)
   set(OPT_TOOL $<TARGET_FILE:opt>)
+  set(EXTRACT_TOOL $<TARGET_FILE:llvm-extract>)
   libomptarget_say("Building DeviceRTL. Using clang from in-tree build")
 else()
   libomptarget_say("Not building DeviceRTL. No appropriate clang found")
   return()
 endif()
 
 set(devicertl_base_directory ${CMAKE_CURRENT_SOURCE_DIR})
 set(include_directory ${devicertl_base_directory}/include)
set(src_files
   ${source_directory}/Tasking.cpp
   ${source_directory}/Utils.cpp
   ${source_directory}/Workshare.cpp
 )
 
 set(clang_opt_flags -O3 -mllvm -openmp-opt-disable -DSHARED_SCRATCHPAD_SIZE=512)
 set(link_opt_flags  -O3        -openmp-opt-disable)
 set(link_export_flag -passes=internalize -internalize-public-api-file=${source_directory}/exports)
+set(link_extract_flag --func='__keep_alive' --delete)
 
 # Prepend -I to each list element
 set (LIBOMPTARGET_LLVM_INCLUDE_DIRS_DEVICERTL "${LIBOMPTARGET_LLVM_INCLUDE_DIRS}")
 list(TRANSFORM LIBOMPTARGET_LLVM_INCLUDE_DIRS_DEVICERTL PREPEND "-I")
 
 # Set flags for LLVM Bitcode compilation.
 set(bc_flags -c -emit-llvm -std=c++17 -fvisibility=hidden
               ${clang_opt_flags} --offload-device-only
function(compileDeviceRTLLibrary target_cpu target_name target_triple)
 
   add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/${bclib_name}
       COMMAND ${OPT_TOOL} ${link_opt_flags} ${CMAKE_CURRENT_BINARY_DIR}/internalized_${bclib_name}
                       -o ${CMAKE_CURRENT_BINARY_DIR}/${bclib_name}
       DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/internalized_${bclib_name}
       COMMENT "Optimizing LLVM bitcode ${bclib_name}"
   )
 
+  add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/extracted_${bclib_name}
+      COMMAND ${EXTRACT_TOOL} ${link_extract_flag} ${CMAKE_CURRENT_BINARY_DIR}/${bclib_name}
+                      -o ${CMAKE_CURRENT_BINARY_DIR}/extracted_${bclib_name}
+      DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/${bclib_name}
+      COMMENT "Extracting LLVM bitcode ${bclib_name}"
+  )
+
   # Package the bitcode in the bitcode and embed it in an ELF for the static library
   add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/packaged_${bclib_name}
       COMMAND ${PACKAGER_TOOL} -o ${CMAKE_CURRENT_BINARY_DIR}/packaged_${bclib_name}
-        "--image=file=${CMAKE_CURRENT_BINARY_DIR}/${bclib_name},triple=${target_triple},arch=${target_cpu},kind=openmp"
-      DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/${bclib_name}
+        "--image=file=${CMAKE_CURRENT_BINARY_DIR}/extracted_${bclib_name},triple=${target_triple},arch=${target_cpu},kind=openmp"
+      DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/extracted_${bclib_name}
       COMMENT "Packaging LLVM offloading binary ${bclib_name}.out"
   )
 
   set(output_name "${CMAKE_CURRENT_BINARY_DIR}/devicertl-${target_name}-${target_cpu}.o")
   add_custom_command(OUTPUT ${output_name}
     COMMAND ${CLANG_TOOL} --std=c++17 -c -nostdlib
             -Xclang -fembed-offload-object=${CMAKE_CURRENT_BINARY_DIR}/packaged_${bclib_name}
             -o ${output_name}
add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/linked_${bclib_name}
       DEPENDS llvm-link
       APPEND)
   endif()
   if("${OPT_TOOL}" STREQUAL "$<TARGET_FILE:opt>")
     add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/${bclib_name}
       DEPENDS opt
       APPEND)
   endif()
+  if("${EXTRACT_TOOL}" STREQUAL "$<TARGET_FILE:llvm-extract>")
+    add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/extracted_${bclib_name}
+      DEPENDS opt
+      APPEND)
+  endif()
   if("${PACKAGER_TOOL}" STREQUAL "$<TARGET_FILE:clang-offload-packager>")
     add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/packaged_${bclib_name}
       DEPENDS clang-offload-packager
       APPEND)
   endif()
   if("${CLANG_TOOL}" STREQUAL "$<TARGET_FILE:clang>")
     add_custom_command(OUTPUT ${output_name}
       DEPENDS clang

openmp/libomptarget/DeviceRTL/include/Types.h

 #define THREAD_LOCAL(NAME)                                                     \
   [[clang::address_space(5)]] NAME [[clang::loader_uninitialized]]
 
 // TODO: clang should use address space 4 for omp_const_mem_alloc, maybe it
 //       does?
 #define CONSTANT(NAME)                                                         \
   [[clang::address_space(4)]] NAME [[clang::loader_uninitialized]]
 
-// Attribute to keep alive certain definition for the bitcode library.
-#ifdef LIBOMPTARGET_BC_TARGET
-#define KEEP_ALIVE __attribute__((used, retain))
-#else
-#define KEEP_ALIVE
-#endif
-
 ///}
 
 #endif

openmp/libomptarget/DeviceRTL/src/Mapping.cpp

 ///}
 
 /// Execution mode
 ///
 ///{
 
 // TODO: This is a workaround for initialization coming from kernels outside of
 //       the TU. We will need to solve this more correctly in the future.
-int __attribute__((weak)) KEEP_ALIVE SHARED(IsSPMDMode);
+int __attribute__((weak)) SHARED(IsSPMDMode);
 
 void mapping::init(bool IsSPMD) {
   if (mapping::isInitialThreadInLevel0(IsSPMD))
     IsSPMDMode = IsSPMD;
 }
 
 bool mapping::isSPMDMode() { return IsSPMDMode; }
 

openmp/libomptarget/DeviceRTL/src/Utils.cpp

 #include "Debug.h"
 #include "Interface.h"
 #include "Mapping.h"
 
 #pragma omp begin declare target device_type(nohost)
 
 using namespace _OMP;
 
-namespace _OMP {
+extern "C" __attribute__((weak)) int IsSPMDMode;
+
 /// Helper to keep code alive without introducing a performance penalty.
-__attribute__((weak, optnone, cold)) KEEP_ALIVE void keepAlive() {
+extern "C" __attribute__((weak, optnone, cold, used, retain)) void
+__keep_alive() {
   __kmpc_get_hardware_thread_id_in_block();
   __kmpc_get_hardware_num_threads_in_block();
   __kmpc_get_warp_size();
-  __kmpc_barrier_simple_spmd(nullptr, 0);
-  __kmpc_barrier_simple_generic(nullptr, 0);
+  __kmpc_barrier_simple_spmd(nullptr, IsSPMDMode);
+  __kmpc_barrier_simple_generic(nullptr, IsSPMDMode);
 }
-} // namespace _OMP
 
 namespace impl {
 
 bool isSharedMemPtr(const void *Ptr) { return false; }
 void Unpack(uint64_t Val, uint32_t *LowBits, uint32_t *HighBits);
 uint64_t Pack(uint32_t LowBits, uint32_t HighBits);
 
 /// AMDGCN Implementation

openmp/libomptarget/DeviceRTL/src/exports

 omp_*
 *llvm_*
 __kmpc_*
 
+__keep_alive
+IsSPMDMode
+
 memcmp
 printf
 __assert_fail