Skip to content

Commit 4c69624

Browse files
kbenziekbenzie
authored
Merge pull request #1532 from AllanZyne/review/yang/local_accessor
[DeviceSanitizer] Check out-of-bounds on sycl::local_accessor
2 parents ea00936 + f04b44c commit 4c69624

File tree

5 files changed

+330
-56
lines changed

5 files changed

+330
-56
lines changed

‎source/loader/layers/sanitizer/asan_interceptor.cpp

Lines changed: 98 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,6 @@ namespace ur_sanitizer_layer {
2222

2323
namespace {
2424

25-
constexpr auto kSPIR_DeviceSanitizerReportMem = "__DeviceSanitizerReportMem";
26-
2725
uptr MemToShadow_CPU(uptr USM_SHADOW_BASE, uptr UPtr) {
2826
return USM_SHADOW_BASE + (UPtr >> 3);
2927
}
@@ -348,11 +346,14 @@ ur_result_t SanitizerInterceptor::releaseMemory(ur_context_handle_t Context,
348346

349347
ur_result_t SanitizerInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
350348
ur_queue_handle_t Queue,
351-
LaunchInfo &LaunchInfo) {
349+
USMLaunchInfo &LaunchInfo) {
352350
auto Context = GetContext(Queue);
353351
auto Device = GetDevice(Queue);
354352
auto ContextInfo = getContextInfo(Context);
355353
auto DeviceInfo = getDeviceInfo(Device);
354+
auto KernelInfo = getKernelInfo(Kernel);
355+
356+
UR_CALL(LaunchInfo.updateKernelInfo(*KernelInfo.get()));
356357

357358
ManagedQueue InternalQueue(Context, Device);
358359
if (!InternalQueue) {
@@ -370,23 +371,12 @@ ur_result_t SanitizerInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
370371

371372
ur_result_t SanitizerInterceptor::postLaunchKernel(ur_kernel_handle_t Kernel,
372373
ur_queue_handle_t Queue,
373-
ur_event_handle_t &Event,
374-
LaunchInfo &LaunchInfo) {
375-
auto Program = GetProgram(Kernel);
376-
ur_event_handle_t ReadEvent{};
377-
378-
// If kernel has defined SPIR_DeviceSanitizerReportMem, then we try to read it
379-
// to host, but it's okay that it isn't defined
374+
USMLaunchInfo &LaunchInfo) {
380375
// FIXME: We must use block operation here, until we support urEventSetCallback
381-
auto Result = context.urDdiTable.Enqueue.pfnDeviceGlobalVariableRead(
382-
Queue, Program, kSPIR_DeviceSanitizerReportMem, true,
383-
sizeof(LaunchInfo.SPIR_DeviceSanitizerReportMem), 0,
384-
&LaunchInfo.SPIR_DeviceSanitizerReportMem, 1, &Event, &ReadEvent);
376+
auto Result = context.urDdiTable.Queue.pfnFinish(Queue);
385377

386378
if (Result == UR_RESULT_SUCCESS) {
387-
Event = ReadEvent;
388-
389-
const auto &AH = LaunchInfo.SPIR_DeviceSanitizerReportMem;
379+
const auto &AH = LaunchInfo.Data->SanitizerReport;
390380
if (!AH.Flag) {
391381
return UR_RESULT_SUCCESS;
392382
}
@@ -627,13 +617,44 @@ ur_result_t SanitizerInterceptor::eraseDevice(ur_device_handle_t Device) {
627617
return UR_RESULT_SUCCESS;
628618
}
629619

620+
ur_result_t SanitizerInterceptor::insertKernel(ur_kernel_handle_t Kernel) {
621+
std::scoped_lock<ur_shared_mutex> Guard(m_KernelMapMutex);
622+
if (m_KernelMap.find(Kernel) != m_KernelMap.end()) {
623+
return UR_RESULT_SUCCESS;
624+
}
625+
m_KernelMap.emplace(Kernel, std::make_shared<KernelInfo>(Kernel));
626+
return UR_RESULT_SUCCESS;
627+
}
628+
629+
ur_result_t SanitizerInterceptor::eraseKernel(ur_kernel_handle_t Kernel) {
630+
std::scoped_lock<ur_shared_mutex> Guard(m_KernelMapMutex);
631+
assert(m_KernelMap.find(Kernel) != m_KernelMap.end());
632+
m_KernelMap.erase(Kernel);
633+
return UR_RESULT_SUCCESS;
634+
}
635+
630636
ur_result_t SanitizerInterceptor::prepareLaunch(
631637
ur_context_handle_t Context, std::shared_ptr<DeviceInfo> &DeviceInfo,
632638
ur_queue_handle_t Queue, ur_kernel_handle_t Kernel,
633-
LaunchInfo &LaunchInfo) {
639+
USMLaunchInfo &LaunchInfo) {
634640
auto Program = GetProgram(Kernel);
635641

636642
do {
643+
// Set launch info argument
644+
auto ArgNums = GetKernelNumArgs(Kernel);
645+
if (ArgNums) {
646+
context.logger.debug(
647+
"launch_info {} (numLocalArgs={}, localArgs={})",
648+
(void *)LaunchInfo.Data, LaunchInfo.Data->NumLocalArgs,
649+
(void *)LaunchInfo.Data->LocalArgs);
650+
ur_result_t URes = context.urDdiTable.Kernel.pfnSetArgPointer(
651+
Kernel, ArgNums - 1, nullptr, &LaunchInfo.Data);
652+
if (URes != UR_RESULT_SUCCESS) {
653+
context.logger.error("Failed to set launch info: {}", URes);
654+
return URes;
655+
}
656+
}
657+
637658
// Write global variable to program
638659
auto EnqueueWriteGlobal = [Queue, Program](const char *Name,
639660
const void *Value,
@@ -723,15 +744,17 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
723744
"LocalShadowMemorySize={})",
724745
NumWG, LocalMemorySize, LocalShadowMemorySize);
725746

726-
UR_CALL(EnqueueAllocateDevice(LocalShadowMemorySize,
727-
LaunchInfo.LocalShadowOffset));
747+
UR_CALL(EnqueueAllocateDevice(
748+
LocalShadowMemorySize, LaunchInfo.Data->LocalShadowOffset));
728749

729-
LaunchInfo.LocalShadowOffsetEnd =
730-
LaunchInfo.LocalShadowOffset + LocalShadowMemorySize - 1;
750+
LaunchInfo.Data->LocalShadowOffsetEnd =
751+
LaunchInfo.Data->LocalShadowOffset + LocalShadowMemorySize -
752+
1;
731753

732-
context.logger.info("ShadowMemory(Local, {} - {})",
733-
(void *)LaunchInfo.LocalShadowOffset,
734-
(void *)LaunchInfo.LocalShadowOffsetEnd);
754+
context.logger.info(
755+
"ShadowMemory(Local, {} - {})",
756+
(void *)LaunchInfo.Data->LocalShadowOffset,
757+
(void *)LaunchInfo.Data->LocalShadowOffsetEnd);
735758
}
736759
}
737760
} while (false);
@@ -749,15 +772,61 @@ SanitizerInterceptor::findAllocInfoByAddress(uptr Address) {
749772
return --It;
750773
}
751774

752-
LaunchInfo::~LaunchInfo() {
775+
ur_result_t USMLaunchInfo::initialize() {
776+
UR_CALL(context.urDdiTable.Context.pfnRetain(Context));
777+
UR_CALL(context.urDdiTable.Device.pfnRetain(Device));
778+
UR_CALL(context.urDdiTable.USM.pfnSharedAlloc(
779+
Context, Device, nullptr, nullptr, sizeof(LaunchInfo), (void **)&Data));
780+
*Data = LaunchInfo{};
781+
return UR_RESULT_SUCCESS;
782+
}
783+
784+
ur_result_t USMLaunchInfo::updateKernelInfo(const KernelInfo &KI) {
785+
auto NumArgs = KI.LocalArgs.size();
786+
if (NumArgs) {
787+
Data->NumLocalArgs = NumArgs;
788+
UR_CALL(context.urDdiTable.USM.pfnSharedAlloc(
789+
Context, Device, nullptr, nullptr, sizeof(LocalArgsInfo) * NumArgs,
790+
(void **)&Data->LocalArgs));
791+
uint32_t i = 0;
792+
for (auto [ArgIndex, ArgInfo] : KI.LocalArgs) {
793+
Data->LocalArgs[i++] = ArgInfo;
794+
context.logger.debug(
795+
"local_args (argIndex={}, size={}, sizeWithRZ={})", ArgIndex,
796+
ArgInfo.Size, ArgInfo.SizeWithRedZone);
797+
}
798+
}
799+
return UR_RESULT_SUCCESS;
800+
}
801+
802+
USMLaunchInfo::~USMLaunchInfo() {
753803
[[maybe_unused]] ur_result_t Result;
754-
if (LocalShadowOffset) {
755-
Result =
756-
context.urDdiTable.USM.pfnFree(Context, (void *)LocalShadowOffset);
804+
if (Data) {
805+
auto Type = GetDeviceType(Device);
806+
if (Type == DeviceType::GPU_PVC) {
807+
if (Data->PrivateShadowOffset) {
808+
Result = context.urDdiTable.USM.pfnFree(
809+
Context, (void *)Data->PrivateShadowOffset);
810+
assert(Result == UR_RESULT_SUCCESS);
811+
}
812+
if (Data->LocalShadowOffset) {
813+
Result = context.urDdiTable.USM.pfnFree(
814+
Context, (void *)Data->LocalShadowOffset);
815+
assert(Result == UR_RESULT_SUCCESS);
816+
}
817+
}
818+
if (Data->LocalArgs) {
819+
Result = context.urDdiTable.USM.pfnFree(Context,
820+
(void *)Data->LocalArgs);
821+
assert(Result == UR_RESULT_SUCCESS);
822+
}
823+
Result = context.urDdiTable.USM.pfnFree(Context, (void *)Data);
757824
assert(Result == UR_RESULT_SUCCESS);
758825
}
759826
Result = context.urDdiTable.Context.pfnRelease(Context);
760827
assert(Result == UR_RESULT_SUCCESS);
828+
Result = context.urDdiTable.Device.pfnRelease(Device);
829+
assert(Result == UR_RESULT_SUCCESS);
761830
}
762831

763832
} // namespace ur_sanitizer_layer

‎source/loader/layers/sanitizer/asan_interceptor.hpp

Lines changed: 53 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,26 @@ struct QueueInfo {
7979
}
8080
};
8181

82+
struct KernelInfo {
83+
ur_kernel_handle_t Handle;
84+
85+
ur_shared_mutex Mutex;
86+
// Need preserve the order of local arguments
87+
std::map<uint32_t, LocalArgsInfo> LocalArgs;
88+
89+
explicit KernelInfo(ur_kernel_handle_t Kernel) : Handle(Kernel) {
90+
[[maybe_unused]] auto Result =
91+
context.urDdiTable.Kernel.pfnRetain(Kernel);
92+
assert(Result == UR_RESULT_SUCCESS);
93+
}
94+
95+
~KernelInfo() {
96+
[[maybe_unused]] auto Result =
97+
context.urDdiTable.Kernel.pfnRelease(Handle);
98+
assert(Result == UR_RESULT_SUCCESS);
99+
}
100+
};
101+
82102
struct ContextInfo {
83103
ur_context_handle_t Handle;
84104

@@ -107,31 +127,30 @@ struct ContextInfo {
107127
}
108128
};
109129

110-
struct LaunchInfo {
111-
uptr LocalShadowOffset = 0;
112-
uptr LocalShadowOffsetEnd = 0;
113-
DeviceSanitizerReport SPIR_DeviceSanitizerReportMem;
130+
struct USMLaunchInfo {
131+
LaunchInfo *Data;
114132

115133
ur_context_handle_t Context = nullptr;
134+
ur_device_handle_t Device = nullptr;
116135
const size_t *GlobalWorkSize = nullptr;
117136
const size_t *GlobalWorkOffset = nullptr;
118137
std::vector<size_t> LocalWorkSize;
119138
uint32_t WorkDim = 0;
120139

121-
LaunchInfo(ur_context_handle_t Context, const size_t *GlobalWorkSize,
122-
const size_t *LocalWorkSize, const size_t *GlobalWorkOffset,
123-
uint32_t WorkDim)
124-
: Context(Context), GlobalWorkSize(GlobalWorkSize),
140+
USMLaunchInfo(ur_context_handle_t Context, ur_device_handle_t Device,
141+
const size_t *GlobalWorkSize, const size_t *LocalWorkSize,
142+
const size_t *GlobalWorkOffset, uint32_t WorkDim)
143+
: Context(Context), Device(Device), GlobalWorkSize(GlobalWorkSize),
125144
GlobalWorkOffset(GlobalWorkOffset), WorkDim(WorkDim) {
126-
[[maybe_unused]] auto Result =
127-
context.urDdiTable.Context.pfnRetain(Context);
128-
assert(Result == UR_RESULT_SUCCESS);
129145
if (LocalWorkSize) {
130146
this->LocalWorkSize =
131147
std::vector<size_t>(LocalWorkSize, LocalWorkSize + WorkDim);
132148
}
133149
}
134-
~LaunchInfo();
150+
~USMLaunchInfo();
151+
152+
ur_result_t initialize();
153+
ur_result_t updateKernelInfo(const KernelInfo &KI);
135154
};
136155

137156
struct DeviceGlobalInfo {
@@ -158,12 +177,11 @@ class SanitizerInterceptor {
158177

159178
ur_result_t preLaunchKernel(ur_kernel_handle_t Kernel,
160179
ur_queue_handle_t Queue,
161-
LaunchInfo &LaunchInfo);
180+
USMLaunchInfo &LaunchInfo);
162181

163182
ur_result_t postLaunchKernel(ur_kernel_handle_t Kernel,
164183
ur_queue_handle_t Queue,
165-
ur_event_handle_t &Event,
166-
LaunchInfo &LaunchInfo);
184+
USMLaunchInfo &LaunchInfo);
167185

168186
ur_result_t insertContext(ur_context_handle_t Context,
169187
std::shared_ptr<ContextInfo> &CI);
@@ -173,6 +191,9 @@ class SanitizerInterceptor {
173191
std::shared_ptr<DeviceInfo> &CI);
174192
ur_result_t eraseDevice(ur_device_handle_t Device);
175193

194+
ur_result_t insertKernel(ur_kernel_handle_t Kernel);
195+
ur_result_t eraseKernel(ur_kernel_handle_t Kernel);
196+
176197
std::optional<AllocationIterator> findAllocInfoByAddress(uptr Address);
177198

178199
std::shared_ptr<ContextInfo> getContextInfo(ur_context_handle_t Context) {
@@ -181,6 +202,18 @@ class SanitizerInterceptor {
181202
return m_ContextMap[Context];
182203
}
183204

205+
std::shared_ptr<DeviceInfo> getDeviceInfo(ur_device_handle_t Device) {
206+
std::shared_lock<ur_shared_mutex> Guard(m_DeviceMapMutex);
207+
assert(m_DeviceMap.find(Device) != m_DeviceMap.end());
208+
return m_DeviceMap[Device];
209+
}
210+
211+
std::shared_ptr<KernelInfo> getKernelInfo(ur_kernel_handle_t Kernel) {
212+
std::shared_lock<ur_shared_mutex> Guard(m_KernelMapMutex);
213+
assert(m_KernelMap.find(Kernel) != m_KernelMap.end());
214+
return m_KernelMap[Kernel];
215+
}
216+
184217
private:
185218
ur_result_t updateShadowMemory(std::shared_ptr<ContextInfo> &ContextInfo,
186219
std::shared_ptr<DeviceInfo> &DeviceInfo,
@@ -195,26 +228,23 @@ class SanitizerInterceptor {
195228
std::shared_ptr<DeviceInfo> &DeviceInfo,
196229
ur_queue_handle_t Queue,
197230
ur_kernel_handle_t Kernel,
198-
LaunchInfo &LaunchInfo);
231+
USMLaunchInfo &LaunchInfo);
199232

200233
ur_result_t allocShadowMemory(ur_context_handle_t Context,
201234
std::shared_ptr<DeviceInfo> &DeviceInfo);
202235

203-
std::shared_ptr<DeviceInfo> getDeviceInfo(ur_device_handle_t Device) {
204-
std::shared_lock<ur_shared_mutex> Guard(m_DeviceMapMutex);
205-
assert(m_DeviceMap.find(Device) != m_DeviceMap.end());
206-
return m_DeviceMap[Device];
207-
}
208-
209236
private:
210237
std::unordered_map<ur_context_handle_t, std::shared_ptr<ContextInfo>>
211238
m_ContextMap;
212239
ur_shared_mutex m_ContextMapMutex;
213-
214240
std::unordered_map<ur_device_handle_t, std::shared_ptr<DeviceInfo>>
215241
m_DeviceMap;
216242
ur_shared_mutex m_DeviceMapMutex;
217243

244+
std::unordered_map<ur_kernel_handle_t, std::shared_ptr<KernelInfo>>
245+
m_KernelMap;
246+
ur_shared_mutex m_KernelMapMutex;
247+
218248
/// Assumption: all USM chunks are allocated in one VA
219249
AllocationMap m_AllocationMap;
220250
ur_shared_mutex m_AllocationMapMutex;

‎source/loader/layers/sanitizer/asan_libdevice.hpp

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,23 @@ struct DeviceSanitizerReport {
6262
bool IsRecover = false;
6363
};
6464

65+
struct LocalArgsInfo {
66+
uint64_t Size = 0;
67+
uint64_t SizeWithRedZone = 0;
68+
};
69+
70+
struct LaunchInfo {
71+
uintptr_t PrivateShadowOffset =
72+
0; // don't move this field, we use it in AddressSanitizerPass
73+
74+
uintptr_t LocalShadowOffset = 0;
75+
uintptr_t LocalShadowOffsetEnd = 0;
76+
DeviceSanitizerReport SanitizerReport;
77+
78+
uint32_t NumLocalArgs = 0;
79+
LocalArgsInfo *LocalArgs = nullptr; // ordered by ArgIndex
80+
};
81+
6582
constexpr unsigned ASAN_SHADOW_SCALE = 3;
6683
constexpr unsigned ASAN_SHADOW_GRANULARITY = 1ULL << ASAN_SHADOW_SCALE;
6784

0 commit comments

Comments
 (0)