Skip to content

[DeviceSanitizer] Add ASAN interception for SPIRV 2D block load/store #22041

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.

Sign up for GitHub

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

Merged
KornevNikita merged 3 commits into from
May 20, 2026
Merged

[DeviceSanitizer] Add ASAN interception for SPIRV 2D block load/store #22041

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
225513a
[DeviceSanitizer] Add ASAN interception for SPIRV 2D block load/store
zhaomaosu May 18, 2026
156ecbb
Fix format
zhaomaosu May 18, 2026
45cb0f9
update tests
zhaomaosu May 18, 2026
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
66 changes: 66 additions & 0 deletions libdevice/sanitizer/asan_rtl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1029,4 +1029,70 @@ ASAN_MEMMOVE(1)
ASAN_MEMMOVE(3)
ASAN_MEMMOVE(4)

static void Block2DCheck(uptr surface_base, uptr block_ptr, int element_size,
int block_width, int block_height, int block_count,
int surface_pitch, int coord_x, int coord_y,
bool is_load, const char __SYCL_CONSTANT__ *file,
uint32_t line, const char __SYCL_CONSTANT__ *func) {
// Per SPV_INTEL_2d_block_io spec:
// element_size: bytes per element
// block_width: elements per row
// surface_pitch: bytes between rows
// coord_x: element offset, coord_y: row offset
int row_width_bytes = block_width * element_size * block_count;
size_t block_size = (size_t)row_width_bytes * block_height;

{
DebugInfo debug{block_ptr, /*as=*/4, block_size, is_load, file, func, line};
if (auto poisoned_addr =
IsRegionPoisoned(block_ptr, 4, block_size, &debug)) {
ReportAccessError(poisoned_addr, 4, false, &debug);
return;
}
}

uptr start_addr = surface_base + (uptr)coord_y * surface_pitch +
(uptr)coord_x * element_size;
for (int row = 0; row < block_height; row++) {
uptr row_addr = start_addr + (uptr)row * surface_pitch;
DebugInfo debug{row_addr, /*as=*/4, (size_t)row_width_bytes, !is_load, file,
func, line};
if (auto poisoned_addr =
IsRegionPoisoned(row_addr, 4, (size_t)row_width_bytes, &debug)) {
ReportAccessError(poisoned_addr, 4, false, &debug);
return;
}
}
}

DEVICE_EXTERN_C_NOINLINE void __asan_block2d_load_check(
__attribute__((address_space(4))) const char *src_base_ptr, char *dst_ptr,
int element_size, int block_width, int block_height, int block_count,
int surface_width, int surface_height, int surface_pitch, int coord_x,
int coord_y, const char __SYCL_CONSTANT__ *file, uint32_t line,
const char __SYCL_CONSTANT__ *func) {
if (!__AsanLaunchInfo)
return;
if (__spirv_BuiltInSubgroupLocalInvocationId() != 0)
return;
Block2DCheck((uptr)src_base_ptr, (uptr)dst_ptr, element_size, block_width,
block_height, block_count, surface_pitch, coord_x, coord_y,
/*is_load=*/true, file, line, func);
}

DEVICE_EXTERN_C_NOINLINE void __asan_block2d_store_check(
__attribute__((address_space(4))) const char *dst_base_ptr, char *src_ptr,
int element_size, int block_width, int block_height, int block_count,
int surface_width, int surface_height, int surface_pitch, int coord_x,
int coord_y, const char __SYCL_CONSTANT__ *file, uint32_t line,
const char __SYCL_CONSTANT__ *func) {
if (!__AsanLaunchInfo)
return;
if (__spirv_BuiltInSubgroupLocalInvocationId() != 0)
return;
Block2DCheck((uptr)dst_base_ptr, (uptr)src_ptr, element_size, block_width,
block_height, block_count, surface_pitch, coord_x, coord_y,
/*is_load=*/false, file, line, func);
}

#endif // __SPIR__ || __SPIRV__
93 changes: 90 additions & 3 deletions llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -900,6 +900,7 @@ struct AddressSanitizer {
Value *SizeArgument, uint32_t Exp,
RuntimeCallInserter &RTCI);
void instrumentMemIntrinsic(MemIntrinsic *MI, RuntimeCallInserter &RTCI);
void instrumentBlock2DOp(CallBase *CB, RuntimeCallInserter &RTCI);
Value *memToShadow(Value *Shadow, IRBuilder<> &IRB,
uint32_t AddressSpace = kSpirOffloadPrivateAS);
bool suppressInstrumentationSiteForDebug(int &Instrumented);
Expand Down Expand Up @@ -977,6 +978,7 @@ struct AddressSanitizer {
FunctionCallee AsanMemcpyAS[kNumberOfAddressSpace][kNumberOfAddressSpace],
AsanMemmoveAS[kNumberOfAddressSpace][kNumberOfAddressSpace],
AsanMemsetAS[kNumberOfAddressSpace];
FunctionCallee AsanSGBlock2DLoadCheck, AsanSGBlock2DStoreCheck;
Value *LocalDynamicShadow = nullptr;
const StackSafetyGlobalInfo *SSGI;
DenseMap<const AllocaInst *, bool> ProcessedAllocas;
Expand Down Expand Up @@ -1999,6 +2001,50 @@ void AddressSanitizer::instrumentMemIntrinsic(MemIntrinsic *MI,
MI->eraseFromParent();
}

void AddressSanitizer::instrumentBlock2DOp(CallBase *CB,
RuntimeCallInserter &RTCI) {
InstrumentationIRBuilder IRB(CB);
Function *Callee = CB->getCalledFunction();
StringRef Name = Callee->getName();
bool IsStore = Name.contains("__spirv_Subgroup2DBlockStore");

// Per SPIRVBuiltins.td, both load and store have 10 args:
// Load: (elem_size, width, height, count, src_base_ptr, surf_w, surf_h,
// surf_pitch, <2xi32> coord, dst_ptr)
// Store: (elem_size, width, height, count, src_ptr, dst_base_ptr, surf_w,
// surf_h, surf_pitch, <2xi32> coord)
Value *BasePtr = IsStore ? CB->getArgOperand(5) : CB->getArgOperand(4);
Value *BlockPtr = IsStore ? CB->getArgOperand(4) : CB->getArgOperand(9);
Value *ElemSize = CB->getArgOperand(0);
Value *BlockWidth = CB->getArgOperand(1);
Value *BlockHeight = CB->getArgOperand(2);
Value *BlockCount = CB->getArgOperand(3);
Value *SurfWidth = IsStore ? CB->getArgOperand(6) : CB->getArgOperand(5);
Value *SurfHeight = IsStore ? CB->getArgOperand(7) : CB->getArgOperand(6);
Value *SurfPitch = IsStore ? CB->getArgOperand(8) : CB->getArgOperand(7);
Value *Coord = IsStore ? CB->getArgOperand(9) : CB->getArgOperand(8);

Value *CoordX = IRB.CreateExtractElement(Coord, uint64_t(0));
Value *CoordY = IRB.CreateExtractElement(Coord, uint64_t(1));

SmallVector<Value *, 14> Args;
Args.push_back(BasePtr);
Args.push_back(BlockPtr);
Args.push_back(ElemSize);
Args.push_back(BlockWidth);
Args.push_back(BlockHeight);
Args.push_back(BlockCount);
Args.push_back(SurfWidth);
Args.push_back(SurfHeight);
Args.push_back(SurfPitch);
Args.push_back(CoordX);
Args.push_back(CoordY);
AppendDebugInfoToArgs(CB, Args);

RTCI.createRuntimeCall(
IRB, IsStore ? AsanSGBlock2DStoreCheck : AsanSGBlock2DLoadCheck, Args);
}

/// Check if we want (and can) handle this alloca.
bool AddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
auto [It, Inserted] = ProcessedAllocas.try_emplace(&AI);
Expand Down Expand Up @@ -3857,6 +3903,32 @@ void AddressSanitizer::initializeCallbacks(const TargetLibraryInfo *TLI) {
IRB.getInt32Ty(), Int8PtrTy);
}
}

// 2D block load/store checks
// __asan_block2d_{load,store}_check(
// ptr base_ptr, // surface base address (generic AS 4)
// ptr block_ptr, // dst for load, src for store (generic AS 4)
// i32 element_size, // bytes per element
// i32 block_width, // elements per row
// i32 block_height, // rows per block
// i32 block_count, // number of blocks
// i32 surface_width, // surface width in bytes
// i32 surface_height, // surface height in rows
// i32 surface_pitch, // bytes between rows
// i32 coord_x, // element offset
// i32 coord_y, // row offset
// ptr file, i32 line, ptr func)
auto *GenericPtrTy = PointerType::get(*C, kSpirOffloadGenericAS);
auto GetBlock2DCheckFunc = [&](StringRef Name) {
return M.getOrInsertFunction(
Name, IRB.getVoidTy(), GenericPtrTy, GenericPtrTy, IRB.getInt32Ty(),
IRB.getInt32Ty(), IRB.getInt32Ty(), IRB.getInt32Ty(),
IRB.getInt32Ty(), IRB.getInt32Ty(), IRB.getInt32Ty(),
IRB.getInt32Ty(), IRB.getInt32Ty(), Int8PtrTy, IRB.getInt32Ty(),
Int8PtrTy);
};
AsanSGBlock2DLoadCheck = GetBlock2DCheckFunc("__asan_block2d_load_check");
AsanSGBlock2DStoreCheck = GetBlock2DCheckFunc("__asan_block2d_store_check");
}

const std::string MemIntrinCallbackPrefix =
Expand Down Expand Up @@ -4057,6 +4129,7 @@ bool AddressSanitizer::instrumentFunction(Function &F,
SmallPtrSet<Value *, 16> TempsToInstrument;
SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
SmallVector<CallBase *, 8> Block2DToInstrument;
SmallVector<Instruction *, 8> NoReturnCalls;
SmallVector<BasicBlock *, 16> AllBlocks;
SmallVector<Instruction *, 16> PointerComparisonsOrSubtracts;
Expand Down Expand Up @@ -4103,9 +4176,18 @@ bool AddressSanitizer::instrumentFunction(Function &F,
NumInsnsPerBB++;
} else {
if (auto *CB = dyn_cast<CallBase>(&Inst)) {
// On device side, the only non return cases should be *.trap or
// assert, and none of these cases need to be handles.
if (!TargetTriple.isSPIROrSPIRV()) {
if (TargetTriple.isSPIROrSPIRV()) {
if (Function *F = CB->getCalledFunction()) {
StringRef Name = F->getName();
if (Name.contains("__spirv_Subgroup2DBlockLoad") ||
Name.contains("__spirv_Subgroup2DBlockStore")) {
Block2DToInstrument.push_back(CB);
NumInsnsPerBB++;
}
}
} else {
// On device side, the only non return cases should be *.trap or
// assert, and none of these cases need to be handles.
// A call inside BB.
TempsToInstrument.clear();
if (CB->doesNotReturn())
Expand Down Expand Up @@ -4138,6 +4220,11 @@ bool AddressSanitizer::instrumentFunction(Function &F,
instrumentMemIntrinsic(Inst, RTCI);
FunctionModified = true;
}
for (auto *CB : Block2DToInstrument) {
if (!suppressInstrumentationSiteForDebug(NumInstrumented))
instrumentBlock2DOp(CB, RTCI);
FunctionModified = true;
}

FunctionStackPoisoner FSP(F, *this, RTCI);
bool ChangedStack = FSP.runOnFunction();
Expand Down
103 changes: 103 additions & 0 deletions llvm/test/Instrumentation/AddressSanitizer/SPIRV/block2d_load_store.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 -asan-stack=0 -asan-globals=0 -asan-constructor-kind=none -asan-use-after-return=never -S | FileCheck %s

target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
target triple = "spir64-unknown-unknown"

declare dso_local spir_func void @_Z32__spirv_Subgroup2DBlockLoadINTELiiiiPKviiiDv2_iPv(i32 noundef, i32 noundef, i32 noundef, i32 noundef, ptr addrspace(4) noundef, i32 noundef, i32 noundef, i32 noundef, <2 x i32> noundef, ptr addrspace(4) noundef) local_unnamed_addr

declare dso_local spir_func void @_Z33__spirv_Subgroup2DBlockStoreINTELiiiiPKviPviiiDv2_i(i32 noundef, i32 noundef, i32 noundef, i32 noundef, ptr addrspace(4) noundef, ptr addrspace(4) noundef, i32 noundef, i32 noundef, i32 noundef, <2 x i32> noundef) local_unnamed_addr

; Test: Function without sanitize_address should NOT be instrumented
define weak_odr dso_local spir_kernel void @test_block2d_no_sanitize(ptr addrspace(4) %base, ptr addrspace(4) %dst) {
; CHECK-LABEL: define weak_odr dso_local spir_kernel void @test_block2d_no_sanitize(
; CHECK-SAME: ptr addrspace(4) [[BASE:%.*]], ptr addrspace(4) [[DST:%.*]]) {
; CHECK-NEXT: [[ENTRY:.*:]]
; CHECK-NEXT: store ptr addrspace(1) null, ptr addrspace(3) @__AsanLaunchInfo, align 8
; CHECK-NEXT: [[__PRIVATE_BASE:%.*]] = alloca i8, align 1
; CHECK-NEXT: call void @__asan_set_private_base(ptr [[__PRIVATE_BASE]])
; CHECK-NEXT: [[COORD:%.*]] = insertelement <2 x i32> poison, i32 0, i64 0
; CHECK-NEXT: [[COORD1:%.*]] = insertelement <2 x i32> [[COORD]], i32 0, i64 1
; CHECK-NEXT: call spir_func void @_Z32__spirv_Subgroup2DBlockLoadINTELiiiiPKviiiDv2_iPv(i32 16, i32 32, i32 32, i32 1, ptr addrspace(4) [[BASE]], i32 2047, i32 1023, i32 2047, <2 x i32> [[COORD1]], ptr addrspace(4) [[DST]])
; CHECK-NEXT: ret void
;
entry:
%coord = insertelement <2 x i32> poison, i32 0, i64 0
%coord1 = insertelement <2 x i32> %coord, i32 0, i64 1
call spir_func void @_Z32__spirv_Subgroup2DBlockLoadINTELiiiiPKviiiDv2_iPv(i32 16, i32 32, i32 32, i32 1, ptr addrspace(4) %base, i32 2047, i32 1023, i32 2047, <2 x i32> %coord1, ptr addrspace(4) %dst)
ret void
}

; Test: 2D block load should be instrumented with __asan_block2d_load_check
; Function Attrs: sanitize_address
define weak_odr dso_local spir_kernel void @test_block2d_load(ptr addrspace(4) %base, ptr addrspace(4) %dst) #0 {
; CHECK-LABEL: define weak_odr dso_local spir_kernel void @test_block2d_load(
; CHECK-SAME: ptr addrspace(4) [[BASE:%.*]], ptr addrspace(4) [[DST:%.*]], ptr addrspace(1) noundef [[__ASAN_LAUNCH:%.*]]) #[[ATTR1:[0-9]+]] {
; CHECK-NEXT: [[ENTRY:.*:]]
; CHECK-NEXT: store ptr addrspace(1) [[__ASAN_LAUNCH]], ptr addrspace(3) @__AsanLaunchInfo, align 8
; CHECK-NEXT: [[__PRIVATE_BASE:%.*]] = alloca i8, align 1
; CHECK-NEXT: call void @__asan_set_private_base(ptr [[__PRIVATE_BASE]])
; CHECK-NEXT: [[COORD:%.*]] = insertelement <2 x i32> poison, i32 0, i64 0
; CHECK-NEXT: [[COORD1:%.*]] = insertelement <2 x i32> [[COORD]], i32 0, i64 1
; CHECK-NEXT: [[TMP0:%.*]] = extractelement <2 x i32> [[COORD1]], i64 0
; CHECK-NEXT: [[TMP1:%.*]] = extractelement <2 x i32> [[COORD1]], i64 1
; CHECK-NEXT: call void @__asan_block2d_load_check(ptr addrspace(4) [[BASE]], ptr addrspace(4) [[DST]], i32 16, i32 32, i32 32, i32 1, i32 2047, i32 1023, i32 2047, i32 [[TMP0]], i32 [[TMP1]], ptr addrspace(2) null, i32 0, ptr addrspace(2) @__asan_kernel)
; CHECK-NEXT: call spir_func void @_Z32__spirv_Subgroup2DBlockLoadINTELiiiiPKviiiDv2_iPv(i32 16, i32 32, i32 32, i32 1, ptr addrspace(4) [[BASE]], i32 2047, i32 1023, i32 2047, <2 x i32> [[COORD1]], ptr addrspace(4) [[DST]])
; CHECK-NEXT: ret void
;
entry:
%coord = insertelement <2 x i32> poison, i32 0, i64 0
%coord1 = insertelement <2 x i32> %coord, i32 0, i64 1
call spir_func void @_Z32__spirv_Subgroup2DBlockLoadINTELiiiiPKviiiDv2_iPv(i32 16, i32 32, i32 32, i32 1, ptr addrspace(4) %base, i32 2047, i32 1023, i32 2047, <2 x i32> %coord1, ptr addrspace(4) %dst)
ret void
}

; Test: 2D block store should be instrumented with __asan_block2d_store_check
; Function Attrs: sanitize_address
define weak_odr dso_local spir_kernel void @test_block2d_store(ptr addrspace(4) %src, ptr addrspace(4) %base) #0 {
; CHECK-LABEL: define weak_odr dso_local spir_kernel void @test_block2d_store(
; CHECK-SAME: ptr addrspace(4) [[SRC:%.*]], ptr addrspace(4) [[BASE:%.*]], ptr addrspace(1) noundef [[__ASAN_LAUNCH:%.*]]) #[[ATTR1]] {
; CHECK-NEXT: [[ENTRY:.*:]]
; CHECK-NEXT: store ptr addrspace(1) [[__ASAN_LAUNCH]], ptr addrspace(3) @__AsanLaunchInfo, align 8
; CHECK-NEXT: [[__PRIVATE_BASE:%.*]] = alloca i8, align 1
; CHECK-NEXT: call void @__asan_set_private_base(ptr [[__PRIVATE_BASE]])
; CHECK-NEXT: [[COORD:%.*]] = insertelement <2 x i32> poison, i32 64, i64 0
; CHECK-NEXT: [[COORD1:%.*]] = insertelement <2 x i32> [[COORD]], i32 128, i64 1
; CHECK-NEXT: [[TMP0:%.*]] = extractelement <2 x i32> [[COORD1]], i64 0
; CHECK-NEXT: [[TMP1:%.*]] = extractelement <2 x i32> [[COORD1]], i64 1
; CHECK-NEXT: call void @__asan_block2d_store_check(ptr addrspace(4) [[BASE]], ptr addrspace(4) [[SRC]], i32 16, i32 32, i32 8, i32 1, i32 2047, i32 1023, i32 2047, i32 [[TMP0]], i32 [[TMP1]], ptr addrspace(2) null, i32 0, ptr addrspace(2) @__asan_kernel.1)
; CHECK-NEXT: call spir_func void @_Z33__spirv_Subgroup2DBlockStoreINTELiiiiPKviPviiiDv2_i(i32 16, i32 32, i32 8, i32 1, ptr addrspace(4) [[SRC]], ptr addrspace(4) [[BASE]], i32 2047, i32 1023, i32 2047, <2 x i32> [[COORD1]])
; CHECK-NEXT: ret void
;
entry:
%coord = insertelement <2 x i32> poison, i32 64, i64 0
%coord1 = insertelement <2 x i32> %coord, i32 128, i64 1
call spir_func void @_Z33__spirv_Subgroup2DBlockStoreINTELiiiiPKviPviiiDv2_i(i32 16, i32 32, i32 8, i32 1, ptr addrspace(4) %src, ptr addrspace(4) %base, i32 2047, i32 1023, i32 2047, <2 x i32> %coord1)
ret void
}

; Test: 2D block load with non-zero coords extracts x and y correctly
; Function Attrs: sanitize_address
define weak_odr dso_local spir_kernel void @test_block2d_load_nonzero_coord(ptr addrspace(4) %base, ptr addrspace(4) %dst) #0 {
; CHECK-LABEL: define weak_odr dso_local spir_kernel void @test_block2d_load_nonzero_coord(
; CHECK-SAME: ptr addrspace(4) [[BASE:%.*]], ptr addrspace(4) [[DST:%.*]], ptr addrspace(1) noundef [[__ASAN_LAUNCH:%.*]]) #[[ATTR1]] {
; CHECK-NEXT: [[ENTRY:.*:]]
; CHECK-NEXT: store ptr addrspace(1) [[__ASAN_LAUNCH]], ptr addrspace(3) @__AsanLaunchInfo, align 8
; CHECK-NEXT: [[__PRIVATE_BASE:%.*]] = alloca i8, align 1
; CHECK-NEXT: call void @__asan_set_private_base(ptr [[__PRIVATE_BASE]])
; CHECK-NEXT: [[COORD:%.*]] = insertelement <2 x i32> poison, i32 16, i64 0
; CHECK-NEXT: [[COORD1:%.*]] = insertelement <2 x i32> [[COORD]], i32 32, i64 1
; CHECK-NEXT: [[TMP0:%.*]] = extractelement <2 x i32> [[COORD1]], i64 0
; CHECK-NEXT: [[TMP1:%.*]] = extractelement <2 x i32> [[COORD1]], i64 1
; CHECK-NEXT: call void @__asan_block2d_load_check(ptr addrspace(4) [[BASE]], ptr addrspace(4) [[DST]], i32 2, i32 16, i32 32, i32 2, i32 1023, i32 511, i32 1023, i32 [[TMP0]], i32 [[TMP1]], ptr addrspace(2) null, i32 0, ptr addrspace(2) @__asan_kernel.2)
; CHECK-NEXT: call spir_func void @_Z32__spirv_Subgroup2DBlockLoadINTELiiiiPKviiiDv2_iPv(i32 2, i32 16, i32 32, i32 2, ptr addrspace(4) [[BASE]], i32 1023, i32 511, i32 1023, <2 x i32> [[COORD1]], ptr addrspace(4) [[DST]])
; CHECK-NEXT: ret void
;
entry:
%coord = insertelement <2 x i32> poison, i32 16, i64 0
%coord1 = insertelement <2 x i32> %coord, i32 32, i64 1
call spir_func void @_Z32__spirv_Subgroup2DBlockLoadINTELiiiiPKviiiDv2_iPv(i32 2, i32 16, i32 32, i32 2, ptr addrspace(4) %base, i32 1023, i32 511, i32 1023, <2 x i32> %coord1, ptr addrspace(4) %dst)
ret void
}

attributes #0 = { sanitize_address }
Loading
Loading