1 /* Copyright 2016 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 !TENSORFLOW_USE_SYCL 17 #error This file must only be included when building TensorFlow with SYCL support 18 #endif 19 20 #ifndef TENSORFLOW_CORE_KERNELS_CWISE_OPS_SYCL_COMMON_H_ 21 #define TENSORFLOW_CORE_KERNELS_CWISE_OPS_SYCL_COMMON_H_ 22 23 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 24 25 #include "tensorflow/core/framework/register_types.h" 26 #include "tensorflow/core/kernels/cwise_ops.h" 27 #include "tensorflow/core/platform/types.h" 28 29 namespace tensorflow { 30 namespace functor { 31 32 typedef Eigen::SyclDevice SYCLDevice; 33 34 template <typename OUT, typename RHS> 35 void Assign(const SYCLDevice& d, OUT out, RHS rhs) { 36 out.device(d) = rhs; 37 } 38 39 // Partial specialization of UnaryFunctor<Device=SYCLDevice, Functor>. 40 template <typename Functor> 41 struct UnaryFunctor<SYCLDevice, Functor> { 42 void operator()(const SYCLDevice& 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=SYCLDevice, Functor>. 49 template <typename Functor, int NDIMS, bool has_errors> 50 struct BinaryFunctor<SYCLDevice, Functor, NDIMS, has_errors> { 51 void operator()(const SYCLDevice& 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(To32Bit(in1), typename Functor::func()); 56 } 57 58 void Left(const SYCLDevice& d, typename Functor::tout_type out, 59 typename Functor::tscalar_type scalar, 60 typename Functor::tin_type in, bool* error) { 61 typedef typename Functor::func Binary; 62 constexpr int NumDims = Functor::tin_type::NumDimensions; 63 static_assert(NumDims == 1, "Unexpected size"); 64 Eigen::Sizes<1> scalar_dim; 65 out.device(d) = scalar.reshape(scalar_dim) 66 .broadcast(in.dimensions()) 67 .binaryExpr(in, Binary()); 68 } 69 70 void Right(const SYCLDevice& d, typename Functor::tout_type out, 71 typename Functor::tin_type in, 72 typename Functor::tscalar_type scalar, bool* error) { 73 typedef typename Functor::func Binary; 74 constexpr int NumDims = Functor::tin_type::NumDimensions; 75 static_assert(NumDims == 1, "Unexpected size"); 76 Eigen::Sizes<1> scalar_dim; 77 out.device(d) = in.binaryExpr( 78 scalar.reshape(scalar_dim).broadcast(in.dimensions()), Binary()); 79 } 80 81 void BCast(const SYCLDevice& d, 82 typename TTypes<typename Functor::out_type, NDIMS>::Tensor out, 83 typename TTypes<typename Functor::in_type, NDIMS>::ConstTensor in0, 84 typename Eigen::array<Eigen::DenseIndex, NDIMS> bcast0, 85 typename TTypes<typename Functor::in_type, NDIMS>::ConstTensor in1, 86 typename Eigen::array<Eigen::DenseIndex, NDIMS> bcast1, 87 bool* error) { 88 typedef typename Functor::in_type T; 89 typename Functor::func func; 90 if ((NDIMS == 2) && Functor::use_bcast_optimization && 91 use_bcast_optimization<T>::value) { 92 const bool bcast0_all_one = AllOne<NDIMS>(bcast0); 93 const bool bcast1_all_one = AllOne<NDIMS>(bcast1); 94 if (bcast0_all_one && !bcast1_all_one) { 95 To32Bit(out).device(d) = 96 To32Bit(in0).binaryExpr(To32Bit(in1).broadcast(bcast1), func); 97 return; 98 } 99 if (!bcast0_all_one && bcast1_all_one) { 100 To32Bit(out).device(d) = 101 To32Bit(in0).broadcast(bcast0).binaryExpr(To32Bit(in1), func); 102 return; 103 } 104 } 105 To32Bit(out).device(d) = To32Bit(in0).broadcast(bcast0).binaryExpr( 106 To32Bit(in1).broadcast(bcast1), func); 107 } 108 }; 109 110 // Macros to explicitly instantiate kernels on GPU for multiple types 111 // (T0, T1, etc.) for UnaryFunctor (e.g., functor::sqrt). 112 #define DEFINE_UNARY1(F, T) template struct UnaryFunctor<SYCLDevice, F<T> > 113 #define DEFINE_UNARY2(F, T0, T1) \ 114 DEFINE_UNARY1(F, T0); \ 115 DEFINE_UNARY1(F, T1) 116 #define DEFINE_UNARY3(F, T0, T1, T2) \ 117 DEFINE_UNARY2(F, T0, T1); \ 118 DEFINE_UNARY1(F, T2) 119 #define DEFINE_UNARY4(F, T0, T1, T2, T3) \ 120 DEFINE_UNARY2(F, T0, T1); \ 121 DEFINE_UNARY2(F, T2, T3) 122 #define DEFINE_UNARY5(F, T0, T1, T2, T3, T4) \ 123 DEFINE_UNARY2(F, T0, T1); \ 124 DEFINE_UNARY3(F, T2, T3, T4) 125 126 // Macros to explicitly instantiate kernels on GPU for multiple types 127 // (T0, T1, etc.) for BinaryFunctor. 128 #define DEFINE_BINARY1(F, T) \ 129 template struct BinaryFunctor<SYCLDevice, F<T>, 1>; \ 130 template struct BinaryFunctor<SYCLDevice, F<T>, 2>; \ 131 template struct BinaryFunctor<SYCLDevice, F<T>, 3> 132 #define DEFINE_BINARY2(F, T0, T1) \ 133 DEFINE_BINARY1(F, T0); \ 134 DEFINE_BINARY1(F, T1) 135 #define DEFINE_BINARY3(F, T0, T1, T2) \ 136 DEFINE_BINARY2(F, T0, T1); \ 137 DEFINE_BINARY1(F, T2) 138 #define DEFINE_BINARY4(F, T0, T1, T2, T3) \ 139 DEFINE_BINARY2(F, T0, T1); \ 140 DEFINE_BINARY2(F, T2, T3) 141 #define DEFINE_BINARY5(F, T0, T1, T2, T3, T4) \ 142 DEFINE_BINARY2(F, T0, T1); \ 143 DEFINE_BINARY3(F, T2, T3, T4) 144 #define DEFINE_BINARY6(F, T0, T1, T2, T3, T4, T5) \ 145 DEFINE_BINARY3(F, T0, T1, T2); \ 146 DEFINE_BINARY3(F, T3, T4, T5) 147 #define DEFINE_BINARY7(F, T0, T1, T2, T3, T4, T5, T6) \ 148 DEFINE_BINARY3(F, T0, T1, T2); \ 149 DEFINE_BINARY4(F, T3, T4, T5, T6) 150 #define DEFINE_BINARY8(F, T0, T1, T2, T3, T4, T5, T6, T7) \ 151 DEFINE_BINARY4(F, T0, T1, T2, T3); \ 152 DEFINE_BINARY4(F, T4, T5, T6, T7) 153 #define DEFINE_BINARY9(F, T0, T1, T2, T3, T4, T5, T6, T7, T8) \ 154 DEFINE_BINARY4(F, T0, T1, T2, T3); \ 155 DEFINE_BINARY5(F, T4, T5, T6, T7, T8) 156 #define DEFINE_BINARY10(F, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9) \ 157 DEFINE_BINARY5(F, T0, T1, T2, T3, T4); \ 158 DEFINE_BINARY5(F, T5, T6, T7, T8, T9) 159 160 } // end namespace functor 161 } // end namespace tensorflow 162 163 #endif // TENSORFLOW_CORE_KERNELS_CWISE_OPS_SYCL_COMMON_H_ 164
