Skip to content

Commit

Permalink
[DeviceASAN] Implement asan_load/store for different address space (i…
Browse files Browse the repository at this point in the history
  • Loading branch information
zhaomaosu authored Nov 4, 2024
1 parent 691cb90 commit 644a763
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 26 deletions.
38 changes: 26 additions & 12 deletions libdevice/sanitizer_utils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -668,10 +668,10 @@ constexpr size_t AlignMask(size_t n) { return n - 1; }
/// ASAN Load/Store Report Built-ins
///

#define ASAN_REPORT_ERROR(type, is_write, size) \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size( \
uptr addr, uint32_t as, const char __SYCL_CONSTANT__ *file, \
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
#define ASAN_REPORT_ERROR_BASE(type, is_write, size, as) \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size##_as##as( \
uptr addr, const char __SYCL_CONSTANT__ *file, uint32_t line, \
const char __SYCL_CONSTANT__ *func) { \
if (addr & AlignMask(size)) { \
__asan_report_misalign_error(addr, as, size, is_write, addr, file, line, \
func); \
Expand All @@ -681,9 +681,9 @@ constexpr size_t AlignMask(size_t n) { return n - 1; }
func); \
} \
} \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size##_noabort( \
uptr addr, uint32_t as, const char __SYCL_CONSTANT__ *file, \
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##size##_as##as##_noabort( \
uptr addr, const char __SYCL_CONSTANT__ *file, uint32_t line, \
const char __SYCL_CONSTANT__ *func) { \
if (addr & AlignMask(size)) { \
__asan_report_misalign_error(addr, as, size, is_write, addr, file, line, \
func, true); \
Expand All @@ -694,6 +694,13 @@ constexpr size_t AlignMask(size_t n) { return n - 1; }
} \
}

#define ASAN_REPORT_ERROR(type, is_write, size) \
ASAN_REPORT_ERROR_BASE(type, is_write, size, 0) \
ASAN_REPORT_ERROR_BASE(type, is_write, size, 1) \
ASAN_REPORT_ERROR_BASE(type, is_write, size, 2) \
ASAN_REPORT_ERROR_BASE(type, is_write, size, 3) \
ASAN_REPORT_ERROR_BASE(type, is_write, size, 4)

ASAN_REPORT_ERROR(load, false, 1)
ASAN_REPORT_ERROR(load, false, 2)
ASAN_REPORT_ERROR(load, false, 4)
Expand All @@ -705,24 +712,31 @@ ASAN_REPORT_ERROR(store, true, 4)
ASAN_REPORT_ERROR(store, true, 8)
ASAN_REPORT_ERROR(store, true, 16)

#define ASAN_REPORT_ERROR_N(type, is_write) \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N( \
uptr addr, size_t size, uint32_t as, const char __SYCL_CONSTANT__ *file, \
#define ASAN_REPORT_ERROR_N_BASE(type, is_write, as) \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N_as##as( \
uptr addr, size_t size, const char __SYCL_CONSTANT__ *file, \
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
if (auto poisoned_addr = __asan_region_is_poisoned(addr, as, size)) { \
__asan_report_access_error(addr, as, size, is_write, poisoned_addr, \
file, line, func); \
} \
} \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N_noabort( \
uptr addr, size_t size, uint32_t as, const char __SYCL_CONSTANT__ *file, \
DEVICE_EXTERN_C_NOINLINE void __asan_##type##N_as##as##_noabort( \
uptr addr, size_t size, const char __SYCL_CONSTANT__ *file, \
uint32_t line, const char __SYCL_CONSTANT__ *func) { \
if (auto poisoned_addr = __asan_region_is_poisoned(addr, as, size)) { \
__asan_report_access_error(addr, as, size, is_write, poisoned_addr, \
file, line, func, true); \
} \
}

#define ASAN_REPORT_ERROR_N(type, is_write) \
ASAN_REPORT_ERROR_N_BASE(type, is_write, 0) \
ASAN_REPORT_ERROR_N_BASE(type, is_write, 1) \
ASAN_REPORT_ERROR_N_BASE(type, is_write, 2) \
ASAN_REPORT_ERROR_N_BASE(type, is_write, 3) \
ASAN_REPORT_ERROR_N_BASE(type, is_write, 4)

ASAN_REPORT_ERROR_N(load, false)
ASAN_REPORT_ERROR_N(store, true)

Expand Down
54 changes: 41 additions & 13 deletions llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,7 @@ const char kAsanMemToShadow[] = "__asan_mem_to_shadow";

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

static const uint64_t kAllocaRzSize = 32;

Expand Down Expand Up @@ -897,10 +898,13 @@ struct AddressSanitizer {
// These arrays is indexed by AccessIsWrite, Experiment and log2(AccessSize).
FunctionCallee AsanErrorCallback[2][2][kNumberOfAccessSizes];
FunctionCallee AsanMemoryAccessCallback[2][2][kNumberOfAccessSizes];
FunctionCallee AsanMemoryAccessCallbackAS[2][2][kNumberOfAccessSizes]
[kNumberOfAddressSpace];

// These arrays is indexed by AccessIsWrite and Experiment.
FunctionCallee AsanErrorCallbackSized[2][2];
FunctionCallee AsanMemoryAccessCallbackSized[2][2];
FunctionCallee AsanMemoryAccessCallbackSizedAS[2][2][kNumberOfAddressSpace];

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

CallInst *NewCI = CallInst::Create(NewF, Args, CI->getName(), CI);
CallInst *NewCI =
CallInst::Create(NewF, Args, CI->getName(), CI->getIterator());
NewCI->setCallingConv(CI->getCallingConv());
NewCI->setAttributes(CI->getAttributes());
NewCI->setDebugLoc(CI->getDebugLoc());
Expand Down Expand Up @@ -1547,7 +1552,7 @@ static bool isJointMatrixAccess(Value *V) {
for (Value *Op : CI->args()) {
if (auto *AI = dyn_cast<AllocaInst>(Op->stripInBoundsOffsets()))
if (auto *TargetTy = getTargetExtType(AI->getAllocatedType()))
return TargetTy->getName().startswith("spirv.") &&
return TargetTy->getName().starts_with("spirv.") &&
TargetTy->getName().contains("Matrix");
}
}
Expand Down Expand Up @@ -1623,11 +1628,6 @@ void AddressSanitizer::AppendDebugInfoToArgs(Instruction *InsertBefore,
PointerType *ConstASPtrTy =
Type::getInt8Ty(C)->getPointerTo(kSpirOffloadConstantAS);

// Address Space
Type *PtrTy = cast<PointerType>(Addr->getType()->getScalarType());
Args.push_back(
ConstantInt::get(Type::getInt32Ty(C), PtrTy->getPointerAddressSpace()));

// File & Line
if (Loc) {
llvm::SmallString<128> Source = Loc->getDirectory();
Expand Down Expand Up @@ -2322,10 +2322,13 @@ void AddressSanitizer::instrumentAddress(Instruction *OrigIns,
if (Exp == 0) {
if (TargetTriple.isSPIROrSPIRV()) {
SmallVector<Value *, 5> Args;
auto AS = cast<PointerType>(Addr->getType()->getScalarType())
->getPointerAddressSpace();
Args.push_back(AddrLong);
AppendDebugInfoToArgs(InsertBefore, Addr, Args);
RTCI.createRuntimeCall(
IRB, AsanMemoryAccessCallback[IsWrite][0][AccessSizeIndex], Args);
IRB, AsanMemoryAccessCallbackAS[IsWrite][0][AccessSizeIndex][AS],
Args);
} else {
RTCI.createRuntimeCall(
IRB, AsanMemoryAccessCallback[IsWrite][0][AccessSizeIndex], AddrLong);
Expand Down Expand Up @@ -2404,11 +2407,13 @@ void AddressSanitizer::instrumentUnusualSizeOrAlignment(
if (Exp == 0) {
if (TargetTriple.isSPIROrSPIRV()) {
SmallVector<Value *, 6> Args;
auto AS = cast<PointerType>(Addr->getType()->getScalarType())
->getPointerAddressSpace();
Args.push_back(AddrLong);
Args.push_back(Size);
AppendDebugInfoToArgs(InsertBefore, Addr, Args);
RTCI.createRuntimeCall(IRB, AsanMemoryAccessCallbackSized[IsWrite][0],
Args);
RTCI.createRuntimeCall(
IRB, AsanMemoryAccessCallbackSizedAS[IsWrite][0][AS], Args);
} else {
RTCI.createRuntimeCall(IRB, AsanMemoryAccessCallbackSized[IsWrite][0],
{AddrLong, Size});
Expand Down Expand Up @@ -3286,7 +3291,6 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {

// __asan_loadX/__asan_storeX(
// ...
// int32_t as, // Address Space
// char* file,
// unsigned int line,
// char* func
Expand All @@ -3295,15 +3299,39 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {
auto *Int8PtrTy =
Type::getInt8Ty(*C)->getPointerTo(kSpirOffloadConstantAS);

Args1.push_back(Type::getInt32Ty(*C)); // address_space
Args1.push_back(Int8PtrTy); // file
Args1.push_back(Type::getInt32Ty(*C)); // line
Args1.push_back(Int8PtrTy); // func

Args2.push_back(Type::getInt32Ty(*C)); // address_space
Args2.push_back(Int8PtrTy); // file
Args2.push_back(Type::getInt32Ty(*C)); // line
Args2.push_back(Int8PtrTy); // func

for (size_t AddressSpaceIndex = 0;
AddressSpaceIndex < kNumberOfAddressSpace; AddressSpaceIndex++) {
AsanMemoryAccessCallbackSizedAS
[AccessIsWrite][Exp][AddressSpaceIndex] = M.getOrInsertFunction(
ClMemoryAccessCallbackPrefix + ExpStr + TypeStr + "N" +
"_as" + itostr(AddressSpaceIndex) + EndingStr,
FunctionType::get(IRB.getVoidTy(), Args2, false), AL2);

for (size_t AccessSizeIndex = 0;
AccessSizeIndex < kNumberOfAccessSizes; AccessSizeIndex++) {
const std::string Suffix = TypeStr +
itostr(1ULL << AccessSizeIndex) + "_as" +
itostr(AddressSpaceIndex);
AsanMemoryAccessCallbackAS[AccessIsWrite][Exp][AccessSizeIndex]
[AddressSpaceIndex] =
M.getOrInsertFunction(
ClMemoryAccessCallbackPrefix +
ExpStr + Suffix + EndingStr,
FunctionType::get(IRB.getVoidTy(),
Args1, false),
AL1);
}
}

continue;
}
AsanErrorCallbackSized[AccessIsWrite][Exp] = M.getOrInsertFunction(
kAsanReportErrorTemplate + ExpStr + TypeStr + "_n" + EndingStr,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ define spir_kernel void @sycl_kernel(ptr addrspace(1) %p) #0 {
entry:
%0 = load i32, ptr addrspace(1) %p, align 4
; CHECK-NOT: store ptr addrspace(1) %__asan_launch, ptr addrspace(3) @__AsanLaunchInfo, align 8
; CHECK-NOT: call void @__asan_load4
; CHECK-NOT: call void @__asan_load4_as1
ret void
}

Expand Down

0 comments on commit 644a763

Please sign in to comment.