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 "llvm/Support/DynamicLibrary.h"
15 #include "llvm/Support/Error.h"
16 #include <cstdint>
17 #include <cstdio>
18 #include <memory>
19
20 #if DYNAMIC_CUDA
21 typedef enum cudaError_enum {
22 CUDA_SUCCESS = 0,
23 CUDA_ERROR_NO_DEVICE = 100,
24 } CUresult;
25
26 typedef enum CUdevice_attribute_enum {
27 CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR = 75,
28 CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR = 76,
29 } CUdevice_attribute;
30
31 typedef uint32_t CUdevice;
32
33 CUresult (*cuInit)(unsigned int);
34 CUresult (*cuDeviceGetCount)(int *);
35 CUresult (*cuGetErrorString)(CUresult, const char **);
36 CUresult (*cuDeviceGet)(CUdevice *, int);
37 CUresult (*cuDeviceGetAttribute)(int *, CUdevice_attribute, CUdevice);
38
39 constexpr const char *DynamicCudaPath = "libcuda.so";
40
loadCUDA()41 llvm::Error loadCUDA() {
42 std::string ErrMsg;
43 auto DynlibHandle = std::make_unique<llvm::sys::DynamicLibrary>(
44 llvm::sys::DynamicLibrary::getPermanentLibrary(DynamicCudaPath, &ErrMsg));
45 if (!DynlibHandle->isValid()) {
46 return llvm::createStringError(llvm::inconvertibleErrorCode(),
47 "Failed to 'dlopen' %s", DynamicCudaPath);
48 }
49 #define DYNAMIC_INIT(SYMBOL) \
50 { \
51 void *SymbolPtr = DynlibHandle->getAddressOfSymbol(#SYMBOL); \
52 if (!SymbolPtr) \
53 return llvm::createStringError(llvm::inconvertibleErrorCode(), \
54 "Failed to 'dlsym' " #SYMBOL); \
55 SYMBOL = reinterpret_cast<decltype(SYMBOL)>(SymbolPtr); \
56 }
57 DYNAMIC_INIT(cuInit);
58 DYNAMIC_INIT(cuDeviceGetCount);
59 DYNAMIC_INIT(cuGetErrorString);
60 DYNAMIC_INIT(cuDeviceGet);
61 DYNAMIC_INIT(cuDeviceGetAttribute);
62 #undef DYNAMIC_INIT
63 return llvm::Error::success();
64 }
65 #else
66
67 #include "cuda.h"
loadCUDA()68 llvm::Error loadCUDA() { return llvm::Error::success(); }
69
70 #endif
71
handleError(CUresult Err)72 static int handleError(CUresult Err) {
73 const char *ErrStr = nullptr;
74 CUresult Result = cuGetErrorString(Err, &ErrStr);
75 if (Result != CUDA_SUCCESS)
76 return 1;
77 fprintf(stderr, "CUDA error: %s\n", ErrStr);
78 return 1;
79 }
80
main(int argc,char * argv[])81 int main(int argc, char *argv[]) {
82 // Attempt to load the NVPTX driver runtime.
83 if (llvm::Error Err = loadCUDA()) {
84 logAllUnhandledErrors(std::move(Err), llvm::errs());
85 return 1;
86 }
87
88 if (CUresult Err = cuInit(0)) {
89 if (Err == CUDA_ERROR_NO_DEVICE)
90 return 0;
91 else
92 return handleError(Err);
93 }
94
95 int Count = 0;
96 if (CUresult Err = cuDeviceGetCount(&Count))
97 return handleError(Err);
98 if (Count == 0)
99 return 0;
100 for (int DeviceId = 0; DeviceId < Count; ++DeviceId) {
101 CUdevice Device;
102 if (CUresult Err = cuDeviceGet(&Device, DeviceId))
103 return handleError(Err);
104
105 int32_t Major, Minor;
106 if (CUresult Err = cuDeviceGetAttribute(
107 &Major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, Device))
108 return handleError(Err);
109 if (CUresult Err = cuDeviceGetAttribute(
110 &Minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, Device))
111 return handleError(Err);
112
113 printf("sm_%d%d\n", Major, Minor);
114 }
115 return 0;
116 }
117