1 //===------------------ Client.h - Client Implementation ------------------===// 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 // gRPC Client for the remote plugin. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #ifndef LLVM_OPENMP_LIBOMPTARGET_PLUGINS_REMOTE_SRC_CLIENT_H 14 #define LLVM_OPENMP_LIBOMPTARGET_PLUGINS_REMOTE_SRC_CLIENT_H 15 16 #include "Utils.h" 17 #include "omptarget.h" 18 #include <google/protobuf/arena.h> 19 #include <grpcpp/grpcpp.h> 20 #include <grpcpp/security/credentials.h> 21 #include <grpcpp/support/channel_arguments.h> 22 #include <memory> 23 #include <mutex> 24 #include <numeric> 25 26 using grpc::Channel; 27 using openmp::libomptarget::remote::RemoteOffload; 28 using namespace RemoteOffloading; 29 30 using namespace google; 31 32 class RemoteOffloadClient { 33 const int Timeout; 34 35 int DebugLevel; 36 uint64_t MaxSize; 37 int64_t BlockSize; 38 39 std::unique_ptr<RemoteOffload::Stub> Stub; 40 std::unique_ptr<protobuf::Arena> Arena; 41 42 std::unique_ptr<std::mutex> ArenaAllocatorLock; 43 44 std::map<int32_t, std::unordered_map<void *, void *>> RemoteEntries; 45 std::map<int32_t, std::unique_ptr<__tgt_target_table>> DevicesToTables; 46 47 template <typename Fn1, typename Fn2, typename TReturn> 48 auto remoteCall(Fn1 Preprocess, Fn2 Postprocess, TReturn ErrorValue, 49 bool Timeout = true); 50 51 public: 52 RemoteOffloadClient(std::shared_ptr<Channel> Channel, int Timeout, 53 uint64_t MaxSize, int64_t BlockSize) 54 : Timeout(Timeout), MaxSize(MaxSize), BlockSize(BlockSize), 55 Stub(RemoteOffload::NewStub(Channel)) { 56 DebugLevel = getDebugLevel(); 57 Arena = std::make_unique<protobuf::Arena>(); 58 ArenaAllocatorLock = std::make_unique<std::mutex>(); 59 } 60 61 RemoteOffloadClient(RemoteOffloadClient &&C) = default; 62 63 ~RemoteOffloadClient() { 64 for (auto &TableIt : DevicesToTables) 65 freeTargetTable(TableIt.second.get()); 66 } 67 68 int32_t shutdown(void); 69 70 int32_t registerLib(__tgt_bin_desc *Desc); 71 int32_t unregisterLib(__tgt_bin_desc *Desc); 72 73 int32_t isValidBinary(__tgt_device_image *Image); 74 int32_t getNumberOfDevices(); 75 76 int32_t initDevice(int32_t DeviceId); 77 int32_t initRequires(int64_t RequiresFlags); 78 79 __tgt_target_table *loadBinary(int32_t DeviceId, __tgt_device_image *Image); 80 int64_t synchronize(int32_t DeviceId, __tgt_async_info *AsyncInfo); 81 int32_t isDataExchangeable(int32_t SrcDevId, int32_t DstDevId); 82 83 void *dataAlloc(int32_t DeviceId, int64_t Size, void *HstPtr); 84 int32_t dataDelete(int32_t DeviceId, void *TgtPtr); 85 86 int32_t dataSubmitAsync(int32_t DeviceId, void *TgtPtr, void *HstPtr, 87 int64_t Size, __tgt_async_info *AsyncInfo); 88 int32_t dataRetrieveAsync(int32_t DeviceId, void *HstPtr, void *TgtPtr, 89 int64_t Size, __tgt_async_info *AsyncInfo); 90 91 int32_t dataExchangeAsync(int32_t SrcDevId, void *SrcPtr, int32_t DstDevId, 92 void *DstPtr, int64_t Size, 93 __tgt_async_info *AsyncInfo); 94 95 int32_t runTargetRegionAsync(int32_t DeviceId, void *TgtEntryPtr, 96 void **TgtArgs, ptrdiff_t *TgtOffsets, 97 int32_t ArgNum, __tgt_async_info *AsyncInfo); 98 99 int32_t runTargetTeamRegionAsync(int32_t DeviceId, void *TgtEntryPtr, 100 void **TgtArgs, ptrdiff_t *TgtOffsets, 101 int32_t ArgNum, int32_t TeamNum, 102 int32_t ThreadLimit, uint64_t LoopTripCount, 103 __tgt_async_info *AsyncInfo); 104 }; 105 106 class RemoteClientManager { 107 private: 108 std::vector<std::string> Addresses; 109 std::vector<RemoteOffloadClient> Clients; 110 std::vector<int> Devices; 111 112 std::pair<int32_t, int32_t> mapDeviceId(int32_t DeviceId); 113 int DebugLevel; 114 115 public: 116 RemoteClientManager(std::vector<std::string> Addresses, int Timeout, 117 uint64_t MaxSize, int64_t BlockSize) 118 : Addresses(Addresses) { 119 grpc::ChannelArguments ChArgs; 120 ChArgs.SetMaxReceiveMessageSize(-1); 121 DebugLevel = getDebugLevel(); 122 for (auto Address : Addresses) { 123 Clients.push_back(RemoteOffloadClient( 124 grpc::CreateChannel(Address, grpc::InsecureChannelCredentials()), 125 Timeout, MaxSize, BlockSize)); 126 } 127 } 128 129 int32_t shutdown(void); 130 131 int32_t registerLib(__tgt_bin_desc *Desc); 132 int32_t unregisterLib(__tgt_bin_desc *Desc); 133 134 int32_t isValidBinary(__tgt_device_image *Image); 135 int32_t getNumberOfDevices(); 136 137 int32_t initDevice(int32_t DeviceId); 138 int32_t initRequires(int64_t RequiresFlags); 139 140 __tgt_target_table *loadBinary(int32_t DeviceId, __tgt_device_image *Image); 141 int64_t synchronize(int32_t DeviceId, __tgt_async_info *AsyncInfo); 142 int32_t isDataExchangeable(int32_t SrcDevId, int32_t DstDevId); 143 144 void *dataAlloc(int32_t DeviceId, int64_t Size, void *HstPtr); 145 int32_t dataDelete(int32_t DeviceId, void *TgtPtr); 146 147 int32_t dataSubmitAsync(int32_t DeviceId, void *TgtPtr, void *HstPtr, 148 int64_t Size, __tgt_async_info *AsyncInfo); 149 int32_t dataRetrieveAsync(int32_t DeviceId, void *HstPtr, void *TgtPtr, 150 int64_t Size, __tgt_async_info *AsyncInfo); 151 152 int32_t dataExchangeAsync(int32_t SrcDevId, void *SrcPtr, int32_t DstDevId, 153 void *DstPtr, int64_t Size, 154 __tgt_async_info *AsyncInfo); 155 156 int32_t runTargetRegionAsync(int32_t DeviceId, void *TgtEntryPtr, 157 void **TgtArgs, ptrdiff_t *TgtOffsets, 158 int32_t ArgNum, __tgt_async_info *AsyncInfo); 159 160 int32_t runTargetTeamRegionAsync(int32_t DeviceId, void *TgtEntryPtr, 161 void **TgtArgs, ptrdiff_t *TgtOffsets, 162 int32_t ArgNum, int32_t TeamNum, 163 int32_t ThreadLimit, uint64_t LoopTripCount, 164 __tgt_async_info *AsyncInfo); 165 }; 166 167 #endif 168