78 lines · cpp
1//===- AMDGPUArchByKFD.cpp - list AMDGPU installed ------*- 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 AMD GPUs installed in10// system using the Linux sysfs interface for the AMD KFD driver. This file does11// not respect ROCR_VISIBLE_DEVICES like the ROCm environment would.12//13//===----------------------------------------------------------------------===//14 15#include "llvm/Support/FileSystem.h"16#include "llvm/Support/LineIterator.h"17#include "llvm/Support/MemoryBuffer.h"18#include "llvm/Support/Path.h"19#include <memory>20 21using namespace llvm;22 23constexpr static const char *KFD_SYSFS_NODE_PATH =24 "/sys/devices/virtual/kfd/kfd/topology/nodes";25 26// See the ROCm implementation for how this is handled.27// https://github.com/ROCm/ROCT-Thunk-Interface/blob/master/src/libhsakmt.h#L12628constexpr static long getMajor(long Ver) { return (Ver / 10000) % 100; }29constexpr static long getMinor(long Ver) { return (Ver / 100) % 100; }30constexpr static long getStep(long Ver) { return Ver % 100; }31 32int printGPUsByKFD() {33 SmallVector<std::pair<long, long>> Devices;34 std::error_code EC;35 for (sys::fs::directory_iterator Begin(KFD_SYSFS_NODE_PATH, EC), End;36 Begin != End; Begin.increment(EC)) {37 if (EC)38 return 1;39 40 long Node = 0;41 if (sys::path::stem(Begin->path()).consumeInteger(10, Node))42 return 1;43 44 SmallString<0> Path(Begin->path());45 sys::path::append(Path, "properties");46 47 ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =48 MemoryBuffer::getFileOrSTDIN(Path);49 if (std::error_code EC = BufferOrErr.getError())50 return 1;51 52 long GFXVersion = 0;53 for (line_iterator Lines(**BufferOrErr, false); !Lines.is_at_end();54 ++Lines) {55 StringRef Line(*Lines);56 if (Line.consume_front("gfx_target_version")) {57 if (Line.drop_while([](char C) { return std::isspace(C); })58 .consumeInteger(10, GFXVersion))59 return 1;60 break;61 }62 }63 64 // If this is zero the node is a CPU.65 if (GFXVersion == 0)66 continue;67 Devices.emplace_back(Node, GFXVersion);68 }69 70 // Sort the devices by their node to make sure it prints in order.71 llvm::sort(Devices, [](auto &L, auto &R) { return L.first < R.first; });72 for (const auto &[Node, GFXVersion] : Devices)73 std::fprintf(stdout, "gfx%ld%ld%lx\n", getMajor(GFXVersion),74 getMinor(GFXVersion), getStep(GFXVersion));75 76 return 0;77}78