Skip to content

Commit 70a38ea

Browse files
skc7skc7
committed
[AMDGPU] Enable flag to choose instrumentation of LDS.
[AMDGPU] Update the MD initializer and donot replace uses of MD global.
1 parent 65dd67a commit 70a38ea

File tree

6 files changed

+224
-394
lines changed

6 files changed

+224
-394
lines changed

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

Lines changed: 47 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,11 @@ static cl::opt<bool>
232232
cl::desc("instrument byval call arguments"), cl::Hidden,
233233
cl::init(true));
234234

235+
static cl::opt<bool>
236+
ClInstrumentAMDGPULDS("asan-instrument-amdgpu-lds",
237+
cl::desc("instrument amdgpu LDS accesses"),
238+
cl::Hidden, cl::init(true));
239+
235240
static cl::opt<bool> ClAlwaysSlowPath(
236241
"asan-always-slow-path",
237242
cl::desc("use instrumentation with slow path for all accesses"), cl::Hidden,
@@ -1293,8 +1298,8 @@ static GlobalVariable *getKernelSwLDSBaseGlobal(Module &M) {
12931298
static void updateLDSSizeFnAttr(Function *Func, uint32_t Offset,
12941299
bool UsesDynLDS) {
12951300
if (Offset != 0) {
1296-
std::string Buffer;
1297-
raw_string_ostream SS{Buffer};
1301+
SmallString<256> Buffer;
1302+
raw_svector_ostream SS(Buffer);
12981303
SS << format("%u", Offset);
12991304
if (UsesDynLDS)
13001305
SS << format(",%u", Offset);
@@ -1312,33 +1317,30 @@ static void recordLDSAbsoluteAddress(Module &M, GlobalVariable *GV,
13121317
MDNode::get(Ctx, {MinC, MaxC}));
13131318
}
13141319

1315-
static void UpdateSwLDSMetadataWithRedzoneInfo(Function &F, int Scale) {
1320+
/// Update SwLDS Metadata global initializer with redzone info.
1321+
static SmallVector<std::pair<uint32_t, uint32_t>, 64>
1322+
UpdateSwLDSMetadataWithRedzoneInfo(Function &F, int Scale) {
13161323
Module *M = F.getParent();
13171324
GlobalVariable *SwLDSMetadataGlobal = getKernelSwLDSMetadataGlobal(*M, F);
13181325
GlobalVariable *SwLDSGlobal = getKernelSwLDSGlobal(*M, F);
13191326
if (!SwLDSMetadataGlobal || !SwLDSGlobal)
1320-
return;
1327+
return {};
13211328

13221329
LLVMContext &Ctx = M->getContext();
13231330
Type *Int32Ty = Type::getInt32Ty(Ctx);
1324-
1331+
SmallVector<std::pair<uint32_t, uint32_t>, 64> RedzoneOffsetAndSizeVector;
13251332
Constant *MdInit = SwLDSMetadataGlobal->getInitializer();
1326-
Align MdAlign = Align(SwLDSMetadataGlobal->getAlign().valueOrOne());
13271333
Align LDSAlign = Align(SwLDSGlobal->getAlign().valueOrOne());
13281334

13291335
StructType *MDStructType =
13301336
cast<StructType>(SwLDSMetadataGlobal->getValueType());
1331-
assert(MDStructType);
13321337
unsigned NumStructs = MDStructType->getNumElements();
1333-
1334-
std::vector<Type *> Items;
13351338
std::vector<Constant *> Initializers;
13361339
uint32_t MallocSize = 0;
1337-
//{GV.start, Align(GV.size + Redzone.size), Redzone.start, Redzone.size}
1338-
StructType *LDSItemTy = StructType::create(
1339-
Ctx, {Int32Ty, Int32Ty, Int32Ty, Int32Ty, Int32Ty}, "");
1340+
StructType *LDSItemTy =
1341+
cast<StructType>(MDStructType->getStructElementType(0));
1342+
13401343
for (unsigned i = 0; i < NumStructs; i++) {
1341-
Items.push_back(LDSItemTy);
13421344
ConstantStruct *member =
13431345
dyn_cast<ConstantStruct>(MdInit->getAggregateElement(i));
13441346
Constant *NewInitItem;
@@ -1353,91 +1355,48 @@ static void UpdateSwLDSMetadataWithRedzoneInfo(Function &F, int Scale) {
13531355
const uint64_t RightRedzoneSize =
13541356
getRedzoneSizeForGlobal(Scale, GlobalSizeValue);
13551357
MallocSize += GlobalSizeValue;
1356-
Constant *NewItemRedzoneStartOffset =
1357-
ConstantInt::get(Int32Ty, MallocSize);
1358+
RedzoneOffsetAndSizeVector.emplace_back(MallocSize, RightRedzoneSize);
13581359
MallocSize += RightRedzoneSize;
1359-
Constant *NewItemRedzoneSize =
1360-
ConstantInt::get(Int32Ty, RightRedzoneSize);
1361-
13621360
unsigned NewItemAlignGlobalPlusRedzoneSize =
13631361
alignTo(GlobalSizeValue + RightRedzoneSize, LDSAlign);
13641362
Constant *NewItemAlignGlobalPlusRedzoneSizeConst =
13651363
ConstantInt::get(Int32Ty, NewItemAlignGlobalPlusRedzoneSize);
13661364
NewInitItem = ConstantStruct::get(
13671365
LDSItemTy, {NewItemStartOffset, NewItemGlobalSizeConst,
1368-
NewItemAlignGlobalPlusRedzoneSizeConst,
1369-
NewItemRedzoneStartOffset, NewItemRedzoneSize});
1366+
NewItemAlignGlobalPlusRedzoneSizeConst});
13701367
MallocSize = alignTo(MallocSize, LDSAlign);
13711368
} else {
13721369
Constant *CurrMallocSize = ConstantInt::get(Int32Ty, MallocSize);
13731370
Constant *zero = ConstantInt::get(Int32Ty, 0);
1374-
NewInitItem = ConstantStruct::get(
1375-
LDSItemTy, {CurrMallocSize, zero, zero, zero, zero});
1371+
NewInitItem =
1372+
ConstantStruct::get(LDSItemTy, {CurrMallocSize, zero, zero});
1373+
RedzoneOffsetAndSizeVector.emplace_back(0, 0);
13761374
}
13771375
} else {
13781376
Constant *CurrMallocSize = ConstantInt::get(Int32Ty, MallocSize);
13791377
Constant *zero = ConstantInt::get(Int32Ty, 0);
1380-
NewInitItem = ConstantStruct::get(
1381-
LDSItemTy, {CurrMallocSize, zero, zero, zero, zero});
1378+
NewInitItem =
1379+
ConstantStruct::get(LDSItemTy, {CurrMallocSize, zero, zero});
1380+
RedzoneOffsetAndSizeVector.emplace_back(0, 0);
13821381
}
13831382
Initializers.push_back(NewInitItem);
13841383
}
13851384
GlobalVariable *SwDynLDS = getKernelSwDynLDSGlobal(*M, F);
1386-
bool usesDynLDS = SwDynLDS ? true : false;
1385+
bool usesDynLDS = SwDynLDS != nullptr;
13871386
updateLDSSizeFnAttr(&F, MallocSize, usesDynLDS);
13881387
if (usesDynLDS)
13891388
recordLDSAbsoluteAddress(*M, SwDynLDS, MallocSize);
13901389

1391-
StructType *MetadataStructType = StructType::create(Ctx, Items, "");
1392-
1393-
GlobalVariable *NewSwLDSMetadataGlobal = new GlobalVariable(
1394-
*M, MetadataStructType, false, GlobalValue::InternalLinkage,
1395-
PoisonValue::get(MetadataStructType), "", nullptr,
1396-
GlobalValue::NotThreadLocal, 1, false);
1397-
Constant *Data = ConstantStruct::get(MetadataStructType, Initializers);
1398-
NewSwLDSMetadataGlobal->setInitializer(Data);
1399-
NewSwLDSMetadataGlobal->setAlignment(MdAlign);
1400-
GlobalValue::SanitizerMetadata MD;
1401-
MD.NoAddress = true;
1402-
NewSwLDSMetadataGlobal->setSanitizerMetadata(MD);
1403-
1404-
for (Use &U : make_early_inc_range(SwLDSMetadataGlobal->uses())) {
1405-
if (GEPOperator *GEP = dyn_cast<GEPOperator>(U.getUser())) {
1406-
SmallVector<Constant *> Indices;
1407-
for (Use &Idx : GEP->indices()) {
1408-
Indices.push_back(cast<Constant>(Idx));
1409-
}
1410-
Constant *NewGEP = ConstantExpr::getGetElementPtr(
1411-
MetadataStructType, NewSwLDSMetadataGlobal, Indices, true);
1412-
GEP->replaceAllUsesWith(NewGEP);
1413-
} else if (LoadInst *Load = dyn_cast<LoadInst>(U.getUser())) {
1414-
Constant *zero = ConstantInt::get(Int32Ty, 0);
1415-
SmallVector<Constant *> Indices{zero, zero, zero};
1416-
Constant *NewGEP = ConstantExpr::getGetElementPtr(
1417-
MetadataStructType, NewSwLDSMetadataGlobal, Indices, true);
1418-
IRBuilder<> IRB(Load);
1419-
LoadInst *NewLoad = IRB.CreateLoad(Load->getType(), NewGEP);
1420-
Load->replaceAllUsesWith(NewLoad);
1421-
Load->eraseFromParent();
1422-
} else if (StoreInst *Store = dyn_cast<StoreInst>(U.getUser())) {
1423-
Constant *zero = ConstantInt::get(Int32Ty, 0);
1424-
SmallVector<Constant *> Indices{zero, zero, zero};
1425-
Constant *NewGEP = ConstantExpr::getGetElementPtr(
1426-
MetadataStructType, NewSwLDSMetadataGlobal, Indices, true);
1427-
IRBuilder<> IRB(Store);
1428-
StoreInst *NewStore = IRB.CreateStore(Store->getValueOperand(), NewGEP);
1429-
Store->replaceAllUsesWith(NewStore);
1430-
Store->eraseFromParent();
1431-
} else
1432-
report_fatal_error("AMDGPU Sw LDS Metadata User instruction not handled");
1433-
}
1434-
SwLDSMetadataGlobal->replaceAllUsesWith(NewSwLDSMetadataGlobal);
1435-
NewSwLDSMetadataGlobal->takeName(SwLDSMetadataGlobal);
1436-
SwLDSMetadataGlobal->eraseFromParent();
1437-
return;
1390+
Constant *Data = ConstantStruct::get(MDStructType, Initializers);
1391+
SwLDSMetadataGlobal->setInitializer(Data);
1392+
return RedzoneOffsetAndSizeVector;
14381393
}
14391394

1440-
static void poisonRedzonesForSwLDS(Function &F) {
1395+
/// Poison redzone regions using the redzone size and offset info.
1396+
static void
1397+
poisonRedzonesForSwLDS(Function &F,
1398+
SmallVector<std::pair<uint32_t, uint32_t>, 64>
1399+
&RedzoneOffsetAndSizeVector) {
14411400
Module *M = F.getParent();
14421401
GlobalVariable *SwLDSGlobal = getKernelSwLDSGlobal(*M, F);
14431402
GlobalVariable *SwLDSMetadataGlobal = getKernelSwLDSMetadataGlobal(*M, F);
@@ -1466,10 +1425,10 @@ static void poisonRedzonesForSwLDS(Function &F) {
14661425

14671426
StructType *MDStructType =
14681427
cast<StructType>(SwLDSMetadataGlobal->getValueType());
1469-
assert(MDStructType);
14701428
unsigned NumStructs = MDStructType->getNumElements();
14711429
Value *StoreMallocPointer = SI->getValueOperand();
14721430

1431+
assert(RedzoneOffsetAndSizeVector.size() == NumStructs);
14731432
for (unsigned i = 0; i < NumStructs; i++) {
14741433
ConstantStruct *member =
14751434
dyn_cast<ConstantStruct>(MdInit->getAggregateElement(i));
@@ -1484,35 +1443,28 @@ static void poisonRedzonesForSwLDS(Function &F) {
14841443
continue;
14851444
IRBuilder<> IRB(SI);
14861445
IRB.SetInsertPoint(SI->getNextNode());
1446+
auto &RedzonePair = RedzoneOffsetAndSizeVector[i];
1447+
uint64_t RedzoneOffset = RedzonePair.first;
1448+
uint64_t RedzoneSize = RedzonePair.second;
14871449

1488-
auto *GEPForOffset = IRB.CreateInBoundsGEP(
1489-
MDStructType, SwLDSMetadataGlobal,
1490-
{IRB.getInt32(0), IRB.getInt32(i), IRB.getInt32(3)});
1491-
1492-
auto *GEPForSize = IRB.CreateInBoundsGEP(
1493-
MDStructType, SwLDSMetadataGlobal,
1494-
{IRB.getInt32(0), IRB.getInt32(i), IRB.getInt32(4)});
1495-
1496-
Value *RedzoneOffset = IRB.CreateLoad(IRB.getInt32Ty(), GEPForOffset);
1497-
RedzoneOffset = IRB.CreateZExt(RedzoneOffset, IRB.getInt64Ty());
14981450
Value *RedzoneAddrOffset = IRB.CreateInBoundsGEP(
1499-
IRB.getInt8Ty(), StoreMallocPointer, {RedzoneOffset});
1451+
IRB.getInt8Ty(), StoreMallocPointer, {IRB.getInt64(RedzoneOffset)});
15001452
Value *RedzoneAddress =
15011453
IRB.CreatePtrToInt(RedzoneAddrOffset, IRB.getInt64Ty());
1502-
Value *RedzoneSize = IRB.CreateLoad(IRB.getInt32Ty(), GEPForSize);
1503-
RedzoneSize = IRB.CreateZExt(RedzoneSize, IRB.getInt64Ty());
1504-
IRB.CreateCall(AsanPoisonRegion, {RedzoneAddress, RedzoneSize});
1454+
IRB.CreateCall(AsanPoisonRegion,
1455+
{RedzoneAddress, IRB.getInt64(RedzoneSize)});
15051456
}
15061457
}
1507-
return;
15081458
}
15091459

1460+
/// Update SwLDS Metadata global initializer with redzone info.
1461+
/// Poison redzone regions using the redzone size and offset info.
15101462
static void preProcessAMDGPULDSAccesses(Module &M, int Scale) {
15111463
for (Function &F : M) {
1512-
UpdateSwLDSMetadataWithRedzoneInfo(F, Scale);
1513-
poisonRedzonesForSwLDS(F);
1464+
auto RedzoneOffsetAndSizeVector =
1465+
UpdateSwLDSMetadataWithRedzoneInfo(F, Scale);
1466+
poisonRedzonesForSwLDS(F, RedzoneOffsetAndSizeVector);
15141467
}
1515-
return;
15161468
}
15171469

15181470
AddressSanitizerPass::AddressSanitizerPass(
@@ -1527,7 +1479,7 @@ PreservedAnalyses AddressSanitizerPass::run(Module &M,
15271479
ModuleAnalysisManager &MAM) {
15281480
Triple TargetTriple = Triple(M.getTargetTriple());
15291481

1530-
if (TargetTriple.isAMDGPU()) {
1482+
if (TargetTriple.isAMDGPU() && ClInstrumentAMDGPULDS) {
15311483
unsigned LongSize = M.getDataLayout().getPointerSizeInBits();
15321484
ShadowMapping Mapping = getShadowMapping(TargetTriple, LongSize, false);
15331485
preProcessAMDGPULDSAccesses(M, Mapping.Scale);
@@ -2147,7 +2099,7 @@ void AddressSanitizer::instrumentAddress(Instruction *OrigIns,
21472099
}
21482100

21492101
Value *AddrLong;
2150-
if (TargetTriple.isAMDGCN()) {
2102+
if (TargetTriple.isAMDGPU() && ClInstrumentAMDGPULDS) {
21512103
Type *PtrTy = cast<PointerType>(Addr->getType()->getScalarType());
21522104
if (PtrTy->getPointerAddressSpace() == 3) {
21532105
Module *M = IRB.GetInsertBlock()->getParent()->getParent();
@@ -2168,6 +2120,7 @@ void AddressSanitizer::instrumentAddress(Instruction *OrigIns,
21682120
SwLDS = IRB.CreateIntToPtr(IRB.getInt32(0), IRB.getPtrTy(3));
21692121
}
21702122
}
2123+
assert(SwLDS && "Invalid AMDGPU Sw LDS base ptr");
21712124
Value *PtrToInt = IRB.CreatePtrToInt(Addr, IRB.getInt32Ty());
21722125
Value *LoadMallocPtr = IRB.CreateLoad(IRB.getPtrTy(1), SwLDS);
21732126
Value *GEP =

0 commit comments

Comments
 (0)