Skip to content

Commit 8bbdcde

Browse files
AllanZyneAllanZyne
authored
[DeviceASAN] Fix sycl::group_local_memory (#17769)
1 parent c237313 commit 8bbdcde

File tree

4 files changed

+256
-125
lines changed

4 files changed

+256
-125
lines changed

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

Lines changed: 196 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -839,8 +839,6 @@ struct AddressSanitizer {
839839
bool maybeInsertAsanInitAtFunctionEntry(Function &F);
840840
bool maybeInsertDynamicShadowAtFunctionEntry(Function &F);
841841
void markEscapedLocalAllocas(Function &F);
842-
void instrumentSyclStaticLocalMemory(CallInst *CI,
843-
ArrayRef<Instruction *> RetVec);
844842
bool instrumentSyclDynamicLocalMemory(Function &F,
845843
ArrayRef<Instruction *> RetVec);
846844
void instrumentInitAsanLaunchInfo(Function &F, const TargetLibraryInfo *TLI);
@@ -889,8 +887,6 @@ struct AddressSanitizer {
889887
ShadowMapping Mapping;
890888
FunctionCallee AsanHandleNoReturnFunc;
891889
FunctionCallee AsanPtrCmpFunction, AsanPtrSubFunction;
892-
FunctionCallee AsanSetShadowStaticLocalFunc;
893-
FunctionCallee AsanUnpoisonShadowStaticLocalFunc;
894890
FunctionCallee AsanSetShadowDynamicLocalFunc;
895891
FunctionCallee AsanUnpoisonShadowDynamicLocalFunc;
896892
Constant *AsanShadowGlobal;
@@ -972,6 +968,10 @@ class ModuleAddressSanitizer {
972968
void initializeCallbacks();
973969

974970
void instrumentDeviceGlobal(IRBuilder<> &IRB);
971+
void instrumentSyclStaticLocalMemory(IRBuilder<> &IRB);
972+
void initializeRetVecMap(Function *F);
973+
void initializeKernelCallerMap(Function *F);
974+
975975
void instrumentGlobals(IRBuilder<> &IRB, bool *CtorComdat);
976976
void InstrumentGlobalsCOFF(IRBuilder<> &IRB,
977977
ArrayRef<GlobalVariable *> ExtendedGlobals,
@@ -1030,10 +1030,15 @@ class ModuleAddressSanitizer {
10301030
FunctionCallee AsanUnregisterImageGlobals;
10311031
FunctionCallee AsanRegisterElfGlobals;
10321032
FunctionCallee AsanUnregisterElfGlobals;
1033+
FunctionCallee AsanSetShadowStaticLocalFunc;
1034+
FunctionCallee AsanUnpoisonShadowStaticLocalFunc;
10331035

10341036
Function *AsanCtorFunction = nullptr;
10351037
Function *AsanDtorFunction = nullptr;
10361038
GlobalVariable *ModuleName = nullptr;
1039+
1040+
DenseMap<Function *, SmallVector<Instruction *, 8>> KernelToRetVecMap;
1041+
DenseMap<Function *, DenseSet<Function *>> FuncToKernelCallerMap;
10371042
};
10381043

10391044
// Stack poisoning does not play well with exception handling.
@@ -1661,6 +1666,9 @@ static bool isUnsupportedDeviceGlobal(GlobalVariable *G) {
16611666
if (G->getName().starts_with("__Asan"))
16621667
return true;
16631668

1669+
if (G->getAddressSpace() == kSpirOffloadLocalAS)
1670+
return true;
1671+
16641672
Attribute Attr = G->getAttribute("sycl-device-image-scope");
16651673
return (!Attr.isStringAttribute() || Attr.getValueAsString() == "false");
16661674
}
@@ -1765,68 +1773,6 @@ Value *AddressSanitizer::memToShadow(Value *Shadow, IRBuilder<> &IRB,
17651773
return IRB.CreateAdd(Shadow, ShadowBase);
17661774
}
17671775

1768-
static uint64_t getSizeAndRedzoneSizeForLocal(uint64_t Size,
1769-
uint64_t Granularity,
1770-
uint64_t Alignment) {
1771-
uint64_t Res = 0;
1772-
if (Size <= 4)
1773-
Res = 16;
1774-
else if (Size <= 16)
1775-
Res = 32;
1776-
else if (Size <= 128)
1777-
Res = Size + 32;
1778-
else if (Size <= 512)
1779-
Res = Size + 64;
1780-
else if (Size <= 4096)
1781-
Res = Size + 128;
1782-
else
1783-
Res = Size + 256;
1784-
return alignTo(std::max(Res, 2 * Granularity), Alignment);
1785-
}
1786-
1787-
// Instument static local memory
1788-
void AddressSanitizer::instrumentSyclStaticLocalMemory(
1789-
CallInst *CI, ArrayRef<Instruction *> RetVec) {
1790-
InstrumentationIRBuilder IRB(CI->getNextNode());
1791-
auto *Size = cast<ConstantInt>(CI->getArgOperand(0));
1792-
auto *Alignment = cast<ConstantInt>(CI->getArgOperand(1));
1793-
1794-
const auto Granularity = 1 << Mapping.Scale;
1795-
// The base address of local memory needs to align to granularity
1796-
const auto Align = alignTo(Alignment->getZExtValue(), Granularity);
1797-
1798-
auto *SizeWithRedZone = ConstantInt::get(
1799-
IntptrTy, getSizeAndRedzoneSizeForLocal(Size->getZExtValue(), Granularity,
1800-
Alignment->getZExtValue()));
1801-
1802-
auto *NewCI =
1803-
IRB.CreateCall(CI->getCalledFunction(),
1804-
{SizeWithRedZone, ConstantInt::get(IntptrTy, Align)});
1805-
1806-
// __asan_set_shadow_static_local(
1807-
// uptr beg,
1808-
// size_t size,
1809-
// size_t size_with_redzone,
1810-
// )
1811-
auto LocalAddr = IRB.CreatePointerCast(NewCI, IntptrTy);
1812-
IRB.CreateCall(AsanSetShadowStaticLocalFunc,
1813-
{LocalAddr, Size, SizeWithRedZone});
1814-
1815-
// __asan_unpoison_shadow_static_local(
1816-
// uptr beg,
1817-
// size_t size,
1818-
// size_t size_with_redzone,
1819-
// )
1820-
for (Instruction *Ret : RetVec) {
1821-
IRBuilder<> IRBRet(Ret);
1822-
IRBRet.CreateCall(AsanUnpoisonShadowStaticLocalFunc,
1823-
{LocalAddr, Size, SizeWithRedZone});
1824-
}
1825-
1826-
CI->replaceAllUsesWith(NewCI);
1827-
CI->eraseFromParent();
1828-
}
1829-
18301776
// Instument dynamic local memory
18311777
bool AddressSanitizer::instrumentSyclDynamicLocalMemory(
18321778
Function &F, ArrayRef<Instruction *> RetVec) {
@@ -2810,6 +2756,24 @@ void ModuleAddressSanitizer::initializeCallbacks() {
28102756
AsanUnregisterElfGlobals =
28112757
M.getOrInsertFunction(kAsanUnregisterElfGlobalsName, IRB.getVoidTy(),
28122758
IntptrTy, IntptrTy, IntptrTy);
2759+
2760+
// __asan_set_shadow_static_local(
2761+
// uptr ptr,
2762+
// size_t size,
2763+
// size_t size_with_redzone
2764+
// )
2765+
AsanSetShadowStaticLocalFunc =
2766+
M.getOrInsertFunction("__asan_set_shadow_static_local", IRB.getVoidTy(),
2767+
IntptrTy, IntptrTy, IntptrTy);
2768+
2769+
// __asan_unpoison_shadow_static_local(
2770+
// uptr ptr,
2771+
// size_t size,
2772+
// size_t size_with_redzone
2773+
// )
2774+
AsanUnpoisonShadowStaticLocalFunc =
2775+
M.getOrInsertFunction("__asan_unpoison_shadow_static_local",
2776+
IRB.getVoidTy(), IntptrTy, IntptrTy, IntptrTy);
28132777
}
28142778

28152779
// Put the metadata and the instrumented global in the same group. This ensures
@@ -2949,6 +2913,164 @@ void ModuleAddressSanitizer::instrumentDeviceGlobal(IRBuilder<> &IRB) {
29492913
G->eraseFromParent();
29502914
}
29512915

2916+
static void getFunctionsOfUser(User *User, DenseSet<Function *> &Functions) {
2917+
if (Instruction *Inst = dyn_cast<Instruction>(User)) {
2918+
Functions.insert(Inst->getFunction());
2919+
} else if (ConstantExpr *CE = dyn_cast<ConstantExpr>(User)) {
2920+
for (auto *U : CE->users())
2921+
getFunctionsOfUser(U, Functions);
2922+
}
2923+
}
2924+
2925+
void ModuleAddressSanitizer::initializeRetVecMap(Function *F) {
2926+
if (KernelToRetVecMap.find(F) != KernelToRetVecMap.end())
2927+
return;
2928+
2929+
SmallVector<Instruction *, 8> RetVec;
2930+
for (auto &BB : *F) {
2931+
for (auto &Inst : BB) {
2932+
if (ReturnInst *RI = dyn_cast<ReturnInst>(&Inst)) {
2933+
if (CallInst *CI = RI->getParent()->getTerminatingMustTailCall())
2934+
RetVec.push_back(CI);
2935+
else
2936+
RetVec.push_back(RI);
2937+
} else if (ResumeInst *RI = dyn_cast<ResumeInst>(&Inst)) {
2938+
RetVec.push_back(RI);
2939+
} else if (CleanupReturnInst *CRI = dyn_cast<CleanupReturnInst>(&Inst)) {
2940+
RetVec.push_back(CRI);
2941+
}
2942+
}
2943+
}
2944+
2945+
KernelToRetVecMap[F] = std::move(RetVec);
2946+
}
2947+
2948+
void ModuleAddressSanitizer::initializeKernelCallerMap(Function *F) {
2949+
if (FuncToKernelCallerMap.find(F) != FuncToKernelCallerMap.end())
2950+
return;
2951+
2952+
for (auto *U : F->users()) {
2953+
if (Instruction *Inst = dyn_cast<Instruction>(U)) {
2954+
Function *Caller = Inst->getFunction();
2955+
if (Caller->getCallingConv() == CallingConv::SPIR_KERNEL) {
2956+
FuncToKernelCallerMap[F].insert(Caller);
2957+
continue;
2958+
}
2959+
initializeKernelCallerMap(Caller);
2960+
FuncToKernelCallerMap[F].insert(FuncToKernelCallerMap[Caller].begin(),
2961+
FuncToKernelCallerMap[Caller].end());
2962+
}
2963+
}
2964+
}
2965+
2966+
// Instument static local memory
2967+
void ModuleAddressSanitizer::instrumentSyclStaticLocalMemory(IRBuilder<> &IRB) {
2968+
auto &DL = M.getDataLayout();
2969+
SmallVector<GlobalVariable *, 8> GlobalsToRemove;
2970+
SmallVector<GlobalVariable *, 8> LocalGlobals;
2971+
2972+
Type *IntptrTy = M.getDataLayout().getIntPtrType(*C, kSpirOffloadGlobalAS);
2973+
2974+
// Step1. Create a new global variable with enough space for a redzone.
2975+
for (auto &G : M.globals()) {
2976+
if (G.getAddressSpace() != kSpirOffloadLocalAS)
2977+
continue;
2978+
if (G.getName().starts_with("__Asan"))
2979+
continue;
2980+
2981+
Type *Ty = G.getValueType();
2982+
const uint64_t SizeInBytes = DL.getTypeAllocSize(Ty);
2983+
const uint64_t RightRedzoneSize = getRedzoneSizeForGlobal(SizeInBytes);
2984+
Type *RightRedZoneTy = ArrayType::get(IRB.getInt8Ty(), RightRedzoneSize);
2985+
StructType *NewTy = StructType::get(Ty, RightRedZoneTy);
2986+
Constant *NewInitializer =
2987+
G.hasInitializer()
2988+
? ConstantStruct::get(NewTy, G.getInitializer(),
2989+
Constant::getNullValue(RightRedZoneTy))
2990+
: nullptr;
2991+
2992+
GlobalVariable *NewGlobal = new GlobalVariable(
2993+
M, NewTy, G.isConstant(), G.getLinkage(), NewInitializer, "", &G,
2994+
G.getThreadLocalMode(), G.getAddressSpace());
2995+
NewGlobal->copyAttributesFrom(&G);
2996+
NewGlobal->setComdat(G.getComdat());
2997+
NewGlobal->setAlignment(Align(getMinRedzoneSizeForGlobal()));
2998+
NewGlobal->copyMetadata(&G, 0);
2999+
3000+
Value *Indices2[2];
3001+
Indices2[0] = IRB.getInt32(0);
3002+
Indices2[1] = IRB.getInt32(0);
3003+
3004+
G.replaceAllUsesWith(
3005+
ConstantExpr::getGetElementPtr(NewTy, NewGlobal, Indices2, true));
3006+
NewGlobal->takeName(&G);
3007+
GlobalsToRemove.push_back(&G);
3008+
LocalGlobals.push_back(NewGlobal);
3009+
}
3010+
3011+
if (GlobalsToRemove.empty())
3012+
return;
3013+
3014+
for (auto *G : GlobalsToRemove)
3015+
G->eraseFromParent();
3016+
3017+
// Step2. Instrument initialization functions on kernel
3018+
DenseMap<Function *, Instruction *> FuncToLaunchInfoMap;
3019+
auto Instrument = [&](GlobalVariable *G, Function *F) {
3020+
StructType *Type = cast<StructType>(G->getValueType());
3021+
const uint64_t Size = DL.getTypeAllocSize(Type->getElementType(0));
3022+
const uint64_t SizeWithRedZone = DL.getTypeAllocSize(Type);
3023+
3024+
// Poison shadow of static local memory
3025+
if (FuncToLaunchInfoMap.find(F) == FuncToLaunchInfoMap.end()) {
3026+
for (auto &Inst : F->getEntryBlock()) {
3027+
auto *SI = dyn_cast<StoreInst>(&Inst);
3028+
if (SI && (SI->getPointerOperand()->getName() == "__AsanLaunchInfo")) {
3029+
FuncToLaunchInfoMap[F] = &Inst;
3030+
break;
3031+
}
3032+
}
3033+
}
3034+
assert(FuncToLaunchInfoMap.find(F) != FuncToLaunchInfoMap.end() &&
3035+
"All spir kernels should be instrumented.");
3036+
3037+
IRBuilder<> Builder(FuncToLaunchInfoMap[F]->getNextNode());
3038+
Builder.CreateCall(AsanSetShadowStaticLocalFunc,
3039+
{Builder.CreatePointerCast(G, IntptrTy),
3040+
ConstantInt::get(IntptrTy, Size),
3041+
ConstantInt::get(IntptrTy, SizeWithRedZone)});
3042+
3043+
// Unpoison shadow of static local memory, required by CPU device
3044+
initializeRetVecMap(F);
3045+
for (auto *RI : KernelToRetVecMap[F]) {
3046+
IRBuilder<> Builder(RI);
3047+
Builder.CreateCall(AsanUnpoisonShadowStaticLocalFunc,
3048+
{Builder.CreatePointerCast(G, IntptrTy),
3049+
ConstantInt::get(IntptrTy, Size),
3050+
ConstantInt::get(IntptrTy, SizeWithRedZone)});
3051+
}
3052+
};
3053+
3054+
// We only instrument on spir_kernel, because local variables are
3055+
// kind of global variable
3056+
for (auto *G : LocalGlobals) {
3057+
DenseSet<Function *> InstrumentedFunc;
3058+
for (auto *User : G->users())
3059+
getFunctionsOfUser(User, InstrumentedFunc);
3060+
for (Function *F : InstrumentedFunc) {
3061+
if (F->getCallingConv() == CallingConv::SPIR_KERNEL) {
3062+
Instrument(G, F);
3063+
continue;
3064+
}
3065+
// Get root spir_kernel of spir_func
3066+
initializeKernelCallerMap(F);
3067+
for (Function *Kernel : FuncToKernelCallerMap[F])
3068+
if (!InstrumentedFunc.contains(Kernel))
3069+
Instrument(G, Kernel);
3070+
}
3071+
}
3072+
}
3073+
29523074
void ModuleAddressSanitizer::InstrumentGlobalsCOFF(
29533075
IRBuilder<> &IRB, ArrayRef<GlobalVariable *> ExtendedGlobals,
29543076
ArrayRef<Constant *> MetadataInitializers) {
@@ -3412,12 +3534,10 @@ bool ModuleAddressSanitizer::instrumentModule() {
34123534
}
34133535

34143536
if (TargetTriple.isSPIROrSPIRV()) {
3415-
// Add module metadata "device.sanitizer" for sycl-post-link
3416-
LLVMContext &Ctx = M.getContext();
3417-
auto *MD = M.getOrInsertNamedMetadata("device.sanitizer");
3418-
Metadata *MDVals[] = {MDString::get(Ctx, "asan")};
3419-
MD->addOperand(MDNode::get(Ctx, MDVals));
3420-
3537+
if (ClSpirOffloadLocals) {
3538+
IRBuilder<> IRB(*C);
3539+
instrumentSyclStaticLocalMemory(IRB);
3540+
}
34213541
if (ClDeviceGlobals) {
34223542
IRBuilder<> IRB(*C);
34233543
instrumentDeviceGlobal(IRB);
@@ -3564,23 +3684,6 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {
35643684
ArrayType::get(IRB.getInt8Ty(), 0));
35653685

35663686
if (TargetTriple.isSPIROrSPIRV()) {
3567-
// __asan_set_shadow_static_local(
3568-
// uptr ptr,
3569-
// size_t size,
3570-
// size_t size_with_redzone
3571-
// )
3572-
AsanSetShadowStaticLocalFunc =
3573-
M.getOrInsertFunction("__asan_set_shadow_static_local", IRB.getVoidTy(),
3574-
IntptrTy, IntptrTy, IntptrTy);
3575-
3576-
// __asan_unpoison_shadow_static_local(
3577-
// uptr ptr,
3578-
// size_t size,
3579-
// )
3580-
AsanUnpoisonShadowStaticLocalFunc =
3581-
M.getOrInsertFunction("__asan_unpoison_shadow_static_local",
3582-
IRB.getVoidTy(), IntptrTy, IntptrTy, IntptrTy);
3583-
35843687
// __asan_set_shadow_dynamic_local(
35853688
// uptr ptr,
35863689
// uint32_t num_args
@@ -3747,7 +3850,6 @@ bool AddressSanitizer::instrumentFunction(Function &F,
37473850
SmallVector<Instruction *, 8> NoReturnCalls;
37483851
SmallVector<BasicBlock *, 16> AllBlocks;
37493852
SmallVector<Instruction *, 16> PointerComparisonsOrSubtracts;
3750-
SmallVector<CallInst *, 8> SyclAllocateLocalMemoryCalls;
37513853

37523854
// Fill the set of memory operations to instrument.
37533855
for (auto &BB : F) {
@@ -3800,16 +3902,8 @@ bool AddressSanitizer::instrumentFunction(Function &F,
38003902
NoReturnCalls.push_back(CB);
38013903
}
38023904
}
3803-
if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
3804-
if (TargetTriple.isSPIROrSPIRV() && CI->getCalledFunction() &&
3805-
CI->getCalledFunction()->getCallingConv() ==
3806-
llvm::CallingConv::SPIR_FUNC &&
3807-
CI->getCalledFunction()->getName() ==
3808-
"__sycl_allocateLocalMemory")
3809-
SyclAllocateLocalMemoryCalls.push_back(CI);
3810-
else
3811-
maybeMarkSanitizerLibraryCallNoBuiltin(CI, TLI);
3812-
}
3905+
if (CallInst *CI = dyn_cast<CallInst>(&Inst))
3906+
maybeMarkSanitizerLibraryCallNoBuiltin(CI, TLI);
38133907
}
38143908
if (NumInsnsPerBB >= ClMaxInsnsToInstrumentPerBB) break;
38153909
}
@@ -3855,13 +3949,9 @@ bool AddressSanitizer::instrumentFunction(Function &F,
38553949
if (ChangedStack || !NoReturnCalls.empty())
38563950
FunctionModified = true;
38573951

3858-
// We need to instrument dynamic/static local arguments after stack poisoner
3952+
// We need to instrument dynamic local arguments after stack poisoner
38593953
if (TargetTriple.isSPIROrSPIRV()) {
3860-
for (auto *CI : SyclAllocateLocalMemoryCalls) {
3861-
instrumentSyclStaticLocalMemory(CI, FSP.RetVec);
3862-
FunctionModified = true;
3863-
}
3864-
if (F.getCallingConv() == CallingConv::SPIR_KERNEL) {
3954+
if (ClSpirOffloadLocals && F.getCallingConv() == CallingConv::SPIR_KERNEL) {
38653955
FunctionModified |= instrumentSyclDynamicLocalMemory(F, FSP.RetVec);
38663956
}
38673957
}

0 commit comments

Comments
 (0)