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#include "sanitizer_offload.h"
10
11#include <dlfcn.h>
12
13#include "sanitizer_common.h"
14#include "sanitizer_libc.h"
15#include "sanitizer_mutex.h"
16#include "sanitizer_offload_rpc.h"
17
18namespace __sanitizer {
19namespace {
20
21// Generic helpers to iterate HSA agents and pools.
22template <typename ElemTy, typename IterFuncTy, typename CallbackTy>
23hsa_status_t Iterate(IterFuncTy Func, CallbackTy Cb) {
24 auto L = [](ElemTy Elem, void* Data) -> hsa_status_t {
25 return (*static_cast<CallbackTy*>(Data))(Elem);
26 };
27 return Func(L, &Cb);
28}
29
30template <typename ElemTy, typename IterFuncTy, typename ArgTy,
31 typename CallbackTy>
32hsa_status_t Iterate(IterFuncTy Func, ArgTy Arg, CallbackTy Cb) {
33 auto L = [](ElemTy Elem, void* Data) -> hsa_status_t {
34 return (*static_cast<CallbackTy*>(Data))(Elem);
35 };
36 return Func(Arg, L, &Cb);
37}
38
39template <typename Elem1Ty, typename Elem2Ty, typename IterFuncTy,
40 typename ArgTy, typename CallbackTy>
41hsa_status_t Iterate(IterFuncTy Func, ArgTy Arg, CallbackTy Cb) {
42 auto L = [](Elem1Ty A, Elem2Ty B, void* Data) -> hsa_status_t {
43 return (*static_cast<CallbackTy*>(Data))(A, B);
44 };
45 return Func(Arg, L, &Cb);
46}
47
48void CheckHsa(hsa_status_t S) {
49 if (S == HSA_STATUS_SUCCESS)
50 return;
51 Report(format: "ERROR: %s: HSA query failed\n", SanitizerToolName);
52 Die();
53}
54
55} // namespace
56
57Offload Offload::Ctx;
58
59Offload& Offload::Get() { return Ctx; }
60
61template <typename Cb>
62void Offload::ForEachAgentObject(hsa_executable_t Exec, Cb F) {
63 Iterate<hsa_executable_t, hsa_loaded_code_object_t>(
64 Loader.IterateLoadedCodeObjects, Exec,
65 [&](hsa_executable_t, hsa_loaded_code_object_t Obj) {
66 u32 Kind = 0;
67 if (Loader.GetCodeObjectInfo(
68 Obj, HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_KIND, &Kind) !=
69 HSA_STATUS_SUCCESS ||
70 Kind != HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_KIND_AGENT)
71 return HSA_STATUS_SUCCESS;
72 hsa_agent_t Agent{};
73 if (Loader.GetCodeObjectInfo(
74 Obj, HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_AGENT,
75 &Agent) != HSA_STATUS_SUCCESS ||
76 !Agent.handle)
77 return HSA_STATUS_SUCCESS;
78 F(Obj, Agent);
79 return HSA_STATUS_SUCCESS;
80 });
81}
82
83// Fetch all the real HSA library calls we need.
84bool Offload::Resolve() {
85#define SANITIZER_HSA_RESOLVE(Name) \
86 Api.Name = reinterpret_cast<decltype(&::Name)>(dlsym(RTLD_NEXT, #Name)); \
87 if (!Api.Name) \
88 return false;
89 SANITIZER_HSA_FUNCTIONS(SANITIZER_HSA_RESOLVE)
90#undef SANITIZER_HSA_RESOLVE
91 return true;
92}
93
94// Iterate the topology to discover agents. If this fails the interceptors are
95// disabled.
96bool Offload::Discover() {
97 DeviceList.clear();
98 HostDevice = {};
99
100 auto Info = [&](hsa_agent_t Agent, u32 Attr, void* Out) {
101 return Api.hsa_agent_get_info(Agent, static_cast<hsa_agent_info_t>(Attr),
102 Out) == HSA_STATUS_SUCCESS;
103 };
104 auto Pool = [&](hsa_agent_t Agent) {
105 hsa_amd_memory_pool_t Found{};
106 Iterate<hsa_amd_memory_pool_t>(
107 Func: Api.hsa_amd_agent_iterate_memory_pools, Arg: Agent,
108 Cb: [&](hsa_amd_memory_pool_t Mem) {
109 hsa_amd_segment_t Seg;
110 u32 Flags = 0;
111 if (Api.hsa_amd_memory_pool_get_info(Mem,
112 HSA_AMD_MEMORY_POOL_INFO_SEGMENT,
113 &Seg) != HSA_STATUS_SUCCESS ||
114 Seg != HSA_AMD_SEGMENT_GLOBAL)
115 return HSA_STATUS_SUCCESS;
116 if (Api.hsa_amd_memory_pool_get_info(
117 Mem, HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, &Flags) !=
118 HSA_STATUS_SUCCESS)
119 return HSA_STATUS_SUCCESS;
120 if (Flags & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED)
121 Found = Mem;
122 return HSA_STATUS_SUCCESS;
123 });
124 return Found;
125 };
126
127 CheckHsa(S: Iterate<hsa_agent_t>(Func: Api.hsa_iterate_agents, Cb: [&](hsa_agent_t Agent) {
128 hsa_device_type_t Type;
129 if (hsa_status_t S =
130 Api.hsa_agent_get_info(Agent, HSA_AGENT_INFO_DEVICE, &Type))
131 return S;
132 Device D = {};
133 D.Agent = Agent;
134 D.Pool = Pool(Agent);
135 if (Type == HSA_DEVICE_TYPE_CPU && !HostDevice.Agent.handle) {
136 HostDevice = D;
137 } else if (Type == HSA_DEVICE_TYPE_GPU) {
138 u32 CUs = 0, WavesPerCU = 0;
139 if (!Info(Agent, HSA_AGENT_INFO_WAVEFRONT_SIZE, &D.Lanes) || !D.Lanes ||
140 !Info(Agent, HSA_AMD_AGENT_INFO_COMPUTE_UNIT_COUNT, &CUs) ||
141 !Info(Agent, HSA_AMD_AGENT_INFO_MAX_WAVES_PER_CU, &WavesPerCU) ||
142 !CUs || !WavesPerCU)
143 CheckHsa(S: HSA_STATUS_ERROR);
144 // The RPC interface is deliberately sized to the hardware parallel
145 // limits of the device to make deadlock impossible.
146 D.MaxWaves = CUs * WavesPerCU;
147 DeviceList.push_back(element: D);
148 }
149 return HSA_STATUS_SUCCESS;
150 }));
151 return HostDevice.Agent.handle && !DeviceList.empty() &&
152 HostDevice.Pool.handle;
153}
154
155// Initialize the HSA loader extension used to manage host addresses.
156bool Offload::BindLoader() {
157 Loader = {};
158 CheckHsa(S: Api.hsa_system_get_major_extension_table(HSA_EXTENSION_AMD_LOADER, 1,
159 sizeof(Loader), &Loader));
160 return Loader.QueryHostAddress && Loader.IterateLoadedCodeObjects &&
161 Loader.GetCodeObjectInfo;
162}
163
164bool Offload::ExecutableInfo(hsa_loaded_code_object_t Obj,
165 hsa_ven_amd_loader_loaded_code_object_info_t Attr,
166 u64* Out) {
167 *Out = 0;
168 return Loader.GetCodeObjectInfo(Obj, Attr, Out) == HSA_STATUS_SUCCESS;
169}
170
171bool Offload::Init() {
172 Lock Life(&LifetimeMtx);
173 Lock L(&OffloadMtx);
174 if (++Refs > 1)
175 return Ready();
176 if (!Resolve()) {
177 Report(format: "ERROR: %s: cannot resolve HSA\n", SanitizerToolName);
178 Die();
179 }
180 if (!Discover() || !BindLoader()) {
181 --Refs;
182 return false;
183 }
184 VReport(1, "%s: device reporting on %zu GPU(s)\n", SanitizerToolName,
185 DeviceList.size());
186 atomic_store(a: &Active, v: 1, mo: memory_order_release);
187 return true;
188}
189
190bool Offload::Release() {
191 if (!Refs || --Refs)
192 return false;
193 atomic_store(a: &Active, v: 0, mo: memory_order_release);
194 return true;
195}
196
197void Offload::Teardown() {
198 UntrackImages();
199 Loader = {};
200 DeviceList.clear();
201 HostDevice = {};
202}
203
204void Offload::Shutdown() {
205 Lock Life(&LifetimeMtx);
206 bool Last;
207 {
208 Lock L(&OffloadMtx);
209 Last = Release();
210 }
211 if (!Last)
212 return;
213 OffloadRpc::Stop(O&: *this);
214 Lock L(&OffloadMtx);
215 Teardown();
216}
217
218bool Offload::Ready() const {
219 return atomic_load(a: &Active, mo: memory_order_acquire) != 0;
220}
221
222void Offload::RegisterHandler(Handler Fn) { OffloadRpc::RegisterHandler(Fn); }
223
224// Record every executable we come across for symbolization and address lookup.
225void Offload::TrackExecutable(hsa_executable_t Exec) {
226 {
227 Lock L(&OffloadMtx);
228 if (!Ready())
229 return;
230 ForEachAgentObject(Exec, F: [&](hsa_loaded_code_object_t Obj, hsa_agent_t) {
231 u64 LoadBase = 0, LoadSize = 0;
232 if (!ExecutableInfo(Obj,
233 Attr: HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_LOAD_BASE,
234 Out: &LoadBase) ||
235 !ExecutableInfo(Obj,
236 Attr: HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_LOAD_SIZE,
237 Out: &LoadSize) ||
238 !LoadBase || !LoadSize)
239 return;
240
241 u64 StorageType = 0, StorageBase = 0, StorageSize = 0;
242 const void* Storage = nullptr;
243 if (ExecutableInfo(
244 Obj,
245 Attr: HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_CODE_OBJECT_STORAGE_TYPE,
246 Out: &StorageType) &&
247 StorageType == HSA_VEN_AMD_LOADER_CODE_OBJECT_STORAGE_TYPE_MEMORY &&
248 ExecutableInfo(
249 Obj,
250 Attr: HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_CODE_OBJECT_STORAGE_MEMORY_BASE,
251 Out: &StorageBase) &&
252 ExecutableInfo(
253 Obj,
254 Attr: HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_CODE_OBJECT_STORAGE_MEMORY_SIZE,
255 Out: &StorageSize) &&
256 StorageBase && StorageSize)
257 Storage = reinterpret_cast<const void*>(StorageBase);
258
259 TrackImage(LoadBase: (uptr)LoadBase, LoadSize: (uptr)LoadSize, Storage, StorageSize: (uptr)StorageSize);
260 });
261 }
262 OffloadRpc::Start(O&: *this, Exec);
263}
264
265void Offload::UntrackExecutable(hsa_executable_t Exec) {
266 OffloadRpc::Flush();
267 Lock L(&OffloadMtx);
268 ForEachAgentObject(Exec, F: [&](hsa_loaded_code_object_t Obj, hsa_agent_t) {
269 u64 LoadBase = 0;
270 if (ExecutableInfo(Obj,
271 Attr: HSA_VEN_AMD_LOADER_LOADED_CODE_OBJECT_INFO_LOAD_BASE,
272 Out: &LoadBase) &&
273 LoadBase)
274 UntrackImage(LoadBase: (uptr)LoadBase);
275 });
276}
277
278// Allocate coherent 'fine-grained' memory for host and device communication.
279bool Offload::Alloc(const Device& D, uptr Bytes, void** Out) {
280 void* P = nullptr;
281 if (!D.Pool.handle ||
282 Api.hsa_amd_memory_pool_allocate(D.Pool, Bytes, 0, &P) !=
283 HSA_STATUS_SUCCESS ||
284 !P)
285 return false;
286 InternalMmapVector<hsa_agent_t> All;
287 All.push_back(element: HostDevice.Agent);
288 for (uptr I = 0; I < DeviceList.size(); ++I)
289 All.push_back(element: DeviceList[I].Agent);
290 if (Api.hsa_amd_agents_allow_access(All.size(), All.data(), nullptr, P) !=
291 HSA_STATUS_SUCCESS) {
292 Api.hsa_amd_memory_pool_free(P);
293 return false;
294 }
295 *Out = P;
296 return true;
297}
298
299void Offload::Free(void* P) { Api.hsa_amd_memory_pool_free(P); }
300
301bool Offload::Copy(void* Dst, const void* Src, uptr N) {
302 return Api.hsa_memory_copy(Dst, Src, N) == HSA_STATUS_SUCCESS;
303}
304
305// The runtime knows the original host address of a device pointer that is
306// located inside one of the loaded segments, accesses read-only data we need.
307const void* Offload::HostPointer(uptr DeviceAddr) {
308 Lock L(&OffloadMtx);
309 if (!Ready())
310 return nullptr;
311 if (!DeviceAddr)
312 return nullptr;
313 const void* HostAddr = nullptr;
314 if (Loader.QueryHostAddress(reinterpret_cast<const void*>(DeviceAddr),
315 &HostAddr) != HSA_STATUS_SUCCESS)
316 return nullptr;
317 return HostAddr;
318}
319
320bool Offload::Lookup(hsa_executable_t Exec, const char* Name, hsa_agent_t Agent,
321 u64* Addr) {
322 hsa_executable_symbol_t Symbol;
323 if (Api.hsa_executable_get_symbol_by_name(Exec, Name, &Agent, &Symbol) !=
324 HSA_STATUS_SUCCESS)
325 return false;
326 *Addr = 0;
327 return Api.hsa_executable_symbol_get_info(
328 Symbol, HSA_EXECUTABLE_SYMBOL_INFO_VARIABLE_ADDRESS, Addr) ==
329 HSA_STATUS_SUCCESS &&
330 *Addr;
331}
332
333bool Offload::CreateSignal(hsa_signal_t* Out) {
334 return Api.hsa_amd_signal_create(0, 0, nullptr, 0, Out) == HSA_STATUS_SUCCESS;
335}
336
337void Offload::DestroySignal(hsa_signal_t Sig) { Api.hsa_signal_destroy(Sig); }
338
339void Offload::WaitSignal(hsa_signal_t Sig) {
340 Api.hsa_signal_wait_scacquire(Sig, HSA_SIGNAL_CONDITION_NE, 0, UINT64_MAX,
341 HSA_WAIT_STATE_BLOCKED);
342}
343
344void Offload::StoreSignal(hsa_signal_t Sig) {
345 Api.hsa_signal_store_screlease(Sig, 1);
346}
347
348} // namespace __sanitizer
349