Skip to content

Commit f2af85f

Browse files
AllanZyneAllanZyne
authored
[DeviceSanitizer] Support "-fsanitize-ignorelist=" to disable sanitizing on some of kernels (#2055)
1 parent 49a3d6c commit f2af85f

File tree

4 files changed

+165
-48
lines changed

4 files changed

+165
-48
lines changed

‎source/loader/layers/sanitizer/asan/asan_ddi.cpp‎

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,12 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
5252
return UR_RESULT_SUCCESS;
5353
}
5454

55+
bool isInstrumentedKernel(ur_kernel_handle_t hKernel) {
56+
auto hProgram = GetProgram(hKernel);
57+
auto PI = getAsanInterceptor()->getProgramInfo(hProgram);
58+
return PI->isKernelInstrumented(hKernel);
59+
}
60+
5561
} // namespace
5662

5763
///////////////////////////////////////////////////////////////////////////////
@@ -307,7 +313,7 @@ __urdlllocal ur_result_t UR_APICALL urProgramBuild(
307313

308314
UR_CALL(pfnProgramBuild(hContext, hProgram, pOptions));
309315

310-
UR_CALL(getAsanInterceptor()->registerProgram(hContext, hProgram));
316+
UR_CALL(getAsanInterceptor()->registerProgram(hProgram));
311317

312318
return UR_RESULT_SUCCESS;
313319
}
@@ -331,8 +337,7 @@ __urdlllocal ur_result_t UR_APICALL urProgramBuildExp(
331337
getContext()->logger.debug("==== urProgramBuildExp");
332338

333339
UR_CALL(pfnBuildExp(hProgram, numDevices, phDevices, pOptions));
334-
UR_CALL(
335-
getAsanInterceptor()->registerProgram(GetContext(hProgram), hProgram));
340+
UR_CALL(getAsanInterceptor()->registerProgram(hProgram));
336341

337342
return UR_RESULT_SUCCESS;
338343
}
@@ -359,7 +364,7 @@ __urdlllocal ur_result_t UR_APICALL urProgramLink(
359364

360365
UR_CALL(pfnProgramLink(hContext, count, phPrograms, pOptions, phProgram));
361366

362-
UR_CALL(getAsanInterceptor()->registerProgram(hContext, *phProgram));
367+
UR_CALL(getAsanInterceptor()->registerProgram(*phProgram));
363368

364369
return UR_RESULT_SUCCESS;
365370
}
@@ -390,7 +395,7 @@ ur_result_t UR_APICALL urProgramLinkExp(
390395
UR_CALL(pfnProgramLinkExp(hContext, numDevices, phDevices, count,
391396
phPrograms, pOptions, phProgram));
392397

393-
UR_CALL(getAsanInterceptor()->registerProgram(hContext, *phProgram));
398+
UR_CALL(getAsanInterceptor()->registerProgram(*phProgram));
394399

395400
return UR_RESULT_SUCCESS;
396401
}
@@ -460,7 +465,13 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueKernelLaunch(
460465

461466
getContext()->logger.debug("==== urEnqueueKernelLaunch");
462467

463-
USMLaunchInfo LaunchInfo(GetContext(hQueue), GetDevice(hQueue),
468+
if (!isInstrumentedKernel(hKernel)) {
469+
return pfnKernelLaunch(hQueue, hKernel, workDim, pGlobalWorkOffset,
470+
pGlobalWorkSize, pLocalWorkSize,
471+
numEventsInWaitList, phEventWaitList, phEvent);
472+
}
473+
474+
USMLaunchInfo LaunchInfo(GetContext(hKernel), GetDevice(hQueue),
464475
pGlobalWorkSize, pLocalWorkSize, pGlobalWorkOffset,
465476
workDim);
466477
UR_CALL(LaunchInfo.initialize());
@@ -1351,7 +1362,9 @@ __urdlllocal ur_result_t UR_APICALL urKernelCreate(
13511362
getContext()->logger.debug("==== urKernelCreate");
13521363

13531364
UR_CALL(pfnCreate(hProgram, pKernelName, phKernel));
1354-
UR_CALL(getAsanInterceptor()->insertKernel(*phKernel));
1365+
if (isInstrumentedKernel(*phKernel)) {
1366+
UR_CALL(getAsanInterceptor()->insertKernel(*phKernel));
1367+
}
13551368

13561369
return UR_RESULT_SUCCESS;
13571370
}
@@ -1372,8 +1385,9 @@ __urdlllocal ur_result_t UR_APICALL urKernelRetain(
13721385
UR_CALL(pfnRetain(hKernel));
13731386

13741387
auto KernelInfo = getAsanInterceptor()->getKernelInfo(hKernel);
1375-
UR_ASSERT(KernelInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
1376-
KernelInfo->RefCount++;
1388+
if (KernelInfo) {
1389+
KernelInfo->RefCount++;
1390+
}
13771391

13781392
return UR_RESULT_SUCCESS;
13791393
}
@@ -1393,9 +1407,10 @@ __urdlllocal ur_result_t urKernelRelease(
13931407
UR_CALL(pfnRelease(hKernel));
13941408

13951409
auto KernelInfo = getAsanInterceptor()->getKernelInfo(hKernel);
1396-
UR_ASSERT(KernelInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
1397-
if (--KernelInfo->RefCount == 0) {
1398-
UR_CALL(getAsanInterceptor()->eraseKernel(hKernel));
1410+
if (KernelInfo) {
1411+
if (--KernelInfo->RefCount == 0) {
1412+
UR_CALL(getAsanInterceptor()->eraseKernel(hKernel));
1413+
}
13991414
}
14001415

14011416
return UR_RESULT_SUCCESS;
@@ -1421,10 +1436,11 @@ __urdlllocal ur_result_t UR_APICALL urKernelSetArgValue(
14211436
getContext()->logger.debug("==== urKernelSetArgValue");
14221437

14231438
std::shared_ptr<MemBuffer> MemBuffer;
1439+
std::shared_ptr<KernelInfo> KernelInfo;
14241440
if (argSize == sizeof(ur_mem_handle_t) &&
14251441
(MemBuffer = getAsanInterceptor()->getMemBuffer(
1426-
*ur_cast<const ur_mem_handle_t *>(pArgValue)))) {
1427-
auto KernelInfo = getAsanInterceptor()->getKernelInfo(hKernel);
1442+
*ur_cast<const ur_mem_handle_t *>(pArgValue))) &&
1443+
(KernelInfo = getAsanInterceptor()->getKernelInfo(hKernel))) {
14281444
std::scoped_lock<ur_shared_mutex> Guard(KernelInfo->Mutex);
14291445
KernelInfo->BufferArgs[argIndex] = std::move(MemBuffer);
14301446
} else {
@@ -1452,8 +1468,10 @@ __urdlllocal ur_result_t UR_APICALL urKernelSetArgMemObj(
14521468

14531469
getContext()->logger.debug("==== urKernelSetArgMemObj");
14541470

1455-
if (auto MemBuffer = getAsanInterceptor()->getMemBuffer(hArgValue)) {
1456-
auto KernelInfo = getAsanInterceptor()->getKernelInfo(hKernel);
1471+
std::shared_ptr<MemBuffer> MemBuffer;
1472+
std::shared_ptr<KernelInfo> KernelInfo;
1473+
if ((MemBuffer = getAsanInterceptor()->getMemBuffer(hArgValue)) &&
1474+
(KernelInfo = getAsanInterceptor()->getKernelInfo(hKernel))) {
14571475
std::scoped_lock<ur_shared_mutex> Guard(KernelInfo->Mutex);
14581476
KernelInfo->BufferArgs[argIndex] = std::move(MemBuffer);
14591477
} else {
@@ -1483,8 +1501,7 @@ __urdlllocal ur_result_t UR_APICALL urKernelSetArgLocal(
14831501
"==== urKernelSetArgLocal (argIndex={}, argSize={})", argIndex,
14841502
argSize);
14851503

1486-
{
1487-
auto KI = getAsanInterceptor()->getKernelInfo(hKernel);
1504+
if (auto KI = getAsanInterceptor()->getKernelInfo(hKernel)) {
14881505
std::scoped_lock<ur_shared_mutex> Guard(KI->Mutex);
14891506
// TODO: get local variable alignment
14901507
auto argSizeWithRZ = GetSizeAndRedzoneSizeForLocal(
@@ -1520,8 +1537,9 @@ __urdlllocal ur_result_t UR_APICALL urKernelSetArgPointer(
15201537
"==== urKernelSetArgPointer (argIndex={}, pArgValue={})", argIndex,
15211538
pArgValue);
15221539

1523-
if (getAsanInterceptor()->getOptions().DetectKernelArguments) {
1524-
auto KI = getAsanInterceptor()->getKernelInfo(hKernel);
1540+
std::shared_ptr<KernelInfo> KI;
1541+
if (getAsanInterceptor()->getOptions().DetectKernelArguments &&
1542+
(KI = getAsanInterceptor()->getKernelInfo(hKernel))) {
15251543
std::scoped_lock<ur_shared_mutex> Guard(KI->Mutex);
15261544
KI->PointerArgs[argIndex] = {pArgValue, GetCurrentBacktrace()};
15271545
}

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

Lines changed: 109 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,7 @@ ur_result_t AsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
255255
auto ContextInfo = getContextInfo(Context);
256256
auto DeviceInfo = getDeviceInfo(Device);
257257
auto KernelInfo = getKernelInfo(Kernel);
258+
assert(KernelInfo && "Kernel should be instrumented");
258259

259260
UR_CALL(LaunchInfo.updateKernelInfo(*KernelInfo.get()));
260261

@@ -412,10 +413,107 @@ AsanInterceptor::updateShadowMemory(std::shared_ptr<ContextInfo> &ContextInfo,
412413
return UR_RESULT_SUCCESS;
413414
}
414415

415-
ur_result_t AsanInterceptor::registerProgram(ur_context_handle_t Context,
416-
ur_program_handle_t Program) {
416+
ur_result_t AsanInterceptor::registerProgram(ur_program_handle_t Program) {
417+
ur_result_t Result = UR_RESULT_SUCCESS;
418+
419+
getContext()->logger.info("registerSpirKernels");
420+
Result = registerSpirKernels(Program);
421+
if (Result != UR_RESULT_SUCCESS) {
422+
return Result;
423+
}
424+
425+
getContext()->logger.info("registerDeviceGlobals");
426+
Result = registerDeviceGlobals(Program);
427+
if (Result != UR_RESULT_SUCCESS) {
428+
return Result;
429+
}
430+
431+
return Result;
432+
}
433+
434+
ur_result_t AsanInterceptor::unregisterProgram(ur_program_handle_t Program) {
435+
auto ProgramInfo = getProgramInfo(Program);
436+
437+
for (auto AI : ProgramInfo->AllocInfoForGlobals) {
438+
UR_CALL(getDeviceInfo(AI->Device)->Shadow->ReleaseShadow(AI));
439+
m_AllocationMap.erase(AI->AllocBegin);
440+
}
441+
ProgramInfo->AllocInfoForGlobals.clear();
442+
443+
ProgramInfo->InstrumentedKernels.clear();
444+
445+
return UR_RESULT_SUCCESS;
446+
}
447+
448+
ur_result_t AsanInterceptor::registerSpirKernels(ur_program_handle_t Program) {
449+
auto Context = GetContext(Program);
417450
std::vector<ur_device_handle_t> Devices = GetDevices(Program);
418451

452+
for (auto Device : Devices) {
453+
size_t MetadataSize;
454+
void *MetadataPtr;
455+
ur_result_t Result =
456+
getContext()->urDdiTable.Program.pfnGetGlobalVariablePointer(
457+
Device, Program, kSPIR_AsanSpirKernelMetadata, &MetadataSize,
458+
&MetadataPtr);
459+
if (Result != UR_RESULT_SUCCESS) {
460+
getContext()->logger.error(
461+
"Can't get the pointer of <{}> under device {}: {}",
462+
kSPIR_AsanSpirKernelMetadata, (void *)Device, Result);
463+
return Result;
464+
}
465+
466+
const uint64_t NumOfSpirKernel = MetadataSize / sizeof(SpirKernelInfo);
467+
assert((MetadataSize % sizeof(SpirKernelInfo) == 0) &&
468+
"SpirKernelMetadata size is not correct");
469+
470+
ManagedQueue Queue(Context, Device);
471+
472+
std::vector<SpirKernelInfo> SKInfo(NumOfSpirKernel);
473+
Result = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy(
474+
Queue, true, &SKInfo[0], MetadataPtr,
475+
sizeof(SpirKernelInfo) * NumOfSpirKernel, 0, nullptr, nullptr);
476+
if (Result != UR_RESULT_SUCCESS) {
477+
getContext()->logger.error("Can't read the value of <{}>: {}",
478+
kSPIR_AsanSpirKernelMetadata, Result);
479+
return Result;
480+
}
481+
482+
auto PI = getProgramInfo(Program);
483+
for (const auto &SKI : SKInfo) {
484+
if (SKI.Size == 0) {
485+
continue;
486+
}
487+
std::vector<char> KernelNameV(SKI.Size);
488+
Result = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy(
489+
Queue, true, KernelNameV.data(), (void *)SKI.KernelName,
490+
sizeof(char) * SKI.Size, 0, nullptr, nullptr);
491+
if (Result != UR_RESULT_SUCCESS) {
492+
getContext()->logger.error("Can't read kernel name: {}",
493+
Result);
494+
return Result;
495+
}
496+
497+
std::string KernelName =
498+
std::string(KernelNameV.begin(), KernelNameV.end());
499+
500+
getContext()->logger.info(
501+
"SpirKernel(name='{}', isInstrumented={})", KernelName, true);
502+
503+
PI->InstrumentedKernels.insert(KernelName);
504+
}
505+
getContext()->logger.info("Number of sanitized kernel: {}",
506+
PI->InstrumentedKernels.size());
507+
}
508+
509+
return UR_RESULT_SUCCESS;
510+
}
511+
512+
ur_result_t
513+
AsanInterceptor::registerDeviceGlobals(ur_program_handle_t Program) {
514+
std::vector<ur_device_handle_t> Devices = GetDevices(Program);
515+
assert(Devices.size() != 0 && "No devices in registerDeviceGlobals");
516+
auto Context = GetContext(Program);
419517
auto ContextInfo = getContextInfo(Context);
420518
auto ProgramInfo = getProgramInfo(Program);
421519

@@ -461,33 +559,16 @@ ur_result_t AsanInterceptor::registerProgram(ur_context_handle_t Context,
461559
{}});
462560

463561
ContextInfo->insertAllocInfo({Device}, AI);
562+
ProgramInfo->AllocInfoForGlobals.emplace(AI);
464563

465-
{
466-
std::scoped_lock<ur_shared_mutex, ur_shared_mutex> Guard(
467-
m_AllocationMapMutex, ProgramInfo->Mutex);
468-
ProgramInfo->AllocInfoForGlobals.emplace(AI);
469-
m_AllocationMap.emplace(AI->AllocBegin, std::move(AI));
470-
}
564+
std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
565+
m_AllocationMap.emplace(AI->AllocBegin, std::move(AI));
471566
}
472567
}
473568

474569
return UR_RESULT_SUCCESS;
475570
}
476571

477-
ur_result_t AsanInterceptor::unregisterProgram(ur_program_handle_t Program) {
478-
auto ProgramInfo = getProgramInfo(Program);
479-
480-
std::scoped_lock<ur_shared_mutex, ur_shared_mutex> Guard(
481-
m_AllocationMapMutex, ProgramInfo->Mutex);
482-
for (auto AI : ProgramInfo->AllocInfoForGlobals) {
483-
UR_CALL(getDeviceInfo(AI->Device)->Shadow->ReleaseShadow(AI));
484-
m_AllocationMap.erase(AI->AllocBegin);
485-
}
486-
ProgramInfo->AllocInfoForGlobals.clear();
487-
488-
return UR_RESULT_SUCCESS;
489-
}
490-
491572
ur_result_t AsanInterceptor::insertContext(ur_context_handle_t Context,
492573
std::shared_ptr<ContextInfo> &CI) {
493574
std::scoped_lock<ur_shared_mutex> Guard(m_ContextMapMutex);
@@ -611,6 +692,7 @@ ur_result_t AsanInterceptor::prepareLaunch(
611692

612693
do {
613694
auto KernelInfo = getKernelInfo(Kernel);
695+
assert(KernelInfo && "Kernel should be instrumented");
614696

615697
// Validate pointer arguments
616698
if (getOptions().DetectKernelArguments) {
@@ -837,6 +919,11 @@ AsanInterceptor::findAllocInfoByContext(ur_context_handle_t Context) {
837919
return AllocInfos;
838920
}
839921

922+
bool ProgramInfo::isKernelInstrumented(ur_kernel_handle_t Kernel) const {
923+
const auto Name = GetKernelName(Kernel);
924+
return InstrumentedKernels.find(Name) != InstrumentedKernels.end();
925+
}
926+
840927
ContextInfo::~ContextInfo() {
841928
Stats.Print(Handle);
842929

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

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -111,9 +111,9 @@ struct ProgramInfo {
111111
ur_program_handle_t Handle;
112112
std::atomic<int32_t> RefCount = 1;
113113

114-
// lock this mutex if following fields are accessed
115-
ur_shared_mutex Mutex;
114+
// Program is built only once, so we don't need to lock it
116115
std::unordered_set<std::shared_ptr<AllocInfo>> AllocInfoForGlobals;
116+
std::unordered_set<std::string> InstrumentedKernels;
117117

118118
explicit ProgramInfo(ur_program_handle_t Program) : Handle(Program) {
119119
[[maybe_unused]] auto Result =
@@ -126,6 +126,8 @@ struct ProgramInfo {
126126
getContext()->urDdiTable.Program.pfnRelease(Handle);
127127
assert(Result == UR_RESULT_SUCCESS);
128128
}
129+
130+
bool isKernelInstrumented(ur_kernel_handle_t Kernel) const;
129131
};
130132

131133
struct ContextInfo {
@@ -187,6 +189,11 @@ struct DeviceGlobalInfo {
187189
uptr Addr;
188190
};
189191

192+
struct SpirKernelInfo {
193+
uptr KernelName;
194+
uptr Size;
195+
};
196+
190197
class AsanInterceptor {
191198
public:
192199
explicit AsanInterceptor();
@@ -200,8 +207,7 @@ class AsanInterceptor {
200207
AllocType Type, void **ResultPtr);
201208
ur_result_t releaseMemory(ur_context_handle_t Context, void *Ptr);
202209

203-
ur_result_t registerProgram(ur_context_handle_t Context,
204-
ur_program_handle_t Program);
210+
ur_result_t registerProgram(ur_program_handle_t Program);
205211

206212
ur_result_t unregisterProgram(ur_program_handle_t Program);
207213

@@ -266,8 +272,10 @@ class AsanInterceptor {
266272

267273
std::shared_ptr<KernelInfo> getKernelInfo(ur_kernel_handle_t Kernel) {
268274
std::shared_lock<ur_shared_mutex> Guard(m_KernelMapMutex);
269-
assert(m_KernelMap.find(Kernel) != m_KernelMap.end());
270-
return m_KernelMap[Kernel];
275+
if (m_KernelMap.find(Kernel) != m_KernelMap.end()) {
276+
return m_KernelMap[Kernel];
277+
}
278+
return nullptr;
271279
}
272280

273281
const AsanOptions &getOptions() { return m_Options; }
@@ -298,6 +306,9 @@ class AsanInterceptor {
298306
ur_result_t allocShadowMemory(ur_context_handle_t Context,
299307
std::shared_ptr<DeviceInfo> &DeviceInfo);
300308

309+
ur_result_t registerDeviceGlobals(ur_program_handle_t Program);
310+
ur_result_t registerSpirKernels(ur_program_handle_t Program);
311+
301312
private:
302313
std::unordered_map<ur_context_handle_t, std::shared_ptr<ContextInfo>>
303314
m_ContextMap;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ const int kPrivateMidRedzoneMagic = (char)0xf2;
9999
const int kPrivateRightRedzoneMagic = (char)0xf3;
100100

101101
constexpr auto kSPIR_AsanDeviceGlobalMetadata = "__AsanDeviceGlobalMetadata";
102+
constexpr auto kSPIR_AsanSpirKernelMetadata = "__AsanKernelMetadata";
102103

103104
#if !defined(__SPIR__) && !defined(__SPIRV__)
104105
} // namespace ur_sanitizer_layer

0 commit comments

Comments
 (0)