119 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a tool for detecting name of CUDA gpus installed in the10// 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 llvm::sys::DynamicLibrary::getPermanentLibrary(DynamicCudaPath, &ErrMsg));50 if (!DynlibHandle->isValid()) {51 return llvm::createStringError(llvm::inconvertibleErrorCode(),52 "Failed to 'dlopen' %s", 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_INIT68 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, "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(std::move(Err), llvm::errs());85 else86 consumeError(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 else94 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("sm_%d%d\n", Major, Minor);116 }117 return 0;118}119