309 lines · cpp
1//===----------------------------------------------------------------------===//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/// \file10/// This file contains the implementation of helpers and non-template member11/// functions for the DeviceContext class.12///13//===----------------------------------------------------------------------===//14 15#include "mathtest/DeviceContext.hpp"16 17#include "mathtest/ErrorHandling.hpp"18 19#include "llvm/ADT/STLExtras.h"20#include "llvm/ADT/SetVector.h"21#include "llvm/ADT/SmallString.h"22#include "llvm/ADT/StringExtras.h"23#include "llvm/ADT/StringRef.h"24#include "llvm/ADT/Twine.h"25#include "llvm/Support/Error.h"26#include "llvm/Support/ErrorHandling.h"27#include "llvm/Support/ErrorOr.h"28#include "llvm/Support/MemoryBuffer.h"29#include "llvm/Support/Path.h"30 31#include <OffloadAPI.h>32#include <cstddef>33#include <cstdint>34#include <memory>35#include <optional>36#include <string>37#include <system_error>38#include <vector>39 40using namespace mathtest;41 42//===----------------------------------------------------------------------===//43// Helpers44//===----------------------------------------------------------------------===//45 46namespace {47 48// The static 'Wrapper' instance ensures olInit() is called once at program49// startup and olShutDown() is called once at program termination50struct OffloadInitWrapper {51 OffloadInitWrapper() { OL_CHECK(olInit()); }52 ~OffloadInitWrapper() { OL_CHECK(olShutDown()); }53};54static OffloadInitWrapper Wrapper{};55 56[[nodiscard]] std::string getDeviceName(ol_device_handle_t DeviceHandle) {57 std::size_t PropSize = 0;58 OL_CHECK(olGetDeviceInfoSize(DeviceHandle, OL_DEVICE_INFO_PRODUCT_NAME,59 &PropSize));60 61 if (PropSize == 0)62 return "";63 64 std::string PropValue(PropSize, '\0');65 OL_CHECK(olGetDeviceInfo(DeviceHandle, OL_DEVICE_INFO_PRODUCT_NAME, PropSize,66 PropValue.data()));67 PropValue.pop_back(); // Remove the null terminator68 69 return PropValue;70}71 72[[nodiscard]] ol_platform_handle_t73getDevicePlatform(ol_device_handle_t DeviceHandle) noexcept {74 ol_platform_handle_t PlatformHandle = nullptr;75 OL_CHECK(olGetDeviceInfo(DeviceHandle, OL_DEVICE_INFO_PLATFORM,76 sizeof(PlatformHandle), &PlatformHandle));77 return PlatformHandle;78}79 80[[nodiscard]] std::string getPlatformName(ol_platform_handle_t PlatformHandle) {81 std::size_t PropSize = 0;82 OL_CHECK(83 olGetPlatformInfoSize(PlatformHandle, OL_PLATFORM_INFO_NAME, &PropSize));84 85 if (PropSize == 0)86 return "";87 88 std::string PropValue(PropSize, '\0');89 OL_CHECK(olGetPlatformInfo(PlatformHandle, OL_PLATFORM_INFO_NAME, PropSize,90 PropValue.data()));91 PropValue.pop_back(); // Remove the null terminator92 93 return PropValue;94}95 96[[nodiscard]] ol_platform_backend_t97getPlatformBackend(ol_platform_handle_t PlatformHandle) noexcept {98 ol_platform_backend_t Backend = OL_PLATFORM_BACKEND_UNKNOWN;99 OL_CHECK(olGetPlatformInfo(PlatformHandle, OL_PLATFORM_INFO_BACKEND,100 sizeof(Backend), &Backend));101 return Backend;102}103 104struct Device {105 ol_device_handle_t Handle;106 std::string Name;107 std::string Platform;108 ol_platform_backend_t Backend;109};110 111const std::vector<Device> &getDevices() {112 // Thread-safe initialization of a static local variable113 static auto Devices = []() {114 std::vector<Device> TmpDevices;115 116 // Discovers all devices that are not the host117 const auto *const ResultFromIterate = olIterateDevices(118 [](ol_device_handle_t DeviceHandle, void *Data) {119 ol_platform_handle_t PlatformHandle = getDevicePlatform(DeviceHandle);120 ol_platform_backend_t Backend = getPlatformBackend(PlatformHandle);121 122 if (Backend != OL_PLATFORM_BACKEND_HOST) {123 auto Name = getDeviceName(DeviceHandle);124 auto Platform = getPlatformName(PlatformHandle);125 126 static_cast<std::vector<Device> *>(Data)->push_back(127 {DeviceHandle, Name, Platform, Backend});128 }129 130 return true;131 },132 &TmpDevices);133 134 OL_CHECK(ResultFromIterate);135 136 return TmpDevices;137 }();138 139 return Devices;140}141} // namespace142 143const llvm::SetVector<llvm::StringRef> &mathtest::getPlatforms() {144 // Thread-safe initialization of a static local variable145 static auto Platforms = []() {146 llvm::SetVector<llvm::StringRef> TmpPlatforms;147 148 for (const auto &Device : getDevices())149 TmpPlatforms.insert(Device.Platform);150 151 return TmpPlatforms;152 }();153 154 return Platforms;155}156 157void detail::allocManagedMemory(ol_device_handle_t DeviceHandle,158 std::size_t Size,159 void **AllocationOut) noexcept {160 OL_CHECK(161 olMemAlloc(DeviceHandle, OL_ALLOC_TYPE_MANAGED, Size, AllocationOut));162}163 164//===----------------------------------------------------------------------===//165// DeviceContext166//===----------------------------------------------------------------------===//167 168DeviceContext::DeviceContext(std::size_t GlobalDeviceId)169 : GlobalDeviceId(GlobalDeviceId), DeviceHandle(nullptr) {170 const auto &Devices = getDevices();171 172 if (GlobalDeviceId >= Devices.size())173 FATAL_ERROR("Invalid GlobalDeviceId: " + llvm::Twine(GlobalDeviceId) +174 ", but the number of available devices is " +175 llvm::Twine(Devices.size()));176 177 DeviceHandle = Devices[GlobalDeviceId].Handle;178}179 180DeviceContext::DeviceContext(llvm::StringRef Platform, std::size_t DeviceId)181 : DeviceHandle(nullptr) {182 const auto &Platforms = getPlatforms();183 184 if (!llvm::any_of(Platforms, [&](llvm::StringRef CurrentPlatform) {185 return CurrentPlatform.equals_insensitive(Platform);186 }))187 FATAL_ERROR("There is no platform that matches with '" +188 llvm::Twine(Platform) +189 "'. Available platforms are: " + llvm::join(Platforms, ", "));190 191 const auto &Devices = getDevices();192 193 std::optional<std::size_t> FoundGlobalDeviceId;194 std::size_t MatchCount = 0;195 196 for (std::size_t Index = 0; Index < Devices.size(); ++Index) {197 if (Platform.equals_insensitive(Devices[Index].Platform)) {198 if (MatchCount == DeviceId) {199 FoundGlobalDeviceId = Index;200 break;201 }202 MatchCount++;203 }204 }205 206 if (!FoundGlobalDeviceId)207 FATAL_ERROR("Invalid DeviceId: " + llvm::Twine(DeviceId) +208 ", but the number of available devices on '" + Platform +209 "' is " + llvm::Twine(MatchCount));210 211 GlobalDeviceId = *FoundGlobalDeviceId;212 DeviceHandle = Devices[GlobalDeviceId].Handle;213}214 215[[nodiscard]] llvm::Expected<std::shared_ptr<DeviceImage>>216DeviceContext::loadBinary(llvm::StringRef Directory,217 llvm::StringRef BinaryName) const {218 auto Backend = getDevices()[GlobalDeviceId].Backend;219 llvm::StringRef Extension;220 221 switch (Backend) {222 case OL_PLATFORM_BACKEND_AMDGPU:223 Extension = ".amdgpu.bin";224 break;225 case OL_PLATFORM_BACKEND_CUDA:226 Extension = ".nvptx64.bin";227 break;228 default:229 return llvm::createStringError(230 "Unsupported backend to infer binary extension");231 }232 233 llvm::SmallString<128> FullPath(Directory);234 llvm::sys::path::append(FullPath, llvm::Twine(BinaryName) + Extension);235 236 llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> FileOrErr =237 llvm::MemoryBuffer::getFile(FullPath);238 239 if (std::error_code ErrorCode = FileOrErr.getError())240 return llvm::createStringError(241 llvm::Twine("Failed to read device binary file '") + FullPath +242 "': " + ErrorCode.message());243 244 std::unique_ptr<llvm::MemoryBuffer> &BinaryData = *FileOrErr;245 246 ol_program_handle_t ProgramHandle = nullptr;247 const ol_result_t OlResult =248 olCreateProgram(DeviceHandle, BinaryData->getBufferStart(),249 BinaryData->getBufferSize(), &ProgramHandle);250 251 if (OlResult != OL_SUCCESS) {252 llvm::StringRef Details =253 OlResult->Details ? OlResult->Details : "No details provided";254 255 // clang-format off256 return llvm::createStringError(257 llvm::Twine(Details) +258 " (code " + llvm::Twine(OlResult->Code) + ")");259 // clang-format on260 }261 262 return std::shared_ptr<DeviceImage>(263 new DeviceImage(DeviceHandle, ProgramHandle));264}265 266[[nodiscard]] llvm::Expected<ol_symbol_handle_t>267DeviceContext::getKernelHandle(ol_program_handle_t ProgramHandle,268 llvm::StringRef KernelName) const noexcept {269 ol_symbol_handle_t Handle = nullptr;270 llvm::SmallString<32> NameBuffer(KernelName);271 272 const ol_result_t OlResult = olGetSymbol(ProgramHandle, NameBuffer.c_str(),273 OL_SYMBOL_KIND_KERNEL, &Handle);274 275 if (OlResult != OL_SUCCESS) {276 llvm::StringRef Details =277 OlResult->Details ? OlResult->Details : "No details provided";278 279 // clang-format off280 return llvm::createStringError(281 llvm::Twine(Details) +282 " (code " + llvm::Twine(OlResult->Code) + ")");283 // clang-format on284 }285 286 return Handle;287}288 289void DeviceContext::launchKernelImpl(290 ol_symbol_handle_t KernelHandle, uint32_t NumGroups, uint32_t GroupSize,291 const void *KernelArgs, std::size_t KernelArgsSize) const noexcept {292 ol_kernel_launch_size_args_t LaunchSizeArgs;293 LaunchSizeArgs.Dimensions = 1;294 LaunchSizeArgs.NumGroups = {NumGroups, 1, 1};295 LaunchSizeArgs.GroupSize = {GroupSize, 1, 1};296 LaunchSizeArgs.DynSharedMemory = 0;297 298 OL_CHECK(olLaunchKernel(nullptr, DeviceHandle, KernelHandle, KernelArgs,299 KernelArgsSize, &LaunchSizeArgs));300}301 302[[nodiscard]] llvm::StringRef DeviceContext::getName() const noexcept {303 return getDevices()[GlobalDeviceId].Name;304}305 306[[nodiscard]] llvm::StringRef DeviceContext::getPlatform() const noexcept {307 return getDevices()[GlobalDeviceId].Platform;308}309