[libc] Add code for detecting NVIDIA GPUs as well

[jhuber6]

Summary

Recently the nvptx-arch tool was added to address the lack of a
similar tool for detecting locally installed NVIDIA GPUs. This patch
adds similar functionality for the already existing amdgpu-arch tool
in libc. These will be used to run tests on the user's system in the
future. On a system with both GPUs installed we will just go with the
first GPU detected. In the future we may want to configure the tests to
run on both of them.

[sivachandra]

This comment has been deleted.

libc/cmake/modules/prepare_libc_gpu_build.cmake

53

How can one pick between AMD or Nvidia GPUs?

[jhuber6]

libc/cmake/modules/prepare_libc_gpu_build.cmake

53

I wasn't sure if it warranted a special case right now. Since this is a contingency for a system with both an Nvidia and AMD GPU installed, which is uncommon but not unheard of. Right now I wasn't sure how much effort it would be to support testing multiple builds at the same time, so I figured it would be easier to just stick with one for now.

I could maybe make an override flag. Maybe that would also be useful for ensuring the tests get built even without a GPU installed (as would be the case if someone build LLVM on a cluster's home node and then switched to a compute node).

[sivachandra]

libc/cmake/modules/prepare_libc_gpu_build.cmake

53

Maybe that would also be useful for ensuring the tests get built even without a GPU installed (as would be the case if someone build LLVM on a cluster's home node and then switched to a compute node).

You don't need to address this in this patch but I had this exact use case in mind.

[jhuber6]

libc/cmake/modules/prepare_libc_gpu_build.cmake

53

Yes, I think it would be best to focus on getting the tests up and running first. Right now my plans are

• Implement a new 'loader' like https://reviews.llvm.org/D139839. Right now I'll probably implement it via the OpenMP runtime to make it more common
• Add a very simple integration test that ensures we can call a main function on the GPU.
• Start working on an RPC mechanism to enable printf and exit like features and then use those to implement the unit test suite.

Once we have that we can probably start thinking about wrapping the existing unit tests into integration tests as well. Do you have any suggestions for how that support should look? I see there have been a few changes in libc since I last looked so I want to make sure I'm up-to-date.

Basically the next patches would be like

libc/startup/gpu/amdgpu/start.cpp
libc/startup/gpu/nvptx/start.cpp

Then

libc/utils/gpu/Loader/openmp/Loader.cpp

and finally

libc/utils/gpu/RPC/Server.h

libc/cmake/modules/prepare_libc_gpu_build.cmake

message(FATAL_ERROR "Cannot build libc for GPU. CMake compiler "
                       "'${CMAKE_CXX_COMPILER_ID} ${CMAKE_CXX_COMPILER_VERSION}' "
                       " is not `Clang ${req_ver}.")
 endif()
 if(NOT LLVM_LIBC_FULL_BUILD)
   message(FATAL_ERROR "LLVM_LIBC_FULL_BUILD must be enabled to build libc for "
                       "GPU.")
 endif()
 
-# Identify any locally installed GPUs to use for testing.
+# Identify any locally installed AMD GPUs on the system to use for testing.
 find_program(LIBC_AMDGPU_ARCH
              NAMES amdgpu-arch
              PATHS ${LLVM_BINARY_DIR}/bin /opt/rocm/llvm/bin/)
 if(LIBC_AMDGPU_ARCH)
   execute_process(COMMAND ${LIBC_AMDGPU_ARCH}
                   OUTPUT_VARIABLE LIBC_AMDGPU_ARCH_OUTPUT
                   OUTPUT_STRIP_TRAILING_WHITESPACE)
   string(FIND "${LIBC_AMDGPU_ARCH_OUTPUT}" "\n" first_arch_string)
   string(SUBSTRING "${LIBC_AMDGPU_ARCH_OUTPUT}" 0 ${first_arch_string}
          arch_string)
   if(arch_string)
     set(LIBC_GPU_TARGET_ARCHITECTURE_IS_AMDGPU TRUE)
     set(LIBC_GPU_TARGET_TRIPLE "amdgcn-amd-amdhsa")
     set(LIBC_GPU_TARGET_ARCHITECTURE "${arch_string}")
   endif()
 endif()
-# TODO: Check for Nvidia GPUs.
+
+if(LIBC_GPU_TARGET_ARCHITECTURE_IS_AMDGPU)
+  message(STATUS "Found an installed AMD GPU on the system with target "
+                 "architecture ${LIBC_GPU_TARGET_ARCHITECTURE} ")
+  return()
+endif()
+
+# Identify any locally installed NVIDIA GPUs on the system to use for testing.
+find_program(LIBC_NVPTX_ARCH
+             NAMES nvptx-arch
+             PATHS ${LLVM_BINARY_DIR}/bin)
+if(LIBC_NVPTX_ARCH)
+  execute_process(COMMAND ${LIBC_NVPTX_ARCH}
+                  OUTPUT_VARIABLE LIBC_NVPTX_ARCH_OUTPUT
+                  OUTPUT_STRIP_TRAILING_WHITESPACE)
+  string(FIND "${LIBC_NVPTX_ARCH_OUTPUT}" "\n" first_arch_string)
+  string(SUBSTRING "${LIBC_NVPTX_ARCH_OUTPUT}" 0 ${first_arch_string}
+         arch_string)
+  if(arch_string)
+    set(LIBC_GPU_TARGET_ARCHITECTURE_IS_NVPTX TRUE)
+    set(LIBC_GPU_TARGET_TRIPLE "nvptx64-nvidia-cuda")
+    set(LIBC_GPU_TARGET_ARCHITECTURE "${arch_string}")
+  endif()
+endif()
+
+if(LIBC_GPU_TARGET_ARCHITECTURE_IS_NVPTX)
+  message(STATUS "Found an installed NVIDIA GPU on the system with target "
+                 "architecture ${LIBC_GPU_TARGET_ARCHITECTURE} ")
+  return()
+endif()