Skip to content

Commit 5d227ee

Browse files
ampandey-AMDampandey-1995
authored and
zGoldthorpe
committed
[ASan] Intercept rocr api's 'hsa_amd_vmem_address_reserve_align' and … (llvm#3870)
…(llvm#3208) 'hsa_vmem_address_free'. Implement interception of 'hsa_amd_vmem_address_reserve_align' and 'hsa_vmem_address_free' so as to support ASan overflow errors for memory allocated via 'hipMallocManaged'. [Ticket: SWDEV-483895] --------- Co-authored-by: Amit Pandey <[email protected]>
1 parent c75c589 commit 5d227ee

File tree

6 files changed

+132
-23
lines changed

6 files changed

+132
-23
lines changed

‎compiler-rt/lib/asan/asan_allocator.cpp‎

Lines changed: 60 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1294,6 +1294,9 @@ DECLARE_REAL(hsa_status_t, hsa_amd_ipc_memory_attach,
12941294
const hsa_amd_ipc_memory_t *handle, size_t len, uint32_t num_agents,
12951295
const hsa_agent_t *mapping_agents, void **mapped_ptr)
12961296
DECLARE_REAL(hsa_status_t, hsa_amd_ipc_memory_detach, void *mapped_ptr)
1297+
DECLARE_REAL(hsa_status_t, hsa_amd_vmem_address_reserve_align, void** ptr,
1298+
size_t size, uint64_t address, uint64_t alignment, uint64_t flags)
1299+
DECLARE_REAL(hsa_status_t, hsa_amd_vmem_address_free, void* ptr, size_t size);
12971300

12981301
namespace __asan {
12991302

@@ -1322,21 +1325,17 @@ hsa_status_t asan_hsa_amd_memory_pool_free(
13221325
if (p) {
13231326
instance.Deallocate(ptr, 0, 0, stack, FROM_MALLOC);
13241327
return HSA_STATUS_SUCCESS;
1325-
} else {
1326-
return REAL(hsa_amd_memory_pool_free)(ptr);
13271328
}
1329+
return REAL(hsa_amd_memory_pool_free)(ptr);
13281330
}
13291331

13301332
hsa_status_t asan_hsa_amd_agents_allow_access(
13311333
uint32_t num_agents, const hsa_agent_t *agents, const uint32_t *flags,
13321334
const void *ptr,
13331335
BufferedStackTrace *stack) {
13341336
void *p = get_allocator().GetBlockBegin(ptr);
1335-
if (p) {
1336-
return REAL(hsa_amd_agents_allow_access)(num_agents, agents, flags, p);
1337-
} else {
1338-
return REAL(hsa_amd_agents_allow_access)(num_agents, agents, flags, ptr);
1339-
}
1337+
return REAL(hsa_amd_agents_allow_access)(num_agents, agents, flags,
1338+
p ? p : ptr);
13401339
}
13411340

13421341
// For asan allocator, kMetadataSize is 0 and maximum redzone size is 2048. This
@@ -1384,5 +1383,59 @@ hsa_status_t asan_hsa_amd_ipc_memory_detach(void *mapped_ptr) {
13841383
reinterpret_cast<void *>(reinterpret_cast<uptr>(mapped_ptr) - kPageSize_);
13851384
return REAL(hsa_amd_ipc_memory_detach)(mapped_ptr_);
13861385
}
1386+
1387+
hsa_status_t asan_hsa_amd_vmem_address_reserve_align(
1388+
void** ptr, size_t size, uint64_t address, uint64_t alignment,
1389+
uint64_t flags, BufferedStackTrace* stack) {
1390+
// Bypass the tracking for a fixed address since it cannot be supported.
1391+
// Reasons:
1392+
// 1. Address may not meet the alignment/page-size requirement.
1393+
// 2. Requested range overlaps an existing reserved/mapped range.
1394+
// 3. Insufficient VA space to honor that exact placement.
1395+
if (address)
1396+
return REAL(hsa_amd_vmem_address_reserve_align)(ptr, size, address,
1397+
alignment, flags);
1398+
1399+
if (alignment < kPageSize_)
1400+
alignment = kPageSize_;
1401+
1402+
if (UNLIKELY(!IsPowerOfTwo(alignment))) {
1403+
errno = errno_EINVAL;
1404+
return HSA_STATUS_ERROR_INVALID_ARGUMENT;
1405+
}
1406+
1407+
AmdgpuAllocationInfo aa_info;
1408+
aa_info.alloc_func =
1409+
reinterpret_cast<void*>(asan_hsa_amd_vmem_address_reserve_align);
1410+
aa_info.memory_pool = {0};
1411+
aa_info.size = size;
1412+
aa_info.flags64 = flags;
1413+
aa_info.address = 0;
1414+
aa_info.alignment = alignment;
1415+
aa_info.ptr = nullptr;
1416+
SetErrnoOnNull(*ptr = instance.Allocate(size, alignment, stack, FROM_MALLOC,
1417+
false, &aa_info));
1418+
1419+
return aa_info.status;
1420+
}
1421+
1422+
hsa_status_t asan_hsa_amd_vmem_address_free(void* ptr, size_t size,
1423+
BufferedStackTrace* stack) {
1424+
if (UNLIKELY(!IsAligned(reinterpret_cast<uptr>(ptr), kPageSize_))) {
1425+
errno = errno_EINVAL;
1426+
return HSA_STATUS_ERROR_INVALID_ARGUMENT;
1427+
}
1428+
if (size == 0) {
1429+
errno = errno_EINVAL;
1430+
return HSA_STATUS_ERROR_INVALID_ARGUMENT;
1431+
}
1432+
1433+
void* p = get_allocator().GetBlockBegin(ptr);
1434+
if (p) {
1435+
instance.Deallocate(ptr, 0, 0, stack, FROM_MALLOC);
1436+
return HSA_STATUS_SUCCESS;
1437+
}
1438+
return REAL(hsa_amd_vmem_address_free)(ptr, size);
1439+
}
13871440
} // namespace __asan
13881441
#endif

‎compiler-rt/lib/asan/asan_allocator.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,13 @@ hsa_status_t asan_hsa_amd_ipc_memory_attach(
319319
const hsa_agent_t* mapping_agents, void** mapped_ptr);
320320
hsa_status_t asan_hsa_amd_ipc_memory_detach(
321321
void* mapped_ptr);
322+
hsa_status_t asan_hsa_amd_vmem_address_reserve_align(void** ptr, size_t size,
323+
uint64_t address,
324+
uint64_t alignment,
325+
uint64_t flags,
326+
BufferedStackTrace* stack);
327+
hsa_status_t asan_hsa_amd_vmem_address_free(void* ptr, size_t size,
328+
BufferedStackTrace* stack);
322329
} // namespace __asan
323330
#endif
324331

‎compiler-rt/lib/asan/asan_interceptors.cpp‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -897,6 +897,22 @@ INTERCEPTOR(hsa_status_t, hsa_amd_ipc_memory_detach, void* mapped_ptr) {
897897
return asan_hsa_amd_ipc_memory_detach(mapped_ptr);
898898
}
899899

900+
INTERCEPTOR(hsa_status_t, hsa_amd_vmem_address_reserve_align, void** ptr,
901+
size_t size, uint64_t address, uint64_t alignment, uint64_t flags) {
902+
AsanInitFromRtl();
903+
ENSURE_HSA_INITED();
904+
GET_STACK_TRACE_MALLOC;
905+
return asan_hsa_amd_vmem_address_reserve_align(ptr, size, address, alignment,
906+
flags, &stack);
907+
}
908+
909+
INTERCEPTOR(hsa_status_t, hsa_amd_vmem_address_free, void* ptr, size_t size) {
910+
AsanInitFromRtl();
911+
ENSURE_HSA_INITED();
912+
GET_STACK_TRACE_FREE;
913+
return asan_hsa_amd_vmem_address_free(ptr, size, &stack);
914+
}
915+
900916
void InitializeAmdgpuInterceptors() {
901917
ASAN_INTERCEPT_FUNC(hsa_memory_copy);
902918
ASAN_INTERCEPT_FUNC(hsa_amd_memory_pool_allocate);
@@ -909,6 +925,8 @@ void InitializeAmdgpuInterceptors() {
909925
ASAN_INTERCEPT_FUNC(hsa_amd_ipc_memory_create);
910926
ASAN_INTERCEPT_FUNC(hsa_amd_ipc_memory_attach);
911927
ASAN_INTERCEPT_FUNC(hsa_amd_ipc_memory_detach);
928+
ASAN_INTERCEPT_FUNC(hsa_amd_vmem_address_reserve_align);
929+
ASAN_INTERCEPT_FUNC(hsa_amd_vmem_address_free);
912930
}
913931

914932
void ENSURE_HSA_INITED() {

‎compiler-rt/lib/sanitizer_common/sanitizer_allocator_amdgpu.cpp‎

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,11 @@ struct HsaMemoryFunctions {
2222
void *(*alloc)(size_t),
2323
uint32_t *num_agents_accessible,
2424
hsa_agent_t **accessible);
25+
hsa_status_t (*vmem_address_reserve_align)(void** ptr, size_t size,
26+
uint64_t address,
27+
uint64_t alignment,
28+
uint64_t flags);
29+
hsa_status_t (*vmem_address_free)(void* ptr, size_t size);
2530
};
2631

2732
static HsaMemoryFunctions hsa_amd;
@@ -37,31 +42,49 @@ bool AmdgpuMemFuncs::Init() {
3742
RTLD_NEXT, "hsa_amd_memory_pool_free");
3843
hsa_amd.pointer_info = (decltype(hsa_amd.pointer_info))dlsym(
3944
RTLD_NEXT, "hsa_amd_pointer_info");
45+
hsa_amd.vmem_address_reserve_align =
46+
(decltype(hsa_amd.vmem_address_reserve_align))dlsym(
47+
RTLD_NEXT, "hsa_amd_vmem_address_reserve_align");
48+
hsa_amd.vmem_address_free = (decltype(hsa_amd.vmem_address_free))dlsym(
49+
RTLD_NEXT, "hsa_amd_vmem_address_free");
4050
if (!hsa_amd.memory_pool_allocate || !hsa_amd.memory_pool_free ||
41-
!hsa_amd.pointer_info)
51+
!hsa_amd.pointer_info || !hsa_amd.vmem_address_reserve_align ||
52+
!hsa_amd.vmem_address_free)
4253
return false;
43-
else
44-
return true;
54+
return true;
4555
}
4656

4757
void *AmdgpuMemFuncs::Allocate(uptr size, uptr alignment,
4858
DeviceAllocationInfo *da_info) {
4959
AmdgpuAllocationInfo *aa_info =
5060
reinterpret_cast<AmdgpuAllocationInfo *>(da_info);
51-
52-
aa_info->status = hsa_amd.memory_pool_allocate(aa_info->memory_pool, size,
53-
aa_info->flags, &aa_info->ptr);
61+
if (!aa_info->memory_pool.handle) {
62+
aa_info->status = hsa_amd.vmem_address_reserve_align(
63+
&aa_info->ptr, size, aa_info->address, aa_info->alignment,
64+
aa_info->flags64);
65+
} else {
66+
aa_info->status = hsa_amd.memory_pool_allocate(
67+
aa_info->memory_pool, size, aa_info->flags, &aa_info->ptr);
68+
}
5469
if (aa_info->status != HSA_STATUS_SUCCESS)
5570
return nullptr;
5671

5772
return aa_info->ptr;
5873
}
5974

6075
void AmdgpuMemFuncs::Deallocate(void *p) {
61-
UNUSED hsa_status_t status = hsa_amd.memory_pool_free(p);
76+
DevicePointerInfo DevPtrInfo;
77+
if (AmdgpuMemFuncs::GetPointerInfo(reinterpret_cast<uptr>(p), &DevPtrInfo)) {
78+
if (DevPtrInfo.type == HSA_EXT_POINTER_TYPE_HSA) {
79+
UNUSED hsa_status_t status = hsa_amd.memory_pool_free(p);
80+
} else if (DevPtrInfo.type == HSA_EXT_POINTER_TYPE_RESERVED_ADDR) {
81+
UNUSED hsa_status_t status =
82+
hsa_amd.vmem_address_free(p, DevPtrInfo.map_size);
83+
}
84+
}
6285
}
6386

64-
bool AmdgpuMemFuncs::GetPointerInfo(uptr ptr, DevivePointerInfo *ptr_info) {
87+
bool AmdgpuMemFuncs::GetPointerInfo(uptr ptr, DevicePointerInfo* ptr_info) {
6588
hsa_amd_pointer_info_t info;
6689
info.size = sizeof(hsa_amd_pointer_info_t);
6790
hsa_status_t status =
@@ -70,8 +93,12 @@ bool AmdgpuMemFuncs::GetPointerInfo(uptr ptr, DevivePointerInfo *ptr_info) {
7093
if (status != HSA_STATUS_SUCCESS)
7194
return false;
7295

73-
ptr_info->map_beg = reinterpret_cast<uptr>(info.agentBaseAddress);
96+
if (info.type == HSA_EXT_POINTER_TYPE_RESERVED_ADDR)
97+
ptr_info->map_beg = reinterpret_cast<uptr>(info.hostBaseAddress);
98+
else if (info.type == HSA_EXT_POINTER_TYPE_HSA)
99+
ptr_info->map_beg = reinterpret_cast<uptr>(info.agentBaseAddress);
74100
ptr_info->map_size = info.sizeInBytes;
101+
ptr_info->type = reinterpret_cast<hsa_amd_pointer_type_t>(info.type);
75102

76103
return true;
77104
}

‎compiler-rt/lib/sanitizer_common/sanitizer_allocator_amdgpu.h‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ class AmdgpuMemFuncs {
2020
static void *Allocate(uptr size, uptr alignment,
2121
DeviceAllocationInfo *da_info);
2222
static void Deallocate(void *p);
23-
static bool GetPointerInfo(uptr ptr, DevivePointerInfo *ptr_info);
23+
static bool GetPointerInfo(uptr ptr, DevicePointerInfo* ptr_info);
2424
static uptr GetPageSize();
2525
};
2626

@@ -32,8 +32,11 @@ struct AmdgpuAllocationInfo : public DeviceAllocationInfo {
3232
hsa_status_t status;
3333
void *alloc_func;
3434
hsa_amd_memory_pool_t memory_pool;
35-
size_t size;
36-
uint32_t flags;
35+
u64 alignment;
36+
u64 address;
37+
u64 flags64;
38+
usize size;
39+
u32 flags;
3740
void *ptr;
3841
};
3942
#endif // SANITIZER_AMDGPU

‎compiler-rt/lib/sanitizer_common/sanitizer_allocator_device.h‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@ struct DeviceAllocationInfo {
3131
DeviceAllocationType type_;
3232
};
3333

34-
struct DevivePointerInfo {
34+
struct DevicePointerInfo {
35+
u64 type;
3536
uptr map_beg;
3637
uptr map_size;
3738
};
@@ -165,7 +166,7 @@ class DeviceAllocatorT {
165166
: nullptr;
166167
}
167168

168-
void *GetBlockBegin(const void *ptr) const {
169+
void* GetBlockBegin(const void* ptr) const {
169170
Header header;
170171
if (!mem_funcs_inited_) return nullptr;
171172
uptr p = reinterpret_cast<uptr>(ptr);
@@ -182,7 +183,7 @@ class DeviceAllocatorT {
182183
if (!nearest_chunk)
183184
return nullptr;
184185
if (p != nearest_chunk) {
185-
Header *h = GetHeader(nearest_chunk, &header);
186+
Header* h = GetHeader(nearest_chunk, &header);
186187
CHECK_GE(nearest_chunk, h->map_beg);
187188
CHECK_LT(nearest_chunk, h->map_beg + h->map_size);
188189
CHECK_LE(nearest_chunk, p);
@@ -297,7 +298,7 @@ class DeviceAllocatorT {
297298
return mem_funcs_inited_;
298299
}
299300

300-
typedef DevivePointerInfo Header;
301+
typedef DevicePointerInfo Header;
301302

302303
Header *GetHeaderAnyPointer(uptr p, Header* h) const {
303304
CHECK(IsAligned(p, page_size_));

0 commit comments

Comments
 (0)