From 644a7631b27c792d96a62c0e0af8a9b29975e157 Mon Sep 17 00:00:00 2001 From: Maosu Zhao Date: Mon, 4 Nov 2024 22:33:22 +0800 Subject: [PATCH] [DeviceASAN] Implement asan_load/store for different address space (#15936) --- libdevice/sanitizer_utils.cpp | 38 ++++++++----- .../Instrumentation/AddressSanitizer.cpp | 54 ++++++++++++++----- .../AddressSanitizer/SPIRV/sycl_esimd.ll | 2 +- 3 files changed, 68 insertions(+), 26 deletions(-) diff --git a/libdevice/sanitizer_utils.cpp b/libdevice/sanitizer_utils.cpp index e5d5ac5afdc30..097cd97a4f706 100644 --- a/libdevice/sanitizer_utils.cpp +++ b/libdevice/sanitizer_utils.cpp @@ -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); \ @@ -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); \ @@ -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) @@ -705,17 +712,17 @@ 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, \ @@ -723,6 +730,13 @@ ASAN_REPORT_ERROR(store, true, 16) } \ } +#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) diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp index 7ade5db5046a6..338d8d0ad717c 100644 --- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp @@ -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; @@ -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; @@ -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()); @@ -1547,7 +1552,7 @@ static bool isJointMatrixAccess(Value *V) { for (Value *Op : CI->args()) { if (auto *AI = dyn_cast(Op->stripInBoundsOffsets())) if (auto *TargetTy = getTargetExtType(AI->getAllocatedType())) - return TargetTy->getName().startswith("spirv.") && + return TargetTy->getName().starts_with("spirv.") && TargetTy->getName().contains("Matrix"); } } @@ -1623,11 +1628,6 @@ void AddressSanitizer::AppendDebugInfoToArgs(Instruction *InsertBefore, PointerType *ConstASPtrTy = Type::getInt8Ty(C)->getPointerTo(kSpirOffloadConstantAS); - // Address Space - Type *PtrTy = cast(Addr->getType()->getScalarType()); - Args.push_back( - ConstantInt::get(Type::getInt32Ty(C), PtrTy->getPointerAddressSpace())); - // File & Line if (Loc) { llvm::SmallString<128> Source = Loc->getDirectory(); @@ -2322,10 +2322,13 @@ void AddressSanitizer::instrumentAddress(Instruction *OrigIns, if (Exp == 0) { if (TargetTriple.isSPIROrSPIRV()) { SmallVector Args; + auto AS = cast(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); @@ -2404,11 +2407,13 @@ void AddressSanitizer::instrumentUnusualSizeOrAlignment( if (Exp == 0) { if (TargetTriple.isSPIROrSPIRV()) { SmallVector Args; + auto AS = cast(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}); @@ -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 @@ -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, diff --git a/llvm/test/Instrumentation/AddressSanitizer/SPIRV/sycl_esimd.ll b/llvm/test/Instrumentation/AddressSanitizer/SPIRV/sycl_esimd.ll index 77b73fc5b67dd..ccb3de1092b4c 100644 --- a/llvm/test/Instrumentation/AddressSanitizer/SPIRV/sycl_esimd.ll +++ b/llvm/test/Instrumentation/AddressSanitizer/SPIRV/sycl_esimd.ll @@ -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 }