diff --git a/unsupported/Eigen/src/Tensor/TensorReverse.h b/unsupported/Eigen/src/Tensor/TensorReverse.h index ece85ca1b..158eac4a7 100644 --- a/unsupported/Eigen/src/Tensor/TensorReverse.h +++ b/unsupported/Eigen/src/Tensor/TensorReverse.h @@ -191,15 +191,28 @@ struct TensorEvaluator, Device EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const { eigen_assert(index + PacketSize - 1 < dimensions().TotalSize()); - // TODO(ndjaitly): write a better packing routine that uses - // local structure. + // Fast path: when the whole packet stays inside a single inner-most + // slice of the input, replace PacketSize coeff() calls with one packet + // load (plus a preverse when the inner dim is reversed). + constexpr int inner_dim = (static_cast(Layout) == static_cast(ColMajor)) ? 0 : NumDims - 1; + const Index inner_size = m_dimensions[inner_dim]; + const Index inner_pos = index - (index / inner_size) * inner_size; + if (inner_pos + PacketSize <= inner_size) { + if (m_reverse[inner_dim]) { + const Index input_index = reverseIndex(index + PacketSize - 1); + return internal::preverse(m_impl.template packet(input_index)); + } + return m_impl.template packet(reverseIndex(index)); + } + + // Slow path: the packet crosses an inner-slice boundary, so the + // contiguous-load trick does not apply. EIGEN_ALIGN_MAX std::remove_const_t values[PacketSize]; EIGEN_UNROLL_LOOP for (int i = 0; i < PacketSize; ++i) { values[i] = coeff(index + i); } - PacketReturnType rslt = internal::pload(values); - return rslt; + return internal::pload(values); } EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const { @@ -322,7 +335,9 @@ struct TensorEvaluator, Device compute_cost += 2 * TensorOpCost::AddCost(); } } - return m_impl.costPerCoeff(vectorized) + TensorOpCost(0, 0, compute_cost, false /* vectorized */, PacketSize); + // The inner-slice fast path runs the per-coeff index math once per packet, + // so the amortized compute cost matches the vectorized convention. + return m_impl.costPerCoeff(vectorized) + TensorOpCost(0, 0, compute_cost, vectorized, PacketSize); } EIGEN_DEVICE_FUNC typename Storage::Type data() const { return NULL; } diff --git a/unsupported/Eigen/src/Tensor/TensorRoll.h b/unsupported/Eigen/src/Tensor/TensorRoll.h index e4b5181c0..2184b4d0b 100644 --- a/unsupported/Eigen/src/Tensor/TensorRoll.h +++ b/unsupported/Eigen/src/Tensor/TensorRoll.h @@ -189,13 +189,28 @@ struct TensorEvaluator, Device> { template EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketReturnType packet(Index index) const { eigen_assert(index + PacketSize - 1 < dimensions().TotalSize()); + + // Fast path: when the entire packet stays inside one inner-most slice + // of both the output and the rolled input (no modular wrap on the + // inner dim), the PacketSize coeff() calls collapse to a single + // contiguous packet load from the underlying tensor. + constexpr int inner_dim = (static_cast(Layout) == static_cast(ColMajor)) ? 0 : NumDims - 1; + const Index inner_size = m_dimensions[inner_dim]; + const Index inner_pos = index - (index / inner_size) * inner_size; + if (inner_pos + PacketSize <= inner_size) { + const Index rolled_inner_pos = roll(inner_pos, m_rolls[inner_dim], inner_size); + if (rolled_inner_pos + PacketSize <= inner_size) { + return m_impl.template packet(rollIndex(index)); + } + } + + // Slow path: the packet straddles a slice boundary on either side. EIGEN_ALIGN_MAX std::remove_const_t values[PacketSize]; EIGEN_UNROLL_LOOP for (int i = 0; i < PacketSize; ++i) { values[i] = coeff(index + i); } - PacketReturnType rslt = internal::pload(values); - return rslt; + return internal::pload(values); } EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const { @@ -273,7 +288,9 @@ struct TensorEvaluator, Device> { for (int i = 0; i < NumDims; ++i) { compute_cost += 2 * TensorOpCost::AddCost(); } - return m_impl.costPerCoeff(vectorized) + TensorOpCost(0, 0, compute_cost, false /* vectorized */, PacketSize); + // The inner-slice fast path runs the per-coeff index math once per packet, + // so the amortized compute cost matches the vectorized convention. + return m_impl.costPerCoeff(vectorized) + TensorOpCost(0, 0, compute_cost, vectorized, PacketSize); } EIGEN_DEVICE_FUNC typename Storage::Type data() const { return nullptr; } diff --git a/unsupported/benchmarks/Tensor/CMakeLists.txt b/unsupported/benchmarks/Tensor/CMakeLists.txt index c0a7c9c21..301ce3950 100644 --- a/unsupported/benchmarks/Tensor/CMakeLists.txt +++ b/unsupported/benchmarks/Tensor/CMakeLists.txt @@ -7,3 +7,5 @@ eigen_add_benchmark(bench_tensor_fft bench_tensor_fft.cpp) eigen_add_benchmark(bench_morphing bench_morphing.cpp) eigen_add_benchmark(bench_coefficient_wise bench_coefficient_wise.cpp) eigen_add_benchmark(bench_image_patch bench_image_patch.cpp) +eigen_add_benchmark(bench_reverse bench_reverse.cpp) +eigen_add_benchmark(bench_roll bench_roll.cpp) diff --git a/unsupported/benchmarks/Tensor/bench_reverse.cpp b/unsupported/benchmarks/Tensor/bench_reverse.cpp new file mode 100644 index 000000000..852df6837 --- /dev/null +++ b/unsupported/benchmarks/Tensor/bench_reverse.cpp @@ -0,0 +1,103 @@ +// Benchmarks for Eigen TensorReverse. + +#include +#include + +using namespace Eigen; + +typedef float Scalar; + +// --- Reverse only the inner-most (contiguous) dimension. SIMD preverse case. --- +static void BM_Reverse_Inner(benchmark::State& state) { + const int M = state.range(0); + const int N = state.range(1); + + Tensor A(M, N); + A.setRandom(); + + array dim_rev = {true, false}; + + for (auto _ : state) { + Tensor B = A.reverse(dim_rev); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(M) * N * sizeof(Scalar)); +} + +// --- Reverse only an outer dimension. Inner dim stays contiguous. --- +static void BM_Reverse_Outer(benchmark::State& state) { + const int M = state.range(0); + const int N = state.range(1); + + Tensor A(M, N); + A.setRandom(); + + array dim_rev = {false, true}; + + for (auto _ : state) { + Tensor B = A.reverse(dim_rev); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(M) * N * sizeof(Scalar)); +} + +// --- Reverse every dimension. --- +static void BM_Reverse_All(benchmark::State& state) { + const int M = state.range(0); + const int N = state.range(1); + + Tensor A(M, N); + A.setRandom(); + + array dim_rev = {true, true}; + + for (auto _ : state) { + Tensor B = A.reverse(dim_rev); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(M) * N * sizeof(Scalar)); +} + +// --- 3D reverse with the inner dim reversed (typical CNN-style layout). --- +static void BM_Reverse_3D_Inner(benchmark::State& state) { + const int D0 = state.range(0); + const int D1 = state.range(1); + const int D2 = state.range(2); + + Tensor A(D0, D1, D2); + A.setRandom(); + + array dim_rev = {true, false, false}; + + for (auto _ : state) { + Tensor B = A.reverse(dim_rev); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(D0) * D1 * D2 * sizeof(Scalar)); +} + +// Sweep sizes that span L1 (~32 KB), L2 (~256 KB), and LLC (~MBs) for float +// tensors. Bytes per element = 4, so per-side sizes: +// 64x64 = 16 KB (L1) +// 256x256 = 256 KB (L2) +// 1024x1024 = 4 MB (LLC / DRAM) +static void ReverseSizes(::benchmark::Benchmark* b) { + for (int size : {64, 256, 1024}) { + b->Args({size, size}); + } +} + +static void Reverse3DSizes(::benchmark::Benchmark* b) { + b->Args({32, 32, 32}); // 128 KB + b->Args({64, 64, 64}); // 1 MB + b->Args({128, 128, 128}); // 8 MB +} + +BENCHMARK(BM_Reverse_Inner)->Apply(ReverseSizes); +BENCHMARK(BM_Reverse_Outer)->Apply(ReverseSizes); +BENCHMARK(BM_Reverse_All)->Apply(ReverseSizes); +BENCHMARK(BM_Reverse_3D_Inner)->Apply(Reverse3DSizes); diff --git a/unsupported/benchmarks/Tensor/bench_roll.cpp b/unsupported/benchmarks/Tensor/bench_roll.cpp new file mode 100644 index 000000000..219a8ff86 --- /dev/null +++ b/unsupported/benchmarks/Tensor/bench_roll.cpp @@ -0,0 +1,103 @@ +// Benchmarks for Eigen TensorRoll. + +#include +#include + +using namespace Eigen; + +typedef float Scalar; + +// --- Roll only the inner-most (contiguous) dimension. --- +static void BM_Roll_Inner(benchmark::State& state) { + const int M = state.range(0); + const int N = state.range(1); + const int shift = state.range(2); + + Tensor A(M, N); + A.setRandom(); + + array rolls = {shift, 0}; + + for (auto _ : state) { + Tensor B = A.roll(rolls); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(M) * N * sizeof(Scalar)); +} + +// --- Roll only an outer dimension. Inner dim stays contiguous. --- +static void BM_Roll_Outer(benchmark::State& state) { + const int M = state.range(0); + const int N = state.range(1); + const int shift = state.range(2); + + Tensor A(M, N); + A.setRandom(); + + array rolls = {0, shift}; + + for (auto _ : state) { + Tensor B = A.roll(rolls); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(M) * N * sizeof(Scalar)); +} + +// --- Roll every dimension. --- +static void BM_Roll_All(benchmark::State& state) { + const int M = state.range(0); + const int N = state.range(1); + const int shift = state.range(2); + + Tensor A(M, N); + A.setRandom(); + + array rolls = {shift, shift}; + + for (auto _ : state) { + Tensor B = A.roll(rolls); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(M) * N * sizeof(Scalar)); +} + +// --- 3D roll with the inner dim shifted. --- +static void BM_Roll_3D_Inner(benchmark::State& state) { + const int D0 = state.range(0); + const int D1 = state.range(1); + const int D2 = state.range(2); + + Tensor A(D0, D1, D2); + A.setRandom(); + + array rolls = {D0 / 4, 0, 0}; + + for (auto _ : state) { + Tensor B = A.roll(rolls); + benchmark::DoNotOptimize(B.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(state.iterations() * static_cast(D0) * D1 * D2 * sizeof(Scalar)); +} + +static void RollSizes(::benchmark::Benchmark* b) { + for (int size : {64, 256, 1024}) { + for (int shift : {1, 13}) { + b->Args({size, size, shift}); + } + } +} + +static void Roll3DSizes(::benchmark::Benchmark* b) { + b->Args({32, 32, 32}); + b->Args({64, 64, 64}); + b->Args({128, 128, 128}); +} + +BENCHMARK(BM_Roll_Inner)->Apply(RollSizes); +BENCHMARK(BM_Roll_Outer)->Apply(RollSizes); +BENCHMARK(BM_Roll_All)->Apply(RollSizes); +BENCHMARK(BM_Roll_3D_Inner)->Apply(Roll3DSizes); diff --git a/unsupported/test/tensor_reverse.cpp b/unsupported/test/tensor_reverse.cpp index 37a836a90..d6b5a7cc8 100644 --- a/unsupported/test/tensor_reverse.cpp +++ b/unsupported/test/tensor_reverse.cpp @@ -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 +static void test_packet_reverse() { + using namespace Eigen::internal; + + Tensor tensor(8, 5, 7); + tensor.setRandom(); + + array dim_rev_inner = + (DataLayout == ColMajor) ? array{{true, false, false}} : array{{false, false, true}}; + array dim_rev_outer = + (DataLayout == ColMajor) ? array{{false, false, true}} : array{{true, false, false}}; + array 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; + 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(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()); CALL_SUBTEST(test_simple_reverse()); @@ -176,4 +216,6 @@ EIGEN_DECLARE_TEST(tensor_reverse) { CALL_SUBTEST(test_expr_reverse(true)); CALL_SUBTEST(test_expr_reverse(false)); CALL_SUBTEST(test_expr_reverse(false)); + CALL_SUBTEST(test_packet_reverse()); + CALL_SUBTEST(test_packet_reverse()); } diff --git a/unsupported/test/tensor_roll.cpp b/unsupported/test/tensor_roll.cpp index 0142ba75f..762b48e34 100644 --- a/unsupported/test/tensor_roll.cpp +++ b/unsupported/test/tensor_roll.cpp @@ -146,6 +146,47 @@ static void test_expr_roll(bool LValue) { } } +// Verify that the rvalue evaluator's packet() returns the same lanes as +// coeff() for every offset across a range of shift values, including the +// no-shift, in-slice, and wrap-around cases. +template +static void test_packet_roll() { + using namespace Eigen::internal; + + Tensor tensor(8, 5, 7); + tensor.setRandom(); + + // Cover several shift configurations: zero, in-bounds, and a shift that + // forces a modular wrap-around inside any reasonable packet. + array, 4> roll_configs; + roll_configs[0] = {{0, 0, 0}}; + roll_configs[1] = {{2, 1, 3}}; + roll_configs[2] = {{6, 0, 0}}; // near-end shift on dim 0 + roll_configs[3] = {{0, 0, 6}}; // near-end shift on dim 2 + + for (const auto& rolls : roll_configs) { + auto expr = tensor.roll(rolls); + using Eval = TensorEvaluator; + 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(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_roll) { CALL_SUBTEST(test_simple_roll()); CALL_SUBTEST(test_simple_roll()); @@ -153,4 +194,6 @@ EIGEN_DECLARE_TEST(tensor_roll) { CALL_SUBTEST(test_expr_roll(true)); CALL_SUBTEST(test_expr_roll(false)); CALL_SUBTEST(test_expr_roll(false)); + CALL_SUBTEST(test_packet_roll()); + CALL_SUBTEST(test_packet_roll()); }