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>
Assign(const SYCLDevice & d,OUT out,RHS 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