mirror of
https://gitlab.com/libeigen/eigen.git
synced 2026-04-10 11:34:33 +08:00
Add block evaluation V2 to TensorAsyncExecutor.
Add async evaluation to a number of ops.
This commit is contained in:
@@ -689,15 +689,14 @@ struct TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, M
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_dimensions; }
|
||||
|
||||
EIGEN_STRONG_INLINE
|
||||
#if !defined(EIGEN_HIPCC)
|
||||
// Marking this as EIGEN_DEVICE_FUNC for HIPCC requires also doing the same for all the functions
|
||||
// being called within here, which then leads to proliferation of EIGEN_DEVICE_FUNC markings, one
|
||||
// of which will eventually result in an NVCC error
|
||||
EIGEN_DEVICE_FUNC
|
||||
#endif
|
||||
bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
|
||||
m_impl.evalSubExprsIfNeeded(NULL);
|
||||
|
||||
#if !defined(EIGEN_HIPCC)
|
||||
// Marking this as EIGEN_DEVICE_FUNC for HIPCC requires also doing the same
|
||||
// for all the functions being called within here, which then leads to
|
||||
// proliferation of EIGEN_DEVICE_FUNC markings, one of which will eventually
|
||||
// result in an NVCC error
|
||||
EIGEN_DEVICE_FUNC
|
||||
#endif
|
||||
bool evalSubExprsIfNeededCommon(EvaluatorPointerType data) {
|
||||
// Use the FullReducer if possible.
|
||||
if ((RunningFullReduction && RunningOnSycl) ||(RunningFullReduction &&
|
||||
internal::FullReducer<Self, Op, Device>::HasOptimizedImplementation &&
|
||||
@@ -802,6 +801,34 @@ struct TensorReductionEvaluatorBase<const TensorReductionOp<Op, Dims, ArgType, M
|
||||
return true;
|
||||
}
|
||||
|
||||
#ifdef EIGEN_USE_THREADS
|
||||
template <typename EvalSubExprsCallback>
|
||||
EIGEN_STRONG_INLINE
|
||||
#if !defined(EIGEN_HIPCC)
|
||||
EIGEN_DEVICE_FUNC
|
||||
#endif
|
||||
void
|
||||
evalSubExprsIfNeededAsync(EvaluatorPointerType data,
|
||||
EvalSubExprsCallback done) {
|
||||
m_impl.evalSubExprsIfNeededAsync(NULL, [this, data, done](bool) {
|
||||
done(evalSubExprsIfNeededCommon(data));
|
||||
});
|
||||
}
|
||||
#endif
|
||||
|
||||
EIGEN_STRONG_INLINE
|
||||
#if !defined(EIGEN_HIPCC)
|
||||
// Marking this as EIGEN_DEVICE_FUNC for HIPCC requires also doing the same
|
||||
// for all the functions being called within here, which then leads to
|
||||
// proliferation of EIGEN_DEVICE_FUNC markings, one of which will eventually
|
||||
// result in an NVCC error
|
||||
EIGEN_DEVICE_FUNC
|
||||
#endif
|
||||
bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
|
||||
m_impl.evalSubExprsIfNeeded(NULL);
|
||||
return evalSubExprsIfNeededCommon(data);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
|
||||
m_impl.cleanup();
|
||||
if (m_result) {
|
||||
|
||||
Reference in New Issue
Block a user