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 #ifndef TENSORFLOW_CORE_KERNELS_CAST_OP_IMPL_H_ 17 #define TENSORFLOW_CORE_KERNELS_CAST_OP_IMPL_H_ 18 19 #define EIGEN_USE_THREADS 20 21 #include "tensorflow/core/framework/op_kernel.h" 22 #include "tensorflow/core/kernels/cast_op.h" 23 24 namespace tensorflow { 25 26 namespace functor { 27 28 template <typename O, typename I> 29 struct CastFunctor<Eigen::ThreadPoolDevice, O, I> { 30 void operator()(const Eigen::ThreadPoolDevice& d, typename TTypes<O>::Flat o, 31 typename TTypes<I>::ConstFlat i) { 32 o.device(d) = i.template cast<O>(); 33 } 34 }; 35 36 #ifdef TENSORFLOW_USE_SYCL 37 template <typename O, typename I> 38 struct CastFunctor<Eigen::SyclDevice, O, I> { 39 void operator()(const Eigen::SyclDevice& d, typename TTypes<O>::Flat o, 40 typename TTypes<I>::ConstFlat i) { 41 o.device(d) = i.template cast<O>(); 42 } 43 }; 44 #endif // TENSORFLOW_USE_SYCL 45 46 } // namespace functor 47 48 #define CURRY_TYPES3_NO_HALF(FN, arg0, arg1) \ 49 FN(arg0, arg1, bool); \ 50 FN(arg0, arg1, uint8); \ 51 FN(arg0, arg1, int8); \ 52 FN(arg0, arg1, uint16); \ 53 FN(arg0, arg1, int16); \ 54 FN(arg0, arg1, int32); \ 55 FN(arg0, arg1, int64); \ 56 FN(arg0, arg1, float); \ 57 FN(arg0, arg1, double); \ 58 FN(arg0, arg1, std::complex<float>); \ 59 FN(arg0, arg1, std::complex<double>) 60 61 #define CURRY_TYPES3(FN, arg0, arg1) \ 62 CURRY_TYPES3_NO_HALF(FN, arg0, arg1) \ 63 FN(arg0, arg1, Eigen::half); 64 65 #define CAST_CASE(DEVICE, IN, OUT) \ 66 if (DataTypeToEnum<OUT>::value == dst_dtype) { \ 67 return [](OpKernelContext* ctx, const Tensor& inp, Tensor* out) { \ 68 functor::CastFunctor<DEVICE, OUT, IN> func; \ 69 func(ctx->eigen_device<DEVICE>(), out->flat<OUT>(), inp.flat<IN>()); \ 70 }; \ 71 } 72 73 // The functions below are implemented in the cast_op_impl_*.cc files. 74 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 75 GetCpuCastFromBool(DataType dst_dtype); 76 77 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 78 GetCpuCastFromUint8(DataType dst_dtype); 79 80 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 81 GetCpuCastFromInt8(DataType dst_dtype); 82 83 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 84 GetCpuCastFromUint16(DataType dst_dtype); 85 86 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 87 GetCpuCastFromInt16(DataType dst_dtype); 88 89 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 90 GetCpuCastFromInt32(DataType dst_dtype); 91 92 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 93 GetCpuCastFromInt64(DataType dst_dtype); 94 95 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 96 GetCpuCastFromHalf(DataType dst_dtype); 97 98 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 99 GetCpuCastFromFloat(DataType dst_dtype); 100 101 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 102 GetCpuCastFromDouble(DataType dst_dtype); 103 104 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 105 GetCpuCastFromComplex64(DataType dst_dtype); 106 107 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 108 GetCpuCastFromComplex128(DataType dst_dtype); 109 110 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 111 GetCpuCastFromBfloat(DataType dst_dtype); 112 113 #if GOOGLE_CUDA 114 // Same, for GPU. 115 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 116 GetGpuCastFromBool(DataType dst_dtype); 117 118 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 119 GetGpuCastFromUint8(DataType dst_dtype); 120 121 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 122 GetGpuCastFromInt8(DataType dst_dtype); 123 124 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 125 GetGpuCastFromUint16(DataType dst_dtype); 126 127 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 128 GetGpuCastFromInt16(DataType dst_dtype); 129 130 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 131 GetGpuCastFromInt32(DataType dst_dtype); 132 133 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 134 GetGpuCastFromInt64(DataType dst_dtype); 135 136 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 137 GetGpuCastFromHalf(DataType dst_dtype); 138 139 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 140 GetGpuCastFromFloat(DataType dst_dtype); 141 142 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 143 GetGpuCastFromDouble(DataType dst_dtype); 144 145 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 146 GetGpuCastFromComplex64(DataType dst_dtype); 147 148 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 149 GetGpuCastFromComplex128(DataType dst_dtype); 150 151 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 152 GetGpuCastFromBfloat(DataType dst_dtype); 153 154 #endif // GOOGLE_CUDA 155 156 #ifdef TENSORFLOW_USE_SYCL 157 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 158 GetSyclCastFromBool(DataType dst_dtype); 159 160 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 161 GetSyclCastFromUint8(DataType dst_dtype); 162 163 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 164 GetSyclCastFromUint16(DataType dst_dtype); 165 166 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 167 GetSyclCastFromInt16(DataType dst_dtype); 168 169 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 170 GetSyclCastFromInt32(DataType dst_dtype); 171 172 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 173 GetSyclCastFromInt64(DataType dst_dtype); 174 175 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 176 GetSyclCastFromFloat(DataType dst_dtype); 177 178 std::function<void(OpKernelContext*, const Tensor&, Tensor*)> 179 GetSyclCastFromDouble(DataType dst_dtype); 180 #endif // TENSORFLOW_USE_SYCL 181 182 } // namespace tensorflow 183 184 #endif // TENSORFLOW_CORE_KERNELS_CAST_OP_IMPL_H_ 185
