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