Add ROCM backend

CMakeLists.txt

@@ -15,6 +15,7 @@ option(BUILD_PYTHON "" ON)
 option(BUILD_CPP_API "Option to build OneFlow C++ API (beta)" OFF)
 option(BUILD_RDMA "" OFF)
 option(BUILD_CUDA "" ON)
+option(BUILD_ROCM "" OFF)
 option(BUILD_TESTING "" OFF)
 option(BUILD_GIT_VERSION "" ON)
 option(BUILD_PROFILER "" OFF)
@@ -32,6 +33,16 @@ option(MAYBE_NEED_ERROR_MSG_CHECK "" OFF)
 # https://medium.com/@alasher/colored-c-compiler-output-with-ninja-clang-gcc-10bfe7f2b949
 option(OF_FORCE_COLORED_DIAGNOSTICS "Always produce ANSI-colored diagnostics (GNU/Clang only)." ON)
 
+if (BUILD_CUDA AND BUILD_ROCM)
+  message(FATAL_ERROR "Compile with cuda and rocm enabled simultaneously are not supported right now.")
+endif()
+
+if (BUILD_ROCM)
+  # Search for rocm in common locations
+  list(APPEND CMAKE_PREFIX_PATH /opt/rocm/hip /opt/rocm)
+endif()
+
+
 set(ONEFLOW_CURRENT_VERSION 0.7.0.dev CACHE STRING "")
 if(BUILD_FOR_CI)
   set(ONEFLOW_CURRENT_VERSION ci)

README.md

@@ -1,4 +1,4 @@
-# OneFlow
+# OneFlow 
 
 **OneFlow is a performance-centered and open-source deep learning framework.**
 

cmake/third_party.cmake

@@ -161,6 +161,38 @@ if(BUILD_CUDA)
        ${NCCL_INCLUDE_DIR})
 endif()
 
+
+if (BUILD_ROCM)
+  # Find rocm packages
+  find_package(hip)
+  find_package(hipblas)
+  find_package(hipcub)
+  find_package(hiprand)
+  find_package(rocrand)
+  find_package(miopen)
+  find_package(rccl)
+  add_definitions(-DWITH_ROCM)
+  add_definitions(-D__HIP_PLATFORM_HCC__)
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -D__HIP_PLATFORM_HCC__")
+  list(APPEND oneflow_third_party_libs hip::device)
+  list(APPEND oneflow_third_party_libs roc::hipblas)
+  list(APPEND oneflow_third_party_libs hip::hipcub)
+  list(APPEND oneflow_third_party_libs roc::rocrand)
+  list(APPEND oneflow_third_party_libs hip::hiprand)
+  list(APPEND oneflow_third_party_libs MIOpen)
+  link_directories(/opt/rocm/rccl/lib)
+  list(APPEND oneflow_third_party_libs rccl)
+  list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${HIP_INCLUDE_DIRS} 
+                                               ${HIPBLAS_INCLUDE_DIRS}
+                                               ${HIPCUB_INCLUDE_DIRS}
+                                               "/opt/rocm/hiprand/include"
+                                               "/opt/rocm/rocrand/include"
+                                               ${MIOPEN_INCLUDE_DIRS}
+                                               ${RCCL_INCLUDE_DIRS})
+  message(STATUS "ONEFLOW_THIRD_PARTY_INCLUDE_DIRS: ${ONEFLOW_THIRD_PARTY_INCLUDE_DIRS}")
+endif()
+
+
 if(BUILD_RDMA)
   if(UNIX)
     include(CheckIncludeFiles)

oneflow/core/common/device_type.proto

@@ -6,4 +6,5 @@ enum DeviceType {
   kCPU = 1;
   kCUDA = 2;
   kMockDevice = 3; // pseudo device for test.
+  kROCm = 4;
 }

python/oneflow/framework/tensor_str.py

@@ -349,10 +349,8 @@ def _gen_tensor_str_template(tensor, is_meta):
     if tensor.is_global:
         suffixes.append(f"placement={str(tensor.placement)}")
         suffixes.append(f"sbp={str(tensor.sbp)}")
-    elif tensor.device.type == "cuda":
-        suffixes.append("device='" + str(tensor.device) + "'")
     elif tensor.device.type != "cpu":
-        raise RunTimeError("unknow device type")
+        suffixes.append("device='" + str(tensor.device) + "'")
     if tensor.is_lazy:
         suffixes.append("is_lazy='True'")
 
@@ -366,7 +364,10 @@ def _gen_tensor_str_template(tensor, is_meta):
         tensor_str = "..."
         suffixes.append("size=" + str(tuple(tensor.shape)))
     else:
-        tensor_str = _tensor_str(tensor, indent)
+        if tensor.device.type != "cpu" and tensor.device.type != "cuda":
+            tensor_str = _tensor_str(tensor.detach().to("cpu"), indent)
+        else:
+            tensor_str = _tensor_str(tensor, indent)
 
     suffixes.append("dtype=" + str(tensor.dtype))
     if tensor.grad_fn is not None: