Skip to content

Commit 805abc8

Browse files
AllanZyneAllanZyne
committed
fix
1 parent 16a14c9 commit 805abc8

File tree

2 files changed

+199
-112
lines changed

2 files changed

+199
-112
lines changed

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

Lines changed: 190 additions & 104 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.
@@ -1765,68 +1770,6 @@ Value *AddressSanitizer::memToShadow(Value *Shadow, IRBuilder<> &IRB,
17651770
return IRB.CreateAdd(Shadow, ShadowBase);
17661771
}
17671772

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-
18301773
// Instument dynamic local memory
18311774
bool AddressSanitizer::instrumentSyclDynamicLocalMemory(
18321775
Function &F, ArrayRef<Instruction *> RetVec) {
@@ -2810,6 +2753,24 @@ void ModuleAddressSanitizer::initializeCallbacks() {
28102753
AsanUnregisterElfGlobals =
28112754
M.getOrInsertFunction(kAsanUnregisterElfGlobalsName, IRB.getVoidTy(),
28122755
IntptrTy, IntptrTy, IntptrTy);
2756+
2757+
// __asan_set_shadow_static_local(
2758+
// uptr ptr,
2759+
// size_t size,
2760+
// size_t size_with_redzone
2761+
// )
2762+
AsanSetShadowStaticLocalFunc =
2763+
M.getOrInsertFunction("__asan_set_shadow_static_local", IRB.getVoidTy(),
2764+
IntptrTy, IntptrTy, IntptrTy);
2765+
2766+
// __asan_unpoison_shadow_static_local(
2767+
// uptr ptr,
2768+
// size_t size,
2769+
// size_t size_with_redzone
2770+
// )
2771+
AsanUnpoisonShadowStaticLocalFunc =
2772+
M.getOrInsertFunction("__asan_unpoison_shadow_static_local",
2773+
IRB.getVoidTy(), IntptrTy, IntptrTy, IntptrTy);
28132774
}
28142775

28152776
// Put the metadata and the instrumented global in the same group. This ensures
@@ -2949,6 +2910,163 @@ void ModuleAddressSanitizer::instrumentDeviceGlobal(IRBuilder<> &IRB) {
29492910
G->eraseFromParent();
29502911
}
29512912

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

34143532
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-
3533+
if (ClSpirOffloadLocals) {
3534+
IRBuilder<> IRB(*C);
3535+
instrumentSyclStaticLocalMemory(IRB);
3536+
}
34213537
if (ClDeviceGlobals) {
34223538
IRBuilder<> IRB(*C);
34233539
instrumentDeviceGlobal(IRB);
@@ -3564,23 +3680,6 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {
35643680
ArrayType::get(IRB.getInt8Ty(), 0));
35653681

35663682
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-
35843683
// __asan_set_shadow_dynamic_local(
35853684
// uptr ptr,
35863685
// uint32_t num_args
@@ -3747,7 +3846,6 @@ bool AddressSanitizer::instrumentFunction(Function &F,
37473846
SmallVector<Instruction *, 8> NoReturnCalls;
37483847
SmallVector<BasicBlock *, 16> AllBlocks;
37493848
SmallVector<Instruction *, 16> PointerComparisonsOrSubtracts;
3750-
SmallVector<CallInst *, 8> SyclAllocateLocalMemoryCalls;
37513849

37523850
// Fill the set of memory operations to instrument.
37533851
for (auto &BB : F) {
@@ -3800,16 +3898,8 @@ bool AddressSanitizer::instrumentFunction(Function &F,
38003898
NoReturnCalls.push_back(CB);
38013899
}
38023900
}
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-
}
3901+
if (CallInst *CI = dyn_cast<CallInst>(&Inst))
3902+
maybeMarkSanitizerLibraryCallNoBuiltin(CI, TLI);
38133903
}
38143904
if (NumInsnsPerBB >= ClMaxInsnsToInstrumentPerBB) break;
38153905
}
@@ -3857,10 +3947,6 @@ bool AddressSanitizer::instrumentFunction(Function &F,
38573947

38583948
// We need to instrument dynamic/static local arguments after stack poisoner
38593949
if (TargetTriple.isSPIROrSPIRV()) {
3860-
for (auto *CI : SyclAllocateLocalMemoryCalls) {
3861-
instrumentSyclStaticLocalMemory(CI, FSP.RetVec);
3862-
FunctionModified = true;
3863-
}
38643950
if (F.getCallingConv() == CallingConv::SPIR_KERNEL) {
38653951
FunctionModified |= instrumentSyclDynamicLocalMemory(F, FSP.RetVec);
38663952
}

0 commit comments

Comments
 (0)