12b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// This file is part of Eigen, a lightweight C++ template library 22b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// for linear algebra. 32b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// 42b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com> 52b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// 62b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// This Source Code Form is subject to the terms of the Mozilla 72b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// Public License v. 2.0. If a copy of the MPL was not distributed 82b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 92b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 102b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang#ifndef EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_BLOCKING_H 112b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang#define EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_BLOCKING_H 122b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 132b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 142b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wangnamespace Eigen { 152b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wangnamespace internal { 162b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 172b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wangenum { 182b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang ShardByRow = 0, 192b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang ShardByCol = 1 202b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang}; 212b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 222b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 232b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang// Default Blocking Strategy 242b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wangtemplate <typename LhsMapper, typename RhsMapper, typename Index, int ShardingType=ShardByCol> 252b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wangclass TensorContractionBlocking { 262b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang public: 272b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 282b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang typedef typename LhsMapper::Scalar LhsScalar; 292b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang typedef typename RhsMapper::Scalar RhsScalar; 302b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 312b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang EIGEN_DEVICE_FUNC TensorContractionBlocking(Index k, Index m, Index n, Index num_threads = 1) : 322b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang kc_(k), mc_(m), nc_(n) 332b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang { 342b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang if (ShardingType == ShardByCol) { 352b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang computeProductBlockingSizes<LhsScalar, RhsScalar, 1>(kc_, mc_, nc_, num_threads); 362b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang } 372b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang else { 382b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang computeProductBlockingSizes<LhsScalar, RhsScalar, 1>(kc_, nc_, mc_, num_threads); 392b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang } 402b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang } 412b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 422b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Index kc() const { return kc_; } 432b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Index mc() const { return mc_; } 442b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Index nc() const { return nc_; } 452b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 462b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang private: 472b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang Index kc_; 482b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang Index mc_; 492b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang Index nc_; 502b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang}; 512b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 522b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 532b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang} // end namespace internal 542b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang} // end namespace Eigen 552b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang 562b8756b6f1de65d3f8bffab45be6c44ceb7411fcMiao Wang#endif // EIGEN_CXX11_TENSOR_TENSOR_CONTRACTION_BLOCKING_H 57