1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations 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" 32 namespace tensorflow { 33 namespace functor { 34 35 typedef Eigen::GpuDevice GPUDevice; 36 typedef std::complex<float> complex64; 37 typedef std::complex<double> complex128; 38 39 // Partial specialization of UnaryFunctor<Device=GPUDevice, Functor>. 40 template <typename Functor> 41 struct 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>. 49 template <typename Functor, int NDIMS, bool has_errors> 50 struct 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>. 108 template <typename T> 109 struct 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
