Merge branch 'master' of https://github.com/ROCm-Developer-Tools/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 {
@@ -953,7 +961,7 @@ inline hipDeviceptr_t agent_address_for_symbol(const char* symbolName) {
 
 #if __hcc_workweek__ >= 17481
     size_t byte_cnt = 0u;
-    hipModuleGetGlobal(&r, &byte_cnt, 0, symbolName);
+    ihipModuleGetGlobal(&r, &byte_cnt, 0, symbolName);
 #else
     auto ctx = ihipGetTlsDefaultCtx();
     auto acc = ctx->getDevice()->_acc;
@@ -1098,6 +1106,23 @@ hipError_t hipMemcpyFromSymbolAsync(void* dst, const void* symbolName, size_t co
     return ihipLogStatus(e);
 }
 
+
+hipError_t hipGetSymbolAddress(void** devPtr, const void* symbolName) {
+    HIP_INIT_API(devPtr, symbolName);
+
+    size_t size = 0;
+    return ihipModuleGetGlobal(devPtr, &size, 0, static_cast<const char*>(symbolName));
+}
+
+
+hipError_t hipGetSymbolSize(size_t* size, const void* symbolName) {
+    HIP_INIT_API(size, symbolName);
+
+    void* devPtr = nullptr;
+    return ihipModuleGetGlobal(&devPtr, size, 0, static_cast<const char*>(symbolName));
+}
+
+
 //---
 hipError_t hipMemcpy(void* dst, const void* src, size_t sizeBytes, hipMemcpyKind kind) {
     HIP_INIT_SPECIAL_API((TRACE_MCMD), dst, src, sizeBytes, kind);

src/hip_module.cpp

@@ -472,14 +472,19 @@ hipError_t hipModuleGetGlobal(hipDeviceptr_t* dptr, size_t* bytes, hipModule_t h
                               const char* name) {
     HIP_INIT_API(dptr, bytes, hmod, name);
 
-    if (!dptr || !bytes) return ihipLogStatus(hipErrorInvalidValue);
+    return ihipLogStatus(ihipModuleGetGlobal(dptr, bytes, hmod, name));
+}
 
-    if (!name) return ihipLogStatus(hipErrorNotInitialized);
+hipError_t ihipModuleGetGlobal(hipDeviceptr_t* dptr, size_t* bytes, hipModule_t hmod,
+                               const char* name) {
+    if (!dptr || !bytes) return hipErrorInvalidValue;
+
+    if (!name) return hipErrorNotInitialized;
 
     const auto r = hmod ? read_agent_global_from_module(dptr, bytes, hmod, name)
                         : read_agent_global_from_process(dptr, bytes, name);
 
-    return ihipLogStatus(r);
+    return r;
 }
 
 namespace