Add async evaluation support to TensorPadding/TensorImagePatch/TensorShuffling

This commit is contained in:
Eugene Zhulenev 2019-11-26 11:41:57 -08:00
parent c79b6ffe1f
commit bc66c88255
4 changed files with 35 additions and 2 deletions

View File

@ -176,7 +176,8 @@ struct TensorEvaluator<const TensorEvalToOp<ArgType, MakePointer_>, Device>
/*dst_base=*/m_buffer + desc.offset(), /*dst_base=*/m_buffer + desc.offset(),
/*dst_strides=*/internal::strides<Layout>(m_impl.dimensions())); /*dst_strides=*/internal::strides<Layout>(m_impl.dimensions()));
ArgTensorBlock block = m_impl.blockV2(desc, scratch); ArgTensorBlock block =
m_impl.blockV2(desc, scratch, /*root_of_expr_ast=*/true);
// If block was evaluated into a destination buffer, there is no need to do // If block was evaluated into a destination buffer, there is no need to do
// an assignment. // an assignment.

View File

@ -397,6 +397,14 @@ struct TensorEvaluator<const TensorImagePatchOp<Rows, Cols, ArgType>, Device>
return true; return true;
} }
#ifdef EIGEN_USE_THREADS
template <typename EvalSubExprsCallback>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(
EvaluatorPointerType, EvalSubExprsCallback done) {
m_impl.evalSubExprsIfNeededAsync(nullptr, [done](bool) { done(true); });
}
#endif // EIGEN_USE_THREADS
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() { EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
m_impl.cleanup(); m_impl.cleanup();
} }

View File

@ -155,6 +155,15 @@ struct TensorEvaluator<const TensorPaddingOp<PaddingDimensions, ArgType>, Device
m_impl.evalSubExprsIfNeeded(NULL); m_impl.evalSubExprsIfNeeded(NULL);
return true; return true;
} }
#ifdef EIGEN_USE_THREADS
template <typename EvalSubExprsCallback>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(
EvaluatorPointerType, EvalSubExprsCallback done) {
m_impl.evalSubExprsIfNeededAsync(nullptr, [done](bool) { done(true); });
}
#endif // EIGEN_USE_THREADS
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() { EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
m_impl.cleanup(); m_impl.cleanup();
} }

View File

@ -182,6 +182,15 @@ struct TensorEvaluator<const TensorShufflingOp<Shuffle, ArgType>, Device>
m_impl.evalSubExprsIfNeeded(NULL); m_impl.evalSubExprsIfNeeded(NULL);
return true; return true;
} }
#ifdef EIGEN_USE_THREADS
template <typename EvalSubExprsCallback>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(
EvaluatorPointerType, EvalSubExprsCallback done) {
m_impl.evalSubExprsIfNeededAsync(nullptr, [done](bool) { done(true); });
}
#endif // EIGEN_USE_THREADS
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() { EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
m_impl.cleanup(); m_impl.cleanup();
} }
@ -237,10 +246,16 @@ struct TensorEvaluator<const TensorShufflingOp<Shuffle, ArgType>, Device>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void getResourceRequirements( EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void getResourceRequirements(
std::vector<internal::TensorOpResourceRequirements>* resources) const { std::vector<internal::TensorOpResourceRequirements>* resources) const {
static const int inner_dim =
Layout == static_cast<int>(ColMajor) ? 0 : NumDims - 1;
const bool inner_dim_shuffled = m_shuffle[inner_dim] != inner_dim;
Eigen::Index block_total_size_max = numext::maxi<Eigen::Index>( Eigen::Index block_total_size_max = numext::maxi<Eigen::Index>(
1, m_device.firstLevelCacheSize() / sizeof(Scalar)); 1, m_device.firstLevelCacheSize() / sizeof(Scalar));
resources->push_back(internal::TensorOpResourceRequirements( resources->push_back(internal::TensorOpResourceRequirements(
internal::kUniformAllDims, block_total_size_max)); inner_dim_shuffled ? internal::kUniformAllDims
: internal::kSkewedInnerDims,
block_total_size_max));
} }
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlockV2 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlockV2