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