src/Tensor/TensorContractionBlocking.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_CONTRACTION_BLOCKING_H
0011 #define EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_BLOCKING_H
0012
0013
0014 namespace Eigen {
0015 namespace internal {
0016
0017 enum {
0018   ShardByRow = 0,
0019   ShardByCol = 1
0020 };
0021
0022
0023 // Default Blocking Strategy
0024 template<typename ResScalar, typename LhsScalar, typename RhsScalar, typename StorageIndex, int ShardingType = ShardByCol>
0025 class TensorContractionBlocking {
0026  public:
0027
0028  /*
0029    adding EIGEN_DEVICE_FUNC unconditionally to 'TensorContractionBlocking' constructor in `TensorContractionBlocking.h`
0030      requires adding EIGEN_DEVICE_FUNC to `computeProductBlockingSizes` in `GeneralBlockPanelKernel.h`
0031      which in turn, requires adding EIGEN_DEVICE_FUNC to `evaluateProductBlockingSizesHeuristic` in `GeneralBlockPanelKernel.h`
0032      which in turn, requires adding EIGEN_DEVICE_FUNC to `manage_caching_sizes` in `GeneralBlockPanelKernel.h`
0033      (else HIPCC will error out)
0034
0035    However adding EIGEN_DEVICE_FUNC to `manage_caching_sizes` in `GeneralBlockPanelKernel.h`
0036    results in NVCC erroring out with the following error
0037
0038    ../Eigen/src/Core/products/GeneralBlockPanelKernel.h(57): error #2901:
0039       dynamic initialization is not supported for function-scope static variables within a __device__/__global__ function
0040  */
0041
0042   #if !defined(EIGEN_HIPCC)
0043   EIGEN_DEVICE_FUNC
0044   #endif
0045  TensorContractionBlocking(StorageIndex k, StorageIndex m, StorageIndex n, StorageIndex num_threads = 1) :
0046       kc_(k), mc_(m), nc_(n)
0047   {
0048     if (ShardingType == ShardByCol) {
0049       computeProductBlockingSizes<LhsScalar, RhsScalar, 1>(kc_, mc_, nc_, num_threads);
0050     }
0051     else {
0052       computeProductBlockingSizes<LhsScalar, RhsScalar, 1>(kc_, nc_, mc_, num_threads);
0053     }
0054
0055     const int rhs_packet_size = internal::packet_traits<RhsScalar>::size;
0056     kc_ = (rhs_packet_size <= 8 || kc_ <= rhs_packet_size) ?
0057       kc_ : (kc_ / rhs_packet_size) * rhs_packet_size;
0058   }
0059
0060   EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE StorageIndex kc() const { return kc_; }
0061   EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE StorageIndex mc() const { return mc_; }
0062   EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE StorageIndex nc() const { return nc_; }
0063
0064  private:
0065   StorageIndex kc_;
0066   StorageIndex mc_;
0067   StorageIndex nc_;
0068 };
0069
0070 } // end namespace internal
0071 } // end namespace Eigen
0072
0073 #endif // EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_BLOCKING_H