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

Add block evaluation to TensorEvalTo and fix few small bugs

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

46
unsupported/Eigen/CXX11/src/Tensor/TensorEvalTo.h
View File

@@ -111,22 +111,28 @@ struct TensorEvaluator<const TensorEvalToOp<ArgType, MakePointer_>, Device>
IsAligned = TensorEvaluator<ArgType, Device>::IsAligned,
PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
BlockAccess = true,
BlockAccessV2 = false,
BlockAccessV2 = true,
PreferBlockAccess = false,
Layout = TensorEvaluator<ArgType, Device>::Layout,
CoordAccess = false, // to be implemented
RawAccess = true
};
typedef typename internal::TensorBlock<
CoeffReturnType, Index, internal::traits<ArgType>::NumDimensions, Layout>
TensorBlock;
typedef typename internal::TensorBlockReader<
CoeffReturnType, Index, internal::traits<ArgType>::NumDimensions, Layout>
TensorBlockReader;
static const int NumDims = internal::traits<ArgType>::NumDimensions;
typedef typename internal::TensorBlock<CoeffReturnType, Index, NumDims, Layout> TensorBlock;
typedef typename internal::TensorBlockReader<CoeffReturnType, Index, NumDims, Layout> TensorBlockReader;
//===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
typedef internal::TensorBlockNotImplemented TensorBlockV2;
typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
typedef typename TensorEvaluator<const ArgType, Device>::TensorBlockV2
ArgTensorBlock;
typedef internal::TensorBlockAssignment<
Scalar, NumDims, typename ArgTensorBlock::XprType, Index>
TensorBlockAssignment;
//===--------------------------------------------------------------------===//
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
@@ -164,6 +170,30 @@ struct TensorEvaluator<const TensorEvalToOp<ArgType, MakePointer_>, Device>
m_impl.block(&eval_to_block);
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalBlockV2(
TensorBlockDesc& desc, TensorBlockScratch& scratch) {
// Add `m_buffer` as destination buffer to the block descriptor.
desc.AddDestinationBuffer(
/*dst_base=*/m_buffer + desc.offset(),
/*dst_strides=*/internal::strides<Layout>(m_impl.dimensions()),
/*total_dst_bytes=*/
(internal::array_prod(m_impl.dimensions())
* sizeof(Scalar)));
ArgTensorBlock block = m_impl.blockV2(desc, scratch);
// If block was evaluated into a destination buffer, there is no need to do
// an assignment.
if (block.kind() != internal::TensorBlockKind::kMaterializedInOutput) {
TensorBlockAssignment::Run(
TensorBlockAssignment::target(
desc.dimensions(), internal::strides<Layout>(m_impl.dimensions()),
m_buffer, desc.offset()),
block.expr());
}
block.cleanup();
}
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup() {
m_impl.cleanup();
}

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

@@ -238,7 +238,8 @@ class TensorExecutor<Expression, DefaultDevice, Vectorizable,
typedef TensorBlockMapper<ScalarNoConst, StorageIndex, NumDims, Evaluator::Layout> TensorBlockMapper;
typedef typename TensorBlock::Dimensions TensorBlockDimensions;
typedef internal::TensorBlockDescriptor<NumDims> TensorBlockDesc;
typedef internal::TensorBlockDescriptor<NumDims, StorageIndex>
TensorBlockDesc;
typedef internal::TensorBlockScratchAllocator<DefaultDevice>
TensorBlockScratch;

8
unsupported/Eigen/CXX11/src/Tensor/TensorPadding.h
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()),