Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[DeviceASAN] Implement asan_load/store for different address space #15936

Merged
merged 2 commits into from
Nov 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading