Support rocm

CMakeLists.txt

@@ -20,6 +20,7 @@ endif()
 option(USE_CLANG_FORMAT "" OFF)
 option(BUILD_RDMA "" OFF)
 option(BUILD_CUDA "" ON)
+option(BUILD_HIP "" OFF)
 option(BUILD_TESTING "" OFF)
 option(WITH_XLA "Option to build with XLA" OFF)
 option(WITH_TENSORRT "Option to build with TensorRT" OFF)
@@ -31,6 +32,10 @@ set(RPC_BACKEND "GRPC,LOCAL" CACHE STRING "")
 set(THIRD_PARTY_MIRROR "" CACHE STRING "")
 set(PIP_INDEX_MIRROR "" CACHE STRING "")
 
+if (BUILD_CUDA AND BUILD_HIP)
+  message(FATAL_ERROR "Enable cuda and hip simultaneously are not supported.")
+endif()
+
 if (APPLE)
   set(RPC_BACKEND "LOCAL")
   set(BUILD_CUDA OFF)
@@ -45,6 +50,11 @@ endif()
 
 project(oneflow C CXX)
 
+if (BUILD_HIP)
+  # Search for rocm in common locations
+  list(APPEND CMAKE_PREFIX_PATH /opt/rocm/hip /opt/rocm)
+endif()
+
 set(oneflow_cmake_dir ${PROJECT_SOURCE_DIR}/cmake)
 
 get_filename_component(real_src_dir "${CMAKE_SOURCE_DIR}" REALPATH)

cmake/oneflow.cmake

@@ -266,7 +266,7 @@ add_dependencies(of_pyext_obj of_ccobj)
 if(APPLE)
   set(of_libs -Wl,-force_load ${ONEFLOW_CUDA_LIBS} of_ccobj of_protoobj of_cfgobj)
 elseif(UNIX)
-  set(of_libs -Wl,--whole-archive ${ONEFLOW_CUDA_LIBS} of_ccobj of_protoobj of_cfgobj -Wl,--no-whole-archive -ldl -lrt)
+  set(of_libs -Wl,--whole-archive ${ONEFLOW_CUDA_LIBS} of_ccobj of_protoobj of_cfgobj -Wl,--no-whole-archive -ldl -lrt -ludev)
 elseif(WIN32)
   set(of_libs of_ccobj of_protoobj of_cfgobj)
   set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} /WHOLEARCHIVE:of_ccobj")

cmake/third_party.cmake

@@ -222,6 +222,26 @@ if (BUILD_CUDA)
   )
 endif()
 
+if (BUILD_HIP)
+  # Find hip
+  find_package(hip)
+  find_package(hipblas)
+  find_package(hipcub)
+  find_package(rocrand)
+  find_package(hiprand)
+  find_package(miopen)
+  find_package(rccl)
+  add_definitions(-DWITH_HIP)
+  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)
+endif()
+
 if(BUILD_RDMA)
   if(UNIX)
     include(CheckIncludeFiles)

oneflow/api/python/flags.cpp

@@ -19,11 +19,11 @@ namespace oneflow {
 
 ONEFLOW_API_PYBIND11_MODULE("flags", m) {
   m.def("with_cuda", []() {
-#ifdef WITH_CUDA
+#if defined (WITH_CUDA) || defined (WITH_HIP)
     return true;
 #else
     return false;
-#endif  // WITH_CUDA
+#endif  // defined (WITH_CUDA) || defined (WITH_HIP)
   });
 
   m.def("use_cxx11_abi", []() {

oneflow/core/actor/acc_compute_actor.cpp

@@ -53,7 +53,7 @@ void AccCompActor::Init(const TaskProto& task_proto, int32_t max_acc_cnt) {
   if (GetDeviceType() == DeviceType::kCPU) {
     cpy_func_ = std::bind(Memcpy<DeviceType::kCPU>, _1, _2, _3, _4);
   } else {
-#ifdef WITH_CUDA
+#if defined(WITH_CUDA) || defined(WITH_HIP)
     cpy_func_ = std::bind(Memcpy<DeviceType::kGPU>, _1, _2, _3, _4);
 #else
     UNIMPLEMENTED();
@@ -80,7 +80,7 @@ void AccCompActor::Act() {
       Memcpy<DeviceType::kCPU>(kernel_ctx.device_ctx, out_blob->ForceMutDptr(), in_blob->dptr(),
                                out_blob->ByteSizeOfBlobBody());
     } else if (GetDeviceType() == DeviceType::kGPU) {
-#ifdef WITH_CUDA
+#if defined(WITH_CUDA) || defined(WITH_HIP)
       Memcpy<DeviceType::kGPU>(kernel_ctx.device_ctx, out_blob->ForceMutDptr(), in_blob->dptr(),
                                out_blob->ByteSizeOfBlobBody());
 #else

oneflow/core/actor/copy_hd_actor.cpp

@@ -17,7 +17,7 @@ limitations under the License.
 
 namespace oneflow {
 
-#ifdef WITH_CUDA
+#if defined(WITH_CUDA) || defined(WITH_HIP)
 
 void CopyHdActor::VirtualActorInit(const TaskProto& task_proto) {
   OF_SET_MSG_HANDLER(&CopyHdActor::HandlerNormal);

oneflow/core/actor/copy_hd_actor.h

@@ -20,7 +20,7 @@ limitations under the License.
 
 namespace oneflow {
 
-#ifdef WITH_CUDA
+#if defined(WITH_CUDA) || defined(WITH_HIP)
 
 class CopyHdActor final : public Actor {
  public:

oneflow/core/common/data_type.cpp

@@ -72,7 +72,7 @@ void CheckDataType() {
   static_assert(sizeof(int32_t) == sizeof(int), "sizeof(int32_t) != sizeof(int)");
   static_assert(sizeof(int64_t) == sizeof(long long), "sizeof(int64_t) != sizeof(long long)");
 
-#if defined(WITH_CUDA)
+#if defined(WITH_CUDA) || defined(WITH_HIP)
 
 #define CHECK_DEVICE_FP16(get_val)                              \
   do {                                                          \

oneflow/core/common/data_type.h

@@ -20,6 +20,9 @@ limitations under the License.
 #if defined(WITH_CUDA)
 #include <cuda_fp16.h>
 #endif
+#if defined(WITH_HIP)
+#include <hip/hip_fp16.h>
+#endif
 #include "oneflow/core/common/fp16_data_type.h"
 #include "oneflow/core/common/data_type.pb.h"
 #include "oneflow/core/common/data_type_seq.h"
@@ -29,7 +32,7 @@ limitations under the License.
 
 namespace oneflow {
 
-#if defined(WITH_CUDA)
+#if defined(WITH_CUDA) || defined(WITH_HIP)
 #define DEVICE_TYPE_SEQ                  \
   OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) \
   OF_PP_MAKE_TUPLE_SEQ(DeviceType::kGPU)
@@ -81,7 +84,7 @@ struct GetDataType<T, typename std::enable_if<IsFloat16<T>::value>::type>
 template<DataType type>
 using DataTypeToType = decltype(GetTypeByDataType(std::integral_constant<DataType, type>{}));
 
-#if defined(__CUDACC__)
+#if defined(__CUDACC__) || defined(__HIP_PLATFORM_HCC__)
 #define OF_DEVICE_FUNC __device__ __host__ __forceinline__
 #else
 #define OF_DEVICE_FUNC inline
@@ -198,7 +201,7 @@ struct DevDType {
   typedef T type;
 };
 
-#if defined(WITH_CUDA)
+#if defined(WITH_CUDA) || defined(WITH_HIP)
 template<>
 struct DevDType<DeviceType::kGPU, float16> {
   static_assert(sizeof(float16) == sizeof(half), "sizeof(float16) != sizeof(half)");

oneflow/core/common/data_type_seq.h

@@ -49,7 +49,7 @@ limitations under the License.
 
 #define FLOAT16_DATA_TYPE_SEQ OF_PP_MAKE_TUPLE_SEQ(float16, DataType::kFloat16)
 
-#if defined(WITH_CUDA)
+#if defined(WITH_CUDA) || defined(WITH_HIP)
 #define HALF_DATA_TYPE_SEQ OF_PP_MAKE_TUPLE_SEQ(half, DataType::kFloat16)
 #endif
 

oneflow/core/device/cuda_util.h

@@ -186,13 +186,15 @@ int GetCudaPtxVersion();
 
 #else
 
+#if !defined(WITH_HIP)
 namespace oneflow {
 
 enum class CudaWorkType {};
 
 inline size_t GetCudaWorkTypeSize() { return 0; }
 
 }  // namespace oneflow
+#endif
 
 #endif  // WITH_CUDA
 

oneflow/core/device/device_context.h

@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_CORE_DEVICE_DEVICE_CONTEXT_H_
 
 #include "oneflow/core/device/cuda_util.h"
+#include "oneflow/core/device/hip_util.hip.h"
 #include "oneflow/core/common/auto_registration_factory.h"
 
 namespace oneflow {
@@ -40,6 +41,14 @@ class DeviceCtx {
   virtual const cudnnHandle_t& cudnn_handle() const { UNIMPLEMENTED(); }
 #endif
 
+#ifdef WITH_HIP
+  virtual const hipStream_t& hip_stream() const { UNIMPLEMENTED(); }
+  virtual const hipblasHandle_t& hipblas_pmh_handle() const { UNIMPLEMENTED(); }
+  virtual const hipblasHandle_t& hipblas_pmd_handle() const { UNIMPLEMENTED(); }
+  virtual const hipblasHandle_t& hipblas_tensor_op_math_handle() const { UNIMPLEMENTED(); }
+  virtual const miopenHandle_t& miopen_handle() const { UNIMPLEMENTED(); }
+#endif
+
   virtual void SyncDevice() { UNIMPLEMENTED(); }
   virtual void AddCallBack(std::function<void()>) const { UNIMPLEMENTED(); }
 

oneflow/core/device/hip_device_context.hip.cpp

@@ -0,0 +1,31 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/device/hip_device_context.hip.h"
+#include "oneflow/core/thread/thread_context.h"
+
+namespace oneflow {
+
+#ifdef WITH_HIP
+
+REGISTER_DEVICE_CONTEXT(DeviceType::kGPU, ([](const ThreadCtx& thread_ctx) -> DeviceCtx* {
+                          HipStreamHandle* hip_handle = nullptr;
+                          hip_handle = thread_ctx.g_hip_stream.get();
+                          return new HipDeviceCtx(hip_handle);
+                        }));
+
+#endif  // WITH_HIP
+
+}  // namespace oneflow

oneflow/core/device/hip_device_context.hip.h

@@ -0,0 +1,63 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_CORE_DEVICE_HIP_DEVICE_CONTEXT_H_
+#define ONEFLOW_CORE_DEVICE_HIP_DEVICE_CONTEXT_H_
+
+#include "oneflow/core/kernel/kernel_context.h"
+#include "oneflow/core/device/device_context.h"
+#include "oneflow/core/device/hip_stream_handle.hip.h"
+
+namespace oneflow {
+
+#ifdef WITH_HIP
+
+class HipDeviceCtx : public DeviceCtx {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(HipDeviceCtx);
+  HipDeviceCtx() = delete;
+  ~HipDeviceCtx() override = default;
+
+  explicit HipDeviceCtx(HipStreamHandle* hip_handler) : hip_handler_(hip_handler) {}
+
+  const hipStream_t& hip_stream() const override { return *(hip_handler_->hip_stream()); }
+
+  const hipblasHandle_t& hipblas_pmh_handle() const override {
+    return *(hip_handler_->hipblas_pmh_handle());
+  }
+  const hipblasHandle_t& hipblas_tensor_op_math_handle() const override {
+    return *(hip_handler_->hipblas_tensor_op_math_handle());
+  }
+  const hipblasHandle_t& hipblas_pmd_handle() const override {
+    return *(hip_handler_->hipblas_pmd_handle());
+  }
+
+  const miopenHandle_t& miopen_handle() const override { return *(hip_handler_->miopen_handle()); }
+
+  void SyncDevice() override { OF_HIP_CHECK(hipStreamSynchronize(hip_stream())); }
+
+  void AddCallBack(std::function<void()> callback) const override {
+    hip_handler_->AddCallBack(callback);
+  }
+
+ protected:
+  HipStreamHandle* hip_handler_;
+};
+
+#endif  // WITH_HIP
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_DEVICE_HIP_DEVICE_CONTEXT_H_