[Caffe2] Enabling AMD GPU Backend for Caffe2 (#7566)

* Add hip support for caffe2 core

* Add MIOPEN header/wrapper to caffe2 core

* Add HIP device into caffe2 PB

* top level makefile change for rocm/hip

* makefile scaffolding for AMD/RocM/HIP

* Makefile scafodding for AMD/RocM/HIP; add makefile/utility for HIP files

* caffe2 PB update for AMD/ROCM HIP device

* Add AMD/RocM/Thrust dependency

* HIP threadpool update

* Fix makefile macro

* makefile fix: duplicate test/binary name

* makefile clean-up

* makefile clean-up

* add HIP operator registry

* add utilities for hip device

* Add USE_HIP to config summary

* makefile fix for BUILD_TEST

* merge latest

* Fix indentation

* code clean-up

* Guard builds without HIP and use the same cmake script as PyTorch to find HIP

* Setup rocm environment variables in build.sh (ideally should be done in the docker images)

* setup locale

* set HIP_PLATFORM

* Revert "set HIP_PLATFORM"

This reverts commit 8ec58db.

* continue the build script environment variables mess

* HCC_AMDGPU_TARGET

* Cleanup the mess, has been fixed in the lastest docker images

* Assign protobuf field hip_gpu_id a new field number for backward compatibility

* change name to avoid conflict

* Fix duplicated thread pool flag

* Refactor cmake files to not add hip includes and libs globally

* Fix the wrong usage of environment variables detection in cmake

* Add MIOPEN CNN operators

* Revert "Add MIOPEN CNN operators"

This reverts commit 6e89ad4.

.jenkins/caffe2/build.sh

@@ -131,6 +131,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

CMakeLists.txt

@@ -59,6 +59,7 @@ cmake_dependent_option(
     USE_GLOO "Use Gloo" ON
     "BUILD_CAFFE2" OFF)
 option(USE_GLOO_IBVERBS "Use Gloo IB verbs for distributed support" OFF)  # New option
+option(USE_HIP "Use HIP" ON)
 cmake_dependent_option(
     USE_LEVELDB "Use LEVELDB" ON
     "BUILD_CAFFE2" OFF)
@@ -319,6 +320,11 @@ if (USE_CUDA)
 
 endif()
 
+if (USE_HIP)
+  # TODO: check if we should include other hip dependency libraries
+  # to the interface as well.
+
+endif()
 # Note(jiayq): when building static libraries, all PRIVATE dependencies
 # will also become interface libraries, and as a result if there are any
 # dependency libraries that are not exported, the following install export

aten/CMakeLists.txt

@@ -121,6 +121,10 @@ ENDIF()
 # Find the HIP package, set the HIP paths, load the HIP CMake.
 IF(WITH_ROCM)
   include(LoadHIP)
+  if (NOT PYTORCH_FOUND_HIP)
+    MESSAGE(FATAL_ERROR
+      "Could not find HIP installation")
+  endif()
 ENDIF()
 
 IF(MSVC)

caffe2/CMakeLists.txt

@@ -56,6 +56,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
@@ -71,6 +72,11 @@ if (FALSE)
     message(STATUS "  " ${tmp})
   endforeach()
 
+  message(STATUS "HIP sources: ")
+  foreach(tmp ${Caffe2_HIP_SRCS})
+    message(STATUS "  " ${tmp})
+  endforeach()
+
   message(STATUS "CPU test sources: ")
   foreach(tmp ${Caffe2_CPU_TEST_SRCS})
     message(STATUS "  " ${tmp})
@@ -80,6 +86,11 @@ if (FALSE)
   foreach(tmp ${Caffe2_GPU_TEST_SRCS})
     message(STATUS "  " ${tmp})
   endforeach()
+
+  message(STATUS "HIP test sources: ")
+  foreach(tmp ${Caffe2_HIP_TEST_SRCS})
+    message(STATUS "  " ${tmp})
+  endforeach()
 endif()
 
 # ---[ Generate and install header files.
@@ -184,6 +195,26 @@ if(BUILD_CAFFE2)
     list(APPEND Caffe2_MAIN_LIBS caffe2_gpu_library)
     install(TARGETS caffe2_gpu EXPORT Caffe2Targets DESTINATION lib)
   endif()
+
+  # ---[ HIP library.
+  if(USE_HIP)
+    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()
 
 # ---[ Test binaries.
@@ -193,6 +224,9 @@ if(BUILD_CAFFE2)
     if (USE_CUDA)
       list(APPEND Caffe2_ALL_TEST_SRCS ${Caffe2_GPU_TEST_SRCS})
     endif()
+    if(USE_HIP)
+      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)
@@ -204,6 +238,18 @@ if(BUILD_CAFFE2)
       add_test(NAME ${test_name} COMMAND $<TARGET_FILE:${test_name}>)
       install(TARGETS ${test_name} DESTINATION test)
     endforeach()
+
+    if(USE_HIP)
+      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()
 

caffe2/core/CMakeLists.txt

@@ -12,25 +12,44 @@ set(Caffe2_GPU_SRCS ${Caffe2_GPU_SRCS} ${tmp})
 file(GLOB tmp *_test.cc)
 exclude(Caffe2_GPU_SRCS "${Caffe2_GPU_SRCS}" ${tmp})
 
+# ---[ HIP files
+# ------[ HIP Sources
+file(GLOB tmp *_hip.cc)
+set(Caffe2_HIP_SRCS ${Caffe2_HIP_SRCS} ${tmp})
+# ------[ MIOpen files
+file(GLOB tmp *_miopen.cc)
+set(Caffe2_HIP_SRCS ${Caffe2_HIP_SRCS} ${tmp})
+# exclude test files
+file(GLOB 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 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)