Skip to content

Commit 9a57946

Browse files
zhaomaosuzhaomaosu
committed
[DeviceASAN] Implement asan_load/store for different address space
1 parent 036fff0 commit 9a57946

File tree

3 files changed

+68
-26
lines changed

3 files changed

+68
-26
lines changed

‎libdevice/sanitizer_utils.cpp‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -668,10 +668,10 @@ constexpr size_t AlignMask(size_t n) { return n - 1; }
668668
/// ASAN Load/Store Report Built-ins
669669
///
670670

671-
#define ASAN_REPORT_ERROR(type, is_write, size) \
672-
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size( \
673-
uptr addr, uint32_t as, const char __SYCL_CONSTANT__ *file, \
674-
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
671+
#define ASAN_REPORT_ERROR_BASE(type, is_write, size, as) \
672+
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size##_as##as( \
673+
uptr addr, const char __SYCL_CONSTANT__ *file, uint32_t line, \
674+
const char __SYCL_CONSTANT__ *func) { \
675675
if (addr & AlignMask(size)) { \
676676
__asan_report_misalign_error(addr, as, size, is_write, addr, file, line, \
677677
func); \
@@ -681,9 +681,9 @@ constexpr size_t AlignMask(size_t n) { return n - 1; }
681681
func); \
682682
} \
683683
} \
684-
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size##_noabort( \
685-
uptr addr, uint32_t as, const char __SYCL_CONSTANT__ *file, \
686-
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
684+
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size##_as##as##_noabort( \
685+
uptr addr, const char __SYCL_CONSTANT__ *file, uint32_t line, \
686+
const char __SYCL_CONSTANT__ *func) { \
687687
if (addr & AlignMask(size)) { \
688688
__asan_report_misalign_error(addr, as, size, is_write, addr, file, line, \
689689
func, true); \
@@ -694,6 +694,13 @@ constexpr size_t AlignMask(size_t n) { return n - 1; }
694694
} \
695695
}
696696

697+
#define ASAN_REPORT_ERROR(type, is_write, size) \
698+
ASAN_REPORT_ERROR_BASE(type, is_write, size, 0) \
699+
ASAN_REPORT_ERROR_BASE(type, is_write, size, 1) \
700+
ASAN_REPORT_ERROR_BASE(type, is_write, size, 2) \
701+
ASAN_REPORT_ERROR_BASE(type, is_write, size, 3) \
702+
ASAN_REPORT_ERROR_BASE(type, is_write, size, 4)
703+
697704
ASAN_REPORT_ERROR(load, false, 1)
698705
ASAN_REPORT_ERROR(load, false, 2)
699706
ASAN_REPORT_ERROR(load, false, 4)
@@ -705,24 +712,31 @@ ASAN_REPORT_ERROR(store, true, 4)
705712
ASAN_REPORT_ERROR(store, true, 8)
706713
ASAN_REPORT_ERROR(store, true, 16)
707714

708-
#define ASAN_REPORT_ERROR_N(type, is_write) \
709-
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N( \
710-
uptr addr, size_t size, uint32_t as, const char __SYCL_CONSTANT__ *file, \
715+
#define ASAN_REPORT_ERROR_N_BASE(type, is_write, as) \
716+
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N_as##as( \
717+
uptr addr, size_t size, const char __SYCL_CONSTANT__ *file, \
711718
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
712719
if (auto poisoned_addr = __asan_region_is_poisoned(addr, as, size)) { \
713720
__asan_report_access_error(addr, as, size, is_write, poisoned_addr, \
714721
file, line, func); \
715722
} \
716723
} \
717-
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N_noabort( \
718-
uptr addr, size_t size, uint32_t as, const char __SYCL_CONSTANT__ *file, \
724+
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N_as##as##_noabort( \
725+
uptr addr, size_t size, const char __SYCL_CONSTANT__ *file, \
719726
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
720727
if (auto poisoned_addr = __asan_region_is_poisoned(addr, as, size)) { \
721728
__asan_report_access_error(addr, as, size, is_write, poisoned_addr, \
722729
file, line, func, true); \
723730
} \
724731
}
725732

733+
#define ASAN_REPORT_ERROR_N(type, is_write) \
734+
ASAN_REPORT_ERROR_N_BASE(type, is_write, 0) \
735+
ASAN_REPORT_ERROR_N_BASE(type, is_write, 1) \
736+
ASAN_REPORT_ERROR_N_BASE(type, is_write, 2) \
737+
ASAN_REPORT_ERROR_N_BASE(type, is_write, 3) \
738+
ASAN_REPORT_ERROR_N_BASE(type, is_write, 4)
739+
726740
ASAN_REPORT_ERROR_N(load, false)
727741
ASAN_REPORT_ERROR_N(store, true)
728742

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

Lines changed: 41 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,7 @@ const char kAsanMemToShadow[] = "__asan_mem_to_shadow";
183183

184184
// Accesses sizes are powers of two: 1, 2, 4, 8, 16.
185185
static const size_t kNumberOfAccessSizes = 5;
186+
static const size_t kNumberOfAddressSpace = 5;
186187

187188
static const uint64_t kAllocaRzSize = 32;
188189

@@ -897,10 +898,13 @@ struct AddressSanitizer {
897898
// These arrays is indexed by AccessIsWrite, Experiment and log2(AccessSize).
898899
FunctionCallee AsanErrorCallback[2][2][kNumberOfAccessSizes];
899900
FunctionCallee AsanMemoryAccessCallback[2][2][kNumberOfAccessSizes];
901+
FunctionCallee AsanMemoryAccessCallbackAS[2][2][kNumberOfAccessSizes]
902+
[kNumberOfAddressSpace];
900903

901904
// These arrays is indexed by AccessIsWrite and Experiment.
902905
FunctionCallee AsanErrorCallbackSized[2][2];
903906
FunctionCallee AsanMemoryAccessCallbackSized[2][2];
907+
FunctionCallee AsanMemoryAccessCallbackSizedAS[2][2][kNumberOfAddressSpace];
904908

905909
FunctionCallee AsanMemmove, AsanMemcpy, AsanMemset;
906910
Value *LocalDynamicShadow = nullptr;
@@ -1403,7 +1407,8 @@ static void ExtendSpirKernelArgs(Module &M, FunctionAnalysisManager &FAM) {
14031407
auto *CurF = CI->getFunction();
14041408
Args.push_back(CurF->getArg(CurF->arg_size() - 1));
14051409

1406-
CallInst *NewCI = CallInst::Create(NewF, Args, CI->getName(), CI);
1410+
CallInst *NewCI =
1411+
CallInst::Create(NewF, Args, CI->getName(), CI->getIterator());
14071412
NewCI->setCallingConv(CI->getCallingConv());
14081413
NewCI->setAttributes(CI->getAttributes());
14091414
NewCI->setDebugLoc(CI->getDebugLoc());
@@ -1541,7 +1546,7 @@ static bool isJointMatrixAccess(Value *V) {
15411546
for (Value *Op : CI->args()) {
15421547
if (auto *AI = dyn_cast<AllocaInst>(Op->stripInBoundsOffsets()))
15431548
if (auto *TargetTy = getTargetExtType(AI->getAllocatedType()))
1544-
return TargetTy->getName().startswith("spirv.") &&
1549+
return TargetTy->getName().starts_with("spirv.") &&
15451550
TargetTy->getName().contains("Matrix");
15461551
}
15471552
}
@@ -1617,11 +1622,6 @@ void AddressSanitizer::AppendDebugInfoToArgs(Instruction *InsertBefore,
16171622
PointerType *ConstASPtrTy =
16181623
Type::getInt8Ty(C)->getPointerTo(kSpirOffloadConstantAS);
16191624

1620-
// Address Space
1621-
Type *PtrTy = cast<PointerType>(Addr->getType()->getScalarType());
1622-
Args.push_back(
1623-
ConstantInt::get(Type::getInt32Ty(C), PtrTy->getPointerAddressSpace()));
1624-
16251625
// File & Line
16261626
if (Loc) {
16271627
llvm::SmallString<128> Source = Loc->getDirectory();
@@ -2316,10 +2316,13 @@ void AddressSanitizer::instrumentAddress(Instruction *OrigIns,
23162316
if (Exp == 0) {
23172317
if (TargetTriple.isSPIROrSPIRV()) {
23182318
SmallVector<Value *, 5> Args;
2319+
auto AS = cast<PointerType>(Addr->getType()->getScalarType())
2320+
->getPointerAddressSpace();
23192321
Args.push_back(AddrLong);
23202322
AppendDebugInfoToArgs(InsertBefore, Addr, Args);
23212323
RTCI.createRuntimeCall(
2322-
IRB, AsanMemoryAccessCallback[IsWrite][0][AccessSizeIndex], Args);
2324+
IRB, AsanMemoryAccessCallbackAS[IsWrite][0][AccessSizeIndex][AS],
2325+
Args);
23232326
} else {
23242327
RTCI.createRuntimeCall(
23252328
IRB, AsanMemoryAccessCallback[IsWrite][0][AccessSizeIndex], AddrLong);
@@ -2398,11 +2401,13 @@ void AddressSanitizer::instrumentUnusualSizeOrAlignment(
23982401
if (Exp == 0) {
23992402
if (TargetTriple.isSPIROrSPIRV()) {
24002403
SmallVector<Value *, 6> Args;
2404+
auto AS = cast<PointerType>(Addr->getType()->getScalarType())
2405+
->getPointerAddressSpace();
24012406
Args.push_back(AddrLong);
24022407
Args.push_back(Size);
24032408
AppendDebugInfoToArgs(InsertBefore, Addr, Args);
2404-
RTCI.createRuntimeCall(IRB, AsanMemoryAccessCallbackSized[IsWrite][0],
2405-
Args);
2409+
RTCI.createRuntimeCall(
2410+
IRB, AsanMemoryAccessCallbackSizedAS[IsWrite][0][AS], Args);
24062411
} else {
24072412
RTCI.createRuntimeCall(IRB, AsanMemoryAccessCallbackSized[IsWrite][0],
24082413
{AddrLong, Size});
@@ -3280,7 +3285,6 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {
32803285

32813286
// __asan_loadX/__asan_storeX(
32823287
// ...
3283-
// int32_t as, // Address Space
32843288
// char* file,
32853289
// unsigned int line,
32863290
// char* func
@@ -3289,15 +3293,39 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {
32893293
auto *Int8PtrTy =
32903294
Type::getInt8Ty(*C)->getPointerTo(kSpirOffloadConstantAS);
32913295

3292-
Args1.push_back(Type::getInt32Ty(*C)); // address_space
32933296
Args1.push_back(Int8PtrTy); // file
32943297
Args1.push_back(Type::getInt32Ty(*C)); // line
32953298
Args1.push_back(Int8PtrTy); // func
32963299

3297-
Args2.push_back(Type::getInt32Ty(*C)); // address_space
32983300
Args2.push_back(Int8PtrTy); // file
32993301
Args2.push_back(Type::getInt32Ty(*C)); // line
33003302
Args2.push_back(Int8PtrTy); // func
3303+
3304+
for (size_t AddressSpaceIndex = 0;
3305+
AddressSpaceIndex < kNumberOfAddressSpace; AddressSpaceIndex++) {
3306+
AsanMemoryAccessCallbackSizedAS
3307+
[AccessIsWrite][Exp][AddressSpaceIndex] = M.getOrInsertFunction(
3308+
ClMemoryAccessCallbackPrefix + ExpStr + TypeStr + "N" +
3309+
"_as" + itostr(AddressSpaceIndex) + EndingStr,
3310+
FunctionType::get(IRB.getVoidTy(), Args2, false), AL2);
3311+
3312+
for (size_t AccessSizeIndex = 0;
3313+
AccessSizeIndex < kNumberOfAccessSizes; AccessSizeIndex++) {
3314+
const std::string Suffix = TypeStr +
3315+
itostr(1ULL << AccessSizeIndex) + "_as" +
3316+
itostr(AddressSpaceIndex);
3317+
AsanMemoryAccessCallbackAS[AccessIsWrite][Exp][AccessSizeIndex]
3318+
[AddressSpaceIndex] =
3319+
M.getOrInsertFunction(
3320+
ClMemoryAccessCallbackPrefix +
3321+
ExpStr + Suffix + EndingStr,
3322+
FunctionType::get(IRB.getVoidTy(),
3323+
Args1, false),
3324+
AL1);
3325+
}
3326+
}
3327+
3328+
continue;
33013329
}
33023330
AsanErrorCallbackSized[AccessIsWrite][Exp] = M.getOrInsertFunction(
33033331
kAsanReportErrorTemplate + ExpStr + TypeStr + "_n" + EndingStr,

‎llvm/test/Instrumentation/AddressSanitizer/SPIRV/sycl_esimd.ll‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ define spir_kernel void @sycl_kernel(ptr addrspace(1) %p) #0 {
88
entry:
99
%0 = load i32, ptr addrspace(1) %p, align 4
1010
; CHECK: store ptr addrspace(1) %__asan_launch, ptr addrspace(3) @__AsanLaunchInfo, align 8
11-
; CHECK: call void @__asan_load4
11+
; CHECK: call void @__asan_load4_as1
1212
ret void
1313
}
1414

0 commit comments

Comments
 (0)