Skip to content

Commit

Permalink
Set preferred location to the device for HIP Managed Memory
Browse files Browse the repository at this point in the history
The semantics for specifying different kinds of advice is unclear so I
set it in two stages.
  • Loading branch information
qedawkins authored and github-actions[bot] committed Jan 26, 2024
1 parent b262054 commit ed49a41
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 0 deletions.
1 change: 1 addition & 0 deletions experimental/rocm/dynamic_symbol_tables.h
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ RC_PFN_DECL(hipInit, unsigned int)
RC_PFN_DECL(hipModuleLaunchKernel, hipFunction_t, unsigned int, unsigned int,
unsigned int, unsigned int, unsigned int, unsigned int,
unsigned int, hipStream_t, void **, void **)
RC_PFN_DECL(hipMemAdvise, const void *, size_t, int, int)
RC_PFN_DECL(hipMemset, void *, int, size_t)
RC_PFN_DECL(hipMemsetAsync, void *, int, size_t, hipStream_t)
RC_PFN_DECL(hipMemsetD32Async, void *, int, size_t, hipStream_t)
Expand Down
10 changes: 10 additions & 0 deletions experimental/rocm/rocm_allocator.c
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,16 @@ static iree_status_t iree_hal_rocm_allocator_allocate_buffer(
status = ROCM_RESULT_TO_STATUS(
allocator->context->syms,
hipMallocManaged(&device_ptr, allocation_size, hipMemAttachGlobal));
if (iree_status_is_ok(status)) {
status = ROCM_RESULT_TO_STATUS(
allocator->context->syms,
hipMemAdvise(device_ptr, allocation_size,
hipMemAdviseSetPreferredLocation, allocator->device));
status = ROCM_RESULT_TO_STATUS(
allocator->context->syms,
hipMemAdvise(device_ptr, allocation_size,
hipMemAdviseSetCoarseGrain, allocator->device));
}
if (iree_status_is_ok(status) &&
allocator->supports_concurrent_managed_access) {
// Prefetch the buffer on the GPU device.
Expand Down

0 comments on commit ed49a41

Please sign in to comment.