1//===-- ubsan_offload_hsa_interceptors.cpp ----------------------*- C++ -*-===//
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 <dlfcn.h>
10
11#include "interception/interception.h"
12#include "sanitizer_common/sanitizer_atomic.h"
13#include "sanitizer_common/sanitizer_common.h"
14#include "sanitizer_common/sanitizer_libc.h"
15#include "sanitizer_common/sanitizer_mutex.h"
16#include "sanitizer_common/sanitizer_offload.h"
17#include "sanitizer_common/sanitizer_platform.h"
18#include "ubsan_diag.h"
19#include "ubsan_offload.h"
20
21#if !SANITIZER_LINUX
22#error "Offload UBSan reporting is supported on Linux only"
23#endif
24
25#if SANITIZER_GLIBC
26#pragma weak dlvsym
27#endif
28
29using namespace __sanitizer;
30using namespace __ubsan;
31
32namespace __ubsan {
33
34static StaticSpinMutex InitMutex;
35static atomic_uint8_t Initialized;
36
37void Initialize() {
38 if (LIKELY(atomic_load(&Initialized, memory_order_acquire)))
39 return;
40 SpinMutexLock L(&InitMutex);
41 if (atomic_load(a: &Initialized, mo: memory_order_relaxed))
42 return;
43 SanitizerToolName = "UndefinedBehaviorSanitizer";
44 __ubsan_set_offload_symbolize(
45 Fn: [](uptr PC) { return Offload::Get().Symbolize(PC); });
46 Offload::Get().RegisterHandler(Fn: HandleOffloadReport);
47 Atexit(function: [] { Offload::Get().UntrackImages(); });
48 AddDieCallback(callback: [] { Offload::Get().UntrackImages(); });
49 atomic_store(a: &Initialized, v: 1, mo: memory_order_release);
50}
51
52} // namespace __ubsan
53
54#define UBSAN_HSA_ENTER(name) \
55 Initialize(); \
56 if (UNLIKELY(!REAL(name))) { \
57 INTERCEPT_FUNCTION(name); \
58 if (UNLIKELY(!REAL(name))) { \
59 Report("ERROR: %s: cannot find %s in this process\n", SanitizerToolName, \
60 #name); \
61 Die(); \
62 } \
63 }
64
65#define UBSAN_HSA_FORWARD(name, ...) \
66 UBSAN_HSA_ENTER(name); \
67 if (UNLIKELY(!Offload::Get().Ready())) \
68 return REAL(name)(__VA_ARGS__);
69
70// PPC cannot transparently tail-call an indirect dlsym target for RTLD_NEXT.
71#if !SANITIZER_PPC
72#define UBSAN_HSA_WRAPS(X) \
73 X(hsa_init) \
74 X(hsa_shut_down) \
75 X(hsa_executable_freeze) \
76 X(hsa_executable_destroy)
77
78static void *WrapperFor(const char *Name) {
79#define UBSAN_HSA_WRAP(Fn) \
80 if (!internal_strcmp(Name, #Fn)) \
81 return reinterpret_cast<void *>(Fn);
82 UBSAN_HSA_WRAPS(UBSAN_HSA_WRAP)
83#undef UBSAN_HSA_WRAP
84 return nullptr;
85}
86
87static bool FromHsa(void *P) {
88 Dl_info Info = {};
89 if (!dladdr(address: P, info: &Info) || !Info.dli_fname)
90 return false;
91 return internal_strstr(haystack: Info.dli_fname, SANITIZER_HSA_LIBRARY);
92}
93
94static void BindRealDlsym();
95
96// OpenMP and sometimes HIP access HSA through 'dlsym' so we need to intercept
97// it here if we want to reliably override its definitions.
98INTERCEPTOR(void *, dlsym, void *Handle, const char *Name) {
99 Initialize();
100 BindRealDlsym();
101
102 // This interceptor interferes with the order of 'RTLD_NEXT'. Force a tail
103 // call to bypass this process in the stack.
104 if (Handle == RTLD_NEXT) [[clang::musttail]]
105 return REAL(dlsym)(Handle, Name);
106
107 void *Sym = REAL(dlsym)(Handle, Name);
108 if (!Sym || !Name)
109 return Sym;
110
111 void *Wrapper = WrapperFor(Name);
112 if (!Wrapper || !FromHsa(P: Sym))
113 return Sym;
114 return Wrapper;
115}
116
117static void BindRealDlsym() {
118 if (LIKELY(REAL(dlsym)))
119 return;
120#if SANITIZER_GLIBC
121 static const char *kVers[] = {"GLIBC_2.34", "GLIBC_2.17", "GLIBC_2.2.5",
122 "GLIBC_2.0"};
123 if (dlvsym) {
124 for (const char *Ver : kVers) {
125 if (void *P = dlvsym(RTLD_NEXT, name: "dlsym", version: Ver)) {
126 REAL(dlsym) = reinterpret_cast<decltype(REAL(dlsym))>(P);
127 return;
128 }
129 }
130 }
131#endif
132 Report(format: "ERROR: %s: cannot bind dlsym\n", SanitizerToolName);
133 Die();
134}
135#endif
136
137INTERCEPTOR(hsa_status_t, hsa_init, void) {
138 UBSAN_HSA_ENTER(hsa_init);
139
140 hsa_status_t Status = REAL(hsa_init)();
141 if (Status != HSA_STATUS_SUCCESS)
142 return Status;
143
144 Offload::Get().Init();
145 return Status;
146}
147
148INTERCEPTOR(hsa_status_t, hsa_shut_down, void) {
149 UBSAN_HSA_ENTER(hsa_shut_down);
150
151 Offload::Get().Shutdown();
152 return REAL(hsa_shut_down)();
153}
154
155INTERCEPTOR(hsa_status_t, hsa_executable_freeze, hsa_executable_t Executable,
156 const char *Options) {
157 UBSAN_HSA_FORWARD(hsa_executable_freeze, Executable, Options);
158
159 hsa_status_t Status = REAL(hsa_executable_freeze)(Executable, Options);
160 if (Status == HSA_STATUS_SUCCESS)
161 Offload::Get().TrackExecutable(Exec: Executable);
162 return Status;
163}
164
165INTERCEPTOR(hsa_status_t, hsa_executable_destroy, hsa_executable_t Executable) {
166 UBSAN_HSA_FORWARD(hsa_executable_destroy, Executable);
167
168 Offload::Get().UntrackExecutable(Exec: Executable);
169 return REAL(hsa_executable_destroy)(Executable);
170}
171
172extern "C" void __ubsan_offload_init() { __ubsan::Initialize(); }
173
174#if SANITIZER_CAN_USE_PREINIT_ARRAY
175__attribute__((section(".preinit_array"), used)) static void (
176 *ubsan_offload_preinit)(void) = __ubsan_offload_init;
177#endif
178
179__attribute__((constructor(0))) static void UbsanOffloadDynInit() {
180 __ubsan_offload_init();
181}
182