| 1 | //===- OffloadArch.cpp - list available GPUs ------------------------------===// |
| 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 | #include "clang/Basic/Version.h" |
| 10 | #include "llvm/Support/CommandLine.h" |
| 11 | #include "llvm/Support/Path.h" |
| 12 | |
| 13 | using namespace llvm; |
| 14 | |
| 15 | static cl::opt<bool> Help("h" , cl::desc("Alias for -help" ), cl::Hidden); |
| 16 | |
| 17 | // Mark all our options with this category. |
| 18 | static cl::OptionCategory OffloadArchCategory("offload-arch options" ); |
| 19 | |
| 20 | enum VendorName { |
| 21 | all, |
| 22 | amdgpu, |
| 23 | nvptx, |
| 24 | }; |
| 25 | |
| 26 | static cl::opt<VendorName> |
| 27 | Only("only" , cl::desc("Restrict to vendor:" ), cl::cat(OffloadArchCategory), |
| 28 | cl::init(Val: all), |
| 29 | cl::values(clEnumVal(all, "Print all GPUs (default)" ), |
| 30 | clEnumVal(amdgpu, "Only print AMD GPUs" ), |
| 31 | clEnumVal(nvptx, "Only print NVIDIA GPUs" ))); |
| 32 | |
| 33 | cl::opt<bool> Verbose("verbose" , cl::desc("Enable verbose output" ), |
| 34 | cl::init(Val: false), cl::cat(OffloadArchCategory)); |
| 35 | |
| 36 | static void PrintVersion(raw_ostream &OS) { |
| 37 | OS << clang::getClangToolFullVersion(ToolName: "offload-arch" ) << '\n'; |
| 38 | } |
| 39 | |
| 40 | int printGPUsByKFD(); |
| 41 | int printGPUsByHIP(); |
| 42 | int printGPUsByCUDA(); |
| 43 | |
| 44 | static int printAMD() { |
| 45 | #ifndef _WIN32 |
| 46 | if (!printGPUsByKFD()) |
| 47 | return 0; |
| 48 | #endif |
| 49 | |
| 50 | return printGPUsByHIP(); |
| 51 | } |
| 52 | |
| 53 | static int printNVIDIA() { return printGPUsByCUDA(); } |
| 54 | |
| 55 | int main(int argc, char *argv[]) { |
| 56 | cl::HideUnrelatedOptions(Category&: OffloadArchCategory); |
| 57 | |
| 58 | cl::SetVersionPrinter(PrintVersion); |
| 59 | cl::ParseCommandLineOptions( |
| 60 | argc, argv, |
| 61 | Overview: "A tool to detect the presence of offloading devices on the system. \n\n" |
| 62 | "The tool will output each detected GPU architecture separated by a\n" |
| 63 | "newline character. If multiple GPUs of the same architecture are found\n" |
| 64 | "a string will be printed for each\n" ); |
| 65 | |
| 66 | if (Help) { |
| 67 | cl::PrintHelpMessage(); |
| 68 | return 0; |
| 69 | } |
| 70 | |
| 71 | // If this was invoked from the legacy symlinks provide the same behavior. |
| 72 | bool AMDGPUOnly = Only == VendorName::amdgpu || |
| 73 | sys::path::stem(path: argv[0]).starts_with(Prefix: "amdgpu-arch" ); |
| 74 | bool NVIDIAOnly = Only == VendorName::nvptx || |
| 75 | sys::path::stem(path: argv[0]).starts_with(Prefix: "nvptx-arch" ); |
| 76 | |
| 77 | int NVIDIAResult = 0; |
| 78 | if (!AMDGPUOnly) |
| 79 | NVIDIAResult = printNVIDIA(); |
| 80 | |
| 81 | int AMDResult = 0; |
| 82 | if (!NVIDIAOnly) |
| 83 | AMDResult = printAMD(); |
| 84 | |
| 85 | // We only failed if all cases returned an error. |
| 86 | return AMDResult && NVIDIAResult; |
| 87 | } |
| 88 | |