[Issue]: Releasing virtual address regions fails when several memory regions mapped via hipMemMap are unmaped with one call to hipMemUnmap

[IMbackK]

Problem Description

When multiple memory regions are mapped to a virtual address space via hipMemMap and then unmapped via a single call to hipMemUnmap, hipMemAddressFree will fail as rocr will consider there to still be another user of that region, even when all memory has been unmaped.

in this case hipMemAddressFree will ultimately fail here as it->second.use_count is > 0:
https://github.com/ROCm/ROCR-Runtime/blob/8b6b9c140d0ec01f73b69348eb53206eea5d5e50/runtime/hsa-runtime/core/runtime/runtime.cpp#L3092C7-L3092C27

reproducer:

#include <hip/hip_runtime.h>
#include <stdio.h>

#define HIP_CHECK(fn) { hipError_t err = fn; if(err != hipSuccess){fprintf(stderr, "Error: %s: %s at %d\n", hipGetErrorName(err), hipGetErrorString(err), __LINE__);} }

int main()
{
	size_t granularity;
	hipMemAllocationProp alloc_prop = {};
	alloc_prop.type = hipMemAllocationTypePinned;
	alloc_prop.location.type = hipMemLocationTypeDevice;
	alloc_prop.location.id = 0;
	HIP_CHECK(hipMemGetAllocationGranularity(&granularity, &alloc_prop, hipMemAllocationGranularityRecommended));
	printf("Device recommended granularity %zu\n", granularity);
	
	constexpr size_t maxSize = 1ull << 35; // 32 GB
	hipDeviceptr_t pool_addr = 0;
	HIP_CHECK(hipMemAddressReserve(&pool_addr, maxSize, 0, 0, 0));
	printf("reserved pool at %p\n", pool_addr);
	
	hipMemAllocationProp prop = {};
	prop.type = hipMemAllocationTypePinned;
	prop.location.type = hipMemLocationTypeDevice;
	prop.location.id = 0;
	hipMemGenericAllocationHandle_t handle;
	
	size_t pool_size = 0;
	HIP_CHECK(hipMemCreate(&handle, granularity, &prop, 0));
	HIP_CHECK(hipMemMap(static_cast<char*>(pool_addr) + pool_size, granularity, 0, handle, 0));
	HIP_CHECK(hipMemRelease(handle));
	pool_size += granularity;
	
	HIP_CHECK(hipMemCreate(&handle, granularity, &prop, 0));
	HIP_CHECK(hipMemMap(static_cast<char*>(pool_addr) + pool_size, granularity, 0, handle, 0));
	HIP_CHECK(hipMemRelease(handle));
	pool_size += granularity;
	
	HIP_CHECK(hipMemCreate(&handle, granularity, &prop, 0));
	HIP_CHECK(hipMemMap(static_cast<char*>(pool_addr) + pool_size, granularity, 0, handle, 0));
	HIP_CHECK(hipMemRelease(handle));
	pool_size += granularity;
	
	printf("unmapping %zu at %p\n", pool_size, pool_addr);
	HIP_CHECK(hipMemUnmap(pool_addr, pool_size));
	printf("Freeing virtual space %zu at %p\n", maxSize, pool_addr);
	HIP_CHECK(hipMemAddressFree(pool_addr, maxSize));
	
	return 0;
}

Output:

Device recommended granularity 4096
reserved pool at 0x73415fe00000
unmapping 12288 at 0x73415fe00000
Freeing virtual space 34359738368 at 0x73415fe00000
Error: hipErrorUnknown: unknown error at 46

Meanwhile this will work:

#include <hip/hip_runtime.h>
#include <stdio.h>

#define HIP_CHECK(fn) { hipError_t err = fn; if(err != hipSuccess){fprintf(stderr, "Error: %s: %s at %d\n", hipGetErrorName(err), hipGetErrorString(err), __LINE__);} }

int main()
{
	size_t granularity;
	hipMemAllocationProp alloc_prop = {};
	alloc_prop.type = hipMemAllocationTypePinned;
	alloc_prop.location.type = hipMemLocationTypeDevice;
	alloc_prop.location.id = 0;
	HIP_CHECK(hipMemGetAllocationGranularity(&granularity, &alloc_prop, hipMemAllocationGranularityRecommended));
	printf("Device recommended granularity %zu\n", granularity);
	
	constexpr size_t maxSize = 1ull << 35; // 32 GB
	hipDeviceptr_t pool_addr = 0;
	HIP_CHECK(hipMemAddressReserve(&pool_addr, maxSize, 0, 0, 0));
	printf("reserved pool at %p\n", pool_addr);
	
	hipMemAllocationProp prop = {};
	prop.type = hipMemAllocationTypePinned;
	prop.location.type = hipMemLocationTypeDevice;
	prop.location.id = 0;
	hipMemGenericAllocationHandle_t handle;
	
	size_t pool_size = 0;
	HIP_CHECK(hipMemCreate(&handle, granularity, &prop, 0));
	HIP_CHECK(hipMemMap(static_cast<char*>(pool_addr) + pool_size, granularity, 0, handle, 0));
	HIP_CHECK(hipMemUnmap(static_cast<char*>(pool_addr) + pool_size, granularity));
	HIP_CHECK(hipMemRelease(handle));
	pool_size += granularity;
	
	HIP_CHECK(hipMemCreate(&handle, granularity, &prop, 0));
	HIP_CHECK(hipMemMap(static_cast<char*>(pool_addr) + pool_size, granularity, 0, handle, 0));
	HIP_CHECK(hipMemUnmap(static_cast<char*>(pool_addr) + pool_size, granularity));
	HIP_CHECK(hipMemRelease(handle));
	pool_size += granularity;
	
	HIP_CHECK(hipMemCreate(&handle, granularity, &prop, 0));
	HIP_CHECK(hipMemMap(static_cast<char*>(pool_addr) + pool_size, granularity, 0, handle, 0));
	HIP_CHECK(hipMemUnmap(static_cast<char*>(pool_addr) + pool_size, granularity));
	HIP_CHECK(hipMemRelease(handle));
	pool_size += granularity;
	
	printf("unmapping %zu at %p\n", pool_size, pool_addr);
	//HIP_CHECK(hipMemUnmap(pool_addr, pool_size));
	printf("Freeing virtual space %zu at %p\n", maxSize, pool_addr);
	HIP_CHECK(hipMemAddressFree(pool_addr, maxSize));
	
	return 0;
}

Operating System

Ubuntu 24.04

CPU

Epyc 7552

GPU

MI100, RX 6800 XT

ROCm Version

ROCm 6.3.0

ROCm Component

ROCR-Runtime, clr

Steps to Reproduce

Run the repoducer

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[taylding-amd]

Hi @IMbackK, I was able to reproduce the bug and confirm that the flow works correctly on CUDA. The team is currently working on a fix, and I will keep you updated on the issue. Thank you for bringing this to our attention!

[IMbackK]

Hi @taylding-amd, thank you for the update!

[YangWang92]

Hi @ppanchad-amd, do we have any updates?
We are trying to use this api to support sleep mode on vllm.

Thanks!
Yang

[taylding-amd]

Hi @YangWang92, thank you for reaching out. The team is currently working on a fix, with a target release for ROCm 7.0. I will update the ticket once the fix is ready.

[IMbackK]

My understanding is that ROCm 7.0 is targeted for fall/winter 2025, is that correct?

[taylding-amd]

Yes, ROCm 7.0 is scheduled for release following ROCm 6.5. While we don't have an exact release date at this time, it is expected to be available later this year.