From 89efed29d78e182229a46c017e4312d617523edf Mon Sep 17 00:00:00 2001 From: Rahul Garg Date: Sat, 17 Nov 2018 05:38:35 +0530 Subject: [PATCH] Fix hipHostRegister [ROCm/hip commit: 1a038879a9dc135e6ce033a8df71c4ba69402b7e] --- projects/hip/src/hip_memory.cpp | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/projects/hip/src/hip_memory.cpp b/projects/hip/src/hip_memory.cpp index 7c25b714f8..2688307017 100644 --- a/projects/hip/src/hip_memory.cpp +++ b/projects/hip/src/hip_memory.cpp @@ -909,14 +909,22 @@ hipError_t hipHostRegister(void* hostPtr, size_t sizeBytes, unsigned int flags) } am_status = hc::am_memory_host_lock(device->_acc, hostPtr, sizeBytes, &vecAcc[0], vecAcc.size()); -#if USE_APP_PTR_FOR_CTX - hc::am_memtracker_update(hostPtr, device->_deviceId, flags, ctx); -#else - hc::am_memtracker_update(hostPtr, device->_deviceId, flags); -#endif + if ( am_status == AM_SUCCESS ) { + am_status = hc::am_memtracker_getinfo(&amPointerInfo, hostPtr); - tprintf(DB_MEM, " %s registered ptr=%p and allowed access to %zu peers\n", __func__, - hostPtr, vecAcc.size()); + if ( am_status == AM_SUCCESS ) { + void *devPtr = amPointerInfo._devicePointer; + #if USE_APP_PTR_FOR_CTX + hc::am_memtracker_update(hostPtr, device->_deviceId, flags, ctx); + hc::am_memtracker_update(devPtr, device->_deviceId, flags, ctx); + #else + hc::am_memtracker_update(hostPtr, device->_deviceId, flags); + hc::am_memtracker_update(devPtr, device->_deviceId, flags); + #endif + tprintf(DB_MEM, " %s registered ptr=%p and allowed access to %zu peers\n", __func__, + hostPtr, vecAcc.size()); + }; + }; if (am_status == AM_SUCCESS) { hip_status = hipSuccess; } else {