Adding ROCm support for slice and strided_slice ops

This commit is contained in:
Deven Desai 2019-06-10 15:20:12 +00:00
parent dac4bd7750
commit 09e3127c38
18 changed files with 32 additions and 32 deletions

View File

@ -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.

View File

@ -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

View File

@ -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<CPUDevice, int32, true>);
#undef REGISTER_GPU
#endif // GOOGLE_CUDA
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#ifdef TENSORFLOW_USE_SYCL
#define REGISTER_SYCL(type) \

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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_

View File

@ -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

View File

@ -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

View File

@ -230,7 +230,7 @@ class HandleStridedSliceAssignCase<Device, T, 0> {
// 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<Device, T, 0> {
#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);

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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