pytorch · bddppq · Jun 4, 2018 · Apr 26, 2018 · Apr 26, 2018 · Apr 26, 2018
diff --git a/.jenkins/caffe2/build.sh b/.jenkins/caffe2/build.sh
@@ -136,6 +136,10 @@ case "${BUILD_ENVIRONMENT}" in
     # Ensure the ccache symlink can still find the real nvcc binary.
     export PATH="/usr/local/cuda/bin:$PATH"
     ;;
+  *-rocm*)
+    export LANG=C.UTF-8
+    export LC_ALL=C.UTF-8
+    export HCC_AMDGPU_TARGET=gfx900
 esac
 
 # Try to include Redis support for Linux builds

diff --git a/aten/CMakeLists.txt b/aten/CMakeLists.txt
@@ -119,6 +119,10 @@ add_subdirectory(src/THS)
 # Find the HIP package, set the HIP paths, load the HIP CMake.
 IF(USE_ROCM)
   include(LoadHIP)
+  if (NOT PYTORCH_FOUND_HIP)
+    MESSAGE(FATAL_ERROR
+      "Could not find HIP installation")
+  endif()
 ENDIF()
 
 IF(MSVC)

diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
@@ -84,6 +84,7 @@ endif()
 if (CAFFE2_WHITELISTED_FILES)
   caffe2_do_whitelist(Caffe2_CPU_SRCS CAFFE2_WHITELISTED_FILES)
   caffe2_do_whitelist(Caffe2_GPU_SRCS CAFFE2_WHITELISTED_FILES)
+  caffe2_do_whitelist(Caffe2_HIP_SRCS CAFFE2_WHITELISTED_FILES)
 endif()
 
 # Debug messages - if you want to get a list of source files, enable the
@@ -119,6 +120,16 @@ if (FALSE)
     message(STATUS "  " ${tmp})
   endforeach()
 
+  message(STATUS "HIP sources: ")
+  foreach(tmp ${Caffe2_HIP_SRCS})
+    message(STATUS "  " ${tmp})
+  endforeach()
+
+  message(STATUS "HIP test sources: ")
+  foreach(tmp ${Caffe2_HIP_TEST_SRCS})
+    message(STATUS "  " ${tmp})
+  endforeach()
+
   message(STATUS "ATen CPU test sources: ")
   foreach(tmp ${ATen_CPU_TEST_SRCS})
     message(STATUS "  " ${tmp})
@@ -194,7 +205,7 @@ caffe2_interface_library(caffe2 caffe2_library)
 list(APPEND Caffe2_MAIN_LIBS caffe2_library)
 
 # ---[ CUDA library.
-if(USE_CUDA OR USE_ROCM)
+if(USE_CUDA OR (USE_ROCM AND NOT BUILD_CAFFE2))
   # A hack to deal with cuda library dependencies and modern CMake: the
   # CUDA_ADD_LIBRARY includes a target_link_libraries, and as a result,
   # one cannot use PUBLIC/PRIVATE/INTERFACE for the target anymore. This
@@ -232,6 +243,28 @@ if(USE_CUDA OR USE_ROCM)
   list(APPEND Caffe2_MAIN_LIBS caffe2_gpu_library)
 endif()
 
+# ---[ Caffe2 HIP sources.
+if(BUILD_CAFFE2)
+  if(USE_ROCM)
+    HIP_ADD_LIBRARY(caffe2_hip ${Caffe2_HIP_SRCS})
+    set_target_properties(caffe2_hip PROPERTIES COMPILE_FLAGS ${Caffe2_HIP_CXX_FLAGS})
+
+    target_include_directories(
+      caffe2_hip PUBLIC ${Caffe2_HIP_INCLUDES})
+    target_include_directories(
+      caffe2_hip INTERFACE $<INSTALL_INTERFACE:include>)
+
+    target_link_libraries(caffe2_hip PUBLIC caffe2)
+    target_link_libraries(caffe2_hip PUBLIC ${Caffe2_HIP_DEPENDENCY_LIBS})
+
+    set_target_properties(caffe2_hip PROPERTIES LINKER_LANGUAGE HIP)
+
+    caffe2_interface_library(caffe2_hip caffe2_hip_library)
+    list(APPEND Caffe2_MAIN_LIBS caffe2_hip_library)
+    install(TARGETS caffe2_hip EXPORT Caffe2Targets DESTINATION lib)
+  endif()
+endif()
+
 # ---[ Check if warnings should be errors.
 if ($ENV{WERROR})
   target_compile_options(caffe2 PRIVATE -Werror)
@@ -244,9 +277,12 @@ endif()
 if(BUILD_CAFFE2)
   if (BUILD_TEST)
     set(Caffe2_ALL_TEST_SRCS ${Caffe2_CPU_TEST_SRCS})
-    if (USE_CUDA OR USE_ROCM)
+    if (USE_CUDA)
       list(APPEND Caffe2_ALL_TEST_SRCS ${Caffe2_GPU_TEST_SRCS})
     endif()
+    if(USE_ROCM)
+      list(APPEND Caffe2_ALL_TEST_SRCS ${Caffe2_HIP_TEST_SRCS})
+    endif()
 
     foreach(test_src ${Caffe2_ALL_TEST_SRCS})
       get_filename_component(test_name ${test_src} NAME_WE)
@@ -258,6 +294,18 @@ if(BUILD_CAFFE2)
       add_test(NAME ${test_name} COMMAND $<TARGET_FILE:${test_name}>)
       install(TARGETS ${test_name} DESTINATION test)
     endforeach()
+
+    if(USE_ROCM)
+      foreach(test_src ${Caffe2_HIP_TEST_SRCS})
+        get_filename_component(test_name ${test_src} NAME_WE)
+        set_target_properties(${test_name} PROPERTIES COMPILE_FLAGS ${Caffe2_HIP_CXX_FLAGS})
+        set_target_properties(${test_name} PROPERTIES LINKER_LANGUAGE HIP)
+        if (${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION} GREATER 3.0)
+          target_compile_features(${test_name} PRIVATE cxx_range_for)
+        endif()
+      endforeach()
+    endif()
+
   endif()
 endif()
 
@@ -336,7 +384,7 @@ if(BUILD_CAFFE2)
     endif(WIN32)
     install(TARGETS caffe2_pybind11_state DESTINATION "${PYTHON_LIB_REL_PATH}/caffe2/python")
 
-    if(USE_CUDA OR USE_ROCM)
+    if(USE_CUDA)
       add_library(caffe2_pybind11_state_gpu MODULE ${Caffe2_GPU_PYTHON_SRCS})
       set_target_properties(caffe2_pybind11_state_gpu PROPERTIES COMPILE_FLAGS "-fvisibility=hidden")
       set_target_properties(caffe2_pybind11_state_gpu PROPERTIES PREFIX "" DEBUG_POSTFIX "")
@@ -366,7 +414,7 @@ if(BUILD_CAFFE2)
           COMMAND ${CMAKE_COMMAND} -E copy
           $<TARGET_FILE:caffe2_pybind11_state>
           ${CMAKE_BINARY_DIR}/caffe2/python)
-      if (USE_CUDA OR USE_ROCM)
+      if (USE_CUDA)
         add_dependencies(windows_python_copy_lib caffe2_pybind11_state_gpu)
         add_custom_command(
             TARGET windows_python_copy_lib POST_BUILD

diff --git a/caffe2/core/CMakeLists.txt b/caffe2/core/CMakeLists.txt
@@ -12,25 +12,43 @@ set(Caffe2_GPU_SRCS ${Caffe2_GPU_SRCS} ${tmp})
 file(GLOB tmp *_test.cc)
 exclude(Caffe2_GPU_SRCS "${Caffe2_GPU_SRCS}" ${tmp})
 
+# ---[ HIP files
+file(GLOB_RECURSE tmp *_hip.cc)
+set(Caffe2_HIP_SRCS ${Caffe2_HIP_SRCS} ${tmp})
+# ------[ MIOpen files
+file(GLOB_RECURSE tmp *_miopen.cc)
+set(Caffe2_HIP_SRCS ${Caffe2_HIP_SRCS} ${tmp})
+# exclude test files
+file(GLOB_RECURSE tmp *_test.cc)
+exclude(Caffe2_HIP_SRCS "${Caffe2_HIP_SRCS}" ${tmp})
+
 # ---[ CPU files.
 file(GLOB tmp *.cc)
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${tmp})
 # exclude test files and gpu files
 file(GLOB tmp *_test.cc)
 exclude(Caffe2_CPU_SRCS "${Caffe2_CPU_SRCS}" ${tmp})
 exclude(Caffe2_CPU_SRCS "${Caffe2_CPU_SRCS}" ${Caffe2_GPU_SRCS})
+exclude(Caffe2_CPU_SRCS "${Caffe2_CPU_SRCS}" ${Caffe2_HIP_SRCS})
 
 # ---[ GPU test files
 file(GLOB tmp *_gpu_test.cc)
 set(Caffe2_GPU_TEST_SRCS ${Caffe2_GPU_TEST_SRCS} ${tmp})
 
+# ---[ HIP test files
+file(GLOB_RECURSE tmp *_hip_test.cc)
+set(Caffe2_HIP_TEST_SRCS ${Caffe2_HIP_TEST_SRCS} ${tmp})
+
 # ---[ CPU test files
 file(GLOB tmp *_test.cc)
 set(Caffe2_CPU_TEST_SRCS ${Caffe2_CPU_TEST_SRCS} ${tmp})
 exclude(Caffe2_CPU_TEST_SRCS "${Caffe2_CPU_TEST_SRCS}" ${Caffe2_GPU_TEST_SRCS})
+exclude(Caffe2_CPU_TEST_SRCS "${Caffe2_CPU_TEST_SRCS}" ${Caffe2_HIP_TEST_SRCS})
 
 # ---[ Send the lists to the parent scope.
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} PARENT_SCOPE)
 set(Caffe2_GPU_SRCS ${Caffe2_GPU_SRCS} PARENT_SCOPE)
+set(Caffe2_HIP_SRCS ${Caffe2_HIP_SRCS} PARENT_SCOPE)
 set(Caffe2_CPU_TEST_SRCS ${Caffe2_CPU_TEST_SRCS} PARENT_SCOPE)
 set(Caffe2_GPU_TEST_SRCS ${Caffe2_GPU_TEST_SRCS} PARENT_SCOPE)
+set(Caffe2_HIP_TEST_SRCS ${Caffe2_HIP_TEST_SRCS} PARENT_SCOPE)
diff --git a/caffe2/core/common.cc b/caffe2/core/common.cc
@@ -8,15 +8,24 @@ namespace caffe2 {
 // Do not directly use this variable, but instead use the HasCudaRuntime()
 // function below.
 std::atomic<bool> g_caffe2_has_cuda_linked{false};
+std::atomic<bool> g_caffe2_has_hip_linked{false};
 
 bool HasCudaRuntime() {
   return g_caffe2_has_cuda_linked.load();
 }
 
+bool HasHipRuntime() {
+  return g_caffe2_has_hip_linked.load();
+}
+
 namespace internal {
 void SetCudaRuntimeFlag() {
   g_caffe2_has_cuda_linked.store(true);
 }
+
+void SetHipRuntimeFlag() {
+  g_caffe2_has_hip_linked.store(true);
+}
 } // namespace internal
 
 const std::map<string, string>& GetBuildOptions() {

diff --git a/caffe2/core/common.h b/caffe2/core/common.h
@@ -280,11 +280,13 @@ class SkipIndices<> {
 // as the underlying boolean variable is going to be switched on when one
 // loads libcaffe2_gpu.so.
 bool HasCudaRuntime();
+bool HasHipRuntime();
 namespace internal {
 // Sets the Cuda Runtime flag that is used by HasCudaRuntime(). You should
 // never use this function - it is only used by the Caffe2 gpu code to notify
 // Caffe2 core that cuda runtime has been loaded.
 void SetCudaRuntimeFlag();
+void SetHipRuntimeFlag();
 }
 // Returns which setting Caffe2 was configured and built with (exported from
 // CMake)