wip

Author: q10

fbgemm_gpu/src/memory_utils/memory_utils.cu

@@ -146,20 +146,24 @@ std::tuple<void*, size_t> adjust_to_page_boundaries(void* ptr, size_t size) {
 
 #ifdef USE_ROCM
 using gpuMemLocation = hipMemLocation;
+using gpuMemLocationTypeHost = hipMemLocationTypeHost;
+using gpuMemLocationTypeDevice = hipMemLocationTypeDevice;
 #else
 using gpuMemLocation = cudaMemLocation;
+using gpuMemLocationTypeHost = cudaMemLocationTypeHost;
+using gpuMemLocationTypeDevice = cudaMemLocationTypeDevice;
 #endif
 
 inline gpuMemLocation new_mem_location_from_device(const int device_id) {
   gpuMemLocation deviceLoc;
-  deviceLoc.type = cudaMemLocationTypeDevice;
+  deviceLoc.type = gpuMemLocationTypeDevice;
   deviceLoc.id = device_id;
   return deviceLoc;
 }
 
 inline gpuMemLocation new_mem_location_cpu() {
   gpuMemLocation deviceLoc;
-  deviceLoc.type = cudaMemLocationTypeHost;
+  deviceLoc.type = gpuMemLocationTypeHost;
   deviceLoc.id = cudaCpuDeviceId;
   return deviceLoc;
 }