| 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 | // Defined in lib/OffloadArchOptions.cpp. |
| 16 | extern cl::OptionCategory OffloadArchCategory; |
| 17 | |
| 18 | extern cl::OptionCategory AMDGPUArchByHIPCategory; |
| 19 | |
| 20 | enum VendorName { |
| 21 | all, |
| 22 | amdgpu, |
| 23 | nvptx, |
| 24 | intel, |
| 25 | }; |
| 26 | |
| 27 | static cl::opt<VendorName> |
| 28 | Only("only" , cl::desc("Restrict to vendor:" ), cl::cat(OffloadArchCategory), |
| 29 | cl::init(Val: all), |
| 30 | cl::values(clEnumVal(all, "Print all GPUs (default)" ), |
| 31 | clEnumVal(amdgpu, "Only print AMD GPUs" ), |
| 32 | clEnumVal(nvptx, "Only print NVIDIA GPUs" ), |
| 33 | clEnumVal(intel, "Only print Intel GPUs" ))); |
| 34 | |
| 35 | static void PrintVersion(raw_ostream &OS) { |
| 36 | OS << clang::getClangToolFullVersion(ToolName: "offload-arch" ) << '\n'; |
| 37 | } |
| 38 | |
| 39 | int printGPUsByKFD(); |
| 40 | int printGPUsByHIP(); |
| 41 | int printGPUsByCUDA(); |
| 42 | int printGPUsByLevelZero(); |
| 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 | static int printIntel() { return printGPUsByLevelZero(); } |
| 55 | |
| 56 | const std::array<std::pair<VendorName, function_ref<int()>>, 3> VendorTable{ |
| 57 | ._M_elems: {{VendorName::amdgpu, printAMD}, |
| 58 | {VendorName::nvptx, printNVIDIA}, |
| 59 | {VendorName::intel, printIntel}}}; |
| 60 | |
| 61 | int main(int argc, char *argv[]) { |
| 62 | cl::HideUnrelatedOptions(Categories: {&OffloadArchCategory, &AMDGPUArchByHIPCategory}); |
| 63 | |
| 64 | cl::SetVersionPrinter(PrintVersion); |
| 65 | cl::ParseCommandLineOptions( |
| 66 | argc, argv, |
| 67 | Overview: "A tool to detect the presence of offloading devices on the system. \n\n" |
| 68 | "The tool will output each detected GPU architecture separated by a\n" |
| 69 | "newline character. If multiple GPUs of the same architecture are found\n" |
| 70 | "a string will be printed for each\n" ); |
| 71 | |
| 72 | // Support legacy binaries. |
| 73 | if (sys::path::stem(path: argv[0]).starts_with(Prefix: "amdgpu-arch" )) |
| 74 | Only = VendorName::amdgpu; |
| 75 | if (sys::path::stem(path: argv[0]).starts_with(Prefix: "nvptx-arch" )) |
| 76 | Only = VendorName::nvptx; |
| 77 | |
| 78 | int Result = 1; |
| 79 | for (auto [Name, Func] : VendorTable) { |
| 80 | if (Only == VendorName::all || Only == Name) |
| 81 | Result &= Func(); |
| 82 | } |
| 83 | |
| 84 | return Result; |
| 85 | } |
| 86 | |