| // This file is part of Eigen, a lightweight C++ template library | |
| // for linear algebra. | |
| // | |
| // Copyright (C) 2018 Andy Davis <andydavis@google.com> | |
| // Copyright (C) 2018 Eugene Zhulenev <ezhulenev@google.com> | |
| // | |
| // This Source Code Form is subject to the terms of the Mozilla | |
| // Public License v. 2.0. If a copy of the MPL was not distributed | |
| // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. | |
| #ifndef EIGEN_CXX11_TENSOR_TENSOR_BLOCK_H | |
| #define EIGEN_CXX11_TENSOR_TENSOR_BLOCK_H | |
| namespace Eigen { | |
| namespace internal { | |
| namespace { | |
| // Helper template to choose between ColMajor and RowMajor values. | |
| template <int Layout> | |
| struct cond; | |
| template <> | |
| struct cond<ColMajor> { | |
| template <typename T> | |
| EIGEN_STRONG_INLINE const T& operator()(const T& col, | |
| const T& /*row*/) const { | |
| return col; | |
| } | |
| }; | |
| template <> | |
| struct cond<RowMajor> { | |
| template <typename T> | |
| EIGEN_STRONG_INLINE const T& operator()(const T& /*col*/, | |
| const T& row) const { | |
| return row; | |
| } | |
| }; | |
| } // namespace | |
| /** | |
| * \enum TensorBlockShapeType | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block shape type. | |
| * | |
| * Tensor block shape type defines what are the shape preference for the blocks | |
| * extracted from the larger tensor. | |
| * | |
| * Example: | |
| * | |
| * We want to extract blocks of 100 elements from the large 100x100 tensor: | |
| * - tensor: 100x100 | |
| * - target_block_size: 100 | |
| * | |
| * TensorBlockShapeType: | |
| * - kUniformAllDims: 100 blocks of size 10x10 | |
| * - kSkewedInnerDims: 100 blocks of size 100x1 (or 1x100 depending on a column | |
| * or row major layout) | |
| */ | |
| enum TensorBlockShapeType { | |
| kUniformAllDims, | |
| kSkewedInnerDims | |
| }; | |
| struct TensorOpResourceRequirements { | |
| TensorBlockShapeType block_shape; | |
| Index block_total_size; | |
| // TODO(andydavis) Add 'target_num_threads' to support communication of | |
| // thread-resource requirements. This will allow ops deep in the | |
| // expression tree (like reductions) to communicate resources | |
| // requirements based on local state (like the total number of reductions | |
| // to be computed). | |
| TensorOpResourceRequirements(TensorBlockShapeType shape, | |
| const Index size) | |
| : block_shape(shape), block_total_size(size) {} | |
| }; | |
| // Tries to merge multiple resource requirements. | |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void MergeResourceRequirements( | |
| const std::vector<TensorOpResourceRequirements>& resources, | |
| TensorBlockShapeType* block_shape, Index* block_total_size) { | |
| if (resources.empty()) { | |
| return; | |
| } | |
| // TODO(andydavis) Implement different policies (i.e. revert to a default | |
| // policy if block shapes/sizes conflict). | |
| *block_shape = resources[0].block_shape; | |
| *block_total_size = resources[0].block_total_size; | |
| for (std::vector<TensorOpResourceRequirements>::size_type i = 1; i < resources.size(); ++i) { | |
| if (resources[i].block_shape == kSkewedInnerDims && | |
| *block_shape != kSkewedInnerDims) { | |
| *block_shape = kSkewedInnerDims; | |
| } | |
| *block_total_size = | |
| numext::maxi(*block_total_size, resources[i].block_total_size); | |
| } | |
| } | |
| /** | |
| * \class TensorBlock | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block class. | |
| * | |
| * This class represents a tensor block specified by the index of the | |
| * first block coefficient, and the size of the block in each dimension. | |
| */ | |
| template <typename Scalar, typename StorageIndex, int NumDims, int Layout> | |
| class TensorBlock { | |
| public: | |
| typedef DSizes<StorageIndex, NumDims> Dimensions; | |
| TensorBlock(const StorageIndex first_coeff_index, const Dimensions& block_sizes, | |
| const Dimensions& block_strides, const Dimensions& tensor_strides, | |
| Scalar* data) | |
| : m_first_coeff_index(first_coeff_index), | |
| m_block_sizes(block_sizes), | |
| m_block_strides(block_strides), | |
| m_tensor_strides(tensor_strides), | |
| m_data(data) {} | |
| StorageIndex first_coeff_index() const { return m_first_coeff_index; } | |
| const Dimensions& block_sizes() const { return m_block_sizes; } | |
| const Dimensions& block_strides() const { return m_block_strides; } | |
| const Dimensions& tensor_strides() const { return m_tensor_strides; } | |
| Scalar* data() { return m_data; } | |
| const Scalar* data() const { return m_data; } | |
| private: | |
| StorageIndex m_first_coeff_index; | |
| Dimensions m_block_sizes; | |
| Dimensions m_block_strides; | |
| Dimensions m_tensor_strides; | |
| Scalar* m_data; // Not owned. | |
| }; | |
| template <typename Scalar, typename StorageIndex> | |
| struct TensorBlockCopyOp { | |
| typedef typename packet_traits<Scalar>::type Packet; | |
| enum { | |
| Vectorizable = packet_traits<Scalar>::Vectorizable, | |
| PacketSize = packet_traits<Scalar>::size | |
| }; | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const StorageIndex num_coeff_to_copy, const StorageIndex dst_index, | |
| const StorageIndex dst_stride, Scalar* EIGEN_RESTRICT dst_data, | |
| const StorageIndex src_index, const StorageIndex src_stride, | |
| const Scalar* EIGEN_RESTRICT src_data) { | |
| const Scalar* src = &src_data[src_index]; | |
| Scalar* dst = &dst_data[dst_index]; | |
| if (!Vectorizable) { | |
| for (Index i = 0; i < num_coeff_to_copy; ++i) { | |
| dst[i * dst_stride] = src[i * src_stride]; | |
| } | |
| return; | |
| } | |
| if (src_stride == 1) { | |
| const StorageIndex vectorized_size = (num_coeff_to_copy / PacketSize) * PacketSize; | |
| if (dst_stride == 1) { | |
| // LINEAR | |
| for (StorageIndex i = 0; i < vectorized_size; i += PacketSize) { | |
| Packet p = ploadu<Packet>(src + i); | |
| pstoreu<Scalar, Packet>(dst + i, p); | |
| } | |
| for (StorageIndex i = vectorized_size; i < num_coeff_to_copy; ++i) { | |
| dst[i] = src[i]; | |
| } | |
| } else { | |
| // SCATTER | |
| for (StorageIndex i = 0; i < vectorized_size; i += PacketSize) { | |
| Packet p = ploadu<Packet>(src + i); | |
| pscatter<Scalar, Packet>(dst + i * dst_stride, p, dst_stride); | |
| } | |
| for (StorageIndex i = vectorized_size; i < num_coeff_to_copy; ++i) { | |
| dst[i * dst_stride] = src[i]; | |
| } | |
| } | |
| } else if (src_stride == 0) { | |
| const StorageIndex vectorized_size = (num_coeff_to_copy / PacketSize) * PacketSize; | |
| if (dst_stride == 1) { | |
| // LINEAR | |
| for (StorageIndex i = 0; i < vectorized_size; i += PacketSize) { | |
| Packet p = pload1<Packet>(src); | |
| pstoreu<Scalar, Packet>(dst + i, p); | |
| } | |
| for (StorageIndex i = vectorized_size; i < num_coeff_to_copy; ++i) { | |
| dst[i] = *src; | |
| } | |
| } else { | |
| // SCATTER | |
| for (StorageIndex i = 0; i < vectorized_size; i += PacketSize) { | |
| Packet p = pload1<Packet>(src); | |
| pscatter<Scalar, Packet>(dst + i * dst_stride, p, dst_stride); | |
| } | |
| for (StorageIndex i = vectorized_size; i < num_coeff_to_copy; ++i) { | |
| dst[i * dst_stride] = *src; | |
| } | |
| } | |
| } else { | |
| if (dst_stride == 1) { | |
| // GATHER | |
| const StorageIndex vectorized_size = (num_coeff_to_copy / PacketSize) * PacketSize; | |
| for (StorageIndex i = 0; i < vectorized_size; i += PacketSize) { | |
| Packet p = pgather<Scalar, Packet>(src + i * src_stride, src_stride); | |
| pstoreu<Scalar, Packet>(dst + i, p); | |
| } | |
| for (StorageIndex i = vectorized_size; i < num_coeff_to_copy; ++i) { | |
| dst[i] = src[i * src_stride]; | |
| } | |
| } else { | |
| // RANDOM | |
| for (StorageIndex i = 0; i < num_coeff_to_copy; ++i) { | |
| dst[i * dst_stride] = src[i * src_stride]; | |
| } | |
| } | |
| } | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockIO | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block IO class. | |
| * | |
| * This class is responsible for copying data between a tensor and a tensor | |
| * block. | |
| */ | |
| template <typename Scalar, typename StorageIndex, int NumDims, int Layout, | |
| bool BlockRead> | |
| class TensorBlockIO { | |
| public: | |
| typedef TensorBlock<Scalar, StorageIndex, NumDims, Layout> Block; | |
| typedef TensorBlockCopyOp<Scalar, StorageIndex> BlockCopyOp; | |
| protected: | |
| typedef array<StorageIndex, NumDims> Dimensions; | |
| struct BlockIteratorState { | |
| StorageIndex input_stride; | |
| StorageIndex output_stride; | |
| StorageIndex input_span; | |
| StorageIndex output_span; | |
| StorageIndex size; | |
| StorageIndex count; | |
| BlockIteratorState() | |
| : input_stride(0), | |
| output_stride(0), | |
| input_span(0), | |
| output_span(0), | |
| size(0), | |
| count(0) {} | |
| }; | |
| // Compute how many inner dimensions it's allowed to squeeze when doing IO | |
| // between a tensor and a block. It's safe to squeeze inner dimensions, only | |
| // if they are not reordered. | |
| static int NumSqueezableInnerDims(const Dimensions& tensor_to_block_dim_map) { | |
| int num_squeezable_dims = 0; | |
| if (Layout == ColMajor) { | |
| for (int i = 0; i < NumDims; ++i) { | |
| if (tensor_to_block_dim_map[i] == i) num_squeezable_dims++; | |
| else break; | |
| } | |
| } else { | |
| for (int i = NumDims - 1; i >= 0; --i) { | |
| if (tensor_to_block_dim_map[i] == i) num_squeezable_dims++; | |
| else break; | |
| } | |
| } | |
| return num_squeezable_dims; | |
| } | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Copy( | |
| const Block& block, StorageIndex first_coeff_index, | |
| const Dimensions& tensor_to_block_dim_map, | |
| const Dimensions& tensor_strides, | |
| const Scalar* src_data, | |
| Scalar* dst_data) { | |
| // Do not squeeze reordered inner dimensions. | |
| int num_squeezable_dims = NumSqueezableInnerDims(tensor_to_block_dim_map); | |
| // Find the innermost tensor dimension whose size is not 1. This is the | |
| // effective inner dim. If all dimensions are of size 1, then fallback to | |
| // using the actual innermost dim to avoid out-of-bound access. | |
| StorageIndex num_size_one_inner_dims = 0; | |
| for (int i = 0; i < num_squeezable_dims; ++i) { | |
| const int dim = cond<Layout>()(i, NumDims - i - 1); | |
| if (block.block_sizes()[tensor_to_block_dim_map[dim]] != 1) { | |
| num_size_one_inner_dims = i; | |
| break; | |
| } | |
| } | |
| // Calculate strides and dimensions. | |
| const StorageIndex tensor_stride1_dim = cond<Layout>()( | |
| num_size_one_inner_dims, NumDims - num_size_one_inner_dims - 1); | |
| const StorageIndex block_dim_for_tensor_stride1_dim = | |
| NumDims == 0 ? 1 : tensor_to_block_dim_map[tensor_stride1_dim]; | |
| StorageIndex block_inner_dim_size = | |
| NumDims == 0 ? 1 | |
| : block.block_sizes()[block_dim_for_tensor_stride1_dim]; | |
| // Squeeze multiple inner dims into one for larger inner dim size. | |
| for (Index i = num_size_one_inner_dims + 1; i < num_squeezable_dims; ++i) { | |
| const Index dim = cond<Layout>()(i, NumDims - i - 1); | |
| const StorageIndex block_stride = | |
| block.block_strides()[tensor_to_block_dim_map[dim]]; | |
| if (block_inner_dim_size == block_stride && | |
| block_stride == tensor_strides[dim]) { | |
| block_inner_dim_size *= | |
| block.block_sizes()[tensor_to_block_dim_map[dim]]; | |
| ++num_size_one_inner_dims; | |
| } else { | |
| break; | |
| } | |
| } | |
| StorageIndex inputIndex; | |
| StorageIndex outputIndex; | |
| StorageIndex input_stride; | |
| StorageIndex output_stride; | |
| // Setup strides to read/write along the tensor's stride1 dimension. | |
| if (BlockRead) { | |
| inputIndex = first_coeff_index; | |
| outputIndex = 0; | |
| input_stride = NumDims == 0 ? 1 : tensor_strides[tensor_stride1_dim]; | |
| output_stride = | |
| NumDims == 0 | |
| ? 1 | |
| : block.block_strides()[block_dim_for_tensor_stride1_dim]; | |
| } else { | |
| inputIndex = 0; | |
| outputIndex = first_coeff_index; | |
| input_stride = | |
| NumDims == 0 | |
| ? 1 | |
| : block.block_strides()[block_dim_for_tensor_stride1_dim]; | |
| output_stride = NumDims == 0 ? 1 : tensor_strides[tensor_stride1_dim]; | |
| } | |
| const int at_least_1_dim = NumDims <= 1 ? 1 : NumDims - 1; | |
| array<BlockIteratorState, at_least_1_dim> block_iter_state; | |
| // Initialize block iterator state. Squeeze away any dimension of size 1. | |
| Index num_squeezed_dims = 0; | |
| for (Index i = num_size_one_inner_dims; i < NumDims - 1; ++i) { | |
| const Index dim = cond<Layout>()(i + 1, NumDims - i - 2); | |
| const StorageIndex size = block.block_sizes()[tensor_to_block_dim_map[dim]]; | |
| if (size == 1) { | |
| continue; | |
| } | |
| block_iter_state[num_squeezed_dims].size = size; | |
| if (BlockRead) { | |
| block_iter_state[num_squeezed_dims].input_stride = tensor_strides[dim]; | |
| block_iter_state[num_squeezed_dims].output_stride = | |
| block.block_strides()[tensor_to_block_dim_map[dim]]; | |
| } else { | |
| block_iter_state[num_squeezed_dims].input_stride = | |
| block.block_strides()[tensor_to_block_dim_map[dim]]; | |
| block_iter_state[num_squeezed_dims].output_stride = tensor_strides[dim]; | |
| } | |
| block_iter_state[num_squeezed_dims].input_span = | |
| block_iter_state[num_squeezed_dims].input_stride * | |
| (block_iter_state[num_squeezed_dims].size - 1); | |
| block_iter_state[num_squeezed_dims].output_span = | |
| block_iter_state[num_squeezed_dims].output_stride * | |
| (block_iter_state[num_squeezed_dims].size - 1); | |
| ++num_squeezed_dims; | |
| } | |
| // Iterate copying data from src to dst. | |
| const StorageIndex block_total_size = | |
| NumDims == 0 ? 1 : block.block_sizes().TotalSize(); | |
| for (StorageIndex i = 0; i < block_total_size; i += block_inner_dim_size) { | |
| BlockCopyOp::Run(block_inner_dim_size, outputIndex, output_stride, | |
| dst_data, inputIndex, input_stride, src_data); | |
| // Update index. | |
| for (int j = 0; j < num_squeezed_dims; ++j) { | |
| if (++block_iter_state[j].count < block_iter_state[j].size) { | |
| inputIndex += block_iter_state[j].input_stride; | |
| outputIndex += block_iter_state[j].output_stride; | |
| break; | |
| } | |
| block_iter_state[j].count = 0; | |
| inputIndex -= block_iter_state[j].input_span; | |
| outputIndex -= block_iter_state[j].output_span; | |
| } | |
| } | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockReader | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block reader class. | |
| * | |
| * This class is responsible for reading a tensor block. | |
| * | |
| */ | |
| template <typename Scalar, typename StorageIndex, int NumDims, int Layout> | |
| class TensorBlockReader : public TensorBlockIO<Scalar, StorageIndex, NumDims, | |
| Layout, /*BlockRead=*/true> { | |
| public: | |
| typedef TensorBlock<Scalar, StorageIndex, NumDims, Layout> Block; | |
| typedef TensorBlockIO<Scalar, StorageIndex, NumDims, Layout, /*BlockRead=*/true> Base; | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| Block* block, const Scalar* src_data) { | |
| array<StorageIndex, NumDims> tensor_to_block_dim_map; | |
| for (int i = 0; i < NumDims; ++i) { | |
| tensor_to_block_dim_map[i] = i; | |
| } | |
| Base::Copy(*block, block->first_coeff_index(), tensor_to_block_dim_map, | |
| block->tensor_strides(), src_data, block->data()); | |
| } | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| Block* block, StorageIndex first_coeff_index, | |
| const array<StorageIndex, NumDims>& tensor_to_block_dim_map, | |
| const array<StorageIndex, NumDims>& tensor_strides, const Scalar* src_data) { | |
| Base::Copy(*block, first_coeff_index, tensor_to_block_dim_map, | |
| tensor_strides, src_data, block->data()); | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockWriter | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block writer class. | |
| * | |
| * This class is responsible for writing a tensor block. | |
| * | |
| */ | |
| template <typename Scalar, typename StorageIndex, int NumDims, int Layout> | |
| class TensorBlockWriter : public TensorBlockIO<Scalar, StorageIndex, NumDims, | |
| Layout, /*BlockRead=*/false> { | |
| public: | |
| typedef TensorBlock<Scalar, StorageIndex, NumDims, Layout> Block; | |
| typedef TensorBlockIO<Scalar, StorageIndex, NumDims, Layout, /*BlockRead=*/false> Base; | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const Block& block, Scalar* dst_data) { | |
| array<StorageIndex, NumDims> tensor_to_block_dim_map; | |
| for (int i = 0; i < NumDims; ++i) { | |
| tensor_to_block_dim_map[i] = i; | |
| } | |
| Base::Copy(block, block.first_coeff_index(), tensor_to_block_dim_map, | |
| block.tensor_strides(), block.data(), dst_data); | |
| } | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const Block& block, StorageIndex first_coeff_index, | |
| const array<StorageIndex, NumDims>& tensor_to_block_dim_map, | |
| const array<StorageIndex, NumDims>& tensor_strides, Scalar* dst_data) { | |
| Base::Copy(block, first_coeff_index, tensor_to_block_dim_map, | |
| tensor_strides, block.data(), dst_data); | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockCwiseUnaryOp | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Carries out a cwise binary op on a number of coefficients. | |
| * | |
| * This class reads strided input from the argument, and writes the | |
| * result of the cwise unary op to the strided output array. | |
| * | |
| */ | |
| template <bool Vectorizable> | |
| struct TensorBlockCwiseUnaryOp { | |
| template <typename StorageIndex, typename UnaryFunctor, | |
| typename OutputScalar, typename InputScalar> | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const UnaryFunctor& functor, const StorageIndex num_coeff, | |
| const StorageIndex output_index, const StorageIndex output_stride, | |
| OutputScalar* output_data, const StorageIndex input_index, | |
| const StorageIndex input_stride, const InputScalar* input_data) { | |
| typedef const Array<InputScalar, Dynamic, 1> Input; | |
| typedef Array<OutputScalar, Dynamic, 1> Output; | |
| typedef Map<Input, 0, InnerStride<> > InputMap; | |
| typedef Map<Output, 0, InnerStride<> > OutputMap; | |
| const InputScalar* input_base = &input_data[input_index]; | |
| OutputScalar* output_base = &output_data[output_index]; | |
| const InputMap input(input_base, num_coeff, InnerStride<>(input_stride)); | |
| OutputMap output(output_base, num_coeff, InnerStride<>(output_stride)); | |
| output = CwiseUnaryOp<UnaryFunctor, InputMap>(input, functor); | |
| } | |
| }; | |
| template<> | |
| struct TensorBlockCwiseUnaryOp<true> { | |
| template <typename StorageIndex, typename UnaryFunctor, | |
| typename OutputScalar, typename InputScalar> | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const UnaryFunctor& functor, const StorageIndex num_coeff, | |
| const StorageIndex output_index, const StorageIndex output_stride, | |
| OutputScalar* output_data, const StorageIndex input_index, | |
| const StorageIndex input_stride, const InputScalar* input_data) { | |
| if (input_stride == 1 && output_stride == 1) { | |
| typedef const Array<InputScalar, Dynamic, 1> Input; | |
| typedef Array<OutputScalar, Dynamic, 1> Output; | |
| const Map<Input> input(&input_data[input_index], num_coeff); | |
| Map<Output> output(&output_data[output_index], num_coeff); | |
| output = CwiseUnaryOp<UnaryFunctor, Map<Input> >(input, functor); | |
| } else { | |
| TensorBlockCwiseUnaryOp<false>::Run( | |
| functor, num_coeff, output_index, output_stride, output_data, | |
| input_index, input_stride, input_data); | |
| } | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockCwiseUnaryIO | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block IO class for carrying out cwise unary ops. | |
| * | |
| * This class carries out the unary op on given blocks. | |
| */ | |
| template <typename UnaryFunctor, typename StorageIndex, typename OutputScalar, | |
| int NumDims, int Layout> | |
| struct TensorBlockCwiseUnaryIO { | |
| typedef typename TensorBlock<OutputScalar, StorageIndex, NumDims, | |
| Layout>::Dimensions Dimensions; | |
| typedef TensorBlockCwiseUnaryOp< | |
| packet_traits<OutputScalar>::Vectorizable && | |
| functor_traits<UnaryFunctor>::PacketAccess> | |
| TensorBlockCwiseUnaryOpImpl; | |
| struct BlockIteratorState { | |
| StorageIndex output_stride, output_span; | |
| StorageIndex input_stride, input_span; | |
| StorageIndex size, count; | |
| }; | |
| template <typename InputScalar> | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const UnaryFunctor& functor, const Dimensions& block_sizes, | |
| const Dimensions& block_strides, OutputScalar* output_data, | |
| const array<StorageIndex, NumDims>& input_strides, | |
| const InputScalar* input_data) { | |
| // Find the innermost dimension whose size is not 1. This is the effective | |
| // inner dim. If all dimensions are of size 1, fallback to using the actual | |
| // innermost dim to avoid out-of-bound access. | |
| int num_size_one_inner_dims = 0; | |
| for (int i = 0; i < NumDims; ++i) { | |
| const int dim = cond<Layout>()(i, NumDims - i - 1); | |
| if (block_sizes[dim] != 1) { | |
| num_size_one_inner_dims = i; | |
| break; | |
| } | |
| } | |
| // Calculate strides and dimensions. | |
| const int inner_dim = | |
| NumDims == 0 ? 1 | |
| : cond<Layout>()(num_size_one_inner_dims, | |
| NumDims - num_size_one_inner_dims - 1); | |
| StorageIndex inner_dim_size = NumDims == 0 ? 1 : block_sizes[inner_dim]; | |
| for (int i = num_size_one_inner_dims + 1; i < NumDims; ++i) { | |
| const int dim = cond<Layout>()(i, NumDims - i - 1); | |
| // Merge multiple inner dims into one for larger inner dim size (i.e. | |
| // fewer calls to TensorBlockCwiseUnaryOp::Run()). | |
| if (inner_dim_size == block_strides[dim] && | |
| block_strides[dim] == input_strides[dim]) { | |
| inner_dim_size *= block_sizes[dim]; | |
| ++num_size_one_inner_dims; | |
| } else { | |
| break; | |
| } | |
| } | |
| StorageIndex output_index = 0, input_index = 0; | |
| const StorageIndex output_stride = | |
| NumDims == 0 ? 1 : block_strides[inner_dim]; | |
| const StorageIndex input_stride = | |
| NumDims == 0 ? 1 : input_strides[inner_dim]; | |
| const int at_least_1_dim = NumDims <= 1 ? 1 : NumDims - 1; | |
| array<BlockIteratorState, at_least_1_dim> block_iter_state; | |
| // Initialize block iterator state. Squeeze away any dimension of size 1. | |
| int num_squeezed_dims = 0; | |
| for (int i = num_size_one_inner_dims; i < NumDims - 1; ++i) { | |
| const int dim = cond<Layout>()(i + 1, NumDims - i - 2); | |
| const StorageIndex size = block_sizes[dim]; | |
| if (size == 1) { | |
| continue; | |
| } | |
| BlockIteratorState& state = block_iter_state[num_squeezed_dims]; | |
| state.output_stride = block_strides[dim]; | |
| state.input_stride = input_strides[dim]; | |
| state.size = size; | |
| state.output_span = state.output_stride * (size - 1); | |
| state.input_span = state.input_stride * (size - 1); | |
| state.count = 0; | |
| ++num_squeezed_dims; | |
| } | |
| // Compute cwise unary op. | |
| const StorageIndex block_total_size = | |
| NumDims == 0 ? 1 : block_sizes.TotalSize(); | |
| for (StorageIndex i = 0; i < block_total_size; i += inner_dim_size) { | |
| TensorBlockCwiseUnaryOpImpl::Run(functor, inner_dim_size, output_index, | |
| output_stride, output_data, input_index, | |
| input_stride, input_data); | |
| // Update index. | |
| for (int j = 0; j < num_squeezed_dims; ++j) { | |
| BlockIteratorState& state = block_iter_state[j]; | |
| if (++state.count < state.size) { | |
| output_index += state.output_stride; | |
| input_index += state.input_stride; | |
| break; | |
| } | |
| state.count = 0; | |
| output_index -= state.output_span; | |
| input_index -= state.input_span; | |
| } | |
| } | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockCwiseBinaryOp | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Carries out a cwise binary op on a number of coefficients. | |
| * | |
| * This class reads strided inputs from left and right operands, and writes the | |
| * result of the cwise binary op to the strided output array. | |
| * | |
| */ | |
| template<bool Vectorizable> | |
| struct TensorBlockCwiseBinaryOp { | |
| template <typename StorageIndex, typename BinaryFunctor, typename OutputScalar, | |
| typename LeftScalar, typename RightScalar> | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const BinaryFunctor& functor, const StorageIndex num_coeff, | |
| const StorageIndex output_index, const StorageIndex output_stride, | |
| OutputScalar* output_data, const StorageIndex left_index, | |
| const StorageIndex left_stride, const LeftScalar* left_data, | |
| const StorageIndex right_index, const StorageIndex right_stride, | |
| const RightScalar* right_data) { | |
| typedef const Array<LeftScalar, Dynamic, 1> Lhs; | |
| typedef const Array<RightScalar, Dynamic, 1> Rhs; | |
| typedef Array<OutputScalar, Dynamic, 1> Out; | |
| typedef Map<Lhs, 0, InnerStride<> > LhsMap; | |
| typedef Map<Rhs, 0, InnerStride<> > RhsMap; | |
| typedef Map<Out, 0, InnerStride<> > OutMap; | |
| const LeftScalar* lhs_base = &left_data[left_index]; | |
| const RightScalar* rhs_base = &right_data[right_index]; | |
| OutputScalar* out_base = &output_data[output_index]; | |
| const LhsMap lhs(lhs_base, num_coeff, InnerStride<>(left_stride)); | |
| const RhsMap rhs(rhs_base, num_coeff, InnerStride<>(right_stride)); | |
| OutMap out(out_base, num_coeff, InnerStride<>(output_stride)); | |
| out = CwiseBinaryOp<BinaryFunctor, LhsMap, RhsMap>(lhs, rhs, functor); | |
| } | |
| }; | |
| template<> | |
| struct TensorBlockCwiseBinaryOp<true> { | |
| template <typename StorageIndex, typename BinaryFunctor, typename OutputScalar, | |
| typename LeftScalar, typename RightScalar> | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const BinaryFunctor& functor, const StorageIndex num_coeff, | |
| const StorageIndex output_index, const StorageIndex output_stride, | |
| OutputScalar* output_data, const StorageIndex left_index, | |
| const StorageIndex left_stride, const LeftScalar* left_data, | |
| const StorageIndex right_index, const StorageIndex right_stride, | |
| const RightScalar* right_data) { | |
| if (left_stride == 1 && right_stride == 1 && output_stride == 1) { | |
| typedef const Array<LeftScalar, Dynamic, 1> Lhs; | |
| typedef const Array<RightScalar, Dynamic, 1> Rhs; | |
| typedef Array<OutputScalar, Dynamic, 1> Out; | |
| const LeftScalar* lhs_base = &left_data[left_index]; | |
| const RightScalar* rhs_base = &right_data[right_index]; | |
| OutputScalar* out_base = &output_data[output_index]; | |
| const Map<Lhs> lhs(lhs_base, num_coeff); | |
| const Map<Rhs> rhs(rhs_base, num_coeff); | |
| Map<Out> out(out_base, num_coeff); | |
| out = CwiseBinaryOp<BinaryFunctor, Map<Lhs>, Map<Rhs> >(lhs, rhs, functor); | |
| } else { | |
| TensorBlockCwiseBinaryOp<false>::Run( | |
| functor, num_coeff, output_index, output_stride, output_data, | |
| left_index, left_stride, left_data, right_index, right_stride, | |
| right_data); | |
| } | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockCwiseBinaryIO | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block IO class for carrying out cwise binary ops. | |
| * | |
| * This class carries out the binary op on given blocks. | |
| * | |
| */ | |
| template <typename BinaryFunctor, typename StorageIndex, typename OutputScalar, | |
| int NumDims, int Layout> | |
| struct TensorBlockCwiseBinaryIO { | |
| typedef typename TensorBlock<OutputScalar, StorageIndex, NumDims, Layout>::Dimensions Dimensions; | |
| typedef TensorBlockCwiseBinaryOp< | |
| packet_traits<OutputScalar>::Vectorizable && | |
| functor_traits<BinaryFunctor>::PacketAccess> | |
| TensorBlockCwiseBinaryOpImpl; | |
| struct BlockIteratorState { | |
| StorageIndex output_stride, output_span; | |
| StorageIndex left_stride, left_span; | |
| StorageIndex right_stride, right_span; | |
| StorageIndex size, count; | |
| }; | |
| template <typename LeftScalar, typename RightScalar> | |
| static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void Run( | |
| const BinaryFunctor& functor, const Dimensions& block_sizes, | |
| const Dimensions& block_strides, OutputScalar* output_data, | |
| const array<StorageIndex, NumDims>& left_strides, | |
| const LeftScalar* left_data, | |
| const array<StorageIndex, NumDims>& right_strides, | |
| const RightScalar* right_data) { | |
| // Find the innermost dimension whose size is not 1. This is the effective | |
| // inner dim. If all dimensions are of size 1, fallback to using the actual | |
| // innermost dim to avoid out-of-bound access. | |
| int num_size_one_inner_dims = 0; | |
| for (int i = 0; i < NumDims; ++i) { | |
| const int dim = cond<Layout>()(i, NumDims - i - 1); | |
| if (block_sizes[dim] != 1) { | |
| num_size_one_inner_dims = i; | |
| break; | |
| } | |
| } | |
| // Calculate strides and dimensions. | |
| const int inner_dim = | |
| NumDims == 0 ? 1 | |
| : cond<Layout>()(num_size_one_inner_dims, | |
| NumDims - num_size_one_inner_dims - 1); | |
| StorageIndex inner_dim_size = NumDims == 0 ? 1 : block_sizes[inner_dim]; | |
| for (int i = num_size_one_inner_dims + 1; i < NumDims; ++i) { | |
| const int dim = cond<Layout>()(i, NumDims - i - 1); | |
| // Merge multiple inner dims into one for larger inner dim size (i.e. | |
| // fewer calls to TensorBlockCwiseBinaryOp::Run()). | |
| if (inner_dim_size == block_strides[dim] && | |
| block_strides[dim] == left_strides[dim] && | |
| block_strides[dim] == right_strides[dim]) { | |
| inner_dim_size *= block_sizes[dim]; | |
| ++num_size_one_inner_dims; | |
| } else { | |
| break; | |
| } | |
| } | |
| StorageIndex output_index = 0, left_index = 0, right_index = 0; | |
| const StorageIndex output_stride = | |
| NumDims == 0 ? 1 : block_strides[inner_dim]; | |
| const StorageIndex left_stride = NumDims == 0 ? 1 : left_strides[inner_dim]; | |
| const StorageIndex right_stride = | |
| NumDims == 0 ? 1 : right_strides[inner_dim]; | |
| const int at_least_1_dim = NumDims <= 1 ? 1 : NumDims - 1; | |
| array<BlockIteratorState, at_least_1_dim> block_iter_state; | |
| // Initialize block iterator state. Squeeze away any dimension of size 1. | |
| int num_squeezed_dims = 0; | |
| for (int i = num_size_one_inner_dims; i < NumDims - 1; ++i) { | |
| const int dim = cond<Layout>()(i + 1, NumDims - i - 2); | |
| const StorageIndex size = block_sizes[dim]; | |
| if (size == 1) { | |
| continue; | |
| } | |
| BlockIteratorState& state = block_iter_state[num_squeezed_dims]; | |
| state.output_stride = block_strides[dim]; | |
| state.left_stride = left_strides[dim]; | |
| state.right_stride = right_strides[dim]; | |
| state.size = size; | |
| state.output_span = state.output_stride * (size - 1); | |
| state.left_span = state.left_stride * (size - 1); | |
| state.right_span = state.right_stride * (size - 1); | |
| state.count = 0; | |
| ++num_squeezed_dims; | |
| } | |
| // Compute cwise binary op. | |
| const StorageIndex block_total_size = | |
| NumDims == 0 ? 1 : block_sizes.TotalSize(); | |
| for (StorageIndex i = 0; i < block_total_size; i += inner_dim_size) { | |
| TensorBlockCwiseBinaryOpImpl::Run(functor, inner_dim_size, output_index, | |
| output_stride, output_data, left_index, | |
| left_stride, left_data, right_index, | |
| right_stride, right_data); | |
| // Update index. | |
| for (int j = 0; j < num_squeezed_dims; ++j) { | |
| BlockIteratorState& state = block_iter_state[j]; | |
| if (++state.count < state.size) { | |
| output_index += state.output_stride; | |
| left_index += state.left_stride; | |
| right_index += state.right_stride; | |
| break; | |
| } | |
| state.count = 0; | |
| output_index -= state.output_span; | |
| left_index -= state.left_span; | |
| right_index -= state.right_span; | |
| } | |
| } | |
| } | |
| }; | |
| /** | |
| * \class TensorBlockView | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Read-only view into a block of data. | |
| * | |
| * This class provides read-only access to a block of data in impl. It may need | |
| * to allocate space for holding the intermediate result. | |
| * | |
| */ | |
| template <class ArgType, class Device> | |
| struct TensorBlockView { | |
| typedef TensorEvaluator<ArgType, Device> Impl; | |
| typedef typename Impl::Index StorageIndex; | |
| typedef typename remove_const<typename Impl::Scalar>::type Scalar; | |
| static const int NumDims = array_size<typename Impl::Dimensions>::value; | |
| typedef DSizes<StorageIndex, NumDims> Dimensions; | |
| // Constructs a TensorBlockView for `impl`. `block` is only used for for | |
| // specifying the start offset, shape, and strides of the block. | |
| template <typename OtherTensorBlock> | |
| TensorBlockView(const Device& device, | |
| const TensorEvaluator<ArgType, Device>& impl, | |
| const OtherTensorBlock& block) | |
| : m_device(device), | |
| m_block_sizes(block.block_sizes()), | |
| m_data(NULL), | |
| m_allocated_data(NULL) { | |
| if (Impl::RawAccess && impl.data() != NULL) { | |
| m_data = impl.data() + block.first_coeff_index(); | |
| m_block_strides = block.tensor_strides(); | |
| } else { | |
| // Actually make a copy. | |
| // TODO(wuke): This sometimes put a lot pressure on the heap allocator. | |
| // Consider allowing ops to request additional temporary block memory in | |
| // TensorOpResourceRequirements. | |
| m_allocated_data = static_cast<Scalar*>( | |
| m_device.allocate(m_block_sizes.TotalSize() * sizeof(Scalar))); | |
| m_data = m_allocated_data; | |
| if (NumDims > 0) { | |
| if (static_cast<int>(Impl::Layout) == static_cast<int>(ColMajor)) { | |
| m_block_strides[0] = 1; | |
| for (int i = 1; i < NumDims; ++i) { | |
| m_block_strides[i] = m_block_strides[i - 1] * m_block_sizes[i - 1]; | |
| } | |
| } else { | |
| m_block_strides[NumDims - 1] = 1; | |
| for (int i = NumDims - 2; i >= 0; --i) { | |
| m_block_strides[i] = m_block_strides[i + 1] * m_block_sizes[i + 1]; | |
| } | |
| } | |
| } | |
| TensorBlock<Scalar, StorageIndex, NumDims, Impl::Layout> input_block( | |
| block.first_coeff_index(), m_block_sizes, m_block_strides, | |
| block.tensor_strides(), m_allocated_data); | |
| impl.block(&input_block); | |
| } | |
| } | |
| ~TensorBlockView() { | |
| if (m_allocated_data != NULL) { | |
| m_device.deallocate(m_allocated_data); | |
| } | |
| } | |
| const Dimensions& block_sizes() const { return m_block_sizes; } | |
| const Dimensions& block_strides() const { return m_block_strides; } | |
| const Scalar* data() const { return m_data; } | |
| private: | |
| const Device EIGEN_DEVICE_REF m_device; | |
| Dimensions m_block_sizes, m_block_strides; | |
| const Scalar* m_data; // Not owned. | |
| Scalar* m_allocated_data; // Owned. | |
| }; | |
| /** | |
| * \class TensorBlockMapper | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor block mapper class. | |
| * | |
| * This class is responsible for iterating over the blocks of a tensor. | |
| */ | |
| template <typename Scalar, typename StorageIndex, int NumDims, int Layout> | |
| class TensorBlockMapper { | |
| public: | |
| typedef TensorBlock<Scalar, StorageIndex, NumDims, Layout> Block; | |
| typedef DSizes<StorageIndex, NumDims> Dimensions; | |
| TensorBlockMapper() {} | |
| TensorBlockMapper(const Dimensions& dims, | |
| const TensorBlockShapeType block_shape, | |
| Index min_target_size) | |
| : m_dimensions(dims), | |
| m_block_dim_sizes(BlockDimensions(dims, block_shape, convert_index<StorageIndex>(min_target_size))) { | |
| // Calculate block counts by dimension and total block count. | |
| DSizes<StorageIndex, NumDims> block_count; | |
| for (Index i = 0; i < block_count.rank(); ++i) { | |
| block_count[i] = divup(m_dimensions[i], m_block_dim_sizes[i]); | |
| } | |
| m_total_block_count = array_prod(block_count); | |
| // Calculate block strides (used for enumerating blocks). | |
| if (NumDims > 0) { | |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { | |
| m_block_strides[0] = 1; | |
| m_tensor_strides[0] = 1; | |
| for (int i = 1; i < NumDims; ++i) { | |
| m_block_strides[i] = m_block_strides[i - 1] * block_count[i - 1]; | |
| m_tensor_strides[i] = m_tensor_strides[i - 1] * m_dimensions[i - 1]; | |
| } | |
| } else { | |
| m_block_strides[NumDims - 1] = 1; | |
| m_tensor_strides[NumDims - 1] = 1; | |
| for (int i = NumDims - 2; i >= 0; --i) { | |
| m_block_strides[i] = m_block_strides[i + 1] * block_count[i + 1]; | |
| m_tensor_strides[i] = m_tensor_strides[i + 1] * m_dimensions[i + 1]; | |
| } | |
| } | |
| } | |
| } | |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Block | |
| GetBlockForIndex(StorageIndex block_index, Scalar* data) const { | |
| StorageIndex first_coeff_index = 0; | |
| DSizes<StorageIndex, NumDims> coords; | |
| DSizes<StorageIndex, NumDims> sizes; | |
| DSizes<StorageIndex, NumDims> strides; | |
| if (NumDims > 0) { | |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { | |
| for (int i = NumDims - 1; i > 0; --i) { | |
| const StorageIndex idx = block_index / m_block_strides[i]; | |
| coords[i] = idx * m_block_dim_sizes[i]; | |
| sizes[i] = | |
| numext::mini((m_dimensions[i] - coords[i]), m_block_dim_sizes[i]); | |
| block_index -= idx * m_block_strides[i]; | |
| first_coeff_index += coords[i] * m_tensor_strides[i]; | |
| } | |
| coords[0] = block_index * m_block_dim_sizes[0]; | |
| sizes[0] = | |
| numext::mini((m_dimensions[0] - coords[0]), m_block_dim_sizes[0]); | |
| first_coeff_index += coords[0] * m_tensor_strides[0]; | |
| strides[0] = 1; | |
| for (int i = 1; i < NumDims; ++i) { | |
| strides[i] = strides[i - 1] * sizes[i - 1]; | |
| } | |
| } else { | |
| for (int i = 0; i < NumDims - 1; ++i) { | |
| const StorageIndex idx = block_index / m_block_strides[i]; | |
| coords[i] = idx * m_block_dim_sizes[i]; | |
| sizes[i] = | |
| numext::mini((m_dimensions[i] - coords[i]), m_block_dim_sizes[i]); | |
| block_index -= idx * m_block_strides[i]; | |
| first_coeff_index += coords[i] * m_tensor_strides[i]; | |
| } | |
| coords[NumDims - 1] = block_index * m_block_dim_sizes[NumDims - 1]; | |
| sizes[NumDims - 1] = | |
| numext::mini((m_dimensions[NumDims - 1] - coords[NumDims - 1]), | |
| m_block_dim_sizes[NumDims - 1]); | |
| first_coeff_index += | |
| coords[NumDims - 1] * m_tensor_strides[NumDims - 1]; | |
| strides[NumDims - 1] = 1; | |
| for (int i = NumDims - 2; i >= 0; --i) { | |
| strides[i] = strides[i + 1] * sizes[i + 1]; | |
| } | |
| } | |
| } | |
| return Block(first_coeff_index, sizes, strides, m_tensor_strides, data); | |
| } | |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageIndex total_block_count() const { | |
| return m_total_block_count; | |
| } | |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageIndex | |
| block_dims_total_size() const { | |
| return m_block_dim_sizes.TotalSize(); | |
| } | |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& | |
| block_dim_sizes() const { | |
| return m_block_dim_sizes; | |
| } | |
| private: | |
| static Dimensions BlockDimensions(const Dimensions& tensor_dims, | |
| const TensorBlockShapeType block_shape, | |
| StorageIndex min_target_size) { | |
| min_target_size = numext::maxi<StorageIndex>(1, min_target_size); | |
| // If tensor fully fits into the target size, we'll treat it a single block. | |
| Dimensions block_dim_sizes = tensor_dims; | |
| if (tensor_dims.TotalSize() == 0) { | |
| // Corner case: one of the dimensions is zero. Logic below is too complex | |
| // to handle this case on a general basis, just use unit block size. | |
| // Note: we must not yield blocks with zero dimensions (recipe for | |
| // overflows/underflows, divisions by zero and NaNs later). | |
| for (int i = 0; i < NumDims; ++i) { | |
| block_dim_sizes[i] = 1; | |
| } | |
| } else if (block_dim_sizes.TotalSize() > min_target_size) { | |
| if (block_shape == kUniformAllDims) { | |
| // Tensor will not fit within 'min_target_size' budget: calculate tensor | |
| // block dimension sizes based on "square" dimension size target. | |
| const StorageIndex dim_size_target = convert_index<StorageIndex>( | |
| std::pow(static_cast<float>(min_target_size), | |
| 1.0f / static_cast<float>(block_dim_sizes.rank()))); | |
| for (Index i = 0; i < block_dim_sizes.rank(); ++i) { | |
| // TODO(andydavis) Adjust the inner most 'block_dim_size' to make it | |
| // a multiple of the packet size. Note that reducing | |
| // 'block_dim_size' in this manner can increase the number of | |
| // blocks, and so will amplify any per-block overhead. | |
| block_dim_sizes[i] = numext::mini(dim_size_target, tensor_dims[i]); | |
| } | |
| // Add any un-allocated coefficients to inner dimension(s). | |
| StorageIndex total_size = block_dim_sizes.TotalSize(); | |
| for (int i = 0; i < NumDims; ++i) { | |
| const int dim = cond<Layout>()(i, NumDims - i - 1); | |
| if (block_dim_sizes[dim] < tensor_dims[dim]) { | |
| const StorageIndex total_size_other_dims = | |
| total_size / block_dim_sizes[dim]; | |
| const StorageIndex alloc_avail = | |
| divup<StorageIndex>(min_target_size, total_size_other_dims); | |
| if (alloc_avail == block_dim_sizes[dim]) { | |
| // Insufficient excess coefficients to allocate. | |
| break; | |
| } | |
| block_dim_sizes[dim] = numext::mini(tensor_dims[dim], alloc_avail); | |
| total_size = total_size_other_dims * block_dim_sizes[dim]; | |
| } | |
| } | |
| } else if (block_shape == kSkewedInnerDims) { | |
| StorageIndex coeff_to_allocate = min_target_size; | |
| for (int i = 0; i < NumDims; ++i) { | |
| const int dim = cond<Layout>()(i, NumDims - i - 1); | |
| block_dim_sizes[dim] = | |
| numext::mini(coeff_to_allocate, tensor_dims[dim]); | |
| coeff_to_allocate = divup( | |
| coeff_to_allocate, | |
| numext::maxi(static_cast<StorageIndex>(1), block_dim_sizes[dim])); | |
| } | |
| eigen_assert(coeff_to_allocate == 1); | |
| } else { | |
| eigen_assert(false); // someone added new block shape type | |
| } | |
| } | |
| eigen_assert( | |
| block_dim_sizes.TotalSize() >= | |
| numext::mini<Index>(min_target_size, tensor_dims.TotalSize())); | |
| return block_dim_sizes; | |
| } | |
| Dimensions m_dimensions; | |
| Dimensions m_block_dim_sizes; | |
| Dimensions m_block_strides; | |
| Dimensions m_tensor_strides; | |
| StorageIndex m_total_block_count; | |
| }; | |
| /** | |
| * \class TensorSliceBlockMapper | |
| * \ingroup CXX11_Tensor_Module | |
| * | |
| * \brief Tensor slice block mapper class. | |
| * | |
| * This class is responsible for iterating over the blocks of | |
| * a slice of a tensor. Supports shuffling of the block strides | |
| * for callers that want to reduce strides for dimensions to be | |
| * processed together. | |
| * | |
| */ | |
| template <typename Scalar, typename StorageIndex, int NumDims, int Layout> | |
| class TensorSliceBlockMapper { | |
| public: | |
| typedef TensorBlock<Scalar, StorageIndex, NumDims, Layout> Block; | |
| typedef DSizes<StorageIndex, NumDims> Dimensions; | |
| TensorSliceBlockMapper(const Dimensions& tensor_dims, | |
| const Dimensions& tensor_slice_offsets, | |
| const Dimensions& tensor_slice_extents, | |
| const Dimensions& block_dim_sizes, | |
| const Dimensions& block_stride_order) | |
| : m_tensor_dimensions(tensor_dims), | |
| m_tensor_slice_offsets(tensor_slice_offsets), | |
| m_tensor_slice_extents(tensor_slice_extents), | |
| m_block_dim_sizes(block_dim_sizes), | |
| m_block_stride_order(block_stride_order), | |
| m_total_block_count(1) { | |
| // Calculate block counts by dimension and total block count. | |
| DSizes<StorageIndex, NumDims> block_count; | |
| for (Index i = 0; i < block_count.rank(); ++i) { | |
| block_count[i] = divup(m_tensor_slice_extents[i], m_block_dim_sizes[i]); | |
| } | |
| m_total_block_count = array_prod(block_count); | |
| // Calculate block strides (used for enumerating blocks). | |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { | |
| m_block_strides[0] = 1; | |
| m_tensor_strides[0] = 1; | |
| for (int i = 1; i < NumDims; ++i) { | |
| m_block_strides[i] = m_block_strides[i - 1] * block_count[i - 1]; | |
| m_tensor_strides[i] = | |
| m_tensor_strides[i - 1] * m_tensor_dimensions[i - 1]; | |
| } | |
| } else { | |
| m_block_strides[NumDims - 1] = 1; | |
| m_tensor_strides[NumDims - 1] = 1; | |
| for (int i = NumDims - 2; i >= 0; --i) { | |
| m_block_strides[i] = m_block_strides[i + 1] * block_count[i + 1]; | |
| m_tensor_strides[i] = | |
| m_tensor_strides[i + 1] * m_tensor_dimensions[i + 1]; | |
| } | |
| } | |
| } | |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Block | |
| GetBlockForIndex(StorageIndex block_index, Scalar* data) const { | |
| StorageIndex first_coeff_index = 0; | |
| DSizes<StorageIndex, NumDims> coords; | |
| DSizes<StorageIndex, NumDims> sizes; | |
| DSizes<StorageIndex, NumDims> strides; | |
| if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) { | |
| for (int i = NumDims - 1; i > 0; --i) { | |
| const Index idx = block_index / m_block_strides[i]; | |
| coords[i] = m_tensor_slice_offsets[i] + idx * m_block_dim_sizes[i]; | |
| sizes[i] = numext::mini( | |
| m_tensor_slice_offsets[i] + m_tensor_slice_extents[i] - coords[i], | |
| m_block_dim_sizes[i]); | |
| block_index -= idx * m_block_strides[i]; | |
| first_coeff_index += coords[i] * m_tensor_strides[i]; | |
| } | |
| coords[0] = | |
| m_tensor_slice_offsets[0] + block_index * m_block_dim_sizes[0]; | |
| sizes[0] = numext::mini( | |
| m_tensor_slice_offsets[0] + m_tensor_slice_extents[0] - coords[0], | |
| m_block_dim_sizes[0]); | |
| first_coeff_index += coords[0] * m_tensor_strides[0]; | |
| StorageIndex prev_dim = m_block_stride_order[0]; | |
| strides[prev_dim] = 1; | |
| for (int i = 1; i < NumDims; ++i) { | |
| const StorageIndex curr_dim = m_block_stride_order[i]; | |
| strides[curr_dim] = strides[prev_dim] * sizes[prev_dim]; | |
| prev_dim = curr_dim; | |
| } | |
| } else { | |
| for (int i = 0; i < NumDims - 1; ++i) { | |
| const StorageIndex idx = block_index / m_block_strides[i]; | |
| coords[i] = m_tensor_slice_offsets[i] + idx * m_block_dim_sizes[i]; | |
| sizes[i] = numext::mini( | |
| m_tensor_slice_offsets[i] + m_tensor_slice_extents[i] - coords[i], | |
| m_block_dim_sizes[i]); | |
| block_index -= idx * m_block_strides[i]; | |
| first_coeff_index += coords[i] * m_tensor_strides[i]; | |
| } | |
| coords[NumDims - 1] = m_tensor_slice_offsets[NumDims - 1] + | |
| block_index * m_block_dim_sizes[NumDims - 1]; | |
| sizes[NumDims - 1] = numext::mini( | |
| m_tensor_slice_offsets[NumDims - 1] + | |
| m_tensor_slice_extents[NumDims - 1] - coords[NumDims - 1], | |
| m_block_dim_sizes[NumDims - 1]); | |
| first_coeff_index += coords[NumDims - 1] * m_tensor_strides[NumDims - 1]; | |
| StorageIndex prev_dim = m_block_stride_order[NumDims - 1]; | |
| strides[prev_dim] = 1; | |
| for (int i = NumDims - 2; i >= 0; --i) { | |
| const StorageIndex curr_dim = m_block_stride_order[i]; | |
| strides[curr_dim] = strides[prev_dim] * sizes[prev_dim]; | |
| prev_dim = curr_dim; | |
| } | |
| } | |
| return Block(first_coeff_index, sizes, strides, m_tensor_strides, data); | |
| } | |
| EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE StorageIndex total_block_count() const { | |
| return m_total_block_count; | |
| } | |
| private: | |
| Dimensions m_tensor_dimensions; | |
| Dimensions m_tensor_slice_offsets; | |
| Dimensions m_tensor_slice_extents; | |
| Dimensions m_tensor_strides; | |
| Dimensions m_block_dim_sizes; | |
| Dimensions m_block_stride_order; | |
| Dimensions m_block_strides; | |
| StorageIndex m_total_block_count; | |
| }; | |
| } // namespace internal | |
| } // namespace Eigen | |
| #endif // EIGEN_CXX11_TENSOR_TENSOR_BLOCK_H |