From 9f46967616d8f32bff757e3663c97955f5f7235a Mon Sep 17 00:00:00 2001
From: "Wen-Heng (Jack) Chung" <whchung@gmail.com>
Date: Thu, 30 May 2019 15:12:02 +0000
Subject: [PATCH] Add ROCm support for adjust_hue and adjust_saturation op
---
tensorflow/core/kernels/adjust_hsv_gpu.cu.h | 4 ++--
tensorflow/core/kernels/adjust_hue_op.cc | 4 ++--
tensorflow/core/kernels/adjust_hue_op.h | 4 ++--
tensorflow/core/kernels/adjust_hue_op_gpu.cu.cc | 12 ++++++------
tensorflow/core/kernels/adjust_saturation_op.cc | 4 ++--
tensorflow/core/kernels/adjust_saturation_op.h | 4 ++--
.../core/kernels/adjust_saturation_op_gpu.cu.cc | 12 ++++++------
7 files changed, 22 insertions(+), 22 deletions(-)
diff --git a/tensorflow/core/kernels/adjust_hsv_gpu.cu.h b/tensorflow/core/kernels/adjust_hsv_gpu.cu.h
index d2ccef4ca89..4a43dc55d6d 100644
--- a/tensorflow/core/kernels/adjust_hsv_gpu.cu.h
+++ b/tensorflow/core/kernels/adjust_hsv_gpu.cu.h
@@ -14,7 +14,7 @@ limitations under the License.
#ifndef TENSORFLOW_CORE_KERNELS_ADJUST_HSV_GPU_CU_H_
#define TENSORFLOW_CORE_KERNELS_ADJUST_HSV_GPU_CU_H_
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
@@ -141,5 +141,5 @@ __global__ void adjust_hsv_nhwc(const int64 number_elements,
} // namespace internal
} // namespace tensorflow
-#endif // GOOGLE_CUDA
+#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#endif // TENSORFLOW_CORE_KERNELS_ADJUST_HSV_GPU_CU_H_
diff --git a/tensorflow/core/kernels/adjust_hue_op.cc b/tensorflow/core/kernels/adjust_hue_op.cc
index 06de5ea3fb6..c1993029ac6 100644
--- a/tensorflow/core/kernels/adjust_hue_op.cc
+++ b/tensorflow/core/kernels/adjust_hue_op.cc
@@ -13,7 +13,7 @@ limitations under the License.
==============================================================================*/
#define EIGEN_USE_THREADS
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#endif
@@ -249,7 +249,7 @@ REGISTER_KERNEL_BUILDER(
Name("AdjustHue").Device(DEVICE_CPU).TypeConstraint<float>("T"),
AdjustHueOp<CPUDevice, float>);
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
template <typename T>
class AdjustHueOp<GPUDevice, T> : public AdjustHueOpBase {
public:
diff --git a/tensorflow/core/kernels/adjust_hue_op.h b/tensorflow/core/kernels/adjust_hue_op.h
index 6d6699de3fb..edaf7f538e3 100644
--- a/tensorflow/core/kernels/adjust_hue_op.h
+++ b/tensorflow/core/kernels/adjust_hue_op.h
@@ -14,7 +14,7 @@ limitations under the License.
#ifndef TENSORFLOW_CORE_KERNELS_ADJUST_HUE_OP_H_
#define TENSORFLOW_CORE_KERNELS_ADJUST_HUE_OP_H_
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
@@ -37,5 +37,5 @@ struct AdjustHueGPU {
} // namespace functor
} // namespace tensorflow
-#endif // GOOGLE_CUDA
+#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#endif // TENSORFLOW_CORE_KERNELS_ADJUST_HUE_OP_H_
diff --git a/tensorflow/core/kernels/adjust_hue_op_gpu.cu.cc b/tensorflow/core/kernels/adjust_hue_op_gpu.cu.cc
index d1ac433035c..a9f36fdbc3d 100644
--- a/tensorflow/core/kernels/adjust_hue_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/adjust_hue_op_gpu.cu.cc
@@ -12,7 +12,7 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
@@ -35,10 +35,10 @@ void AdjustHueGPU<T>::operator()(GPUDevice* device,
const int threads_per_block = config.thread_per_block;
const int block_count =
(number_of_elements + threads_per_block - 1) / threads_per_block;
- TF_CHECK_OK(CudaLaunchKernel(internal::adjust_hsv_nhwc<true, false, false, T>,
- block_count, threads_per_block, 0, stream,
- number_of_elements, input, output, delta,
- nullptr, nullptr));
+ TF_CHECK_OK(GpuLaunchKernel(internal::adjust_hsv_nhwc<true, false, false, T>,
+ block_count, threads_per_block, 0, stream,
+ number_of_elements, input, output, delta,
+ nullptr, nullptr));
}
template struct AdjustHueGPU<float>;
@@ -46,4 +46,4 @@ template struct AdjustHueGPU<Eigen::half>;
} // namespace functor
} // namespace tensorflow
-#endif // GOOGLE_CUDA
+#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
diff --git a/tensorflow/core/kernels/adjust_saturation_op.cc b/tensorflow/core/kernels/adjust_saturation_op.cc
index 98264c4a1de..d1fc9d349be 100644
--- a/tensorflow/core/kernels/adjust_saturation_op.cc
+++ b/tensorflow/core/kernels/adjust_saturation_op.cc
@@ -14,7 +14,7 @@ limitations under the License.
==============================================================================*/
#define EIGEN_USE_THREADS
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#endif
@@ -215,7 +215,7 @@ REGISTER_KERNEL_BUILDER(
Name("AdjustSaturation").Device(DEVICE_CPU).TypeConstraint<float>("T"),
AdjustSaturationOp<CPUDevice, float>);
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
template <typename T>
class AdjustSaturationOp<GPUDevice, T> : public AdjustSaturationOpBase {
public:
diff --git a/tensorflow/core/kernels/adjust_saturation_op.h b/tensorflow/core/kernels/adjust_saturation_op.h
index c21ce4e3608..0117a48ead8 100644
--- a/tensorflow/core/kernels/adjust_saturation_op.h
+++ b/tensorflow/core/kernels/adjust_saturation_op.h
@@ -14,7 +14,7 @@ limitations under the License.
#ifndef TENSORFLOW_CORE_KERNELS_ADJUST_SATURATION_OP_H_
#define TENSORFLOW_CORE_KERNELS_ADJUST_SATURATION_OP_H_
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
@@ -37,5 +37,5 @@ struct AdjustSaturationGPU {
} // namespace functor
} // namespace tensorflow
-#endif // GOOGLE_CUDA
+#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#endif // TENSORFLOW_CORE_KERNELS_ADJUST_SATURATION_OP_H_
diff --git a/tensorflow/core/kernels/adjust_saturation_op_gpu.cu.cc b/tensorflow/core/kernels/adjust_saturation_op_gpu.cu.cc
index f9c2806ee7b..749f1381f03 100644
--- a/tensorflow/core/kernels/adjust_saturation_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/adjust_saturation_op_gpu.cu.cc
@@ -12,7 +12,7 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
@@ -36,10 +36,10 @@ void AdjustSaturationGPU<T>::operator()(GPUDevice* device,
const int threads_per_block = config.thread_per_block;
const int block_count =
(number_of_elements + threads_per_block - 1) / threads_per_block;
- TF_CHECK_OK(CudaLaunchKernel(internal::adjust_hsv_nhwc<false, true, false, T>,
- block_count, threads_per_block, 0, stream,
- number_of_elements, input, output, nullptr,
- scale, nullptr));
+ TF_CHECK_OK(GpuLaunchKernel(internal::adjust_hsv_nhwc<false, true, false, T>,
+ block_count, threads_per_block, 0, stream,
+ number_of_elements, input, output, nullptr,
+ scale, nullptr));
}
template struct AdjustSaturationGPU<float>;
@@ -47,4 +47,4 @@ template struct AdjustSaturationGPU<Eigen::half>;
} // namespace functor
} // namespace tensorflow
-#endif // GOOGLE_CUDA
+#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM