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 | |