From 94c85fd4fc873a945868e2c764127d324c33b3ed Mon Sep 17 00:00:00 2001
From: Ben Sander <ben.sander@amd.com>
Date: Fri, 17 Feb 2017 11:53:38 -0600
Subject: [PATCH] Fix copying of registered memory.

Set device properly so copying can recover context.
Enhance test to catch this case.
---
 src/hip_memory.cpp                            |  2 ++
 .../src/runtimeApi/memory/hipHostRegister.cpp | 21 +++++++++++++------
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/hip_memory.cpp b/src/hip_memory.cpp
index df11205344..314dee97dc 100644
--- a/src/hip_memory.cpp
+++ b/src/hip_memory.cpp
@@ -194,6 +194,7 @@ hipError_t hipHostMalloc(void** ptr, size_t sizeBytes, unsigned int flags)
                 if(sizeBytes < 1 && (*ptr == NULL)){
                     hip_status = hipErrorMemoryAllocation;
                 } else {
+                    // TODO - should OR in flags here?
                     hc::am_memtracker_update(*ptr, device->_deviceId, amHostCoherent);
                 }
                 tprintf(DB_MEM, " %s: finegrained system memory ptr=%p\n", __func__, *ptr);
@@ -403,6 +404,7 @@ hipError_t hipHostRegister(void *hostPtr, size_t sizeBytes, unsigned int flags)
                     vecAcc.push_back(ihipGetDevice(i)->_acc);
                 }
                 am_status = hc::am_memory_host_lock(device->_acc, hostPtr, sizeBytes, &vecAcc[0], vecAcc.size());
+                hc::am_memtracker_update(hostPtr, device->_deviceId, flags);
 
                 tprintf(DB_MEM, " %s registered ptr=%p\n", __func__, hostPtr);
                 if(am_status == AM_SUCCESS){
diff --git a/tests/src/runtimeApi/memory/hipHostRegister.cpp b/tests/src/runtimeApi/memory/hipHostRegister.cpp
index 37ee9b1b78..6c81ec0d91 100644
--- a/tests/src/runtimeApi/memory/hipHostRegister.cpp
+++ b/tests/src/runtimeApi/memory/hipHostRegister.cpp
@@ -39,21 +39,30 @@ int main(){
 	const size_t size = N * sizeof(float);
 	A = (float*)malloc(size);
 	HIPCHECK(hipHostRegister(A, size, 0));
+
+
 	for(int i=0;i<N;i++){
 		A[i] = float(1);
 	}
+
+    // Copy to B, this should be optimal pinned malloc copy:
+    float *B;
+    HIPCHECK(hipMalloc(&B, size));
+    HIPCHECK(hipMemcpy(B, A, size, hipMemcpyHostToDevice));
+
+
 	for(int i=0;i<num_devices;i++){
-	HIPCHECK(hipSetDevice(i));
-	HIPCHECK(hipHostGetDevicePointer((void**)&Ad[i], A, 0));
+        HIPCHECK(hipSetDevice(i));
+        HIPCHECK(hipHostGetDevicePointer((void**)&Ad[i], A, 0));
 	}
 
 	for(int i=0;i<num_devices;i++){
-	HIPCHECK(hipSetDevice(i));
-	hipLaunchKernel(HIP_KERNEL_NAME(Inc), dim3(N/512), dim3(512), 0, 0, Ad[i]);
-
-	HIPCHECK(hipDeviceSynchronize());
+        HIPCHECK(hipSetDevice(i));
+        hipLaunchKernel(HIP_KERNEL_NAME(Inc), dim3(N/512), dim3(512), 0, 0, Ad[i]);
 
+        HIPCHECK(hipDeviceSynchronize());
 	}
+
 	HIPASSERT(A[10] == 1.0f + float(num_devices));
 	HIPCHECK(hipHostUnregister(A));
 	passed();