Fix copying of registered memory.

Set device properly so copying can recover context.
Enhance test to catch this case.


[ROCm/hip commit: 94c85fd4fc]

projects/hip/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){

projects/hip/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();