From 09e3127c384c3842f26b4115624e59268b87e229 Mon Sep 17 00:00:00 2001 From: Deven Desai Date: Mon, 10 Jun 2019 15:20:12 +0000 Subject: [PATCH] Adding ROCm support for slice and strided_slice ops --- tensorflow/core/kernels/slice_op.cc | 8 ++++---- tensorflow/core/kernels/slice_op_gpu.cu.cc | 4 ++-- tensorflow/core/kernels/strided_slice_op.cc | 8 ++++---- tensorflow/core/kernels/strided_slice_op_define_grad.cc | 2 +- tensorflow/core/kernels/strided_slice_op_gpu_bool.cu.cc | 4 ++-- .../core/kernels/strided_slice_op_gpu_complex.cu.cc | 4 ++-- tensorflow/core/kernels/strided_slice_op_gpu_impl.h | 4 ++-- tensorflow/core/kernels/strided_slice_op_gpu_int.cu.cc | 4 ++-- .../core/kernels/strided_slice_op_gpu_number_types.cu.cc | 4 ++-- tensorflow/core/kernels/strided_slice_op_impl.h | 6 +++--- tensorflow/core/kernels/strided_slice_op_inst_0.cc | 2 +- tensorflow/core/kernels/strided_slice_op_inst_1.cc | 2 +- tensorflow/core/kernels/strided_slice_op_inst_2.cc | 2 +- tensorflow/core/kernels/strided_slice_op_inst_3.cc | 2 +- tensorflow/core/kernels/strided_slice_op_inst_4.cc | 2 +- tensorflow/core/kernels/strided_slice_op_inst_5.cc | 2 +- tensorflow/core/kernels/strided_slice_op_inst_6.cc | 2 +- tensorflow/core/kernels/strided_slice_op_inst_7.cc | 2 +- 18 files changed, 32 insertions(+), 32 deletions(-) diff --git a/tensorflow/core/kernels/slice_op.cc b/tensorflow/core/kernels/slice_op.cc index 0b0ff95093e..1ac89b07c77 100644 --- a/tensorflow/core/kernels/slice_op.cc +++ b/tensorflow/core/kernels/slice_op.cc @@ -17,9 +17,9 @@ limitations under the License. #define EIGEN_USE_THREADS -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #define EIGEN_USE_GPU -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #include "tensorflow/core/kernels/slice_op.h" @@ -267,7 +267,7 @@ TF_CALL_POD_STRING_TYPES(REGISTER_SLICE); TF_CALL_QUANTIZED_TYPES(REGISTER_SLICE); #undef REGISTER_SLICE -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM // Forward declarations of the functor specializations for GPU. namespace functor { #define DECLARE_GPU_SPEC(T, NDIM) \ @@ -331,7 +331,7 @@ REGISTER_KERNEL_BUILDER(Name("Slice") #undef REGISTER_GPU -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #ifdef TENSORFLOW_USE_SYCL // Forward declarations of the functor specializations for SYCL. diff --git a/tensorflow/core/kernels/slice_op_gpu.cu.cc b/tensorflow/core/kernels/slice_op_gpu.cu.cc index 044948f4065..f76eab04e11 100644 --- a/tensorflow/core/kernels/slice_op_gpu.cu.cc +++ b/tensorflow/core/kernels/slice_op_gpu.cu.cc @@ -13,7 +13,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 @@ -49,4 +49,4 @@ DEFINE_GPU_KERNELS(int64); } // end namespace tensorflow -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc index eb202070042..5d4ee523bdb 100644 --- a/tensorflow/core/kernels/strided_slice_op.cc +++ b/tensorflow/core/kernels/strided_slice_op.cc @@ -18,9 +18,9 @@ limitations under the License. #include "tensorflow/core/lib/core/refcount.h" #define EIGEN_USE_THREADS -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #define EIGEN_USE_GPU -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #include "tensorflow/core/kernels/strided_slice_op.h" @@ -440,7 +440,7 @@ TF_CALL_ALL_TYPES(REGISTER_STRIDED_SLICE); #undef REGISTER_STRIDED_SLICE -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #define REGISTER_GPU(type) \ REGISTER_KERNEL_BUILDER(Name("StridedSlice") \ @@ -536,7 +536,7 @@ REGISTER_KERNEL_BUILDER(Name("TensorStridedSliceUpdate") StridedSliceAssignOp); #undef REGISTER_GPU -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #ifdef TENSORFLOW_USE_SYCL #define REGISTER_SYCL(type) \ diff --git a/tensorflow/core/kernels/strided_slice_op_define_grad.cc b/tensorflow/core/kernels/strided_slice_op_define_grad.cc index eac120d4d37..c6f8985e590 100644 --- a/tensorflow/core/kernels/strided_slice_op_define_grad.cc +++ b/tensorflow/core/kernels/strided_slice_op_define_grad.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 diff --git a/tensorflow/core/kernels/strided_slice_op_gpu_bool.cu.cc b/tensorflow/core/kernels/strided_slice_op_gpu_bool.cu.cc index 8c3f8f2ad30..38b29b74bc6 100644 --- a/tensorflow/core/kernels/strided_slice_op_gpu_bool.cu.cc +++ b/tensorflow/core/kernels/strided_slice_op_gpu_bool.cu.cc @@ -13,7 +13,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 @@ -24,4 +24,4 @@ namespace tensorflow { TF_CALL_bool(DEFINE_GPU_KERNELS); } // end namespace tensorflow -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/strided_slice_op_gpu_complex.cu.cc b/tensorflow/core/kernels/strided_slice_op_gpu_complex.cu.cc index f6951924655..a930b8a3fac 100644 --- a/tensorflow/core/kernels/strided_slice_op_gpu_complex.cu.cc +++ b/tensorflow/core/kernels/strided_slice_op_gpu_complex.cu.cc @@ -13,7 +13,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 @@ -25,4 +25,4 @@ TF_CALL_complex64(DEFINE_GPU_KERNELS); TF_CALL_complex128(DEFINE_GPU_KERNELS); } // end namespace tensorflow -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/strided_slice_op_gpu_impl.h b/tensorflow/core/kernels/strided_slice_op_gpu_impl.h index d70f369ac07..f451fe4bb7a 100644 --- a/tensorflow/core/kernels/strided_slice_op_gpu_impl.h +++ b/tensorflow/core/kernels/strided_slice_op_gpu_impl.h @@ -16,7 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_CORE_KERNELS_STRIDED_SLICE_OP_GPU_IMPL_H_ #define TENSORFLOW_CORE_KERNELS_STRIDED_SLICE_OP_GPU_IMPL_H_ -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #define EIGEN_USE_GPU @@ -56,5 +56,5 @@ typedef Eigen::GpuDevice GPUDevice; } // end namespace tensorflow -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM #endif // TENSORFLOW_CORE_KERNELS_STRIDED_SLICE_OP_GPU_IMPL_H_ diff --git a/tensorflow/core/kernels/strided_slice_op_gpu_int.cu.cc b/tensorflow/core/kernels/strided_slice_op_gpu_int.cu.cc index 07dd0130adc..d86c7b2fe23 100644 --- a/tensorflow/core/kernels/strided_slice_op_gpu_int.cu.cc +++ b/tensorflow/core/kernels/strided_slice_op_gpu_int.cu.cc @@ -13,7 +13,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 @@ -26,4 +26,4 @@ TF_CALL_int32(DEFINE_GPU_KERNELS); TF_CALL_int64(DEFINE_GPU_KERNELS); } // end namespace tensorflow -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/strided_slice_op_gpu_number_types.cu.cc b/tensorflow/core/kernels/strided_slice_op_gpu_number_types.cu.cc index 149886308cd..1725185bbbc 100644 --- a/tensorflow/core/kernels/strided_slice_op_gpu_number_types.cu.cc +++ b/tensorflow/core/kernels/strided_slice_op_gpu_number_types.cu.cc @@ -13,7 +13,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 @@ -24,4 +24,4 @@ namespace tensorflow { TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_KERNELS); } // end namespace tensorflow -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/strided_slice_op_impl.h b/tensorflow/core/kernels/strided_slice_op_impl.h index d9b62d4c754..e7d9a5e129f 100644 --- a/tensorflow/core/kernels/strided_slice_op_impl.h +++ b/tensorflow/core/kernels/strided_slice_op_impl.h @@ -230,7 +230,7 @@ class HandleStridedSliceAssignCase { // Dimension 0 only instantiates some functors. So we only need // to prevent ones defined by PREVENT_INSTANTIATE_DIM0_ONLY -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if STRIDED_SLICE_INSTANTIATE_DIM == 0 #define PREVENT_INSTANTIATE(T, NDIM) PREVENT_INSTANTIATE_DIM0_ONLY(T, NDIM) #else @@ -276,7 +276,7 @@ class HandleStridedSliceAssignCase { #define DECLARE_FOR_N_GPU(T) \ INSTANTIATE(GPUDevice, T, STRIDED_SLICE_INSTANTIATE_DIM) -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM TF_CALL_GPU_PROXY_TYPES(PREVENT_FOR_N_GPU); TF_CALL_complex64(PREVENT_FOR_N_GPU); TF_CALL_complex128(PREVENT_FOR_N_GPU); @@ -288,7 +288,7 @@ TF_CALL_bool(DECLARE_FOR_N_GPU); TF_CALL_int8(DECLARE_FOR_N_GPU); DECLARE_FOR_N_GPU(int32); DECLARE_FOR_N_GPU(int64); -#endif // END GOOGLE_CUDA +#endif // END GOOGLE_CUDA || TENSORFLOW_USE_ROCM TF_CALL_ALL_TYPES(DECLARE_FOR_N_CPU); diff --git a/tensorflow/core/kernels/strided_slice_op_inst_0.cc b/tensorflow/core/kernels/strided_slice_op_inst_0.cc index 48b52442d65..fa9efe813b6 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_0.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_0.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 diff --git a/tensorflow/core/kernels/strided_slice_op_inst_1.cc b/tensorflow/core/kernels/strided_slice_op_inst_1.cc index eac120d4d37..c6f8985e590 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_1.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_1.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 diff --git a/tensorflow/core/kernels/strided_slice_op_inst_2.cc b/tensorflow/core/kernels/strided_slice_op_inst_2.cc index b8810ccc5ba..55810a6fb35 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_2.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_2.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 diff --git a/tensorflow/core/kernels/strided_slice_op_inst_3.cc b/tensorflow/core/kernels/strided_slice_op_inst_3.cc index 67b9581ab3c..b1e5342a95c 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_3.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_3.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 diff --git a/tensorflow/core/kernels/strided_slice_op_inst_4.cc b/tensorflow/core/kernels/strided_slice_op_inst_4.cc index f84f372769c..77553127f23 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_4.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_4.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 diff --git a/tensorflow/core/kernels/strided_slice_op_inst_5.cc b/tensorflow/core/kernels/strided_slice_op_inst_5.cc index 3648a50d392..1b0d6896d6d 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_5.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_5.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 diff --git a/tensorflow/core/kernels/strided_slice_op_inst_6.cc b/tensorflow/core/kernels/strided_slice_op_inst_6.cc index ac97bb861bb..1c2b34f4a5b 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_6.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_6.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 diff --git a/tensorflow/core/kernels/strided_slice_op_inst_7.cc b/tensorflow/core/kernels/strided_slice_op_inst_7.cc index 88b44cfbb18..03f716b7fe5 100644 --- a/tensorflow/core/kernels/strided_slice_op_inst_7.cc +++ b/tensorflow/core/kernels/strided_slice_op_inst_7.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