Skip to content

Commit 04fd002

Browse files
AllanZynekbenzie
AllanZyne
and
kbenzie
authored
[DeviceMSAN] Check use-of-uninitialized value on dynamic local memory (#17180)
Support check use-of-uninitialized value on dynamic local memory, such as "sycl::local_accessor" --------- Co-authored-by: Kenneth Benzie (Benie) <[email protected]>
1 parent b40c543 commit 04fd002

File tree

10 files changed

+306
-15
lines changed

10 files changed

+306
-15
lines changed

‎libdevice/sanitizer/msan_rtl.cpp‎

Lines changed: 71 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ inline uptr __msan_get_shadow_dg2(uptr addr, uint32_t as) {
160160
if (addr < shadow_begin) {
161161
return addr + (shadow_begin - DG2_DEVICE_USM_BEGIN);
162162
} else {
163-
return addr - (DG2_DEVICE_USM_END - shadow_end);
163+
return addr - (DG2_DEVICE_USM_END - shadow_end + 1);
164164
}
165165
}
166166

@@ -176,7 +176,7 @@ inline uptr __msan_get_shadow_pvc(uptr addr, uint32_t as) {
176176
if (addr < shadow_begin) {
177177
return addr + (shadow_begin - PVC_DEVICE_USM_BEGIN);
178178
} else {
179-
return addr - (PVC_DEVICE_USM_END - shadow_end);
179+
return addr - (PVC_DEVICE_USM_END - shadow_end + 1);
180180
}
181181
} else if (as == ADDRESS_SPACE_LOCAL) {
182182
// The size of SLM is 128KB on PVC
@@ -417,4 +417,73 @@ DEVICE_EXTERN_C_INLINE void __msan_barrier() {
417417
__spv::MemorySemanticsMask::WorkgroupMemory);
418418
}
419419

420+
static __SYCL_CONSTANT__ const char __msan_print_local_arg[] =
421+
"[kernel] local_arg(index=%d, size=%d)\n";
422+
423+
static __SYCL_CONSTANT__ const char
424+
__msan_print_set_shadow_dynamic_local_begin[] =
425+
"[kernel] BEGIN __msan_poison_shadow_dynamic_local\n";
426+
static __SYCL_CONSTANT__ const char
427+
__msan_print_set_shadow_dynamic_local_end[] =
428+
"[kernel] END __msan_poison_shadow_dynamic_local\n";
429+
static __SYCL_CONSTANT__ const char __msan_print_report_arg_count_incorrect[] =
430+
"[kernel] ERROR: The number of local args is incorrect, expect %d, actual "
431+
"%d\n";
432+
433+
DEVICE_EXTERN_C_NOINLINE void
434+
__msan_poison_shadow_dynamic_local(uptr ptr, uint32_t num_args) {
435+
if (!GetMsanLaunchInfo)
436+
return;
437+
438+
MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow_dynamic_local_begin));
439+
440+
if (num_args != GetMsanLaunchInfo->NumLocalArgs) {
441+
__spirv_ocl_printf(__msan_print_report_arg_count_incorrect, num_args,
442+
GetMsanLaunchInfo->NumLocalArgs);
443+
return;
444+
}
445+
446+
uptr *args = (uptr *)ptr;
447+
448+
for (uint32_t i = 0; i < num_args; ++i) {
449+
auto *local_arg = &GetMsanLaunchInfo->LocalArgs[i];
450+
MSAN_DEBUG(__spirv_ocl_printf(__msan_print_local_arg, i, local_arg->Size));
451+
452+
__msan_poison_shadow_static_local(args[i], local_arg->Size);
453+
}
454+
455+
MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow_dynamic_local_end));
456+
}
457+
458+
static __SYCL_CONSTANT__ const char
459+
__mem_unpoison_shadow_dynamic_local_begin[] =
460+
"[kernel] BEGIN __msan_unpoison_shadow_dynamic_local\n";
461+
static __SYCL_CONSTANT__ const char __mem_unpoison_shadow_dynamic_local_end[] =
462+
"[kernel] END __msan_unpoison_shadow_dynamic_local\n";
463+
464+
DEVICE_EXTERN_C_NOINLINE void
465+
__msan_unpoison_shadow_dynamic_local(uptr ptr, uint32_t num_args) {
466+
if (!GetMsanLaunchInfo)
467+
return;
468+
469+
MSAN_DEBUG(__spirv_ocl_printf(__mem_unpoison_shadow_dynamic_local_begin));
470+
471+
if (num_args != GetMsanLaunchInfo->NumLocalArgs) {
472+
__spirv_ocl_printf(__msan_print_report_arg_count_incorrect, num_args,
473+
GetMsanLaunchInfo->NumLocalArgs);
474+
return;
475+
}
476+
477+
uptr *args = (uptr *)ptr;
478+
479+
for (uint32_t i = 0; i < num_args; ++i) {
480+
auto *local_arg = &GetMsanLaunchInfo->LocalArgs[i];
481+
MSAN_DEBUG(__spirv_ocl_printf(__msan_print_local_arg, i, local_arg->Size));
482+
483+
__msan_unpoison_shadow_static_local(args[i], local_arg->Size);
484+
}
485+
486+
MSAN_DEBUG(__spirv_ocl_printf(__mem_unpoison_shadow_dynamic_local_end));
487+
}
488+
420489
#endif // __SPIR__ || __SPIRV__

‎llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp‎

Lines changed: 88 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -777,9 +777,11 @@ class MemorySanitizerOnSpirv {
777777
IsSPIRV = TargetTriple.isSPIROrSPIRV();
778778

779779
IntptrTy = DL.getIntPtrType(C);
780+
Int32Ty = Type::getInt32Ty(C);
780781
}
781782

782783
bool instrumentModule();
784+
void instrumentFunction(Function &F);
783785

784786
Constant *getOrCreateGlobalString(StringRef Name, StringRef Value,
785787
unsigned AddressSpace);
@@ -788,6 +790,7 @@ class MemorySanitizerOnSpirv {
788790
void initializeCallbacks();
789791
void instrumentGlobalVariables();
790792
void instrumentStaticLocalMemory();
793+
void instrumentDynamicLocalMemory(Function &F);
791794
void instrumentKernelsMetadata();
792795

793796
void initializeRetVecMap(Function *F);
@@ -799,15 +802,25 @@ class MemorySanitizerOnSpirv {
799802
const DataLayout &DL;
800803
bool IsSPIRV;
801804
Type *IntptrTy;
805+
Type *Int32Ty;
802806

803807
StringMap<GlobalVariable *> GlobalStringMap;
804808

805809
DenseMap<Function *, SmallVector<Instruction *, 8>> KernelToRetVecMap;
806810
DenseMap<Function *, SmallVector<Constant *, 8>> KernelToLocalMemMap;
807811
DenseMap<Function *, DenseSet<Function *>> FuncToKernelCallerMap;
808812

813+
// Make sure that we insert barriers only once per function, and the barrier
814+
// needs to be inserted after all "MsanPoisonShadowStaticLocalFunc" and
815+
// "MsanPoisonShadowDynamicLocalFunc", and before
816+
// "MsanUnpoisonShadowStaticLocalFunc" and
817+
// "MsanUnpoisonShadowDynamicLocalFunc".
818+
DenseMap<Function *, bool> InsertBarrier;
819+
809820
FunctionCallee MsanPoisonShadowStaticLocalFunc;
810821
FunctionCallee MsanUnpoisonShadowStaticLocalFunc;
822+
FunctionCallee MsanPoisonShadowDynamicLocalFunc;
823+
FunctionCallee MsanUnpoisonShadowDynamicLocalFunc;
811824
FunctionCallee MsanBarrierFunc;
812825
};
813826

@@ -874,6 +887,21 @@ void MemorySanitizerOnSpirv::initializeCallbacks() {
874887
M.getOrInsertFunction("__msan_unpoison_shadow_static_local",
875888
IRB.getVoidTy(), IntptrTy, IntptrTy);
876889

890+
// __asan_poison_shadow_dynamic_local(
891+
// uptr ptr,
892+
// uint32_t num_args
893+
// )
894+
MsanPoisonShadowDynamicLocalFunc = M.getOrInsertFunction(
895+
"__msan_poison_shadow_dynamic_local", IRB.getVoidTy(), IntptrTy, Int32Ty);
896+
897+
// __asan_unpoison_shadow_dynamic_local(
898+
// uptr ptr,
899+
// uint32_t num_args
900+
// )
901+
MsanUnpoisonShadowDynamicLocalFunc =
902+
M.getOrInsertFunction("__msan_unpoison_shadow_dynamic_local",
903+
IRB.getVoidTy(), IntptrTy, Int32Ty);
904+
877905
// __msan_barrier()
878906
MsanBarrierFunc = M.getOrInsertFunction("__msan_barrier", IRB.getVoidTy());
879907
}
@@ -951,16 +979,15 @@ void MemorySanitizerOnSpirv::instrumentStaticLocalMemory() {
951979
if (!ClSpirOffloadLocals)
952980
return;
953981

954-
DenseMap<Function *, bool> InsertBarrier;
955-
956-
auto Instrument = [this, &InsertBarrier](GlobalVariable *G, Function *F) {
982+
auto Instrument = [this](GlobalVariable *G, Function *F) {
957983
const uint64_t SizeInBytes = DL.getTypeAllocSize(G->getValueType());
958984

959-
// Poison shadow of static local memory
960985
if (!InsertBarrier[F]) {
961986
IRBuilder<> Builder(&F->getEntryBlock().front());
962987
Builder.CreateCall(MsanBarrierFunc);
963988
}
989+
990+
// Poison shadow of static local memory
964991
IRBuilder<> Builder(&F->getEntryBlock().front());
965992
Builder.CreateCall(MsanPoisonShadowStaticLocalFunc,
966993
{Builder.CreatePointerCast(G, IntptrTy),
@@ -1001,6 +1028,54 @@ void MemorySanitizerOnSpirv::instrumentStaticLocalMemory() {
10011028
}
10021029
}
10031030

1031+
void MemorySanitizerOnSpirv::instrumentDynamicLocalMemory(Function &F) {
1032+
if (!ClSpirOffloadLocals)
1033+
return;
1034+
1035+
// Poison shadow of local memory in kernel argument, required by CPU device
1036+
SmallVector<Argument *> LocalArgs;
1037+
for (auto &Arg : F.args()) {
1038+
Type *PtrTy = dyn_cast<PointerType>(Arg.getType()->getScalarType());
1039+
if (PtrTy && PtrTy->getPointerAddressSpace() == kSpirOffloadLocalAS)
1040+
LocalArgs.push_back(&Arg);
1041+
}
1042+
1043+
if (LocalArgs.empty())
1044+
return;
1045+
1046+
if (!InsertBarrier[&F]) {
1047+
IRBuilder<> Builder(&F.getEntryBlock().front());
1048+
Builder.CreateCall(MsanBarrierFunc);
1049+
}
1050+
1051+
IRBuilder<> IRB(&F.getEntryBlock().front());
1052+
1053+
AllocaInst *ArgsArray = IRB.CreateAlloca(
1054+
IntptrTy, ConstantInt::get(Int32Ty, LocalArgs.size()), "local_args");
1055+
for (size_t i = 0; i < LocalArgs.size(); i++) {
1056+
auto *StoreDest =
1057+
IRB.CreateGEP(IntptrTy, ArgsArray, ConstantInt::get(Int32Ty, i));
1058+
IRB.CreateStore(IRB.CreatePointerCast(LocalArgs[i], IntptrTy), StoreDest);
1059+
}
1060+
1061+
auto *ArgsArrayAddr = IRB.CreatePointerCast(ArgsArray, IntptrTy);
1062+
IRB.CreateCall(MsanPoisonShadowDynamicLocalFunc,
1063+
{ArgsArrayAddr, ConstantInt::get(Int32Ty, LocalArgs.size())});
1064+
1065+
// Unpoison shadow of dynamic local memory, required by CPU device
1066+
initializeRetVecMap(&F);
1067+
for (Instruction *Ret : KernelToRetVecMap[&F]) {
1068+
IRBuilder<> IRBRet(Ret);
1069+
if (!InsertBarrier[&F])
1070+
IRBRet.CreateCall(MsanBarrierFunc);
1071+
IRBRet.CreateCall(
1072+
MsanUnpoisonShadowDynamicLocalFunc,
1073+
{ArgsArrayAddr, ConstantInt::get(Int32Ty, LocalArgs.size())});
1074+
}
1075+
1076+
InsertBarrier[&F] = true;
1077+
}
1078+
10041079
// Instrument __MsanKernelMetadata, which records information of sanitized
10051080
// kernel
10061081
void MemorySanitizerOnSpirv::instrumentKernelsMetadata() {
@@ -1087,6 +1162,14 @@ bool MemorySanitizerOnSpirv::instrumentModule() {
10871162
return true;
10881163
}
10891164

1165+
void MemorySanitizerOnSpirv::instrumentFunction(Function &F) {
1166+
if (!IsSPIRV)
1167+
return;
1168+
1169+
if (F.getCallingConv() == CallingConv::SPIR_KERNEL)
1170+
instrumentDynamicLocalMemory(F);
1171+
}
1172+
10901173
PreservedAnalyses MemorySanitizerPass::run(Module &M,
10911174
ModuleAnalysisManager &AM) {
10921175
// Return early if nosanitize_memory module flag is present for the module.
@@ -1110,6 +1193,7 @@ PreservedAnalyses MemorySanitizerPass::run(Module &M,
11101193
MemorySanitizer Msan(*F.getParent(), MsanSpirv, Options);
11111194
Modified |=
11121195
Msan.sanitizeFunction(F, FAM.getResult<TargetLibraryAnalysis>(F));
1196+
MsanSpirv.instrumentFunction(F);
11131197
}
11141198

11151199
if (!Modified)

‎llvm/test/Instrumentation/MemorySanitizer/SPIRV/instrument_dynamic_local_mem.ll‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-locals=1 -S | FileCheck %s
2+
target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
3+
target triple = "spir64-unknown-unknown"
4+
5+
@WGLocal = internal addrspace(3) global i64 zeroinitializer, align 8
6+
7+
define spir_kernel void @MyKernel(ptr addrspace(3) noundef align 4 %_arg_acc) sanitize_memory {
8+
; CHECK-LABEL: @MyKernel
9+
entry:
10+
; CHECK: %local_args = alloca i64, align 8
11+
; CHECK-NEXT: %0 = getelementptr i64, ptr %local_args, i32 0
12+
; CHECK-NEXT: %1 = ptrtoint ptr addrspace(3) %_arg_acc to i64
13+
; CHECK-NEXT: store i64 %1, ptr %0, align 8
14+
; CHECK-NEXT: %2 = ptrtoint ptr %local_args to i64
15+
; CHECK-NEXT: call void @__msan_poison_shadow_dynamic_local(i64 %2, i32 1)
16+
17+
; CHECK: @__msan_poison_shadow_static_local{{.*}}@WGLocal
18+
; CHECK: @__msan_barrier
19+
store i32 0, ptr addrspace(3) @WGLocal, align 8
20+
; CHECK: @__msan_barrier
21+
; CHECK: @__msan_unpoison_shadow_static_local{{.*}}@WGLocal
22+
23+
; CHECK: call void @__msan_unpoison_shadow_dynamic_local(i64 %2, i32 1)
24+
ret void
25+
}

‎sycl/test-e2e/MemorySanitizer/local/local_accessor.cpp‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// REQUIRES: linux, cpu || (gpu && level_zero)
2+
// RUN: %{build} %device_msan_flags -g -O0 -o %t1.out
3+
// RUN: %{run} not %t1.out 2>&1 | FileCheck %s
4+
// RUN: %{build} %device_msan_flags -g -O1 -o %t2.out
5+
// RUN: %{run} not %t2.out 2>&1 | FileCheck %s
6+
// RUN: %{build} %device_msan_flags -g -O2 -o %t3.out
7+
// RUN: %{run} not %t3.out 2>&1 | FileCheck %s
8+
9+
// XFAIL: spirv-backend && gpu && run-mode
10+
// XFAIL-TRACKER: https://github.com/llvm/llvm-project/issues/122075
11+
12+
#include <sycl/ext/oneapi/group_local_memory.hpp>
13+
#include <sycl/usm.hpp>
14+
15+
constexpr std::size_t global_size = 4;
16+
constexpr std::size_t local_size = 1;
17+
18+
__attribute__((noinline)) int check(int data, sycl::nd_item<1> &item) {
19+
auto ptr =
20+
sycl::ext::oneapi::group_local_memory<int[global_size]>(item.get_group());
21+
auto &ref = *ptr;
22+
return data + ref[0];
23+
}
24+
25+
int main() {
26+
sycl::queue Q;
27+
28+
Q.submit([&](sycl::handler &cgh) {
29+
auto acc = sycl::local_accessor<int>(local_size, cgh);
30+
cgh.parallel_for<class MyKernel>(
31+
sycl::nd_range<1>(global_size, local_size),
32+
[=](sycl::nd_item<1> item) { check(acc[item.get_local_id()], item); });
33+
});
34+
Q.wait();
35+
// CHECK-NOT: [kernel]
36+
// CHECK: DeviceSanitizer: use-of-uninitialized-value
37+
// CHECK: #0 {{.*}} {{.*local_accessor.cpp}}:[[@LINE-5]]
38+
39+
return 0;
40+
}

‎unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1389,6 +1389,31 @@ ur_result_t urKernelSetArgMemObj(
13891389
return UR_RESULT_SUCCESS;
13901390
}
13911391

1392+
///////////////////////////////////////////////////////////////////////////////
1393+
/// @brief Intercept function for urKernelSetArgLocal
1394+
__urdlllocal ur_result_t UR_APICALL urKernelSetArgLocal(
1395+
/// [in] handle of the kernel object
1396+
ur_kernel_handle_t hKernel,
1397+
/// [in] argument index in range [0, num args - 1]
1398+
uint32_t argIndex,
1399+
/// [in] size of the local buffer to be allocated by the runtime
1400+
size_t argSize,
1401+
/// [in][optional] pointer to local buffer properties.
1402+
const ur_kernel_arg_local_properties_t *pProperties) {
1403+
auto pfnSetArgLocal = getContext()->urDdiTable.Kernel.pfnSetArgLocal;
1404+
1405+
getContext()->logger.debug(
1406+
"==== urKernelSetArgLocal (argIndex={}, argSize={})", argIndex, argSize);
1407+
1408+
{
1409+
auto &KI = getMsanInterceptor()->getOrCreateKernelInfo(hKernel);
1410+
std::scoped_lock<ur_shared_mutex> Guard(KI.Mutex);
1411+
KI.LocalArgs[argIndex] = MsanLocalArgsInfo{argSize};
1412+
}
1413+
1414+
return pfnSetArgLocal(hKernel, argIndex, argSize, pProperties);
1415+
}
1416+
13921417
///////////////////////////////////////////////////////////////////////////////
13931418
/// @brief Intercept function for urEnqueueUSMFill
13941419
ur_result_t UR_APICALL urEnqueueUSMFill(
@@ -1738,6 +1763,7 @@ ur_result_t urGetKernelProcAddrTable(
17381763
pDdiTable->pfnRelease = ur_sanitizer_layer::msan::urKernelRelease;
17391764
pDdiTable->pfnSetArgValue = ur_sanitizer_layer::msan::urKernelSetArgValue;
17401765
pDdiTable->pfnSetArgMemObj = ur_sanitizer_layer::msan::urKernelSetArgMemObj;
1766+
pDdiTable->pfnSetArgLocal = ur_sanitizer_layer::msan::urKernelSetArgLocal;
17411767

17421768
return result;
17431769
}

0 commit comments

Comments
 (0)