Manual hipify caffe2/distributed and rocm update (no hcc modules support) (#18088)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18088

Manually hipify the distributed folder

Reviewed By: bddppq

Differential Revision: D14482702

fbshipit-source-id: cc0abdf525b423ab1f18db8010d21e27c6668d36

diff --git a/caffe2/distributed/file_store_handler_op_gpu.cc b/caffe2/distributed/file_store_handler_op_gpu.cc
index 6c13d14..b60b650 100644 (file)
--- a/caffe2/distributed/file_store_handler_op_gpu.cc
+++ b/caffe2/distributed/file_store_handler_op_gpu.cc
@@ -1,11 +1,21 @@
 #include "caffe2/distributed/file_store_handler_op.h"
 
+#ifndef __HIP_PLATFORM_HCC__
 #include <caffe2/core/context_gpu.h>
+#else
+#include <caffe2/core/hip/context_gpu.h>
+#endif
 
 namespace caffe2 {
 
+#ifndef __HIP_PLATFORM_HCC__
 REGISTER_CUDA_OPERATOR(
     FileStoreHandlerCreate,
     FileStoreHandlerCreateOp<CUDAContext>);
+#else
+REGISTER_HIP_OPERATOR(
+    FileStoreHandlerCreate,
+    FileStoreHandlerCreateOp<HIPContext>);
+#endif
 
 } // namespace caffe2

diff --git a/caffe2/distributed/redis_store_handler_op_gpu.cc b/caffe2/distributed/redis_store_handler_op_gpu.cc
index 5a759e5..6ad3a53 100644 (file)
--- a/caffe2/distributed/redis_store_handler_op_gpu.cc
+++ b/caffe2/distributed/redis_store_handler_op_gpu.cc
@@ -1,11 +1,21 @@
 #include "caffe2/distributed/redis_store_handler_op.h"
 
+#ifndef __HIP_PLATFORM_HCC__
 #include <caffe2/core/context_gpu.h>
+#else
+#include <caffe2/core/hip/context_gpu.h>
+#endif
 
 namespace caffe2 {
 
+#ifndef __HIP_PLATFORM_HCC__
 REGISTER_CUDA_OPERATOR(
     RedisStoreHandlerCreate,
     RedisStoreHandlerCreateOp<CUDAContext>);
+#else
+REGISTER_HIP_OPERATOR(
+    RedisStoreHandlerCreate,
+    RedisStoreHandlerCreateOp<HIPContext>);
+#endif
 
 } // namespace caffe2

diff --git a/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py b/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py
index 0172307..31f69ec 100644 (file)
--- a/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py
+++ b/tools/amd_build/pyHIPIFY/cuda_to_hip_mappings.py
@@ -2262,6 +2262,10 @@ PYTORCH_SPECIFIC_MAPPINGS = collections.OrderedDict([
 
 CAFFE2_SPECIFIC_MAPPINGS = collections.OrderedDict([
     ("cuda_stream" , ("hip_stream", API_CAFFE2)),
+    # if the header is a native hip folder (under hip directory),
+    # there is no need to add a hip path to it; the trie in hipify script
+    # takes this mapping order to forbid further replacement
+    ("/hip/" , ("/hip/", API_CAFFE2)),
     ("/context_gpu" , ("/hip/context_gpu", API_CAFFE2)),
     ("/common_gpu"  , ("/hip/common_gpu", API_CAFFE2)),
     ("/mixed_utils" , ("/hip/mixed_utils", API_CAFFE2)),