SWDEV-546293 - hipMemPrefetchAsync_v2 and hipMemAdvise_v2 implementation (#869)

SWDEV-546293 - hipMemPrefetchAsync hipMemAdvise_v2

Please enter the commit message for your changes. Lines starting

[ROCm/clr commit: cbee74a80e]
This commit is contained in:
Betigeri, Sourabh
2025-08-15 22:40:04 -07:00
committed by GitHub
parent ad9eb56dd4
commit 35e48d1eaf
14 changed files with 309 additions and 106 deletions
+152 -77
View File
@@ -29,6 +29,10 @@ namespace hip {
// Forward declaraiton of a function
hipError_t ihipMallocManaged(void** ptr, size_t size, size_t align = 0, bool use_host_ptr = 0);
hipError_t ihipMemPrefetchAsync(const void* dev_ptr, size_t count, hipMemLocation location,
hipStream_t stream);
hipError_t ihipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
hipMemLocation location);
// Make sure HIP defines match ROCclr to avoid double conversion
static_assert(hipCpuDeviceId == amd::CpuDeviceId, "CPU device ID mismatch with ROCclr!");
@@ -83,98 +87,51 @@ hipError_t hipMallocManaged(void** dev_ptr, size_t size, unsigned int flags) {
hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device,
hipStream_t stream) {
HIP_INIT_API(hipMemPrefetchAsync, dev_ptr, count, device, stream);
if ((dev_ptr == nullptr) || (count == 0)) {
HIP_RETURN(hipErrorInvalidValue);
}
getStreamPerThread(stream);
size_t offset = 0;
amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
if ((memObj != nullptr) && (count > (memObj->getSize() - offset))) {
HIP_RETURN(hipErrorInvalidValue);
}
if (device != hipCpuDeviceId && (static_cast<size_t>(device) >= g_devices.size())) {
HIP_RETURN(hipErrorInvalidDevice);
}
hip::Stream* hip_stream = nullptr;
amd::Device* dev = nullptr;
bool cpu_access = false;
if ((memObj == nullptr) && (device != hipCpuDeviceId) &&
(!g_devices[device]->devices()[0]->info().hmmCpuMemoryAccessible_)) {
HIP_RETURN(hipErrorNotSupported);
}
// Pick the specified stream or Null one from the provided device
CHECK_STREAM_CAPTURE_SUPPORTED();
hipMemLocation location;
if (device == hipCpuDeviceId) {
cpu_access = true;
hip_stream = (stream == nullptr || stream == hipStreamLegacy) ?
hip::getCurrentDevice()->NullStream() : hip::getStream(stream);
location.type = hipMemLocationTypeHost;
location.id = hipCpuDeviceId;
} else {
dev = g_devices[device]->devices()[0];
hip_stream = (stream == nullptr || stream == hipStreamLegacy) ?
g_devices[device]->NullStream() : hip::getStream(stream);
location.type = hipMemLocationTypeDevice;
location.id = device;
}
HIP_RETURN(ihipMemPrefetchAsync(dev_ptr, count, location, stream));
}
if (hip_stream == nullptr) {
// ================================================================================================
hipError_t hipMemPrefetchAsync_v2(const void* dev_ptr, size_t count, hipMemLocation location,
unsigned int flags, hipStream_t stream) {
HIP_INIT_API(hipMemPrefetchAsync_v2, dev_ptr, count, location, flags, stream);
CHECK_STREAM_CAPTURE_SUPPORTED();
if (flags != 0) {
HIP_RETURN(hipErrorInvalidValue);
}
amd::Command::EventWaitList waitList;
amd::SvmPrefetchAsyncCommand* command =
new amd::SvmPrefetchAsyncCommand(*hip_stream, waitList, dev_ptr, count, dev, cpu_access);
if (command == nullptr) {
return hipErrorOutOfMemory;
}
command->enqueue();
command->release();
HIP_RETURN(hipSuccess);
HIP_RETURN(ihipMemPrefetchAsync(dev_ptr, count, location, stream));
}
// ================================================================================================
hipError_t hipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice, int device) {
HIP_INIT_API(hipMemAdvise, dev_ptr, count, advice, device);
CHECK_STREAM_CAPTURE_SUPPORTED();
bool isAdviseReadMostly = (advice == hipMemAdviseSetReadMostly) ||
(advice == hipMemAdviseUnsetReadMostly);
if (!isAdviseReadMostly && ((device != hipCpuDeviceId) &&
(static_cast<size_t>(device) >= g_devices.size()))) {
HIP_RETURN(hipErrorInvalidDevice);
hipMemLocation location;
if (device == hipCpuDeviceId) {
location.type = hipMemLocationTypeHost;
location.id = hipCpuDeviceId;
} else {
location.type = hipMemLocationTypeDevice;
location.id = device;
}
if ((dev_ptr == nullptr) || (count == 0)) {
HIP_RETURN(hipErrorInvalidValue);
}
HIP_RETURN(ihipMemAdvise(dev_ptr, count, advice, location));
}
if (!hip::tls.capture_streams_.empty() || !g_captureStreams.empty()) {
HIP_RETURN(hipErrorStreamCaptureUnsupported);
}
size_t offset = 0;
amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
if (memObj && count > (memObj->getSize() - offset)) {
HIP_RETURN(hipErrorInvalidValue);
}
amd::Device* dev = (device == hipCpuDeviceId || isAdviseReadMostly) ?
g_devices[0]->devices()[0] : g_devices[device]->devices()[0];
bool use_cpu = (device == hipCpuDeviceId) ? true : false;
// Set the allocation attributes in AMD HMM
if (!dev->SetSvmAttributes(dev_ptr, count, static_cast<amd::MemoryAdvice>(advice), use_cpu)) {
HIP_RETURN(hipErrorInvalidValue);
}
HIP_RETURN(hipSuccess);
// ================================================================================================
hipError_t hipMemAdvise_v2(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
hipMemLocation location) {
HIP_INIT_API(hipMemAdvise_v2, dev_ptr, count, advice, location);
CHECK_STREAM_CAPTURE_SUPPORTED();
HIP_RETURN(ihipMemAdvise(dev_ptr, count, advice, location));
}
// ================================================================================================
@@ -322,4 +279,122 @@ hipError_t ihipMallocManaged(void** ptr, size_t size, size_t align, bool use_hos
ClPrint(amd::LOG_INFO, amd::LOG_API, "ihipMallocManaged ptr=0x%zx", *ptr);
return hipSuccess;
}
// ================================================================================================
hipError_t ihipMemPrefetchAsync(const void* dev_ptr, size_t count, hipMemLocation location,
hipStream_t stream) {
if ((dev_ptr == nullptr) || (count == 0)) {
return hipErrorInvalidValue;
}
getStreamPerThread(stream);
size_t offset = 0;
amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
if ((memObj != nullptr) && (count > (memObj->getSize() - offset))) {
return hipErrorInvalidValue;
}
// Compute the type of prefetch
const bool isHost = (location.type == hipMemLocationTypeHost);
const bool isHostNuma = (location.type == hipMemLocationTypeHostNuma);
const bool isHostCurrent = (location.type == hipMemLocationTypeHostNumaCurrent);
const bool cpuAccess = isHost || isHostNuma || isHostCurrent;
// Determine the target device index:
// - for host-prefetch and host-current, always use device 0
// - for host-NUMA or device-prefetch, use the provided id
int targetDevice = (isHost || isHostCurrent) ? hipCpuDeviceId : location.id;
amd::Device* dev = nullptr;
if (cpuAccess == false) {
if (static_cast<size_t>(targetDevice) >= g_devices.size()) {
return hipErrorInvalidDevice;
}
dev = g_devices[targetDevice]->devices()[0];
if (memObj == nullptr && !dev->info().hmmCpuMemoryAccessible_) {
return hipErrorNotSupported;
}
}
hip::Stream* hip_stream = nullptr;
// Pick the specified stream or Null one from the provided target device
if (cpuAccess == true) {
hip_stream = (stream == nullptr || stream == hipStreamLegacy)
? hip::getCurrentDevice()->NullStream()
: hip::getStream(stream);
} else {
dev = g_devices[targetDevice]->devices()[0];
hip_stream = (stream == nullptr || stream == hipStreamLegacy)
? g_devices[targetDevice]->NullStream()
: hip::getStream(stream);
}
if (hip_stream == nullptr) {
return hipErrorInvalidValue;
}
amd::Command::EventWaitList waitList;
amd::SvmPrefetchAsyncCommand* command = new amd::SvmPrefetchAsyncCommand(
*hip_stream, waitList, dev_ptr, count, dev, cpuAccess, targetDevice);
if (command == nullptr) {
return hipErrorOutOfMemory;
}
command->enqueue();
command->release();
return hipSuccess;
}
// ================================================================================================
hipError_t ihipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advice,
hipMemLocation location) {
if ((dev_ptr == nullptr) || (count == 0)) {
return hipErrorInvalidValue;
}
if (!hip::tls.capture_streams_.empty() || !g_captureStreams.empty()) {
return hipErrorStreamCaptureUnsupported;
}
// Determine device and CPU access from location
int targetDevice = hipCpuDeviceId;
bool use_cpu = true;
bool isAdviseReadMostly =
(advice == hipMemAdviseSetReadMostly) || (advice == hipMemAdviseUnsetReadMostly);
switch (location.type) {
case hipMemLocationTypeDevice:
targetDevice = location.id;
use_cpu = false;
break;
case hipMemLocationTypeHostNuma:
targetDevice = location.id; // NUMA node ID
use_cpu = true;
break;
case hipMemLocationTypeHost:
case hipMemLocationTypeHostNumaCurrent:
targetDevice = hipCpuDeviceId;
use_cpu = true;
break;
default:
return hipErrorInvalidValue;
}
if (!isAdviseReadMostly && !use_cpu && (static_cast<size_t>(targetDevice) >= g_devices.size())) {
return hipErrorInvalidDevice;
}
size_t offset = 0;
amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
if (memObj && count > (memObj->getSize() - offset)) {
return hipErrorInvalidValue;
}
amd::Device* dev = (use_cpu || isAdviseReadMostly) ? g_devices[0]->devices()[0]
: g_devices[targetDevice]->devices()[0];
// Set the allocation attributes in AMD HMM
if (!dev->SetSvmAttributes(dev_ptr, count, static_cast<amd::MemoryAdvice>(advice), use_cpu,
targetDevice)) {
return hipErrorInvalidValue;
}
return hipSuccess;
}
} //namespace hip