chromium/third_party/eigen3/src/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <[email protected]>
//
// 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_DEVICE_H
#define EIGEN_CXX11_TENSOR_TENSOR_DEVICE_H

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

namespace Eigen {

/** \class TensorDevice
 * \ingroup CXX11_Tensor_Module
 *
 * \brief Pseudo expression providing an operator = that will evaluate its argument
 * on the specified computing 'device' (GPU, thread pool, ...)
 *
 * Example:
 *    C.device(EIGEN_GPU) = A + B;
 *
 * Todo: operator *= and /=.
 */

template <typename ExpressionType, typename DeviceType>
class TensorDevice { … };

/** \class TensorAsyncDevice
 * \ingroup CXX11_Tensor_Module
 *
 * \brief Pseudo expression providing an operator = that will evaluate its
 * argument asynchronously on the specified device. Currently only
 * ThreadPoolDevice implements proper asynchronous execution, while the default
 * and GPU devices just run the expression synchronously and call m_done() on
 * completion..
 *
 * Example:
 *    auto done = []() { ... expression evaluation done ... };
 *    C.device(thread_pool_device, std::move(done)) = A + B;
 */

template <typename ExpressionType, typename DeviceType, typename DoneCallback>
class TensorAsyncDevice { … };

#ifdef EIGEN_USE_THREADS
template <typename ExpressionType, typename DoneCallback>
class TensorAsyncDevice<ExpressionType, ThreadPoolDevice, DoneCallback> {
 public:
  TensorAsyncDevice(const ThreadPoolDevice& device, ExpressionType& expression, DoneCallback done)
      : m_device(device), m_expression(expression), m_done(std::move(done)) {}

  template <typename OtherDerived>
  EIGEN_STRONG_INLINE TensorAsyncDevice& operator=(const OtherDerived& other) {
    typedef TensorAssignOp<ExpressionType, const OtherDerived> Assign;
    typedef internal::TensorAsyncExecutor<const Assign, ThreadPoolDevice, DoneCallback> Executor;

    // WARNING: After assignment 'm_done' callback will be in undefined state.
    Assign assign(m_expression, other);
    Executor::runAsync(assign, m_device, std::move(m_done));

    return *this;
  }

 protected:
  const ThreadPoolDevice& m_device;
  ExpressionType& m_expression;
  DoneCallback m_done;
};
#endif

}  // end namespace Eigen

#endif  // EIGEN_CXX11_TENSOR_TENSOR_DEVICE_H