1//===- NVPTXArch.cpp - list installed NVPTX devies ------*- C++ -*---------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements a tool for detecting name of CUDA gpus installed in the
10// system.
11//
12//===----------------------------------------------------------------------===//
13
14#include "clang/Basic/Version.h"
15#include "llvm/Support/CommandLine.h"
16#include "llvm/Support/DynamicLibrary.h"
17#include "llvm/Support/Error.h"
18#include <cstdint>
19#include <cstdio>
20#include <memory>
21
22using namespace llvm;
23
24typedef enum cudaError_enum {
25 CUDA_SUCCESS = 0,
26 CUDA_ERROR_NO_DEVICE = 100,
27} CUresult;
28
29typedef enum CUdevice_attribute_enum {
30 CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR = 75,
31 CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR = 76,
32} CUdevice_attribute;
33
34typedef uint32_t CUdevice;
35
36CUresult (*cuInit)(unsigned int);
37CUresult (*cuDeviceGetCount)(int *);
38CUresult (*cuGetErrorString)(CUresult, const char **);
39CUresult (*cuDeviceGet)(CUdevice *, int);
40CUresult (*cuDeviceGetAttribute)(int *, CUdevice_attribute, CUdevice);
41
42constexpr const char *DynamicCudaPath = "libcuda.so.1";
43
44llvm::Error loadCUDA() {
45 std::string ErrMsg;
46 auto DynlibHandle = std::make_unique<llvm::sys::DynamicLibrary>(
47 args: llvm::sys::DynamicLibrary::getPermanentLibrary(filename: DynamicCudaPath, errMsg: &ErrMsg));
48 if (!DynlibHandle->isValid()) {
49 return llvm::createStringError(EC: llvm::inconvertibleErrorCode(),
50 Fmt: "Failed to 'dlopen' %s", Vals: DynamicCudaPath);
51 }
52#define DYNAMIC_INIT(SYMBOL) \
53 { \
54 void *SymbolPtr = DynlibHandle->getAddressOfSymbol(#SYMBOL); \
55 if (!SymbolPtr) \
56 return llvm::createStringError(llvm::inconvertibleErrorCode(), \
57 "Failed to 'dlsym' " #SYMBOL); \
58 SYMBOL = reinterpret_cast<decltype(SYMBOL)>(SymbolPtr); \
59 }
60 DYNAMIC_INIT(cuInit);
61 DYNAMIC_INIT(cuDeviceGetCount);
62 DYNAMIC_INIT(cuGetErrorString);
63 DYNAMIC_INIT(cuDeviceGet);
64 DYNAMIC_INIT(cuDeviceGetAttribute);
65#undef DYNAMIC_INIT
66 return llvm::Error::success();
67}
68
69static int handleError(CUresult Err) {
70 const char *ErrStr = nullptr;
71 CUresult Result = cuGetErrorString(Err, &ErrStr);
72 if (Result != CUDA_SUCCESS)
73 return 1;
74 fprintf(stderr, format: "CUDA error: %s\n", ErrStr);
75 return 1;
76}
77
78int printGPUsByCUDA() {
79 // Attempt to load the NVPTX driver runtime.
80 if (llvm::Error Err = loadCUDA()) {
81 logAllUnhandledErrors(E: std::move(Err), OS&: llvm::errs());
82 return 1;
83 }
84
85 if (CUresult Err = cuInit(0)) {
86 if (Err == CUDA_ERROR_NO_DEVICE)
87 return 0;
88 else
89 return handleError(Err);
90 }
91
92 int Count = 0;
93 if (CUresult Err = cuDeviceGetCount(&Count))
94 return handleError(Err);
95 if (Count == 0)
96 return 0;
97 for (int DeviceId = 0; DeviceId < Count; ++DeviceId) {
98 CUdevice Device;
99 if (CUresult Err = cuDeviceGet(&Device, DeviceId))
100 return handleError(Err);
101
102 int32_t Major, Minor;
103 if (CUresult Err = cuDeviceGetAttribute(
104 &Major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, Device))
105 return handleError(Err);
106 if (CUresult Err = cuDeviceGetAttribute(
107 &Minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, Device))
108 return handleError(Err);
109
110 printf(format: "sm_%d%d\n", Major, Minor);
111 }
112 return 0;
113}
114