1 // This file is part of Eigen, a lightweight C++ template library 2 // for linear algebra. 3 // 4 // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog (at) gmail.com> 5 // 6 // This Source Code Form is subject to the terms of the Mozilla 7 // Public License v. 2.0. If a copy of the MPL was not distributed 8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 10 #ifndef EIGEN_CXX11_TENSOR_TENSOR_ASSIGN_H 11 #define EIGEN_CXX11_TENSOR_TENSOR_ASSIGN_H 12 13 namespace Eigen { 14 15 /** \class TensorAssign 16 * \ingroup CXX11_Tensor_Module 17 * 18 * \brief The tensor assignment class. 19 * 20 * This class is represents the assignment of the values resulting from the evaluation of 21 * the rhs expression to the memory locations denoted by the lhs expression. 22 */ 23 namespace internal { 24 template<typename LhsXprType, typename RhsXprType> 25 struct traits<TensorAssignOp<LhsXprType, RhsXprType> > 26 { 27 typedef typename LhsXprType::Scalar Scalar; 28 typedef typename traits<LhsXprType>::StorageKind StorageKind; 29 typedef typename promote_index_type<typename traits<LhsXprType>::Index, 30 typename traits<RhsXprType>::Index>::type Index; 31 typedef typename LhsXprType::Nested LhsNested; 32 typedef typename RhsXprType::Nested RhsNested; 33 typedef typename remove_reference<LhsNested>::type _LhsNested; 34 typedef typename remove_reference<RhsNested>::type _RhsNested; 35 static const std::size_t NumDimensions = internal::traits<LhsXprType>::NumDimensions; 36 static const int Layout = internal::traits<LhsXprType>::Layout; 37 38 enum { 39 Flags = 0 40 }; 41 }; 42 43 template<typename LhsXprType, typename RhsXprType> 44 struct eval<TensorAssignOp<LhsXprType, RhsXprType>, Eigen::Dense> 45 { 46 typedef const TensorAssignOp<LhsXprType, RhsXprType>& type; 47 }; 48 49 template<typename LhsXprType, typename RhsXprType> 50 struct nested<TensorAssignOp<LhsXprType, RhsXprType>, 1, typename eval<TensorAssignOp<LhsXprType, RhsXprType> >::type> 51 { 52 typedef TensorAssignOp<LhsXprType, RhsXprType> type; 53 }; 54 55 } // end namespace internal 56 57 58 59 template<typename LhsXprType, typename RhsXprType> 60 class TensorAssignOp : public TensorBase<TensorAssignOp<LhsXprType, RhsXprType> > 61 { 62 public: 63 typedef typename Eigen::internal::traits<TensorAssignOp>::Scalar Scalar; 64 typedef typename Eigen::NumTraits<Scalar>::Real RealScalar; 65 typedef typename LhsXprType::CoeffReturnType CoeffReturnType; 66 typedef typename Eigen::internal::nested<TensorAssignOp>::type Nested; 67 typedef typename Eigen::internal::traits<TensorAssignOp>::StorageKind StorageKind; 68 typedef typename Eigen::internal::traits<TensorAssignOp>::Index Index; 69 70 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorAssignOp(LhsXprType& lhs, const RhsXprType& rhs) 71 : m_lhs_xpr(lhs), m_rhs_xpr(rhs) {} 72 73 /** \returns the nested expressions */ 74 EIGEN_DEVICE_FUNC 75 typename internal::remove_all<typename LhsXprType::Nested>::type& 76 lhsExpression() const { return *((typename internal::remove_all<typename LhsXprType::Nested>::type*)&m_lhs_xpr); } 77 78 EIGEN_DEVICE_FUNC 79 const typename internal::remove_all<typename RhsXprType::Nested>::type& 80 rhsExpression() const { return m_rhs_xpr; } 81 82 protected: 83 typename internal::remove_all<typename LhsXprType::Nested>::type& m_lhs_xpr; 84 const typename internal::remove_all<typename RhsXprType::Nested>::type& m_rhs_xpr; 85 }; 86 87 88 template<typename LeftArgType, typename RightArgType, typename Device> 89 struct TensorEvaluator<const TensorAssignOp<LeftArgType, RightArgType>, Device> 90 { 91 typedef TensorAssignOp<LeftArgType, RightArgType> XprType; 92 typedef typename XprType::Index Index; 93 typedef typename XprType::Scalar Scalar; 94 typedef typename XprType::CoeffReturnType CoeffReturnType; 95 typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType; 96 typedef typename TensorEvaluator<RightArgType, Device>::Dimensions Dimensions; 97 static const int PacketSize = internal::unpacket_traits<PacketReturnType>::size; 98 99 enum { 100 IsAligned = TensorEvaluator<LeftArgType, Device>::IsAligned & TensorEvaluator<RightArgType, Device>::IsAligned, 101 PacketAccess = TensorEvaluator<LeftArgType, Device>::PacketAccess & TensorEvaluator<RightArgType, Device>::PacketAccess, 102 Layout = TensorEvaluator<LeftArgType, Device>::Layout, 103 RawAccess = TensorEvaluator<LeftArgType, Device>::RawAccess 104 }; 105 106 EIGEN_DEVICE_FUNC TensorEvaluator(const XprType& op, const Device& device) : 107 m_leftImpl(op.lhsExpression(), device), 108 m_rightImpl(op.rhsExpression(), device) 109 { 110 EIGEN_STATIC_ASSERT((static_cast<int>(TensorEvaluator<LeftArgType, Device>::Layout) == static_cast<int>(TensorEvaluator<RightArgType, Device>::Layout)), YOU_MADE_A_PROGRAMMING_MISTAKE); 111 } 112 113 EIGEN_DEVICE_FUNC const Dimensions& dimensions() const 114 { 115 // The dimensions of the lhs and the rhs tensors should be equal to prevent 116 // overflows and ensure the result is fully initialized. 117 // TODO: use left impl instead if right impl dimensions are known at compile time. 118 return m_rightImpl.dimensions(); 119 } 120 121 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(Scalar*) { 122 eigen_assert(dimensions_match(m_leftImpl.dimensions(), m_rightImpl.dimensions())); 123 m_leftImpl.evalSubExprsIfNeeded(NULL); 124 // If the lhs provides raw access to its storage area (i.e. if m_leftImpl.data() returns a non 125 // null value), attempt to evaluate the rhs expression in place. Returns true iff in place 126 // evaluation isn't supported and the caller still needs to manually assign the values generated 127 // by the rhs to the lhs. 128 return m_rightImpl.evalSubExprsIfNeeded(m_leftImpl.data()); 129 } 130 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() { 131 m_leftImpl.cleanup(); 132 m_rightImpl.cleanup(); 133 } 134 135 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalScalar(Index i) { 136 m_leftImpl.coeffRef(i) = m_rightImpl.coeff(i); 137 } 138 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalPacket(Index i) { 139 const int LhsStoreMode = TensorEvaluator<LeftArgType, Device>::IsAligned ? Aligned : Unaligned; 140 const int RhsLoadMode = TensorEvaluator<RightArgType, Device>::IsAligned ? Aligned : Unaligned; 141 m_leftImpl.template writePacket<LhsStoreMode>(i, m_rightImpl.template packet<RhsLoadMode>(i)); 142 } 143 EIGEN_DEVICE_FUNC CoeffReturnType coeff(Index index) const 144 { 145 return m_leftImpl.coeff(index); 146 } 147 template<int LoadMode> 148 EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const 149 { 150 return m_leftImpl.template packet<LoadMode>(index); 151 } 152 153 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost 154 costPerCoeff(bool vectorized) const { 155 // We assume that evalPacket or evalScalar is called to perform the 156 // assignment and account for the cost of the write here, but reduce left 157 // cost by one load because we are using m_leftImpl.coeffRef. 158 TensorOpCost left = m_leftImpl.costPerCoeff(vectorized); 159 return m_rightImpl.costPerCoeff(vectorized) + 160 TensorOpCost( 161 numext::maxi(0.0, left.bytes_loaded() - sizeof(CoeffReturnType)), 162 left.bytes_stored(), left.compute_cycles()) + 163 TensorOpCost(0, sizeof(CoeffReturnType), 0, vectorized, PacketSize); 164 } 165 166 /// required by sycl in order to extract the accessor 167 const TensorEvaluator<LeftArgType, Device>& left_impl() const { return m_leftImpl; } 168 /// required by sycl in order to extract the accessor 169 const TensorEvaluator<RightArgType, Device>& right_impl() const { return m_rightImpl; } 170 171 EIGEN_DEVICE_FUNC CoeffReturnType* data() const { return m_leftImpl.data(); } 172 173 private: 174 TensorEvaluator<LeftArgType, Device> m_leftImpl; 175 TensorEvaluator<RightArgType, Device> m_rightImpl; 176 }; 177 178 } 179 180 181 #endif // EIGEN_CXX11_TENSOR_TENSOR_ASSIGN_H 182
