diff --git a/.gitmodules b/.gitmodules index 7207ef9b6..0777fa968 100644 --- a/.gitmodules +++ b/.gitmodules @@ -11,3 +11,6 @@ [submodule "third_party/libnop"] path = third_party/libnop url = https://github.com/google/libnop.git +[submodule "third_party/hipify"] + path = third_party/hipify + url = https://github.com/ROCmSoftwarePlatform/hipify-torch.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 640054ee3..c31e61436 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,6 +21,13 @@ include(Sanitize) # Misc checks to cope with various compiler modes. include(MiscCheck) +# ROCm related +if (TP_USE_ROCM) + list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/third_party/hipify/cmake") + include(Hipify) + hipify(CUDA_SOURCE_DIR ${PROJECT_SOURCE_DIR}) +endif() + add_subdirectory(tensorpipe) install(EXPORT TensorpipeTargets diff --git a/cmake/Options.cmake b/cmake/Options.cmake index 10e09c94e..4a6da4458 100644 --- a/cmake/Options.cmake +++ b/cmake/Options.cmake @@ -31,6 +31,12 @@ endmacro() # TODO: Default to ON if CUDA available. option(TP_USE_CUDA "Enable support for CUDA tensors" OFF) +option(TP_USE_ROCM "Enable support for ROCM tensors" OFF) + +# if both TP_USE_CUDA and TP_USE_ROCM is set then break +if(TP_USE_CUDA AND TP_USE_ROCM) + message(FATAL_ERROR "Tensorpipe can be built either for CUDA or ROCm, TP_USE_CUDA and TP_USE_ROCM both are set, erroring out!!!!") +endif() # Optional features option(TP_BUILD_BENCHMARK "Build benchmarks" OFF)