1 files changed, 137 insertions, 0 deletions
diff --git a/contrib/llvm-project/clang/tools/nvptx-arch/NVPTXArch.cpp b/contrib/llvm-project/clang/tools/nvptx-arch/NVPTXArch.cpp
new file mode 100644
index 000000000000..71a48657576e
--- /dev/null
+++ b/contrib/llvm-project/clang/tools/nvptx-arch/NVPTXArch.cpp
@@ -0,0 +1,137 @@
+//===- NVPTXArch.cpp - list installed NVPTX devies ------*- C++ -*---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a tool for detecting name of CUDA gpus installed in the
+// system.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Basic/Version.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/DynamicLibrary.h"
+#include "llvm/Support/Error.h"
+#include <cstdint>
+#include <cstdio>
+#include <memory>
+
+using namespace llvm;
+
+static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
+
+static void PrintVersion(raw_ostream &OS) {
+  OS << clang::getClangToolFullVersion("nvptx-arch") << '\n';
+}
+// Mark all our options with this category, everything else (except for -version
+// and -help) will be hidden.
+static cl::OptionCategory NVPTXArchCategory("nvptx-arch options");
+
+typedef enum cudaError_enum {
+  CUDA_SUCCESS = 0,
+  CUDA_ERROR_NO_DEVICE = 100,
+} CUresult;
+
+typedef enum CUdevice_attribute_enum {
+  CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR = 75,
+  CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR = 76,
+} CUdevice_attribute;
+
+typedef uint32_t CUdevice;
+
+CUresult (*cuInit)(unsigned int);
+CUresult (*cuDeviceGetCount)(int *);
+CUresult (*cuGetErrorString)(CUresult, const char **);
+CUresult (*cuDeviceGet)(CUdevice *, int);
+CUresult (*cuDeviceGetAttribute)(int *, CUdevice_attribute, CUdevice);
+
+constexpr const char *DynamicCudaPath = "libcuda.so.1";
+
+llvm::Error loadCUDA() {
+  std::string ErrMsg;
+  auto DynlibHandle = std::make_unique<llvm::sys::DynamicLibrary>(
+      llvm::sys::DynamicLibrary::getPermanentLibrary(DynamicCudaPath, &ErrMsg));
+  if (!DynlibHandle->isValid()) {
+    return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                   "Failed to 'dlopen' %s", DynamicCudaPath);
+  }
+#define DYNAMIC_INIT(SYMBOL)                                                   \
+  {                                                                            \
+    void *SymbolPtr = DynlibHandle->getAddressOfSymbol(#SYMBOL);               \
+    if (!SymbolPtr)                                                            \
+      return llvm::createStringError(llvm::inconvertibleErrorCode(),           \
+                                     "Failed to 'dlsym' " #SYMBOL);            \
+    SYMBOL = reinterpret_cast<decltype(SYMBOL)>(SymbolPtr);                    \
+  }
+  DYNAMIC_INIT(cuInit);
+  DYNAMIC_INIT(cuDeviceGetCount);
+  DYNAMIC_INIT(cuGetErrorString);
+  DYNAMIC_INIT(cuDeviceGet);
+  DYNAMIC_INIT(cuDeviceGetAttribute);
+#undef DYNAMIC_INIT
+  return llvm::Error::success();
+}
+
+static int handleError(CUresult Err) {
+  const char *ErrStr = nullptr;
+  CUresult Result = cuGetErrorString(Err, &ErrStr);
+  if (Result != CUDA_SUCCESS)
+    return 1;
+  fprintf(stderr, "CUDA error: %s\n", ErrStr);
+  return 1;
+}
+
+int main(int argc, char *argv[]) {
+  cl::HideUnrelatedOptions(NVPTXArchCategory);
+
+  cl::SetVersionPrinter(PrintVersion);
+  cl::ParseCommandLineOptions(
+      argc, argv,
+      "A tool to detect the presence of NVIDIA devices on the system. \n\n"
+      "The tool will output each detected GPU architecture separated by a\n"
+      "newline character. If multiple GPUs of the same architecture are found\n"
+      "a string will be printed for each\n");
+
+  if (Help) {
+    cl::PrintHelpMessage();
+    return 0;
+  }
+
+  // Attempt to load the NVPTX driver runtime.
+  if (llvm::Error Err = loadCUDA()) {
+    logAllUnhandledErrors(std::move(Err), llvm::errs());
+    return 1;
+  }
+
+  if (CUresult Err = cuInit(0)) {
+    if (Err == CUDA_ERROR_NO_DEVICE)
+      return 0;
+    else
+      return handleError(Err);
+  }
+
+  int Count = 0;
+  if (CUresult Err = cuDeviceGetCount(&Count))
+    return handleError(Err);
+  if (Count == 0)
+    return 0;
+  for (int DeviceId = 0; DeviceId < Count; ++DeviceId) {
+    CUdevice Device;
+    if (CUresult Err = cuDeviceGet(&Device, DeviceId))
+      return handleError(Err);
+
+    int32_t Major, Minor;
+    if (CUresult Err = cuDeviceGetAttribute(
+            &Major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, Device))
+      return handleError(Err);
+    if (CUresult Err = cuDeviceGetAttribute(
+            &Minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, Device))
+      return handleError(Err);
+
+    printf("sm_%d%d\n", Major, Minor);
+  }
+  return 0;
+}