src/Tensor/TensorChipping.h

0001 // This file is part of Eigen, a lightweight C++ template library
0002 // for linear algebra.
0003 //
0004 // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
0005 //
0006 // This Source Code Form is subject to the terms of the Mozilla
0007 // Public License v. 2.0. If a copy of the MPL was not distributed
0008 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
0009
0010 #ifndef EIGEN_CXX11_TENSOR_TENSOR_CHIPPING_H
0011 #define EIGEN_CXX11_TENSOR_TENSOR_CHIPPING_H
0012
0013 namespace Eigen {
0014
0015 /** \class TensorKChippingReshaping
0016   * \ingroup CXX11_Tensor_Module
0017   *
0018   * \brief A chip is a thin slice, corresponding to a column or a row in a 2-d tensor.
0019   *
0020   *
0021   */
0022
0023 namespace internal {
0024 template<DenseIndex DimId, typename XprType>
0025 struct traits<TensorChippingOp<DimId, XprType> > : public traits<XprType>
0026 {
0027   typedef typename XprType::Scalar Scalar;
0028   typedef traits<XprType> XprTraits;
0029   typedef typename XprTraits::StorageKind StorageKind;
0030   typedef typename XprTraits::Index Index;
0031   typedef typename XprType::Nested Nested;
0032   typedef typename remove_reference<Nested>::type _Nested;
0033   static const int NumDimensions = XprTraits::NumDimensions - 1;
0034   static const int Layout = XprTraits::Layout;
0035   typedef typename XprTraits::PointerType PointerType;
0036 };
0037
0038 template<DenseIndex DimId, typename XprType>
0039 struct eval<TensorChippingOp<DimId, XprType>, Eigen::Dense>
0040 {
0041   typedef const TensorChippingOp<DimId, XprType> EIGEN_DEVICE_REF type;
0042 };
0043
0044 template<DenseIndex DimId, typename XprType>
0045 struct nested<TensorChippingOp<DimId, XprType>, 1, typename eval<TensorChippingOp<DimId, XprType> >::type>
0046 {
0047   typedef TensorChippingOp<DimId, XprType> type;
0048 };
0049
0050 template <DenseIndex DimId>
0051 struct DimensionId
0052 {
0053   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DimensionId(DenseIndex dim) {
0054     EIGEN_UNUSED_VARIABLE(dim);
0055     eigen_assert(dim == DimId);
0056   }
0057   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DenseIndex actualDim() const {
0058     return DimId;
0059   }
0060 };
0061 template <>
0062 struct DimensionId<Dynamic>
0063 {
0064   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DimensionId(DenseIndex dim) : actual_dim(dim) {
0065     eigen_assert(dim >= 0);
0066   }
0067   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DenseIndex actualDim() const {
0068     return actual_dim;
0069   }
0070  private:
0071   const DenseIndex actual_dim;
0072 };
0073
0074
0075 }  // end namespace internal
0076
0077
0078
0079 template<DenseIndex DimId, typename XprType>
0080 class TensorChippingOp : public TensorBase<TensorChippingOp<DimId, XprType> >
0081 {
0082   public:
0083     typedef TensorBase<TensorChippingOp<DimId, XprType> > Base;
0084     typedef typename Eigen::internal::traits<TensorChippingOp>::Scalar Scalar;
0085     typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
0086     typedef typename XprType::CoeffReturnType CoeffReturnType;
0087     typedef typename Eigen::internal::nested<TensorChippingOp>::type Nested;
0088     typedef typename Eigen::internal::traits<TensorChippingOp>::StorageKind StorageKind;
0089     typedef typename Eigen::internal::traits<TensorChippingOp>::Index Index;
0090
0091     EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorChippingOp(const XprType& expr, const Index offset, const Index dim)
0092         : m_xpr(expr), m_offset(offset), m_dim(dim) {
0093     }
0094
0095     EIGEN_DEVICE_FUNC
0096     const Index offset() const { return m_offset; }
0097     EIGEN_DEVICE_FUNC
0098     const Index dim() const { return m_dim.actualDim(); }
0099
0100     EIGEN_DEVICE_FUNC
0101     const typename internal::remove_all<typename XprType::Nested>::type&
0102     expression() const { return m_xpr; }
0103
0104     EIGEN_TENSOR_INHERIT_ASSIGNMENT_OPERATORS(TensorChippingOp)
0105
0106   protected:
0107     typename XprType::Nested m_xpr;
0108     const Index m_offset;
0109     const internal::DimensionId<DimId> m_dim;
0110 };
0111
0112
0113 // Eval as rvalue
0114 template<DenseIndex DimId, typename ArgType, typename Device>
0115 struct TensorEvaluator<const TensorChippingOp<DimId, ArgType>, Device>
0116 {
0117   typedef TensorChippingOp<DimId, ArgType> XprType;
0118   static const int NumInputDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
0119   static const int NumDims = NumInputDims-1;
0120   typedef typename XprType::Index Index;
0121   typedef DSizes<Index, NumDims> Dimensions;
0122   typedef typename XprType::Scalar Scalar;
0123   typedef typename XprType::CoeffReturnType CoeffReturnType;
0124   typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
0125   static const int PacketSize = PacketType<CoeffReturnType, Device>::size;
0126   typedef StorageMemory<CoeffReturnType, Device> Storage;
0127   typedef typename Storage::Type EvaluatorPointerType;
0128
0129   enum {
0130     // Alignment can't be guaranteed at compile time since it depends on the
0131     // slice offsets.
0132     IsAligned         = false,
0133     Layout            = TensorEvaluator<ArgType, Device>::Layout,
0134     PacketAccess      = TensorEvaluator<ArgType, Device>::PacketAccess,
0135     BlockAccess       = TensorEvaluator<ArgType, Device>::BlockAccess,
0136     // Chipping of outer-most dimension is a trivial operation, because we can
0137     // read and write directly from the underlying tensor using single offset.
0138     IsOuterChipping   = (static_cast<int>(Layout) == ColMajor && DimId == NumInputDims - 1) ||
0139                         (static_cast<int>(Layout) == RowMajor && DimId == 0),
0140     // Chipping inner-most dimension.
0141     IsInnerChipping   = (static_cast<int>(Layout) == ColMajor && DimId == 0) ||
0142                         (static_cast<int>(Layout) == RowMajor && DimId == NumInputDims - 1),
0143     // Prefer block access if the underlying expression prefers it, otherwise
0144     // only if chipping is not trivial.
0145     PreferBlockAccess = TensorEvaluator<ArgType, Device>::PreferBlockAccess ||
0146                         !IsOuterChipping,
0147     CoordAccess       = false,  // to be implemented
0148     RawAccess         = false
0149   };
0150
0151   typedef typename internal::remove_const<Scalar>::type ScalarNoConst;
0152
0153   //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
0154   typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
0155   typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
0156
0157   typedef internal::TensorBlockDescriptor<NumInputDims, Index>
0158       ArgTensorBlockDesc;
0159   typedef typename TensorEvaluator<const ArgType, Device>::TensorBlock
0160       ArgTensorBlock;
0161
0162   typedef typename internal::TensorMaterializedBlock<ScalarNoConst, NumDims,
0163                                                      Layout, Index>
0164       TensorBlock;
0165   //===--------------------------------------------------------------------===//
0166
0167   EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
0168       : m_impl(op.expression(), device), m_dim(op.dim()), m_device(device)
0169   {
0170     EIGEN_STATIC_ASSERT((NumInputDims >= 1), YOU_MADE_A_PROGRAMMING_MISTAKE);
0171     eigen_assert(NumInputDims > m_dim.actualDim());
0172
0173     const typename TensorEvaluator<ArgType, Device>::Dimensions& input_dims = m_impl.dimensions();
0174     eigen_assert(op.offset() < input_dims[m_dim.actualDim()]);
0175
0176     int j = 0;
0177     for (int i = 0; i < NumInputDims; ++i) {
0178       if (i != m_dim.actualDim()) {
0179         m_dimensions[j] = input_dims[i];
0180         ++j;
0181       }
0182     }
0183
0184     m_stride = 1;
0185     m_inputStride = 1;
0186     if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
0187       for (int i = 0; i < m_dim.actualDim(); ++i) {
0188         m_stride *= input_dims[i];
0189         m_inputStride *= input_dims[i];
0190       }
0191     } else {
0192       for (int i = NumInputDims-1; i > m_dim.actualDim(); --i) {
0193         m_stride *= input_dims[i];
0194         m_inputStride *= input_dims[i];
0195       }
0196     }
0197     m_inputStride *= input_dims[m_dim.actualDim()];
0198     m_inputOffset = m_stride * op.offset();
0199   }
0200
0201   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
0202
0203   EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType) {
0204     m_impl.evalSubExprsIfNeeded(NULL);
0205     return true;
0206   }
0207
0208   EIGEN_STRONG_INLINE void cleanup() {
0209     m_impl.cleanup();
0210   }
0211
0212   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const
0213   {
0214     return m_impl.coeff(srcCoeff(index));
0215   }
0216
0217   template<int LoadMode>
0218   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const
0219   {
0220     EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
0221     eigen_assert(index+PacketSize-1 < dimensions().TotalSize());
0222
0223     if (isInnerChipping()) {
0224       // m_stride is equal to 1, so let's avoid the integer division.
0225       eigen_assert(m_stride == 1);
0226       Index inputIndex = index * m_inputStride + m_inputOffset;
0227       EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
0228       EIGEN_UNROLL_LOOP
0229       for (int i = 0; i < PacketSize; ++i) {
0230         values[i] = m_impl.coeff(inputIndex);
0231         inputIndex += m_inputStride;
0232       }
0233       PacketReturnType rslt = internal::pload<PacketReturnType>(values);
0234       return rslt;
0235     } else if (isOuterChipping()) {
0236       // m_stride is always greater than index, so let's avoid the integer division.
0237       eigen_assert(m_stride > index);
0238       return m_impl.template packet<LoadMode>(index + m_inputOffset);
0239     } else {
0240       const Index idx = index / m_stride;
0241       const Index rem = index - idx * m_stride;
0242       if (rem + PacketSize <= m_stride) {
0243         Index inputIndex = idx * m_inputStride + m_inputOffset + rem;
0244         return m_impl.template packet<LoadMode>(inputIndex);
0245       } else {
0246         // Cross the stride boundary. Fallback to slow path.
0247         EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
0248        EIGEN_UNROLL_LOOP
0249         for (int i = 0; i < PacketSize; ++i) {
0250           values[i] = coeff(index);
0251           ++index;
0252         }
0253         PacketReturnType rslt = internal::pload<PacketReturnType>(values);
0254         return rslt;
0255       }
0256     }
0257   }
0258
0259   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost
0260   costPerCoeff(bool vectorized) const {
0261     double cost = 0;
0262     if ((static_cast<int>(Layout) == static_cast<int>(ColMajor) &&
0263          m_dim.actualDim() == 0) ||
0264         (static_cast<int>(Layout) == static_cast<int>(RowMajor) &&
0265          m_dim.actualDim() == NumInputDims - 1)) {
0266       cost += TensorOpCost::MulCost<Index>() + TensorOpCost::AddCost<Index>();
0267     } else if ((static_cast<int>(Layout) == static_cast<int>(ColMajor) &&
0268                 m_dim.actualDim() == NumInputDims - 1) ||
0269                (static_cast<int>(Layout) == static_cast<int>(RowMajor) &&
0270                 m_dim.actualDim() == 0)) {
0271       cost += TensorOpCost::AddCost<Index>();
0272     } else {
0273       cost += 3 * TensorOpCost::MulCost<Index>() + TensorOpCost::DivCost<Index>() +
0274               3 * TensorOpCost::AddCost<Index>();
0275     }
0276
0277     return m_impl.costPerCoeff(vectorized) +
0278            TensorOpCost(0, 0, cost, vectorized, PacketSize);
0279   }
0280
0281   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
0282   internal::TensorBlockResourceRequirements getResourceRequirements() const {
0283     const size_t target_size = m_device.lastLevelCacheSize();
0284     return internal::TensorBlockResourceRequirements::merge(
0285         internal::TensorBlockResourceRequirements::skewed<Scalar>(target_size),
0286         m_impl.getResourceRequirements());
0287   }
0288
0289   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock
0290   block(TensorBlockDesc& desc, TensorBlockScratch& scratch,
0291           bool root_of_expr_ast = false) const {
0292     const Index chip_dim = m_dim.actualDim();
0293
0294     DSizes<Index, NumInputDims> input_block_dims;
0295     for (int i = 0; i < NumInputDims; ++i) {
0296       input_block_dims[i]
0297             = i < chip_dim ? desc.dimension(i)
0298             : i > chip_dim ? desc.dimension(i - 1)
0299             : 1;
0300     }
0301
0302     ArgTensorBlockDesc arg_desc(srcCoeff(desc.offset()), input_block_dims);
0303
0304     // Try to reuse destination buffer for materializing argument block.
0305     if (desc.HasDestinationBuffer()) {
0306       DSizes<Index, NumInputDims> arg_destination_strides;
0307       for (int i = 0; i < NumInputDims; ++i) {
0308       arg_destination_strides[i]
0309             = i < chip_dim ? desc.destination().strides()[i]
0310             : i > chip_dim ? desc.destination().strides()[i - 1]
0311             : 0; // for dimensions of size `1` stride should never be used.
0312       }
0313
0314       arg_desc.template AddDestinationBuffer<Layout>(
0315           desc.destination().template data<ScalarNoConst>(),
0316           arg_destination_strides);
0317     }
0318
0319     ArgTensorBlock arg_block = m_impl.block(arg_desc, scratch, root_of_expr_ast);
0320     if (!arg_desc.HasDestinationBuffer()) desc.DropDestinationBuffer();
0321
0322     if (arg_block.data() != NULL) {
0323       // Forward argument block buffer if possible.
0324       return TensorBlock(arg_block.kind(), arg_block.data(),
0325                            desc.dimensions());
0326
0327     } else {
0328       // Assign argument block expression to a buffer.
0329
0330       // Prepare storage for the materialized chipping result.
0331       const typename TensorBlock::Storage block_storage =
0332           TensorBlock::prepareStorage(desc, scratch);
0333
0334       typedef internal::TensorBlockAssignment<
0335           ScalarNoConst, NumInputDims, typename ArgTensorBlock::XprType, Index>
0336           TensorBlockAssignment;
0337
0338       TensorBlockAssignment::Run(
0339           TensorBlockAssignment::target(
0340               arg_desc.dimensions(),
0341               internal::strides<Layout>(arg_desc.dimensions()),
0342               block_storage.data()),
0343           arg_block.expr());
0344
0345       return block_storage.AsTensorMaterializedBlock();
0346     }
0347   }
0348
0349   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename Storage::Type data() const {
0350     typename Storage::Type result = constCast(m_impl.data());
0351     if (isOuterChipping() && result) {
0352       return result + m_inputOffset;
0353     } else {
0354       return NULL;
0355     }
0356   }
0357 #ifdef EIGEN_USE_SYCL
0358   // binding placeholder accessors to a command group handler for SYCL
0359   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void bind(cl::sycl::handler &cgh) const {
0360     m_impl.bind(cgh);
0361   }
0362 #endif
0363
0364  protected:
0365   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index srcCoeff(Index index) const
0366   {
0367     Index inputIndex;
0368     if (isInnerChipping()) {
0369       // m_stride is equal to 1, so let's avoid the integer division.
0370       eigen_assert(m_stride == 1);
0371       inputIndex = index * m_inputStride + m_inputOffset;
0372     } else if (isOuterChipping()) {
0373       // m_stride is always greater than index, so let's avoid the integer
0374       // division.
0375       eigen_assert(m_stride > index);
0376       inputIndex = index + m_inputOffset;
0377     } else {
0378       const Index idx = index / m_stride;
0379       inputIndex = idx * m_inputStride + m_inputOffset;
0380       index -= idx * m_stride;
0381       inputIndex += index;
0382     }
0383     return inputIndex;
0384   }
0385
0386   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool isInnerChipping() const {
0387     return IsInnerChipping ||
0388            (static_cast<int>(Layout) == ColMajor && m_dim.actualDim() == 0) ||
0389            (static_cast<int>(Layout) == RowMajor && m_dim.actualDim() == NumInputDims - 1);
0390   }
0391
0392   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool isOuterChipping() const {
0393     return IsOuterChipping ||
0394            (static_cast<int>(Layout) == ColMajor && m_dim.actualDim() == NumInputDims-1) ||
0395            (static_cast<int>(Layout) == RowMajor && m_dim.actualDim() == 0);
0396   }
0397
0398   Dimensions m_dimensions;
0399   Index m_stride;
0400   Index m_inputOffset;
0401   Index m_inputStride;
0402   TensorEvaluator<ArgType, Device> m_impl;
0403   const internal::DimensionId<DimId> m_dim;
0404   const Device EIGEN_DEVICE_REF m_device;
0405 };
0406
0407
0408 // Eval as lvalue
0409 template<DenseIndex DimId, typename ArgType, typename Device>
0410 struct TensorEvaluator<TensorChippingOp<DimId, ArgType>, Device>
0411   : public TensorEvaluator<const TensorChippingOp<DimId, ArgType>, Device>
0412 {
0413   typedef TensorEvaluator<const TensorChippingOp<DimId, ArgType>, Device> Base;
0414   typedef TensorChippingOp<DimId, ArgType> XprType;
0415   static const int NumInputDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
0416   static const int NumDims = NumInputDims-1;
0417   typedef typename XprType::Index Index;
0418   typedef DSizes<Index, NumDims> Dimensions;
0419   typedef typename XprType::Scalar Scalar;
0420   typedef typename XprType::CoeffReturnType CoeffReturnType;
0421   typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
0422   static const int PacketSize = PacketType<CoeffReturnType, Device>::size;
0423
0424   enum {
0425     IsAligned     = false,
0426     PacketAccess  = TensorEvaluator<ArgType, Device>::PacketAccess,
0427     BlockAccess   = TensorEvaluator<ArgType, Device>::RawAccess,
0428     Layout        = TensorEvaluator<ArgType, Device>::Layout,
0429     RawAccess     = false
0430   };
0431
0432   //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
0433   typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
0434   //===--------------------------------------------------------------------===//
0435
0436   EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
0437     : Base(op, device)
0438     { }
0439
0440   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType& coeffRef(Index index)
0441   {
0442     return this->m_impl.coeffRef(this->srcCoeff(index));
0443   }
0444
0445   template <int StoreMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
0446   void writePacket(Index index, const PacketReturnType& x)
0447   {
0448     EIGEN_STATIC_ASSERT((PacketSize > 1), YOU_MADE_A_PROGRAMMING_MISTAKE)
0449
0450     if (this->isInnerChipping()) {
0451       // m_stride is equal to 1, so let's avoid the integer division.
0452       eigen_assert(this->m_stride == 1);
0453       EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
0454       internal::pstore<CoeffReturnType, PacketReturnType>(values, x);
0455       Index inputIndex = index * this->m_inputStride + this->m_inputOffset;
0456       EIGEN_UNROLL_LOOP
0457       for (int i = 0; i < PacketSize; ++i) {
0458         this->m_impl.coeffRef(inputIndex) = values[i];
0459         inputIndex += this->m_inputStride;
0460       }
0461     } else if (this->isOuterChipping()) {
0462       // m_stride is always greater than index, so let's avoid the integer division.
0463       eigen_assert(this->m_stride > index);
0464       this->m_impl.template writePacket<StoreMode>(index + this->m_inputOffset, x);
0465     } else {
0466       const Index idx = index / this->m_stride;
0467       const Index rem = index - idx * this->m_stride;
0468       if (rem + PacketSize <= this->m_stride) {
0469         const Index inputIndex = idx * this->m_inputStride + this->m_inputOffset + rem;
0470         this->m_impl.template writePacket<StoreMode>(inputIndex, x);
0471       } else {
0472         // Cross stride boundary. Fallback to slow path.
0473         EIGEN_ALIGN_MAX typename internal::remove_const<CoeffReturnType>::type values[PacketSize];
0474         internal::pstore<CoeffReturnType, PacketReturnType>(values, x);
0475         EIGEN_UNROLL_LOOP
0476         for (int i = 0; i < PacketSize; ++i) {
0477           this->coeffRef(index) = values[i];
0478           ++index;
0479         }
0480       }
0481     }
0482   }
0483
0484   template <typename TensorBlock>
0485   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writeBlock(
0486       const TensorBlockDesc& desc, const TensorBlock& block) {
0487     assert(this->m_impl.data() != NULL);
0488
0489     const Index chip_dim = this->m_dim.actualDim();
0490
0491     DSizes<Index, NumInputDims> input_block_dims;
0492     for (int i = 0; i < NumInputDims; ++i) {
0493       input_block_dims[i] = i < chip_dim ? desc.dimension(i)
0494                           : i > chip_dim ? desc.dimension(i - 1)
0495                           : 1;
0496     }
0497
0498     typedef TensorReshapingOp<const DSizes<Index, NumInputDims>,
0499                               const typename TensorBlock::XprType>
0500         TensorBlockExpr;
0501
0502     typedef internal::TensorBlockAssignment<Scalar, NumInputDims,
0503                                             TensorBlockExpr, Index>
0504         TensorBlockAssign;
0505
0506     TensorBlockAssign::Run(
0507         TensorBlockAssign::target(
0508             input_block_dims,
0509             internal::strides<Layout>(this->m_impl.dimensions()),
0510             this->m_impl.data(), this->srcCoeff(desc.offset())),
0511         block.expr().reshape(input_block_dims));
0512   }
0513 };
0514
0515
0516 } // end namespace Eigen
0517
0518 #endif // EIGEN_CXX11_TENSOR_TENSOR_CHIPPING_H