138 lines · plain
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 definition of the DeviceContext class, which serves11/// as the high-level interface to a particular device (GPU).12///13/// This class provides methods for allocating buffers, loading binaries, and14/// getting and launching kernels on the device.15///16//===----------------------------------------------------------------------===//17 18#ifndef MATHTEST_DEVICECONTEXT_HPP19#define MATHTEST_DEVICECONTEXT_HPP20 21#include "mathtest/DeviceResources.hpp"22#include "mathtest/ErrorHandling.hpp"23#include "mathtest/Support.hpp"24 25#include "llvm/ADT/SetVector.h"26#include "llvm/ADT/StringRef.h"27#include "llvm/Support/Error.h"28 29#include <cassert>30#include <cstddef>31#include <cstdint>32#include <memory>33#include <tuple>34#include <type_traits>35#include <utility>36 37namespace mathtest {38 39const llvm::SetVector<llvm::StringRef> &getPlatforms();40 41namespace detail {42 43void allocManagedMemory(ol_device_handle_t DeviceHandle, std::size_t Size,44 void **AllocationOut) noexcept;45} // namespace detail46 47class DeviceContext {48 // For simplicity, the current design of this class doesn't have support for49 // asynchronous operations and all types of memory allocation.50 //51 // Other use cases could benefit from operations like enqueued kernel launch52 // and enqueued memcpy, as well as device and host memory allocation.53 54public:55 explicit DeviceContext(std::size_t GlobalDeviceId = 0);56 57 explicit DeviceContext(llvm::StringRef Platform, std::size_t DeviceId = 0);58 59 template <typename T>60 ManagedBuffer<T> createManagedBuffer(std::size_t Size) const noexcept {61 void *UntypedAddress = nullptr;62 63 detail::allocManagedMemory(DeviceHandle, Size * sizeof(T), &UntypedAddress);64 T *TypedAddress = static_cast<T *>(UntypedAddress);65 66 return ManagedBuffer<T>(TypedAddress, Size);67 }68 69 [[nodiscard]] llvm::Expected<std::shared_ptr<DeviceImage>>70 loadBinary(llvm::StringRef Directory, llvm::StringRef BinaryName) const;71 72 template <typename KernelSignature>73 [[nodiscard]] llvm::Expected<DeviceKernel<KernelSignature>>74 getKernel(const std::shared_ptr<DeviceImage> &Image,75 llvm::StringRef KernelName) const {76 assert(Image && "Image provided to getKernel is null");77 78 if (Image->DeviceHandle != DeviceHandle)79 return llvm::createStringError(80 "Image provided to getKernel was created for a different device");81 82 auto ExpectedHandle = getKernelHandle(Image->Handle, KernelName);83 84 if (!ExpectedHandle)85 return ExpectedHandle.takeError();86 87 return DeviceKernel<KernelSignature>(Image, *ExpectedHandle);88 }89 90 template <typename KernelSignature, typename... ArgTypes>91 void launchKernel(DeviceKernel<KernelSignature> Kernel, uint32_t NumGroups,92 uint32_t GroupSize, ArgTypes &&...Args) const noexcept {93 using ExpectedTypes =94 typename FunctionTypeTraits<KernelSignature>::ArgTypesTuple;95 using ProvidedTypes = std::tuple<std::decay_t<ArgTypes>...>;96 97 static_assert(std::is_same_v<ExpectedTypes, ProvidedTypes>,98 "Argument types provided to launchKernel do not match the "99 "kernel's signature");100 101 if (Kernel.Image->DeviceHandle != DeviceHandle)102 FATAL_ERROR("Kernel provided to launchKernel was created for a different "103 "device");104 105 if constexpr (sizeof...(Args) == 0) {106 launchKernelImpl(Kernel.Handle, NumGroups, GroupSize, nullptr, 0);107 } else {108 auto KernelArgs = makeKernelArgsPack(std::forward<ArgTypes>(Args)...);109 110 static_assert(111 (std::is_trivially_copyable_v<std::decay_t<ArgTypes>> && ...),112 "Argument types provided to launchKernel must be trivially copyable");113 114 launchKernelImpl(Kernel.Handle, NumGroups, GroupSize, &KernelArgs,115 sizeof(KernelArgs));116 }117 }118 119 [[nodiscard]] llvm::StringRef getName() const noexcept;120 121 [[nodiscard]] llvm::StringRef getPlatform() const noexcept;122 123private:124 [[nodiscard]] llvm::Expected<ol_symbol_handle_t>125 getKernelHandle(ol_program_handle_t ProgramHandle,126 llvm::StringRef KernelName) const noexcept;127 128 void launchKernelImpl(ol_symbol_handle_t KernelHandle, uint32_t NumGroups,129 uint32_t GroupSize, const void *KernelArgs,130 std::size_t KernelArgsSize) const noexcept;131 132 std::size_t GlobalDeviceId;133 ol_device_handle_t DeviceHandle;134};135} // namespace mathtest136 137#endif // MATHTEST_DEVICECONTEXT_HPP138