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
17#error This file must only be included when building with Cuda support
18#endif
19
20#ifndef TENSORFLOW_KERNELS_CWISE_OPS_GPU_COMMON_CU_H_
21#define TENSORFLOW_KERNELS_CWISE_OPS_GPU_COMMON_CU_H_
22
23#define EIGEN_USE_GPU
24
25#include <complex>
26
27#include "tensorflow/core/framework/tensor_types.h"
28#include "tensorflow/core/kernels/cwise_ops.h"
29#include "tensorflow/core/platform/types.h"
30
31#include "tensorflow/core/platform/logging.h"
32namespace tensorflow {
33namespace functor {
34
35typedef Eigen::GpuDevice GPUDevice;
36typedef std::complex<float> complex64;
37typedef std::complex<double> complex128;
38
39// Partial specialization of UnaryFunctor<Device=GPUDevice, Functor>.
40template <typename Functor>
41struct UnaryFunctor<GPUDevice, Functor> {
42  void operator()(const GPUDevice& d, typename Functor::tout_type out,
43                  typename Functor::tin_type in) {
44    To32Bit(out).device(d) = To32Bit(in).unaryExpr(typename Functor::func());
45  }
46};
47
48// Partial specialization of BinaryFunctor<Device=GPUDevice, Functor>.
49template <typename Functor, int NDIMS, bool has_errors>
50struct BinaryFunctor<GPUDevice, Functor, NDIMS, has_errors> {
51  void operator()(const GPUDevice& d, typename Functor::tout_type out,
52                  typename Functor::tin_type in0,
53                  typename Functor::tin_type in1, bool* error) {
54    To32Bit(out).device(d) =
55        To32Bit(in0).binaryExpr(in1, typename Functor::func());
56  }
57
58  void Left(const GPUDevice& d, typename Functor::tout_type out,
59            typename Functor::tscalar_type scalar,
60            typename Functor::tin_type in, bool* error) {
61    typedef typename Functor::out_type Tout;
62    typedef typename Functor::in_type Tin;
63    typedef typename Functor::func Binary;
64    typedef typename Eigen::internal::scalar_left<Tout, Tin, Binary> Unary;
65    To32Bit(out).device(d) = To32Bit(in).unaryExpr(Unary(scalar.data()));
66  }
67
68  void Right(const GPUDevice& d, typename Functor::tout_type out,
69             typename Functor::tin_type in,
70             typename Functor::tscalar_type scalar, bool* error) {
71    typedef typename Functor::out_type Tout;
72    typedef typename Functor::in_type Tin;
73    typedef typename Functor::func Binary;
74    typedef typename Eigen::internal::scalar_right<Tout, Tin, Binary> Unary;
75    To32Bit(out).device(d) = To32Bit(in).unaryExpr(Unary(scalar.data()));
76  }
77
78  void BCast(const GPUDevice& d,
79             typename TTypes<typename Functor::out_type, NDIMS>::Tensor out,
80             typename TTypes<typename Functor::in_type, NDIMS>::ConstTensor in0,
81             typename Eigen::array<Eigen::DenseIndex, NDIMS> bcast0,
82             typename TTypes<typename Functor::in_type, NDIMS>::ConstTensor in1,
83             typename Eigen::array<Eigen::DenseIndex, NDIMS> bcast1,
84             bool* error) {
85    typedef typename Functor::in_type T;
86    typename Functor::func func;
87    if ((NDIMS == 2) && Functor::use_bcast_optimization &&
88        use_bcast_optimization<T>::value) {
89      const bool bcast0_all_one = AllOne<NDIMS>(bcast0);
90      const bool bcast1_all_one = AllOne<NDIMS>(bcast1);
91      if (bcast0_all_one && !bcast1_all_one) {
92        To32Bit(out).device(d) =
93            To32Bit(in0).binaryExpr(To32Bit(in1).broadcast(bcast1), func);
94        return;
95      }
96      if (!bcast0_all_one && bcast1_all_one) {
97        To32Bit(out).device(d) =
98            To32Bit(in0).broadcast(bcast0).binaryExpr(To32Bit(in1), func);
99        return;
100      }
101    }
102    To32Bit(out).device(d) = To32Bit(in0).broadcast(bcast0).binaryExpr(
103        To32Bit(in1).broadcast(bcast1), func);
104  }
105};
106
107// Partial specialization of ApproximateEqual<Device=GPUDevice, T>.
108template <typename T>
109struct ApproximateEqual<GPUDevice, T> {
110  void operator()(const GPUDevice& d, typename TTypes<T>::ConstFlat x,
111                  typename TTypes<T>::ConstFlat y, T tolerance,
112                  typename TTypes<bool>::Flat z) {
113    auto diff = x - y;
114    z.device(d) = diff.abs() <= tolerance;
115  }
116};
117
118// Macros to explicitly instantiate kernels on GPU for multiple types
119// (T0, T1, etc.) for UnaryFunctor (e.g., functor::sqrt).
120#define DEFINE_UNARY1(F, T) template struct UnaryFunctor<GPUDevice, F<T> >
121#define DEFINE_UNARY2(F, T0, T1) \
122  DEFINE_UNARY1(F, T0);          \
123  DEFINE_UNARY1(F, T1)
124#define DEFINE_UNARY3(F, T0, T1, T2) \
125  DEFINE_UNARY2(F, T0, T1);          \
126  DEFINE_UNARY1(F, T2)
127#define DEFINE_UNARY4(F, T0, T1, T2, T3) \
128  DEFINE_UNARY2(F, T0, T1);              \
129  DEFINE_UNARY2(F, T2, T3)
130#define DEFINE_UNARY5(F, T0, T1, T2, T3, T4) \
131  DEFINE_UNARY2(F, T0, T1);                  \
132  DEFINE_UNARY3(F, T2, T3, T4)
133#define DEFINE_UNARY6(F, T0, T1, T2, T3, T4, T5) \
134  DEFINE_UNARY2(F, T0, T1);                      \
135  DEFINE_UNARY4(F, T2, T3, T4, T5)
136#define DEFINE_UNARY7(F, T0, T1, T2, T3, T4, T5, T6) \
137  DEFINE_UNARY2(F, T0, T1);                          \
138  DEFINE_UNARY5(F, T2, T3, T4, T5, T6)
139#define DEFINE_UNARY8(F, T0, T1, T2, T3, T4, T5, T6, T7) \
140  DEFINE_UNARY4(F, T0, T1, T2, T3);                      \
141  DEFINE_UNARY4(F, T4, T5, T6, T7)
142
143// Macros to explicitly instantiate kernels on GPU for multiple types
144// (T0, T1, etc.) for BinaryFunctor.
145#define DEFINE_BINARY1(F, T)                         \
146  template struct BinaryFunctor<GPUDevice, F<T>, 1>; \
147  template struct BinaryFunctor<GPUDevice, F<T>, 2>; \
148  template struct BinaryFunctor<GPUDevice, F<T>, 3>; \
149  template struct BinaryFunctor<GPUDevice, F<T>, 4>; \
150  template struct BinaryFunctor<GPUDevice, F<T>, 5>
151#define DEFINE_BINARY2(F, T0, T1) \
152  DEFINE_BINARY1(F, T0);          \
153  DEFINE_BINARY1(F, T1)
154#define DEFINE_BINARY3(F, T0, T1, T2) \
155  DEFINE_BINARY2(F, T0, T1);          \
156  DEFINE_BINARY1(F, T2)
157#define DEFINE_BINARY4(F, T0, T1, T2, T3) \
158  DEFINE_BINARY2(F, T0, T1);              \
159  DEFINE_BINARY2(F, T2, T3)
160#define DEFINE_BINARY5(F, T0, T1, T2, T3, T4) \
161  DEFINE_BINARY2(F, T0, T1);                  \
162  DEFINE_BINARY3(F, T2, T3, T4)
163#define DEFINE_BINARY6(F, T0, T1, T2, T3, T4, T5) \
164  DEFINE_BINARY3(F, T0, T1, T2);                  \
165  DEFINE_BINARY3(F, T3, T4, T5)
166#define DEFINE_BINARY7(F, T0, T1, T2, T3, T4, T5, T6) \
167  DEFINE_BINARY3(F, T0, T1, T2);                      \
168  DEFINE_BINARY4(F, T3, T4, T5, T6)
169#define DEFINE_BINARY8(F, T0, T1, T2, T3, T4, T5, T6, T7) \
170  DEFINE_BINARY4(F, T0, T1, T2, T3);                      \
171  DEFINE_BINARY4(F, T4, T5, T6, T7)
172#define DEFINE_BINARY9(F, T0, T1, T2, T3, T4, T5, T6, T7, T8) \
173  DEFINE_BINARY4(F, T0, T1, T2, T3);                          \
174  DEFINE_BINARY5(F, T4, T5, T6, T7, T8)
175#define DEFINE_BINARY10(F, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9) \
176  DEFINE_BINARY5(F, T0, T1, T2, T3, T4);                           \
177  DEFINE_BINARY5(F, T5, T6, T7, T8, T9)
178#define DEFINE_BINARY11(F, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10) \
179  DEFINE_BINARY5(F, T0, T1, T2, T3, T4);                                \
180  DEFINE_BINARY6(F, T5, T6, T7, T8, T9, T10)
181
182#define DEFINE_APPROXIMATE_EQUAL1(T) \
183  template struct ApproximateEqual<GPUDevice, T>;
184#define DEFINE_APPROXIMATE_EQUAL2(T0, T1) \
185  DEFINE_APPROXIMATE_EQUAL1(T0);          \
186  DEFINE_APPROXIMATE_EQUAL1(T1);
187
188}  // end namespace functor
189}  // end namespace tensorflow
190
191#endif  // TENSORFLOW_KERNELS_CWISE_OPS_GPU_COMMON_CU_H_
192