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:
Rasmus Munk Larsen
2026-04-28 20:20:19 -07:00
co-authored by Rasmus Munk Larsen
parent 8a910e308a
commit 8290a5529b
7 changed files with 333 additions and 8 deletions
+42
View File
@@ -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>());
}