devtools/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2026-04-10 11:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Asynchronous expression evaluation with TensorAsyncDevice

Browse Source
This commit is contained in:
Eugene Zhulenev
2019-08-30 14:49:40 -07:00
parent f6c51d9209
commit 66665e7e76
8 changed files with 414 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
unsupported/Eigen/CXX11/ThreadPool
View File

@@ -43,7 +43,7 @@
#include <mutex>
#include <thread>
#include <functional>
#include <memory>
#include <utility>
#include "src/util/CXX11Meta.h"
#include "src/util/MaxSizeVector.h"

11
unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
View File

@@ -1063,6 +1063,17 @@ class TensorBase : public TensorBase<Derived, ReadOnlyAccessors> {
return TensorDevice<Derived, DeviceType>(dev, derived());
}
#ifdef EIGEN_USE_THREADS
// Select the async device on which to evaluate the expression.
template <typename DeviceType>
typename internal::enable_if<
internal::is_same<DeviceType, ThreadPoolDevice>::value,
TensorAsyncDevice<Derived, DeviceType>>::type
device(const DeviceType& dev, std::function<void()> done) {
return TensorAsyncDevice<Derived, DeviceType>(dev, derived(), std::move(done));
}
#endif // EIGEN_USE_THREADS
protected:
EIGEN_DEVICE_FUNC
EIGEN_STRONG_INLINE Derived& derived() { return *static_cast<Derived*>(this); }

1
unsupported/Eigen/CXX11/src/Tensor/TensorBlock.h
View File

@@ -932,6 +932,7 @@ class TensorBlockMapper {
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)

41
unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h
View File

@@ -63,6 +63,47 @@ template <typename ExpressionType, typename DeviceType> class TensorDevice {
ExpressionType& m_expression;
};
#ifdef EIGEN_USE_THREADS
/** \class TensorAsyncDevice
* \ingroup CXX11_Tensor_Module
*
* \brief Pseudo expression providing an operator = that will evaluate its
* argument asynchronously on the specified device (currently supports only
* ThreadPoolDevice).
*
* Example:
* std::function<void()> done = []() {};
* C.device(EIGEN_THREAD_POOL, std::move(done)) = A + B;
*/
template <typename ExpressionType, typename DeviceType>
class TensorAsyncDevice {
public:
TensorAsyncDevice(const DeviceType& device, ExpressionType& expression,
std::function<void()> 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, DeviceType> 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 DeviceType& m_device;
ExpressionType& m_expression;
std::function<void()> m_done;
};
#endif // EIGEN_USE_THREADS
} // end namespace Eigen
#endif // EIGEN_CXX11_TENSOR_TENSOR_DEVICE_H

271
unsupported/Eigen/CXX11/src/Tensor/TensorExecutor.h
View File

@@ -84,7 +84,7 @@ class TensorExecutor {
EIGEN_DEVICE_FUNC
static EIGEN_STRONG_INLINE void run(const Expression& expr,
const Device& device = Device()) {
const Device& device = Device()) {
TensorEvaluator<Expression, Device> evaluator(expr, device);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(NULL);
if (needs_assign) {
@@ -97,6 +97,14 @@ class TensorExecutor {
}
};
/**
* Default async execution strategy is not implemented. Currently it's only
* available for ThreadPoolDevice (see definition below).
*/
template <typename Expression, typename Device, bool Vectorizable,
bool Tileable>
class TensorAsyncExecutor {};
/**
* Process all the data with a single cpu thread, using vectorized instructions.
*/
@@ -107,8 +115,8 @@ class TensorExecutor<Expression, DefaultDevice, /*Vectorizable*/ true,
typedef typename Expression::Index StorageIndex;
EIGEN_DEVICE_FUNC
static EIGEN_STRONG_INLINE void run(const Expression& expr,
const DefaultDevice& device = DefaultDevice()) {
static EIGEN_STRONG_INLINE void run(
const Expression& expr, const DefaultDevice& device = DefaultDevice()) {
TensorEvaluator<Expression, DefaultDevice> evaluator(expr, device);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(NULL);
if (needs_assign) {
@@ -206,8 +214,81 @@ class TensorExecutor<Expression, DefaultDevice, Vectorizable,
/**
* Multicore strategy: the index space is partitioned and each partition is
* executed on a single core.
*
* (1) TensorExecutor will submit work to the ThreadPoolDevice managed thread
* pool, and will block the caller thread until all tasks are finished.
*
* (2) TensorAsyncExecutor is a non-blocking version, that will submit work to
* the ThreadPoolDevice managed thread pool, and will return immediately.
* It will call 'done' callback after all tasks are finished.
*/
#ifdef EIGEN_USE_THREADS
template <typename TensorBlockMapper>
struct TensorExecutorTilingContext {
typedef typename TensorBlockMapper::Block TensorBlock;
TensorExecutorTilingContext() : buffer(nullptr) {}
TensorExecutorTilingContext(const TensorBlockMapper& b_mapper,
const TensorOpCost& b_cost, void* b_buffer,
size_t b_aligned_size)
: block_mapper(b_mapper),
cost(b_cost),
buffer(b_buffer),
aligned_blocksize(b_aligned_size) {}
template <typename Scalar>
Scalar* GetCurrentThreadBuffer(const ThreadPoolDevice& device) const {
// ThreadPoolDevice::currentThreadId() returns -1 if called from a thread
// not in the thread pool, such as the main thread dispatching Eigen
// expressions.
const int thread_idx = device.currentThreadId();
eigen_assert(thread_idx >= -1 && thread_idx < device.numThreads());
const Index offset = aligned_blocksize * (thread_idx + 1);
return reinterpret_cast<Scalar*>(static_cast<char*>(buffer) + offset);
}
TensorBlockMapper block_mapper; // navigate through blocks
TensorOpCost cost; // cost of computing a single block
void* buffer; // temporary buffer for blocks
size_t aligned_blocksize; // block size after memory alignment
};
// Computes a block evaluation parameters, and allocates temporary memory buffer
// for blocks. See TensorExecutor/TensorAsyncExecutor (Tileable=true) below.
template <typename Evaluator, typename TensorBlockMapper, bool Vectorizable>
TensorExecutorTilingContext<TensorBlockMapper> GetTensorExecutorTilingContext(
const ThreadPoolDevice& device, const Evaluator& evaluator) {
// Prefer blocks skewed toward inner dimension.
TensorBlockShapeType block_shape = kSkewedInnerDims;
Index block_total_size = 0;
// Query expression tree for desired block size/shape.
std::vector<TensorOpResourceRequirements> resources;
evaluator.getResourceRequirements(&resources);
MergeResourceRequirements(resources, &block_shape, &block_total_size);
int num_threads = device.numThreads();
// Estimate minimum block size based on cost.
TensorOpCost cost = evaluator.costPerCoeff(Vectorizable);
double taskSize = TensorCostModel<ThreadPoolDevice>::taskSize(1, cost);
size_t block_size = static_cast<size_t>(1.0 / taskSize);
TensorBlockMapper block_mapper(
typename TensorBlockMapper::Dimensions(evaluator.dimensions()),
block_shape, block_size);
block_size = block_mapper.block_dims_total_size();
const size_t align = numext::maxi(EIGEN_MAX_ALIGN_BYTES, 1);
const size_t aligned_blocksize =
align *
divup<size_t>(block_size * sizeof(typename Evaluator::Scalar), align);
void* buf = device.allocate((num_threads + 1) * aligned_blocksize);
return {block_mapper, cost * block_size, buf, aligned_blocksize};
}
template <typename Evaluator, typename StorageIndex, bool Vectorizable>
struct EvalRange {
static void run(Evaluator* evaluator_in, const StorageIndex firstIdx,
@@ -274,7 +355,7 @@ class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable, Tileable> {
typedef EvalRange<Evaluator, StorageIndex, Vectorizable> EvalRange;
Evaluator evaluator(expr, device);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(NULL);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
if (needs_assign) {
const StorageIndex size = array_prod(evaluator.dimensions());
device.parallelFor(size, evaluator.costPerCoeff(Vectorizable),
@@ -290,18 +371,18 @@ class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable, Tileable> {
template <typename Expression, bool Vectorizable>
class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable, /*Tileable*/ true> {
public:
typedef typename traits<Expression>::Index StorageIndex;
typedef typename traits<Expression>::Scalar Scalar;
typedef typename remove_const<Scalar>::type ScalarNoConst;
typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
typedef typename traits<Expression>::Index StorageIndex;
static const int NumDims = traits<Expression>::NumDimensions;
typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
typedef TensorBlockMapper<ScalarNoConst, StorageIndex, NumDims, Evaluator::Layout> BlockMapper;
typedef TensorExecutorTilingContext<BlockMapper> TilingContext;
static EIGEN_STRONG_INLINE void run(const Expression& expr,
const ThreadPoolDevice& device) {
typedef TensorBlockMapper<ScalarNoConst, StorageIndex, NumDims, Evaluator::Layout> TensorBlockMapper;
Evaluator evaluator(expr, device);
Index total_size = array_prod(evaluator.dimensions());
Index cache_size = device.firstLevelCacheSize() / sizeof(Scalar);
@@ -315,50 +396,152 @@ class TensorExecutor<Expression, ThreadPoolDevice, Vectorizable, /*Tileable*/ tr
return;
}
const bool needs_assign = evaluator.evalSubExprsIfNeeded(NULL);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
if (needs_assign) {
TensorBlockShapeType block_shape = kSkewedInnerDims;
Index block_total_size = 0;
// Query expression tree for desired block size/shape.
std::vector<internal::TensorOpResourceRequirements> resources;
evaluator.getResourceRequirements(&resources);
MergeResourceRequirements(resources, &block_shape, &block_total_size);
int num_threads = device.numThreads();
const TilingContext tiling =
internal::GetTensorExecutorTilingContext<Evaluator, BlockMapper,
Vectorizable>(device, evaluator);
// Estimate minimum block size based on cost.
TensorOpCost cost = evaluator.costPerCoeff(Vectorizable);
double taskSize = TensorCostModel<ThreadPoolDevice>::taskSize(1, cost);
size_t block_size = static_cast<size_t>(1.0 / taskSize);
TensorBlockMapper block_mapper(
typename TensorBlockMapper::Dimensions(evaluator.dimensions()),
block_shape, block_size);
block_size = block_mapper.block_dims_total_size();
const size_t align = numext::maxi(EIGEN_MAX_ALIGN_BYTES, 1);
const size_t aligned_blocksize =
align * divup<size_t>(block_size * sizeof(Scalar), align);
void* buf = device.allocate((num_threads + 1) * aligned_blocksize);
device.parallelFor(
block_mapper.total_block_count(), cost * block_size,
[=, &device, &evaluator, &block_mapper](StorageIndex firstIdx,
StorageIndex lastIdx) {
// currentThreadId() returns -1 if called from a thread not in the
// thread pool, such as the main thread dispatching Eigen
// expressions.
const int thread_idx = device.currentThreadId();
eigen_assert(thread_idx >= -1 && thread_idx < num_threads);
ScalarNoConst* thread_buf = reinterpret_cast<ScalarNoConst*>(
static_cast<char*>(buf) + aligned_blocksize * (thread_idx + 1));
tiling.block_mapper.total_block_count(), tiling.cost,
[=, &device, &evaluator, &tiling](StorageIndex firstIdx,
StorageIndex lastIdx) {
ScalarNoConst* thread_buf =
tiling.template GetCurrentThreadBuffer<ScalarNoConst>(device);
for (StorageIndex i = firstIdx; i < lastIdx; ++i) {
auto block = block_mapper.GetBlockForIndex(i, thread_buf);
auto block = tiling.block_mapper.GetBlockForIndex(i, thread_buf);
evaluator.evalBlock(&block);
}
});
device.deallocate(buf);
device.deallocate(tiling.buffer);
}
evaluator.cleanup();
}
};
template <typename Expression, bool Vectorizable, bool Tileable>
class TensorAsyncExecutor<Expression, ThreadPoolDevice, Vectorizable, Tileable> {
public:
typedef typename Expression::Index StorageIndex;
typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
static EIGEN_STRONG_INLINE void runAsync(const Expression& expr,
const ThreadPoolDevice& device,
std::function<void()> done) {
TensorAsyncExecutorContext* const ctx =
new TensorAsyncExecutorContext(expr, device, std::move(done));
// TODO(ezhulenev): This is a potentially blocking operation. Make it async!
const bool needs_assign = ctx->evaluator.evalSubExprsIfNeeded(nullptr);
typedef EvalRange<Evaluator, StorageIndex, Vectorizable> EvalRange;
if (needs_assign) {
const StorageIndex size = array_prod(ctx->evaluator.dimensions());
device.parallelForAsync(
size, ctx->evaluator.costPerCoeff(Vectorizable),
EvalRange::alignBlockSize,
[ctx](StorageIndex firstIdx, StorageIndex lastIdx) {
EvalRange::run(&ctx->evaluator, firstIdx, lastIdx);
},
[ctx]() { delete ctx; });
}
}
private:
struct TensorAsyncExecutorContext {
TensorAsyncExecutorContext(const Expression& expr,
const ThreadPoolDevice& thread_pool,
std::function<void()> done)
: evaluator(expr, thread_pool), on_done(std::move(done)) {}
~TensorAsyncExecutorContext() {
on_done();
evaluator.cleanup();
}
Evaluator evaluator;
private:
std::function<void()> on_done;
};
};
template <typename Expression, bool Vectorizable>
class TensorAsyncExecutor<Expression, ThreadPoolDevice, Vectorizable, /*Tileable*/ true> {
public:
typedef typename traits<Expression>::Index StorageIndex;
typedef typename traits<Expression>::Scalar Scalar;
typedef typename remove_const<Scalar>::type ScalarNoConst;
static const int NumDims = traits<Expression>::NumDimensions;
typedef TensorEvaluator<Expression, ThreadPoolDevice> Evaluator;
typedef TensorBlockMapper<ScalarNoConst, StorageIndex, NumDims, Evaluator::Layout> BlockMapper;
typedef TensorExecutorTilingContext<BlockMapper> TilingContext;
static EIGEN_STRONG_INLINE void runAsync(const Expression& expr,
const ThreadPoolDevice& device,
std::function<void()> done) {
TensorAsyncExecutorContext* const ctx =
new TensorAsyncExecutorContext(expr, device, std::move(done));
Index total_size = array_prod(ctx->evaluator.dimensions());
Index cache_size = device.firstLevelCacheSize() / sizeof(Scalar);
if (total_size < cache_size &&
!ExpressionHasTensorBroadcastingOp<Expression>::value) {
internal::TensorAsyncExecutor<Expression, ThreadPoolDevice, Vectorizable,
/*Tileable*/ false>::runAsync(
expr, device, [ctx]() { delete ctx; });
return;
}
// TODO(ezhulenev): This is a potentially blocking operation. Make it async!
const bool needs_assign = ctx->evaluator.evalSubExprsIfNeeded(nullptr);
if (needs_assign) {
ctx->tiling =
internal::GetTensorExecutorTilingContext<Evaluator, BlockMapper,
Vectorizable>(device, ctx->evaluator);
device.parallelForAsync(
ctx->tiling.block_mapper.total_block_count(), ctx->tiling.cost,
[ctx](StorageIndex firstIdx, StorageIndex lastIdx) {
ScalarNoConst* thread_buf =
ctx->tiling.template GetCurrentThreadBuffer<ScalarNoConst>(ctx->device);
for (StorageIndex i = firstIdx; i < lastIdx; ++i) {
auto block = ctx->tiling.block_mapper.GetBlockForIndex(i, thread_buf);
ctx->evaluator.evalBlock(&block);
}
},
[ctx]() { delete ctx; });
}
}
private:
struct TensorAsyncExecutorContext {
TensorAsyncExecutorContext(const Expression& expr,
const ThreadPoolDevice& thread_pool,
std::function<void()> done)
: device(thread_pool),
evaluator(expr, thread_pool),
on_done(std::move(done)) {}
~TensorAsyncExecutorContext() {
on_done();
device.deallocate(tiling.buffer);
evaluator.cleanup();
}
const ThreadPoolDevice& device;
Evaluator evaluator;
TilingContext tiling;
private:
std::function<void()> on_done;
};
};
#endif // EIGEN_USE_THREADS
@@ -419,7 +602,7 @@ template <typename Expression, bool Vectorizable, bool Tileable>
EIGEN_STRONG_INLINE void TensorExecutor<Expression, GpuDevice, Vectorizable, Tileable>::run(
const Expression& expr, const GpuDevice& device) {
TensorEvaluator<Expression, GpuDevice> evaluator(expr, device);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(NULL);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
if (needs_assign) {
const int block_size = device.maxGpuThreadsPerBlock();
@@ -517,10 +700,10 @@ struct ExecExprFunctorKernel<Expr, false, Evaluator>
template <typename Expression, bool Vectorizable, bool Tileable>
class TensorExecutor<Expression, Eigen::SyclDevice, Vectorizable, Tileable> {
public:
typedef typename Expression::Index Index;
typedef typename Expression::Index Index;
static EIGEN_STRONG_INLINE void run(const Expression &expr, const Eigen::SyclDevice &dev) {
Eigen::TensorEvaluator<Expression, Eigen::SyclDevice> evaluator(expr, dev);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(NULL);
const bool needs_assign = evaluator.evalSubExprsIfNeeded(nullptr);
if (needs_assign) {
Index range, GRange, tileSize;
Index total_size = ::Eigen::internal::array_prod(evaluator.dimensions());

6
unsupported/Eigen/CXX11/src/Tensor/TensorForwardDeclarations.h
View File

@@ -94,6 +94,7 @@ template<typename XprType, template <class> class MakePointer_ = MakePointer> cl
template<typename XprType> class TensorForcedEvalOp;
template<typename ExpressionType, typename DeviceType> class TensorDevice;
template<typename ExpressionType, typename DeviceType> class TensorAsyncDevice;
template<typename Derived, typename Device> struct TensorEvaluator;
struct NoOpOutputKernel;
@@ -167,6 +168,11 @@ template <typename Expression, typename Device,
bool Tileable = IsTileable<Device, Expression>::value>
class TensorExecutor;
template <typename Expression, typename Device,
bool Vectorizable = IsVectorizable<Device, Expression>::value,
bool Tileable = IsTileable<Device, Expression>::value>
class TensorAsyncExecutor;
} // end namespace internal
} // end namespace Eigen