Add block evaluation to TensorEvalTo and fix few small bugs

This commit is contained in:
Eugene Zhulenev
2019-10-07 15:34:26 -07:00
parent 3afb640b56
commit f74ab8cb8d
4 changed files with 95 additions and 33 deletions

View File

@@ -231,7 +231,11 @@ struct TensorEvaluator<const TensorPaddingOp<PaddingDimensions, ArgType>, Device
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlockV2
blockV2(TensorBlockDesc& desc, TensorBlockScratch& scratch) const {
eigen_assert(m_impl.data() != NULL);
// If one of the dimensions is zero, return empty block view.
if (desc.size() == 0) {
return TensorBlockV2(internal::TensorBlockKind::kView, NULL,
desc.dimensions());
}
// Check if we can reuse `desc` destination, or allocate new scratch buffer.
ScalarNoConst* materialized_output =
@@ -385,6 +389,8 @@ struct TensorEvaluator<const TensorPaddingOp<PaddingDimensions, ArgType>, Device
const Index out = output_offset + output_inner_pad_before_size;
const Index in = input_offset + output_inner_pad_before_size;
eigen_assert(output_inner_copy_size == 0 || m_impl.data() != NULL);
LinCopy::template Run<LinCopy::Kind::Linear>(
typename LinCopy::Dst(out, 1, materialized_output),
typename LinCopy::Src(in, 1, m_impl.data()),