1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
17
18#include "tensorflow/core/kernels/scatter_functor.h"
19#include "tensorflow/core/framework/register_types.h"
20
21namespace tensorflow {
22
23class OpKernelContext;
24typedef Eigen::GpuDevice GPUDevice;
25
26namespace functor {
27
28// Forward declarations of the functor specializations for GPU.
29#define DECLARE_GPU_SPECS_OP(T, Index, op) \
30 template <> \
31 Index ScatterFunctor<GPUDevice, T, Index, op>::operator()( \
32 OpKernelContext* c, const GPUDevice& d, \
33 typename TTypes<T>::Matrix params, \
34 typename TTypes<T>::ConstMatrix updates, \
35 typename TTypes<Index>::ConstFlat indices); \
36 extern template struct ScatterFunctor<GPUDevice, T, Index, op>; \
37 template <> \
38 Index ScatterScalarFunctor<GPUDevice, T, Index, op>::operator()( \
39 OpKernelContext* c, const GPUDevice& d, \
40 typename TTypes<T>::Matrix params, \
41 const typename TTypes<T>::ConstScalar update, \
42 typename TTypes<Index>::ConstFlat indices); \
43 extern template struct ScatterScalarFunctor<GPUDevice, T, Index, op>;
44
45#define DECLARE_GPU_SPECS_INDEX(T, Index) \
46 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ASSIGN); \
47 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::ADD); \
48 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::SUB); \
49 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MUL); \
50 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::DIV); \
51 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MIN); \
52 DECLARE_GPU_SPECS_OP(T, Index, scatter_op::UpdateOp::MAX);
53
54#define DECLARE_GPU_SPECS(T) \
55 DECLARE_GPU_SPECS_INDEX(T, int32); \
56 DECLARE_GPU_SPECS_INDEX(T, int64_t);
57
58TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPECS);
59
60#undef DECLARE_GPU_SPECS
61#undef DECLARE_GPU_SPECS_INDEX
62#undef DECLARE_GPU_SPECS_OP
63
64} // namespace functor
65} // namespace tensorflow
66
67#else
68
69#include "tensorflow/core/kernels/scatter_functor.h"
70
71#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
72