OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2026-01-15 12:15:41 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

[ROCm]: Add backward support for ROCm5.0

Browse Source
This commit is contained in:
Rahul Batra
2022-11-18 02:16:30 +00:00
parent 8b2e0ca8e3
commit 2748adb524
6 changed files with 28 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/xla/service/gpu/nccl_utils.h
View File

@@ -30,8 +30,12 @@ limitations under the License.
// Common place for all collective thunks to include nccl/rccl headers.
#if TENSORFLOW_USE_ROCM
#if (TF_ROCM_VERSION >= 50200)
#include "rocm/include/rccl/rccl.h"
#else
#include "rocm/include/rccl.h"
#endif
#else
#include "third_party/nccl/nccl.h"
#endif

5
tensorflow/compiler/xla/stream_executor/rocm/hipsparse_wrapper.h
View File

@@ -20,7 +20,12 @@ limitations under the License.
#ifndef TENSORFLOW_COMPILER_XLA_STREAM_EXECUTOR_ROCM_HIPSPARSE_WRAPPER_H_
#define TENSORFLOW_COMPILER_XLA_STREAM_EXECUTOR_ROCM_HIPSPARSE_WRAPPER_H_
#if (TF_ROCM_VERSION >= 50200)
#include "rocm/include/hipsparse/hipsparse.h"
#else
#include "rocm/include/hipsparse.h"
#endif
#include "tensorflow/compiler/xla/stream_executor/lib/env.h"
#include "tensorflow/compiler/xla/stream_executor/platform/dso_loader.h"
#include "tensorflow/compiler/xla/stream_executor/platform/port.h"

5
tensorflow/compiler/xla/stream_executor/rocm/rocm_fft.h
View File

@@ -22,7 +22,12 @@ limitations under the License.
#if TENSORFLOW_USE_ROCM
#if (TF_ROCM_VERSION >= 50200)
#include "rocm/include/hipfft/hipfft.h"
#else
#include "rocm/include/hipfft.h"
#endif
#include "rocm/rocm_config.h"
#include "rocm/rocm_config.h"
#endif

6
tensorflow/compiler/xla/stream_executor/rocm/rocsolver_wrapper.h
View File

@@ -20,7 +20,13 @@ limitations under the License.
#ifndef TENSORFLOW_COMPILER_XLA_STREAM_EXECUTOR_ROCM_ROCSOLVER_WRAPPER_H_
#define TENSORFLOW_COMPILER_XLA_STREAM_EXECUTOR_ROCM_ROCSOLVER_WRAPPER_H_
#include "rocm/rocm_config.h"
#if (TF_ROCM_VERSION >= 50200)
#include "rocm/include/rocsolver/rocsolver.h"
#else
#include "rocm/include/rocsolver.h"
#endif
#include "tensorflow/compiler/xla/stream_executor/lib/env.h"
#include "tensorflow/compiler/xla/stream_executor/platform/dso_loader.h"
#include "tensorflow/compiler/xla/stream_executor/platform/port.h"

4
tensorflow/core/kernels/nccl_ops.cc
View File

@@ -20,7 +20,11 @@ limitations under the License.
#if GOOGLE_CUDA
#include "third_party/nccl/nccl.h"
#elif TENSORFLOW_USE_ROCM
#if (TF_ROCM_VERSION >= 50200)
#include "rocm/include/rccl/rccl.h"
#else
#include "rocm/include/rccl.h"
#endif
#endif
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/nccl/nccl_manager.h"

4
tensorflow/core/nccl/nccl_manager.h
View File

@@ -30,7 +30,11 @@ limitations under the License.
#if GOOGLE_CUDA
#include "third_party/nccl/nccl.h"
#elif TENSORFLOW_USE_ROCM
#if (TF_ROCM_VERSION >= 50200)
#include "rocm/include/rccl/rccl.h"
#else
#include "rocm/include/rccl.h"
#endif
#include "tensorflow/core/common_runtime/gpu_device_context.h"
#endif
#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"