Skip to content

Commit 2523098

Browse files
AllanZyneKornevNikita
AllanZyne
authored and
KornevNikita
committed
[DeviceASAN] Fix sycl::group_local_memory (#17769)
1 parent 21f468c commit 2523098

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) {
@@ -2808,6 +2754,24 @@ void ModuleAddressSanitizer::initializeCallbacks() {
28082754
AsanUnregisterElfGlobals =
28092755
M.getOrInsertFunction(kAsanUnregisterElfGlobalsName, IRB.getVoidTy(),
28102756
IntptrTy, IntptrTy, IntptrTy);
2757+
2758+
// __asan_set_shadow_static_local(
2759+
// uptr ptr,
2760+
// size_t size,
2761+
// size_t size_with_redzone
2762+
// )
2763+
AsanSetShadowStaticLocalFunc =
2764+
M.getOrInsertFunction("__asan_set_shadow_static_local", IRB.getVoidTy(),
2765+
IntptrTy, IntptrTy, IntptrTy);
2766+
2767+
// __asan_unpoison_shadow_static_local(
2768+
// uptr ptr,
2769+
// size_t size,
2770+
// size_t size_with_redzone
2771+
// )
2772+
AsanUnpoisonShadowStaticLocalFunc =
2773+
M.getOrInsertFunction("__asan_unpoison_shadow_static_local",
2774+
IRB.getVoidTy(), IntptrTy, IntptrTy, IntptrTy);
28112775
}
28122776

28132777
// Put the metadata and the instrumented global in the same group. This ensures
@@ -2942,6 +2906,164 @@ void ModuleAddressSanitizer::instrumentDeviceGlobal(IRBuilder<> &IRB) {
29422906
G->eraseFromParent();
29432907
}
29442908

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

34073529
if (TargetTriple.isSPIROrSPIRV()) {
3408-
// Add module metadata "device.sanitizer" for sycl-post-link
3409-
LLVMContext &Ctx = M.getContext();
3410-
auto *MD = M.getOrInsertNamedMetadata("device.sanitizer");
3411-
Metadata *MDVals[] = {MDString::get(Ctx, "asan")};
3412-
MD->addOperand(MDNode::get(Ctx, MDVals));
3413-
3530+
if (ClSpirOffloadLocals) {
3531+
IRBuilder<> IRB(*C);
3532+
instrumentSyclStaticLocalMemory(IRB);
3533+
}
34143534
if (ClDeviceGlobals) {
34153535
IRBuilder<> IRB(*C);
34163536
instrumentDeviceGlobal(IRB);
@@ -3557,23 +3677,6 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {
35573677
ArrayType::get(IRB.getInt8Ty(), 0));
35583678

35593679
if (TargetTriple.isSPIROrSPIRV()) {
3560-
// __asan_set_shadow_static_local(
3561-
// uptr ptr,
3562-
// size_t size,
3563-
// size_t size_with_redzone
3564-
// )
3565-
AsanSetShadowStaticLocalFunc =
3566-
M.getOrInsertFunction("__asan_set_shadow_static_local", IRB.getVoidTy(),
3567-
IntptrTy, IntptrTy, IntptrTy);
3568-
3569-
// __asan_unpoison_shadow_static_local(
3570-
// uptr ptr,
3571-
// size_t size,
3572-
// )
3573-
AsanUnpoisonShadowStaticLocalFunc =
3574-
M.getOrInsertFunction("__asan_unpoison_shadow_static_local",
3575-
IRB.getVoidTy(), IntptrTy, IntptrTy, IntptrTy);
3576-
35773680
// __asan_set_shadow_dynamic_local(
35783681
// uptr ptr,
35793682
// uint32_t num_args
@@ -3740,7 +3843,6 @@ bool AddressSanitizer::instrumentFunction(Function &F,
37403843
SmallVector<Instruction *, 8> NoReturnCalls;
37413844
SmallVector<BasicBlock *, 16> AllBlocks;
37423845
SmallVector<Instruction *, 16> PointerComparisonsOrSubtracts;
3743-
SmallVector<CallInst *, 8> SyclAllocateLocalMemoryCalls;
37443846

37453847
// Fill the set of memory operations to instrument.
37463848
for (auto &BB : F) {
@@ -3793,16 +3895,8 @@ bool AddressSanitizer::instrumentFunction(Function &F,
37933895
NoReturnCalls.push_back(CB);
37943896
}
37953897
}
3796-
if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
3797-
if (TargetTriple.isSPIROrSPIRV() && CI->getCalledFunction() &&
3798-
CI->getCalledFunction()->getCallingConv() ==
3799-
llvm::CallingConv::SPIR_FUNC &&
3800-
CI->getCalledFunction()->getName() ==
3801-
"__sycl_allocateLocalMemory")
3802-
SyclAllocateLocalMemoryCalls.push_back(CI);
3803-
else
3804-
maybeMarkSanitizerLibraryCallNoBuiltin(CI, TLI);
3805-
}
3898+
if (CallInst *CI = dyn_cast<CallInst>(&Inst))
3899+
maybeMarkSanitizerLibraryCallNoBuiltin(CI, TLI);
38063900
}
38073901
if (NumInsnsPerBB >= ClMaxInsnsToInstrumentPerBB) break;
38083902
}
@@ -3848,13 +3942,9 @@ bool AddressSanitizer::instrumentFunction(Function &F,
38483942
if (ChangedStack || !NoReturnCalls.empty())
38493943
FunctionModified = true;
38503944

3851-
// We need to instrument dynamic/static local arguments after stack poisoner
3945+
// We need to instrument dynamic local arguments after stack poisoner
38523946
if (TargetTriple.isSPIROrSPIRV()) {
3853-
for (auto *CI : SyclAllocateLocalMemoryCalls) {
3854-
instrumentSyclStaticLocalMemory(CI, FSP.RetVec);
3855-
FunctionModified = true;
3856-
}
3857-
if (F.getCallingConv() == CallingConv::SPIR_KERNEL) {
3947+
if (ClSpirOffloadLocals && F.getCallingConv() == CallingConv::SPIR_KERNEL) {
38583948
FunctionModified |= instrumentSyclDynamicLocalMemory(F, FSP.RetVec);
38593949
}
38603950
}

0 commit comments

Comments
 (0)