mirror of
https://gitlab.com/libeigen/eigen.git
synced 2026-04-10 11:34:33 +08:00
Add block evaluation support to TensorOps
This commit is contained in:
Eugene Zhulenev
@@ -108,16 +108,29 @@ struct TensorEvaluator<const TensorBroadcastingOp<Broadcast, ArgType>, Device>
|
||||
bool isCopy= false, nByOne = false, oneByN = false;
|
||||
|
||||
enum {
|
||||
IsAligned = true,
|
||||
IsAligned = true,
|
||||
PacketAccess = TensorEvaluator<ArgType, Device>::PacketAccess,
|
||||
BlockAccess = false,
|
||||
Layout = TensorEvaluator<ArgType, Device>::Layout,
|
||||
RawAccess = false
|
||||
BlockAccess = TensorEvaluator<ArgType, Device>::BlockAccess,
|
||||
Layout = TensorEvaluator<ArgType, Device>::Layout,
|
||||
RawAccess = false
|
||||
};
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
|
||||
: m_broadcast(op.broadcast()),m_impl(op.expression(), device)
|
||||
{
|
||||
using ScalarNoConst = typename internal::remove_const<Scalar>::type;
|
||||
|
||||
// Block based access to the XprType (input) tensor.
|
||||
using TensorBlock = internal::TensorBlock<ScalarNoConst, Index, NumDims, Layout>;
|
||||
using TensorBlockReader = internal::TensorBlockReader<ScalarNoConst, Index, NumDims, Layout>;
|
||||
// We do block based broadcasting using a a trick with 2x tensor rank and 0
|
||||
// strides. See block method implementation for details.
|
||||
using BroadcastDimensions = DSizes<Index, 2 * NumDims>;
|
||||
using BroadcastTensorBlock = internal::TensorBlock<ScalarNoConst, Index, 2 * NumDims, Layout>;
|
||||
using BroadcastTensorBlockReader = internal::TensorBlockReader<ScalarNoConst, Index, 2 * NumDims, Layout>;
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op,
|
||||
const Device& device)
|
||||
: m_device(device),
|
||||
m_broadcast(op.broadcast()),
|
||||
m_impl(op.expression(), device) {
|
||||
// The broadcasting op doesn't change the rank of the tensor. One can't broadcast a scalar
|
||||
// and store the result in a scalar. Instead one should reshape the scalar into a a N-D
|
||||
// tensor with N >= 1 of 1 element first and then broadcast.
|
||||
@@ -216,8 +229,7 @@ struct TensorEvaluator<const TensorBroadcastingOp<Broadcast, ArgType>, Device>
|
||||
}
|
||||
|
||||
// TODO: attempt to speed this up. The integer divisions and modulo are slow
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeffColMajor(Index index) const
|
||||
{
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index indexColMajor(Index index) const {
|
||||
Index inputIndex = 0;
|
||||
for (int i = NumDims - 1; i > 0; --i) {
|
||||
const Index idx = index / m_outputStrides[i];
|
||||
@@ -243,11 +255,15 @@ struct TensorEvaluator<const TensorBroadcastingOp<Broadcast, ArgType>, Device>
|
||||
inputIndex += (index % m_impl.dimensions()[0]);
|
||||
}
|
||||
}
|
||||
return m_impl.coeff(inputIndex);
|
||||
return inputIndex;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeffRowMajor(Index index) const
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeffColMajor(Index index) const
|
||||
{
|
||||
return m_impl.coeff(indexColMajor(index));
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index indexRowMajor(Index index) const {
|
||||
Index inputIndex = 0;
|
||||
for (int i = 0; i < NumDims - 1; ++i) {
|
||||
const Index idx = index / m_outputStrides[i];
|
||||
@@ -263,17 +279,22 @@ struct TensorEvaluator<const TensorBroadcastingOp<Broadcast, ArgType>, Device>
|
||||
}
|
||||
index -= idx * m_outputStrides[i];
|
||||
}
|
||||
if (internal::index_statically_eq<Broadcast>(NumDims-1, 1)) {
|
||||
eigen_assert(index < m_impl.dimensions()[NumDims-1]);
|
||||
if (internal::index_statically_eq<Broadcast>(NumDims - 1, 1)) {
|
||||
eigen_assert(index < m_impl.dimensions()[NumDims - 1]);
|
||||
inputIndex += index;
|
||||
} else {
|
||||
if (internal::index_statically_eq<InputDimensions>(NumDims-1, 1)) {
|
||||
eigen_assert(index % m_impl.dimensions()[NumDims-1] == 0);
|
||||
if (internal::index_statically_eq<InputDimensions>(NumDims - 1, 1)) {
|
||||
eigen_assert(index % m_impl.dimensions()[NumDims - 1] == 0);
|
||||
} else {
|
||||
inputIndex += (index % m_impl.dimensions()[NumDims-1]);
|
||||
inputIndex += (index % m_impl.dimensions()[NumDims - 1]);
|
||||
}
|
||||
}
|
||||
return m_impl.coeff(inputIndex);
|
||||
return inputIndex;
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeffRowMajor(Index index) const
|
||||
{
|
||||
return m_impl.coeff(indexRowMajor(index));
|
||||
}
|
||||
|
||||
template<int LoadMode>
|
||||
@@ -553,13 +574,291 @@ struct TensorEvaluator<const TensorBroadcastingOp<Broadcast, ArgType>, Device>
|
||||
TensorOpCost(0, 0, compute_cost, vectorized, PacketSize);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void getResourceRequirements(
|
||||
std::vector<internal::TensorOpResourceRequirements>* resources) const {
|
||||
// TODO(wuke): Targeting L1 size is 30% faster than targeting L{-1} on large
|
||||
// tensors. But this might need further tuning.
|
||||
Index l1_cache_scalars = m_device.firstLevelCacheSize() / sizeof(Scalar);
|
||||
Index block_total_size_max = numext::maxi(Index(1), l1_cache_scalars);
|
||||
|
||||
resources->push_back(internal::TensorOpResourceRequirements(
|
||||
internal::TensorBlockShapeType::kSkewedInnerDims,
|
||||
block_total_size_max));
|
||||
|
||||
m_impl.getResourceRequirements(resources);
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void block(
|
||||
TensorBlock* output_block) const {
|
||||
if (NumDims <= 0) {
|
||||
output_block->data()[0] = m_impl.coeff(0);
|
||||
return;
|
||||
}
|
||||
|
||||
// Because we only support kSkewedInnerDims blocking, block size should be
|
||||
// equal to m_dimensions for inner dims, a smaller than m_dimensions[i] size
|
||||
// for the first outer dim, and 1 for other outer dims. This is guaranteed
|
||||
// by MergeResourceRequirements() in TensorBlock.h.
|
||||
const auto& output_block_sizes = output_block->block_sizes();
|
||||
const auto& output_block_strides = output_block->block_strides();
|
||||
|
||||
// Find where outer dims start.
|
||||
int outer_dim_start = 0;
|
||||
Index outer_dim_size = 1, inner_dim_size = 1;
|
||||
for (int i = 0; i < NumDims; ++i) {
|
||||
const int dim = static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? i
|
||||
: NumDims - i - 1;
|
||||
if (i > outer_dim_start) {
|
||||
eigen_assert(output_block_sizes[dim] == 1);
|
||||
} else if (output_block_sizes[dim] != m_dimensions[dim]) {
|
||||
eigen_assert(output_block_sizes[dim] < m_dimensions[dim]);
|
||||
outer_dim_size = output_block_sizes[dim];
|
||||
} else {
|
||||
inner_dim_size *= output_block_sizes[dim];
|
||||
++outer_dim_start;
|
||||
}
|
||||
}
|
||||
|
||||
if (inner_dim_size == 0 || outer_dim_size == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const auto& input_dims = m_impl.dimensions();
|
||||
|
||||
// Pre-fill input_block_sizes, broadcast_block_sizes,
|
||||
// broadcast_block_strides, and broadcast_tensor_strides. Later on we will
|
||||
// only modify the outer_dim_start-th dimension on these arrays.
|
||||
|
||||
// Calculate the input block size for looking into the input.
|
||||
Dimensions input_block_sizes;
|
||||
if (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
|
||||
for (int i = 0; i < outer_dim_start; ++i) {
|
||||
input_block_sizes[i] = input_dims[i];
|
||||
}
|
||||
for (int i = outer_dim_start; i < NumDims; ++i) {
|
||||
input_block_sizes[i] = 1;
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < outer_dim_start; ++i) {
|
||||
input_block_sizes[NumDims - i - 1] = input_dims[NumDims - i - 1];
|
||||
}
|
||||
for (int i = outer_dim_start; i < NumDims; ++i) {
|
||||
input_block_sizes[NumDims - i - 1] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Broadcast with the 0-stride trick: Create 1 extra dim for each
|
||||
// broadcast, set the input stride to 0.
|
||||
//
|
||||
// When ColMajor:
|
||||
// - broadcast_block_sizes is [d_0, b_0, d_1, b_1, ...].
|
||||
//
|
||||
// - broadcast_block_strides is [output_block_strides[0],
|
||||
// output_block_strides[0] * d_0,
|
||||
// output_block_strides[1],
|
||||
// output_block_strides[1] * d_1,
|
||||
// ...].
|
||||
//
|
||||
// - broadcast_tensor_strides is [output_block_strides[0],
|
||||
// 0,
|
||||
// output_block_strides[1],
|
||||
// 0,
|
||||
// ...].
|
||||
BroadcastDimensions broadcast_block_sizes, broadcast_block_strides,
|
||||
broadcast_tensor_strides;
|
||||
|
||||
for (int i = 0; i < outer_dim_start; ++i) {
|
||||
const int dim = static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? i
|
||||
: NumDims - i - 1;
|
||||
const int copy_dim =
|
||||
static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? 2 * i
|
||||
: 2 * NumDims - 2 * i - 1;
|
||||
const int broadcast_dim =
|
||||
static_cast<int>(Layout) == static_cast<int>(ColMajor) ? copy_dim + 1
|
||||
: copy_dim - 1;
|
||||
broadcast_block_sizes[copy_dim] = input_dims[dim];
|
||||
broadcast_block_sizes[broadcast_dim] = m_broadcast[dim];
|
||||
broadcast_block_strides[copy_dim] = output_block_strides[dim];
|
||||
broadcast_block_strides[broadcast_dim] =
|
||||
output_block_strides[dim] * input_dims[dim];
|
||||
broadcast_tensor_strides[copy_dim] = m_inputStrides[dim];
|
||||
broadcast_tensor_strides[broadcast_dim] = 0;
|
||||
}
|
||||
for (int i = 2 * outer_dim_start; i < 2 * NumDims; ++i) {
|
||||
const int dim = static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? i
|
||||
: 2 * NumDims - i - 1;
|
||||
broadcast_block_sizes[dim] = 1;
|
||||
broadcast_block_strides[dim] = 0;
|
||||
broadcast_tensor_strides[dim] = 0;
|
||||
}
|
||||
|
||||
const int outer_dim = static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? outer_dim_start
|
||||
: NumDims - outer_dim_start - 1;
|
||||
|
||||
if (outer_dim_size == 1) {
|
||||
// We just need one block read using the ready-set values above.
|
||||
BroadcastBlock(input_block_sizes, broadcast_block_sizes,
|
||||
broadcast_block_strides, broadcast_tensor_strides, 0,
|
||||
output_block);
|
||||
} else if (input_dims[outer_dim] == 1) {
|
||||
// Broadcast outer_dim_start-th dimension (< NumDims) by outer_dim_size.
|
||||
const int broadcast_outer_dim =
|
||||
static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? 2 * outer_dim_start + 1
|
||||
: 2 * NumDims - 2 * outer_dim_start - 2;
|
||||
broadcast_block_sizes[broadcast_outer_dim] = outer_dim_size;
|
||||
broadcast_tensor_strides[broadcast_outer_dim] = 0;
|
||||
broadcast_block_strides[broadcast_outer_dim] =
|
||||
output_block_strides[outer_dim];
|
||||
BroadcastBlock(input_block_sizes, broadcast_block_sizes,
|
||||
broadcast_block_strides, broadcast_tensor_strides, 0,
|
||||
output_block);
|
||||
} else {
|
||||
// The general case. Let's denote the output block as x[...,
|
||||
// a:a+outer_dim_size, :, ..., :], where a:a+outer_dim_size is a slice on
|
||||
// the outer_dim_start-th dimension (< NumDims). We need to split the
|
||||
// a:a+outer_dim_size into possibly 3 sub-blocks:
|
||||
//
|
||||
// (1) a:b, where b is the smallest multiple of
|
||||
// input_dims[outer_dim_start] in [a, a+outer_dim_size].
|
||||
//
|
||||
// (2) b:c, where c is the largest multiple of input_dims[outer_dim_start]
|
||||
// in [a, a+outer_dim_size].
|
||||
//
|
||||
// (3) c:a+outer_dim_size .
|
||||
//
|
||||
// Or, when b and c do not exist, we just need to process the whole block
|
||||
// together.
|
||||
|
||||
// Find a.
|
||||
const Index outer_dim_left_index =
|
||||
output_block->first_coeff_index() / m_outputStrides[outer_dim];
|
||||
|
||||
// Find b and c.
|
||||
const Index input_outer_dim_size = input_dims[outer_dim];
|
||||
|
||||
// First multiple after a. This is b when <= outer_dim_left_index +
|
||||
// outer_dim_size.
|
||||
const Index first_multiple =
|
||||
divup<Index>(outer_dim_left_index, input_outer_dim_size) *
|
||||
input_outer_dim_size;
|
||||
|
||||
if (first_multiple <= outer_dim_left_index + outer_dim_size) {
|
||||
// b exists, so does c. Find it.
|
||||
const Index last_multiple = (outer_dim_left_index + outer_dim_size) /
|
||||
input_outer_dim_size * input_outer_dim_size;
|
||||
const int copy_outer_dim =
|
||||
static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? 2 * outer_dim_start
|
||||
: 2 * NumDims - 2 * outer_dim_start - 1;
|
||||
const int broadcast_outer_dim =
|
||||
static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? 2 * outer_dim_start + 1
|
||||
: 2 * NumDims - 2 * outer_dim_start - 2;
|
||||
if (first_multiple > outer_dim_left_index) {
|
||||
const Index head_size = first_multiple - outer_dim_left_index;
|
||||
input_block_sizes[outer_dim] = head_size;
|
||||
broadcast_block_sizes[copy_outer_dim] = head_size;
|
||||
broadcast_tensor_strides[copy_outer_dim] = m_inputStrides[outer_dim];
|
||||
broadcast_block_strides[copy_outer_dim] =
|
||||
output_block_strides[outer_dim];
|
||||
broadcast_block_sizes[broadcast_outer_dim] = 1;
|
||||
broadcast_tensor_strides[broadcast_outer_dim] = 0;
|
||||
broadcast_block_strides[broadcast_outer_dim] =
|
||||
output_block_strides[outer_dim] * input_dims[outer_dim];
|
||||
BroadcastBlock(input_block_sizes, broadcast_block_sizes,
|
||||
broadcast_block_strides, broadcast_tensor_strides, 0,
|
||||
output_block);
|
||||
}
|
||||
if (first_multiple < last_multiple) {
|
||||
input_block_sizes[outer_dim] = input_outer_dim_size;
|
||||
broadcast_block_sizes[copy_outer_dim] = input_outer_dim_size;
|
||||
broadcast_tensor_strides[copy_outer_dim] = m_inputStrides[outer_dim];
|
||||
broadcast_block_strides[copy_outer_dim] =
|
||||
output_block_strides[outer_dim];
|
||||
broadcast_block_sizes[broadcast_outer_dim] =
|
||||
(last_multiple - first_multiple) / input_outer_dim_size;
|
||||
broadcast_tensor_strides[broadcast_outer_dim] = 0;
|
||||
broadcast_block_strides[broadcast_outer_dim] =
|
||||
output_block_strides[outer_dim] * input_dims[outer_dim];
|
||||
const Index offset = (first_multiple - outer_dim_left_index) *
|
||||
m_outputStrides[outer_dim];
|
||||
BroadcastBlock(input_block_sizes, broadcast_block_sizes,
|
||||
broadcast_block_strides, broadcast_tensor_strides,
|
||||
offset, output_block);
|
||||
}
|
||||
if (last_multiple < outer_dim_left_index + outer_dim_size) {
|
||||
const Index tail_size =
|
||||
outer_dim_left_index + outer_dim_size - last_multiple;
|
||||
input_block_sizes[outer_dim] = tail_size;
|
||||
broadcast_block_sizes[copy_outer_dim] = tail_size;
|
||||
broadcast_tensor_strides[copy_outer_dim] = m_inputStrides[outer_dim];
|
||||
broadcast_block_strides[copy_outer_dim] =
|
||||
output_block_strides[outer_dim];
|
||||
broadcast_block_sizes[broadcast_outer_dim] = 1;
|
||||
broadcast_tensor_strides[broadcast_outer_dim] = 0;
|
||||
broadcast_block_strides[broadcast_outer_dim] =
|
||||
output_block_strides[outer_dim] * input_dims[outer_dim];
|
||||
const Index offset = (last_multiple - outer_dim_left_index) *
|
||||
m_outputStrides[outer_dim];
|
||||
BroadcastBlock(input_block_sizes, broadcast_block_sizes,
|
||||
broadcast_block_strides, broadcast_tensor_strides,
|
||||
offset, output_block);
|
||||
}
|
||||
} else {
|
||||
// b and c do not exist.
|
||||
const int copy_outer_dim =
|
||||
static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? 2 * outer_dim_start
|
||||
: 2 * NumDims - 2 * outer_dim_start - 1;
|
||||
input_block_sizes[outer_dim] = outer_dim_size;
|
||||
broadcast_block_sizes[copy_outer_dim] = outer_dim_size;
|
||||
broadcast_tensor_strides[copy_outer_dim] = m_inputStrides[outer_dim];
|
||||
broadcast_block_strides[copy_outer_dim] =
|
||||
output_block_strides[outer_dim];
|
||||
BroadcastBlock(input_block_sizes, broadcast_block_sizes,
|
||||
broadcast_block_strides, broadcast_tensor_strides, 0,
|
||||
output_block);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
EIGEN_DEVICE_FUNC typename Eigen::internal::traits<XprType>::PointerType data() const { return NULL; }
|
||||
|
||||
const TensorEvaluator<ArgType, Device>& impl() const { return m_impl; }
|
||||
|
||||
Broadcast functor() const { return m_broadcast; }
|
||||
|
||||
private:
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void BroadcastBlock(
|
||||
const Dimensions& input_block_sizes,
|
||||
const BroadcastDimensions& broadcast_block_sizes,
|
||||
const BroadcastDimensions& broadcast_block_strides,
|
||||
const BroadcastDimensions& broadcast_tensor_strides, Index offset,
|
||||
TensorBlock* output_block) const {
|
||||
TensorBlock input_view_block(
|
||||
static_cast<int>(Layout) == static_cast<int>(ColMajor)
|
||||
? indexColMajor(output_block->first_coeff_index() + offset)
|
||||
: indexRowMajor(output_block->first_coeff_index() + offset),
|
||||
input_block_sizes, Dimensions(m_inputStrides),
|
||||
Dimensions(m_inputStrides), NULL);
|
||||
|
||||
internal::TensorBlockView<ArgType, Device> input_block(m_device, m_impl,
|
||||
input_view_block);
|
||||
BroadcastTensorBlock broadcast_block(
|
||||
0, broadcast_block_sizes, broadcast_block_strides,
|
||||
broadcast_tensor_strides, output_block->data() + offset);
|
||||
|
||||
BroadcastTensorBlockReader::Run(&broadcast_block, input_block.data());
|
||||
}
|
||||
|
||||
protected:
|
||||
const Device& m_device;
|
||||
const Broadcast m_broadcast;
|
||||
Dimensions m_dimensions;
|
||||
array<Index, NumDims> m_outputStrides;
|
||||
|
||||
Reference in New Issue
Block a user