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-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// Host RPC server for device sanitizer reports. One server, many opcode
10// handlers. Prefers liboffload's server when present. Otherwise plants a
11// client on each executable's empty __llvm_rpc_client, sharing one buffer
12// per device.
13//
14//===----------------------------------------------------------------------===//
15
16#include "sanitizer_offload_rpc.h"
17
18#include <dlfcn.h>
19#include <pthread.h>
20
21#include "sanitizer_allocator_internal.h"
22#include "sanitizer_atomic.h"
23#include "sanitizer_common.h"
24#include "sanitizer_libc.h"
25#include "sanitizer_mutex.h"
26#include "sanitizer_offload.h"
27#include "sanitizer_placement_new.h"
28#include "sanitizer_posix.h"
29#include "shared/rpc.h"
30
31namespace __sanitizer {
32namespace {
33
34struct DeviceRpc {
35 void* Buffer;
36 rpc::Server* Server;
37 u32 Ports;
38 u32 Lanes;
39};
40
41struct Rpc {
42 hsa_signal_t Signal{};
43 u64* SignalValue{};
44 u64* SignalMailbox{};
45 u32 SignalEvent{};
46 InternalMmapVectorNoCtor<DeviceRpc> Slots{};
47 InternalMmapVectorNoCtor<Offload::Handler> Handlers{};
48 void* Thread{};
49 bool Liboffload{};
50 bool RegisteredLiboffload{};
51 atomic_uint8_t Halt{};
52 Mutex Mtx{};
53 Mutex HandlerMtx{};
54};
55
56Rpc State;
57
58Rpc& GetRpc() { return State; }
59
60u32 Dispatch(Rpc& R, void* PortPtr, u32 Lanes) {
61 Lock L(&R.HandlerMtx);
62 for (uptr I = 0; I < R.Handlers.size(); ++I) {
63 u32 Status = R.Handlers[I](PortPtr, Lanes);
64 if (Status != rpc::RPC_UNHANDLED_OPCODE)
65 return Status;
66 }
67 return rpc::RPC_UNHANDLED_OPCODE;
68}
69
70bool TryLiboffload(Rpc& R) {
71 using RegisterFn = void (*)(u32 (*)(void*, u32));
72 auto Register = reinterpret_cast<RegisterFn>(
73 dlsym(RTLD_DEFAULT, name: "__tgt_register_rpc_callback"));
74 if (!Register) {
75 R.Liboffload = false;
76 return false;
77 }
78 if (!R.RegisteredLiboffload) {
79 Register([](void* Port, u32 Lanes) -> u32 {
80 return Dispatch(R&: GetRpc(), PortPtr: Port, Lanes);
81 });
82 R.RegisteredLiboffload = true;
83 VReport(1, "%s: device reports through the offload runtime's server\n",
84 SanitizerToolName);
85 }
86 R.Liboffload = true;
87 return true;
88}
89
90void DrainSlot(DeviceRpc& D) {
91 if (!D.Server)
92 return;
93 while (auto Port = D.Server->try_open(lane_size: D.Lanes)) {
94 if (Dispatch(R&: GetRpc(), PortPtr: &*Port, Lanes: D.Lanes) == rpc::RPC_UNHANDLED_OPCODE)
95 VReport(1, "%s: unexpected opcode 0x%x on the report channel\n",
96 SanitizerToolName, Port->get_opcode());
97 }
98}
99
100void DrainAll(Rpc& R) {
101 for (uptr I = 0; I < R.Slots.size(); ++I) DrainSlot(D&: R.Slots[I]);
102}
103
104void FailRpc() {
105 Report(format: "ERROR: %s: failed to start device RPC\n", SanitizerToolName);
106 Die();
107}
108
109void PlantSignal(Rpc& R, void* Buffer) {
110 auto* Bell = reinterpret_cast<rpc::Doorbell*>(static_cast<u8*>(Buffer) +
111 rpc::Server::doorbell_offset());
112 Bell->value = reinterpret_cast<uint64_t*>(R.SignalValue);
113 Bell->mailbox = reinterpret_cast<uint64_t*>(R.SignalMailbox);
114 Bell->event_id = R.SignalEvent;
115}
116
117} // namespace
118
119int internal_pthread_create(void* Th, void* Attr, void* (*Callback)(void*),
120 void* Param) {
121 return pthread_create(newthread: reinterpret_cast<pthread_t*>(Th),
122 attr: reinterpret_cast<const pthread_attr_t*>(Attr), start_routine: Callback,
123 arg: Param);
124}
125
126int internal_pthread_join(void* Th, void** Ret) {
127 return pthread_join(th: reinterpret_cast<pthread_t>(Th), thread_return: Ret);
128}
129
130void* OffloadRpc::ServerLoop(void* Arg) {
131 Offload& O = *static_cast<Offload*>(Arg);
132 Rpc& R = GetRpc();
133 for (;;) {
134 if (!atomic_load_relaxed(a: &R.Halt))
135 O.WaitSignal(Sig: R.Signal);
136 Lock L(&R.Mtx);
137 DrainAll(R);
138 if (atomic_load_relaxed(a: &R.Halt))
139 break;
140 }
141 return nullptr;
142}
143
144void OffloadRpc::RegisterHandler(Offload::Handler Fn) {
145 Rpc& R = GetRpc();
146 Lock L(&R.HandlerMtx);
147 for (uptr I = 0; I < R.Handlers.size(); ++I)
148 if (R.Handlers[I] == Fn)
149 return;
150 R.Handlers.push_back(element: Fn);
151}
152
153void OffloadRpc::Flush() {
154 Rpc& R = GetRpc();
155 Lock L(&R.Mtx);
156 if (R.Liboffload)
157 return;
158 DrainAll(R);
159}
160
161void OffloadRpc::Start(Offload& O, hsa_executable_t Exec) {
162 Rpc& R = GetRpc();
163 Lock L(&R.Mtx);
164 if (!O.Ready() || atomic_load_relaxed(a: &R.Halt))
165 return;
166 if (TryLiboffload(R))
167 return;
168
169 if (!R.Signal.handle) {
170 if (!O.CreateSignal(Out: &R.Signal))
171 FailRpc();
172 // Mirror of ROCr amd_signal_t: KFD interrupt slot used to wake the RPC
173 // server thread.
174 struct AMDSignal {
175 int64_t Kind;
176 int64_t Value;
177 uint64_t EventMailboxPtr;
178 uint32_t EventId;
179 };
180 auto* S = reinterpret_cast<AMDSignal*>(R.Signal.handle);
181 R.SignalValue = reinterpret_cast<u64*>(&S->Value);
182 R.SignalMailbox = reinterpret_cast<u64*>(S->EventMailboxPtr);
183 R.SignalEvent = S->EventId;
184 }
185
186 while (R.Slots.size() < O.Devices().size()) {
187 DeviceRpc Empty = {};
188 R.Slots.push_back(element: Empty);
189 }
190
191 for (uptr I = 0; I < O.Devices().size(); ++I) {
192 const Offload::Device& D = O.Devices()[I];
193 u64 Addr = 0;
194 if (!O.Lookup(Exec, Name: "__llvm_rpc_client", Agent: D.Agent, Addr: &Addr))
195 continue;
196
197 DeviceRpc& Slot = R.Slots[I];
198 if (!Slot.Buffer) {
199 Slot.Lanes = D.Lanes;
200 Slot.Ports = D.MaxWaves;
201 if (Slot.Ports > rpc::MAX_PORT_COUNT)
202 Slot.Ports = rpc::MAX_PORT_COUNT;
203 const uptr Bytes = rpc::Server::allocation_size(lane_size: Slot.Lanes, port_count: Slot.Ports);
204 void* Buffer = nullptr;
205 if (!O.Alloc(D: O.Host(), Bytes, Out: &Buffer) || !Buffer)
206 FailRpc();
207 internal_memset(s: Buffer, c: 0, n: Bytes);
208 PlantSignal(R, Buffer);
209 Slot.Buffer = Buffer;
210 Slot.Server = new (InternalAlloc(size: sizeof(rpc::Server)))
211 rpc::Server(Slot.Ports, Slot.Buffer);
212 VReport(1, "%s: serving device reports on GPU %zu, %u ports, %u lanes\n",
213 SanitizerToolName, I, Slot.Ports, Slot.Lanes);
214 }
215 auto* Client = new (InternalAlloc(size: sizeof(rpc::Client)))
216 rpc::Client(Slot.Ports, Slot.Buffer);
217 bool Installed =
218 O.Copy(Dst: reinterpret_cast<void*>(Addr), Src: Client, N: sizeof(*Client));
219 Client->~Client();
220 InternalFree(p: Client);
221 if (!Installed)
222 FailRpc();
223
224 if (!R.Thread) {
225 if (atomic_load_relaxed(a: &R.Halt) ||
226 !(R.Thread = internal_start_thread(func: ServerLoop, arg: &O)))
227 FailRpc();
228 }
229 }
230}
231
232void OffloadRpc::Stop(Offload& O) {
233 Rpc& R = GetRpc();
234 void* Join = nullptr;
235 {
236 Lock L(&R.Mtx);
237 R.Liboffload = false;
238 atomic_store_relaxed(a: &R.Halt, v: 1);
239 if (R.Thread) {
240 if (R.Signal.handle)
241 O.StoreSignal(Sig: R.Signal);
242 Join = R.Thread;
243 R.Thread = nullptr;
244 }
245 }
246 if (Join)
247 internal_join_thread(th: Join);
248
249 Lock L(&R.Mtx);
250 if (R.Thread)
251 return;
252 for (uptr I = 0; I < R.Slots.size(); ++I) {
253 if (R.Slots[I].Server) {
254 R.Slots[I].Server->~Server();
255 InternalFree(p: R.Slots[I].Server);
256 }
257 if (R.Slots[I].Buffer)
258 O.Free(P: R.Slots[I].Buffer);
259 }
260 R.Slots.clear();
261 if (R.Signal.handle) {
262 O.DestroySignal(Sig: R.Signal);
263 R.Signal = {};
264 R.SignalValue = nullptr;
265 R.SignalMailbox = nullptr;
266 R.SignalEvent = 0;
267 }
268 atomic_store_relaxed(a: &R.Halt, v: 0);
269}
270
271} // namespace __sanitizer
272