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