Vectorize TensorReverse / TensorRoll packet() inner-slice fast path
libeigen/eigen!2478 Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>
This commit is contained in:
co-authored by
Rasmus Munk Larsen
parent
8a910e308a
commit
8290a5529b
@@ -169,6 +169,46 @@ static void test_expr_reverse(bool LValue) {
|
||||
}
|
||||
}
|
||||
|
||||
// Verify that the rvalue evaluator's packet() returns the same lanes as
|
||||
// coeff() at every aligned and unaligned packet offset. This guards against
|
||||
// regressions in the packet implementation that the executor-level tests
|
||||
// (which only compare the assembled result) would not surface.
|
||||
template <int DataLayout>
|
||||
static void test_packet_reverse() {
|
||||
using namespace Eigen::internal;
|
||||
|
||||
Tensor<float, 3, DataLayout> tensor(8, 5, 7);
|
||||
tensor.setRandom();
|
||||
|
||||
array<bool, 3> dim_rev_inner =
|
||||
(DataLayout == ColMajor) ? array<bool, 3>{{true, false, false}} : array<bool, 3>{{false, false, true}};
|
||||
array<bool, 3> dim_rev_outer =
|
||||
(DataLayout == ColMajor) ? array<bool, 3>{{false, false, true}} : array<bool, 3>{{true, false, false}};
|
||||
array<bool, 3> dim_rev_all{{true, true, true}};
|
||||
|
||||
for (const auto& dim_rev : {dim_rev_inner, dim_rev_outer, dim_rev_all}) {
|
||||
auto expr = tensor.reverse(dim_rev);
|
||||
using Eval = TensorEvaluator<const decltype(expr), DefaultDevice>;
|
||||
using Packet = typename Eval::PacketReturnType;
|
||||
constexpr int PacketSize = Eval::PacketSize;
|
||||
|
||||
DefaultDevice device;
|
||||
Eval eval(expr, device);
|
||||
eval.evalSubExprsIfNeeded(nullptr);
|
||||
|
||||
const Index total = tensor.size();
|
||||
EIGEN_ALIGN_MAX float lanes[PacketSize];
|
||||
for (Index offset = 0; offset + PacketSize <= total; ++offset) {
|
||||
Packet p = eval.template packet<Unaligned>(offset);
|
||||
pstoreu(lanes, p);
|
||||
for (int i = 0; i < PacketSize; ++i) {
|
||||
VERIFY_IS_EQUAL(lanes[i], eval.coeff(offset + i));
|
||||
}
|
||||
}
|
||||
eval.cleanup();
|
||||
}
|
||||
}
|
||||
|
||||
EIGEN_DECLARE_TEST(tensor_reverse) {
|
||||
CALL_SUBTEST(test_simple_reverse<ColMajor>());
|
||||
CALL_SUBTEST(test_simple_reverse<RowMajor>());
|
||||
@@ -176,4 +216,6 @@ EIGEN_DECLARE_TEST(tensor_reverse) {
|
||||
CALL_SUBTEST(test_expr_reverse<RowMajor>(true));
|
||||
CALL_SUBTEST(test_expr_reverse<ColMajor>(false));
|
||||
CALL_SUBTEST(test_expr_reverse<RowMajor>(false));
|
||||
CALL_SUBTEST(test_packet_reverse<ColMajor>());
|
||||
CALL_SUBTEST(test_packet_reverse<RowMajor>());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user