[libc-commits] [libc] feedb4a - [libc] Add code for detecting NVIDIA GPUs as well

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2023-01-27T14:06:01-06:00
New Revision: feedb4ad7d305c5864127e2b27ae650e8f04fc58

URL: https://github.com/llvm/llvm-project/commit/feedb4ad7d305c5864127e2b27ae650e8f04fc58
DIFF: https://github.com/llvm/llvm-project/commit/feedb4ad7d305c5864127e2b27ae650e8f04fc58.diff

LOG: [libc] Add code for detecting NVIDIA GPUs as well

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.

Reviewed By: sivachandra, lntue

Differential Revision: https://reviews.llvm.org/D142776

Added: 
    

Modified: 
    libc/cmake/modules/prepare_libc_gpu_build.cmake

Removed: 
    


################################################################################
diff  --git a/libc/cmake/modules/prepare_libc_gpu_build.cmake b/libc/cmake/modules/prepare_libc_gpu_build.cmake
index fa65b9a99826..8dfa0fa5d4c6 100644
--- a/libc/cmake/modules/prepare_libc_gpu_build.cmake
+++ b/libc/cmake/modules/prepare_libc_gpu_build.cmake
@@ -29,7 +29,7 @@ if(NOT LLVM_LIBC_FULL_BUILD)
                       "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/)
@@ -46,4 +46,33 @@ if(LIBC_AMDGPU_ARCH)
     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()