TensorShuffling.h
00001 // This file is part of Eigen, a lightweight C++ template library
00002 // for linear algebra.
00003 //
00004 // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
00005 //
00006 // This Source Code Form is subject to the terms of the Mozilla
00007 // Public License v. 2.0. If a copy of the MPL was not distributed
00008 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
00009 
00010 #ifndef EIGEN_CXX11_TENSOR_TENSOR_SHUFFLING_H
00011 #define EIGEN_CXX11_TENSOR_TENSOR_SHUFFLING_H
00012 
00013 namespace Eigen {
00014 
00022 namespace internal {
00023 template<typename Shuffle, typename XprType>
00024 struct traits<TensorShufflingOp<Shuffle, XprType> > : public traits<XprType>
00025 {
00026   typedef typename XprType::Scalar Scalar;
00027   typedef traits<XprType> XprTraits;
00028   typedef typename XprTraits::StorageKind StorageKind;
00029   typedef typename XprTraits::Index Index;
00030   typedef typename XprType::Nested Nested;
00031   typedef typename remove_reference<Nested>::type _Nested;
00032   static const int NumDimensions = XprTraits::NumDimensions;
00033   static const int Layout = XprTraits::Layout;
00034 };
00035 
00036 template<typename Shuffle, typename XprType>
00037 struct eval<TensorShufflingOp<Shuffle, XprType>, Eigen::Dense>
00038 {
00039   typedef const TensorShufflingOp<Shuffle, XprType>& type;
00040 };
00041 
00042 template<typename Shuffle, typename XprType>
00043 struct nested<TensorShufflingOp<Shuffle, XprType>, 1, typename eval<TensorShufflingOp<Shuffle, XprType> >::type>
00044 {
00045   typedef TensorShufflingOp<Shuffle, XprType> type;
00046 };
00047 
00048 }  // end namespace internal
00049 
00050 
00051 
00052 template<typename Shuffle, typename XprType>
00053 class TensorShufflingOp : public TensorBase<TensorShufflingOp<Shuffle, XprType> >
00054 {
00055   public:
00056   typedef typename Eigen::internal::traits<TensorShufflingOp>::Scalar Scalar;
00057   typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
00058   typedef typename XprType::CoeffReturnType CoeffReturnType;
00059   typedef typename Eigen::internal::nested<TensorShufflingOp>::type Nested;
00060   typedef typename Eigen::internal::traits<TensorShufflingOp>::StorageKind StorageKind;
00061   typedef typename Eigen::internal::traits<TensorShufflingOp>::Index Index;
00062 
00063   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorShufflingOp(const XprType& expr, const Shuffle& shuffle)
00064       : m_xpr(expr), m_shuffle(shuffle) {}
00065 
00066     EIGEN_DEVICE_FUNC
00067     const Shuffle& shufflePermutation() const { return m_shuffle; }
00068 
00069     EIGEN_DEVICE_FUNC
00070     const typename internal::remove_all<typename XprType::Nested>::type&
00071     expression() const { return m_xpr; }
00072 
00073     EIGEN_DEVICE_FUNC
00074     EIGEN_STRONG_INLINE TensorShufflingOp& operator = (const TensorShufflingOp& other)
00075     {
00076       typedef TensorAssignOp<TensorShufflingOp, const TensorShufflingOp> Assign;
00077       Assign assign(*this, other);
00078       internal::TensorExecutor<const Assign, DefaultDevice>::run(assign, DefaultDevice());
00079       return *this;
00080     }
00081 
00082     template<typename OtherDerived>
00083     EIGEN_DEVICE_FUNC
00084     EIGEN_STRONG_INLINE TensorShufflingOp& operator = (const OtherDerived& other)
00085     {
00086       typedef TensorAssignOp<TensorShufflingOp, const OtherDerived> Assign;
00087       Assign assign(*this, other);
00088       internal::TensorExecutor<const Assign, DefaultDevice>::run(assign, DefaultDevice());
00089       return *this;
00090     }
00091 
00092   protected:
00093     typename XprType::Nested m_xpr;
00094     const Shuffle m_shuffle;
00095 };
00096 
00097 
00098 // Eval as rvalue
00099 template<typename Shuffle, typename ArgType, typename Device>
00100 struct TensorEvaluator<const TensorShufflingOp<Shuffle, ArgType>, Device>
00101 {
00102   typedef TensorShufflingOp<Shuffle, ArgType> XprType;
00103   typedef typename XprType::Index Index;
00104   static const int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
00105   typedef DSizes<Index, NumDims> Dimensions;
00106   typedef typename XprType::Scalar Scalar;
00107   typedef typename XprType::CoeffReturnType CoeffReturnType;
00108   typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
00109   static const int PacketSize = internal::unpacket_traits<PacketReturnType>::size;
00110 
00111   enum {
00112     IsAligned = false,
00113     PacketAccess = (internal::packet_traits<Scalar>::size > 1),
00114     Layout = TensorEvaluator<ArgType, Device>::Layout,
00115     CoordAccess = false,  // to be implemented
00116     RawAccess = false
00117   };
00118 
00119   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
00120       : m_impl(op.expression(), device)
00121   {
00122     const typename TensorEvaluator<ArgType, Device>::Dimensions& input_dims = m_impl.dimensions();
00123     const Shuffle& shuffle = op.shufflePermutation();
00124     for (int i = 0; i < NumDims; ++i) {
00125       m_dimensions[i] = input_dims[shuffle[i]];
00126     }
00127 
00128     array<Index, NumDims> inputStrides;
00129 
00130     if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
00131       inputStrides[0] = 1;
00132       m_outputStrides[0] = 1;
00133       for (int i = 1; i < NumDims; ++i) {
00134         inputStrides[i] = inputStrides[i - 1] * input_dims[i - 1];
00135         m_outputStrides[i] = m_outputStrides[i - 1] * m_dimensions[i - 1];
00136       }
00137     } else {
00138       inputStrides[NumDims - 1] = 1;
00139       m_outputStrides[NumDims - 1] = 1;
00140       for (int i = NumDims - 2; i >= 0; --i) {
00141         inputStrides[i] = inputStrides[i + 1] * input_dims[i + 1];
00142         m_outputStrides[i] = m_outputStrides[i + 1] * m_dimensions[i + 1];
00143       }
00144     }
00145 
00146     for (int i = 0; i < NumDims; ++i) {
00147       m_inputStrides[i] = inputStrides[shuffle[i]];
00148     }
00149   }
00150 
00151   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
00152 
00153   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(Scalar* /*data*/) {
00154     m_impl.evalSubExprsIfNeeded(NULL);
00155     return true;
00156   }
00157   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
00158     m_impl.cleanup();
00159   }
00160 
00161   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const
00162   {
00163     return m_impl.coeff(srcCoeff(index));
00164   }
00165 
00166   template<int LoadMode>
00167   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const
00168   {
00169     EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
00170     eigen_assert(index+PacketSize-1 < dimensions().TotalSize());
00171 
00172     EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
00173     for (int i = 0; i < PacketSize; ++i) {
00174       values[i] = coeff(index+i);
00175     }
00176     PacketReturnType rslt = internal::pload<PacketReturnType>(values);
00177     return rslt;
00178   }
00179 
00180   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool vectorized) const {
00181     const double compute_cost = NumDims * (2 * TensorOpCost::AddCost<Index>() +
00182                                            2 * TensorOpCost::MulCost<Index>() +
00183                                            TensorOpCost::DivCost<Index>());
00184     return m_impl.costPerCoeff(vectorized) +
00185            TensorOpCost(0, 0, compute_cost, false /* vectorized */, PacketSize);
00186   }
00187 
00188   EIGEN_DEVICE_FUNC Scalar* data() const { return NULL; }
00189 
00190  protected:
00191   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index srcCoeff(Index index) const {
00192     Index inputIndex = 0;
00193     if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
00194       for (int i = NumDims - 1; i > 0; --i) {
00195         const Index idx = index / m_outputStrides[i];
00196         inputIndex += idx * m_inputStrides[i];
00197         index -= idx * m_outputStrides[i];
00198       }
00199       return inputIndex + index * m_inputStrides[0];
00200     } else {
00201       for (int i = 0; i < NumDims - 1; ++i) {
00202         const Index idx = index / m_outputStrides[i];
00203         inputIndex += idx * m_inputStrides[i];
00204         index -= idx * m_outputStrides[i];
00205       }
00206       return inputIndex + index * m_inputStrides[NumDims - 1];
00207     }
00208   }
00209 
00210   Dimensions m_dimensions;
00211   array<Index, NumDims> m_outputStrides;
00212   array<Index, NumDims> m_inputStrides;
00213   TensorEvaluator<ArgType, Device> m_impl;
00214 };
00215 
00216 
00217 // Eval as lvalue
00218 template<typename Shuffle, typename ArgType, typename Device>
00219 struct TensorEvaluator<TensorShufflingOp<Shuffle, ArgType>, Device>
00220     : public TensorEvaluator<const TensorShufflingOp<Shuffle, ArgType>, Device>
00221 {
00222   typedef TensorEvaluator<const TensorShufflingOp<Shuffle, ArgType>, Device> Base;
00223 
00224   typedef TensorShufflingOp<Shuffle, ArgType> XprType;
00225   typedef typename XprType::Index Index;
00226   static const int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
00227   typedef DSizes<Index, NumDims> Dimensions;
00228   typedef typename XprType::Scalar Scalar;
00229   typedef typename XprType::CoeffReturnType CoeffReturnType;
00230   typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
00231   static const int PacketSize = internal::unpacket_traits<PacketReturnType>::size;
00232 
00233   enum {
00234     IsAligned = false,
00235     PacketAccess = (internal::packet_traits<Scalar>::size > 1),
00236     RawAccess = false
00237   };
00238 
00239   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
00240       : Base(op, device)
00241   { }
00242 
00243   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType& coeffRef(Index index)
00244   {
00245     return this->m_impl.coeffRef(this->srcCoeff(index));
00246   }
00247 
00248   template <int StoreMode> EIGEN_STRONG_INLINE
00249   void writePacket(Index index, const PacketReturnType& x)
00250   {
00251     EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
00252 
00253     EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
00254     internal::pstore<CoeffReturnType, PacketReturnType>(values, x);
00255     for (int i = 0; i < PacketSize; ++i) {
00256       this->coeffRef(index+i) = values[i];
00257     }
00258   }
00259 };
00260 
00261 
00262 } // end namespace Eigen
00263 
00264 #endif // EIGEN_CXX11_TENSOR_TENSOR_SHUFFLING_H
 All Classes Functions Variables Typedefs Enumerator