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_SCAN_OPS_H_ 17 #define TENSORFLOW_CORE_KERNELS_SCAN_OPS_H_ 18 19 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 20 #include "tensorflow/core/framework/tensor_types.h" 21 22 namespace tensorflow { 23 namespace functor { 24 25 typedef Eigen::Index Index; 26 27 // TODO(b/154339590): Needs to be vectorized. 28 template <typename Device, typename Reducer, typename T> 29 struct Scan { operatorScan30 void operator()(const Device& d, typename TTypes<T, 3>::ConstTensor in, 31 typename TTypes<T, 3>::Tensor out, const Reducer& reducer, 32 const bool reverse, const bool exclusive) { 33 // Perform the reverse ops directly with Eigen, which avoids copying the 34 // tensor twice compared to using individual ops. 35 Eigen::array<bool, 3> dims; 36 dims[0] = false; 37 dims[1] = reverse; 38 dims[2] = false; 39 To32Bit(out).device(d) = 40 To32Bit(in).reverse(dims).scan(1, reducer, exclusive).reverse(dims); 41 } 42 }; 43 44 template <typename T> 45 struct LogSumExp { operatorLogSumExp46 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T operator()(const T& a, 47 const T& b) const { 48 auto mi = Eigen::internal::scalar_min_op<T>()(a, b); 49 auto ma = Eigen::internal::scalar_max_op<T>()(a, b); 50 51 auto sub = Eigen::internal::scalar_difference_op<T>(); 52 auto add = Eigen::internal::scalar_sum_op<T>(); 53 auto exp = Eigen::internal::scalar_exp_op<T>(); 54 auto log1p = Eigen::internal::scalar_log1p_op<T>(); 55 auto cmp_lt = 56 Eigen::internal::scalar_cmp_op<T, T, Eigen::internal::cmp_LT>(); 57 58 auto logsumexp = add(log1p(exp(sub(mi, ma))), ma); 59 return cmp_lt(ma, Eigen::NumTraits<T>::lowest()) ? ma : logsumexp; 60 } packetOpLogSumExp61 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T packetOp(const T& a, 62 const T& b) const { 63 auto mi = Eigen::internal::pmin(a, b); 64 auto ma = Eigen::internal::pmax(a, b); 65 using Eigen::internal::padd; 66 using Eigen::internal::pcmp_lt; 67 using Eigen::internal::pexp; 68 using Eigen::internal::plog1p; 69 using Eigen::internal::pset1; 70 using Eigen::internal::psub; 71 72 auto logsumexp = padd(plog1p(pexp(psub(mi, ma))), ma); 73 return pselect(pcmp_lt(ma, pset1(Eigen::NumTraits<T>::lowest())), ma, 74 logsumexp); 75 } 76 }; 77 78 template <typename T> 79 struct LogSumExpReducer { reduceLogSumExpReducer80 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) const { 81 LogSumExp<T> logsumexp; 82 *accum = logsumexp(*accum, t); 83 } 84 85 template <typename Packet> reducePacketLogSumExpReducer86 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reducePacket(const Packet& p, 87 Packet* accum) const { 88 LogSumExp<T> logsumexp; 89 *accum = logsumexp.packetOp(*accum, p); 90 } 91 initializeLogSumExpReducer92 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T initialize() const { 93 return -Eigen::NumTraits<T>::infinity(); 94 } 95 96 template <typename Packet> initializePacketLogSumExpReducer97 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet initializePacket() const { 98 return Eigen::internal::pset1(initialize()); 99 } 100 finalizeLogSumExpReducer101 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalize(const T accum) const { 102 return accum; 103 } 104 105 template <typename Packet> 106 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet finalizePacketLogSumExpReducer107 finalizePacket(const Packet& vaccum) const { 108 return vaccum; 109 } 110 111 template <typename Packet> 112 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T finalizeBothLogSumExpReducer113 finalizeBoth(const T saccum, const Packet& vaccum) const { 114 auto max_reducer = Eigen::internal::MaxReducer<T, Eigen::PropagateNaN>(); 115 auto sum_reducer = Eigen::internal::SumReducer<T>(); 116 auto exp = Eigen::internal::scalar_exp_op<T>(); 117 auto cmp_lt = 118 Eigen::internal::scalar_cmp_op<T, T, Eigen::internal::cmp_LT>(); 119 auto log = Eigen::internal::scalar_log_op<T>(); 120 auto add = Eigen::internal::scalar_sum_op<T>(); 121 122 using Eigen::internal::pexp; 123 using Eigen::internal::psub; 124 125 // `ma = max(x1, ..., xn)` 126 // If the max of all of the `xi` is `-infinity` then the result is 127 // -infinity. If the max is larger than `-infinity` then it's safe to use 128 // for normalization even if the other elements are `-infinity`. 129 // 130 // `logsumexp(x1, ..., xn) = ma + log (exp(x1 - ma) + ... + exp(xn - ma))` 131 auto ma = max_reducer.finalizeBoth(saccum, vaccum); 132 auto logsumexp = add(log(sum_reducer.finalizeBoth( 133 exp(saccum - ma), pexp(psub(vaccum, pset1(ma))))), 134 ma); 135 return cmp_lt(ma, Eigen::NumTraits<T>::lowest()) ? initialize() : logsumexp; 136 } 137 }; 138 139 } // namespace functor 140 } // namespace tensorflow 141 142 #endif // TENSORFLOW_CORE_KERNELS_SCAN_OPS_H_ 143