Add vectorized PrepareDerivatives, test, and benchmark

Author: ProExpertProg

cpp/bench/CMakeLists.txt

@@ -28,6 +28,9 @@ target_link_libraries(bench-hl-filter src-lib benchmark::benchmark_main)
 add_executable(bench-exps exps.cpp)
 target_link_libraries(bench-exps src-lib benchmark::benchmark_main)
 
+add_executable(bench-prepare prepare-derivatives.cpp)
+target_link_libraries(bench-prepare src-lib benchmark::benchmark_main)
+
 if (ENABLE_CILK)
     add_executable(bench-fftw-cilk fftw-cilk.cpp)
     target_link_libraries(bench-fftw-cilk fftw-cpp benchmark::benchmark fftw3_threads)

cpp/bench/prepare-derivatives.cpp

@@ -0,0 +1,35 @@
+#include "PrepareDerivatives.hpp"
+
+#include <benchmark/benchmark.h>
+
+static constexpr auto N_WARMUP_ITERS = 5;
+
+using namespace ahr;
+template <class Prepare> static void BM_Prepare(benchmark::State &state) {
+  Dim const X = state.range(0);
+  Dim const Y = state.range(1);
+
+  Grid grid{1, X, Y};
+  Prepare prepare{grid};
+
+  auto buf = grid.cBufXY();
+  auto buf2 = grid.cBufXY();
+  auto buf3 = grid.cBufXY();
+
+  // Warm-up
+  for (int i = 0; i < N_WARMUP_ITERS; i++) {
+    prepare(buf, {buf2, buf3});
+  }
+
+  for (auto _ : state) {
+    prepare(buf, {buf2, buf3});
+  }
+}
+
+BENCHMARK(BM_Prepare<PrepareDerivatives>)
+    ->ArgsProduct({{2048, 4096, 8192}, {2048, 4096, 8192}})
+    ->Unit(benchmark::kMillisecond);
+
+BENCHMARK(BM_Prepare<PrepareDerivativesVector>)
+    ->ArgsProduct({{2048, 4096, 8192}, {2048, 4096, 8192}})
+    ->Unit(benchmark::kMillisecond);

cpp/lib/PrepareDerivatives.cpp

@@ -1,5 +1,7 @@
 #include "PrepareDerivatives.hpp"
 
+#include <eve/module/core.hpp>
+
 namespace ahr {
 
 void PrepareDerivatives::operator()(View::C_XY const &in, DxDy<View::C_XY> out) const {
@@ -9,4 +11,48 @@ void PrepareDerivatives::operator()(View::C_XY const &in, DxDy<View::C_XY> out)
   });
 }
 
+void PrepareDerivativesVector::operator()(View::C_XY const &in, DxDy<View::C_XY> out) const {
+  eve::logical<VIdx> const even_mask{[](int idx, int) { return idx % 2 == 0; }};
+  VReal const kx_v_init{[](int idx, int) { return Real(idx / 2); }};
+  for (Dim ky = 0; ky < grid.KY; ky += KY_TILE) {
+    // broadcast ky values
+    using TileReal = std::array<VReal, KY_TILE>;
+    TileReal ky_v;
+    for (Dim tile_ky = 0; tile_ky < KY_TILE; ++tile_ky) {
+      ky_v[tile_ky] = ky_(ky + tile_ky);
+    }
+
+    // Initialize kx values
+    VReal kx_v = kx_v_init;
+
+    Dim kx = 0;
+    for (; kx <= grid.KX - C_WIDTH; kx += C_WIDTH, kx_v += C_WIDTH) {
+      auto input = [&](int i) { return (Real *)&in(kx, ky + i); };
+      auto out_dx = [&](int i) { return (Real *)&out.DX(kx, ky + i); };
+      auto out_dy = [&](int i) { return (Real *)&out.DY(kx, ky + i); };
+
+      for (int i = 0; i < KY_TILE; ++i) {
+        auto in_v = VReal{input(i)};
+        // (a + bi) * i -> (-b + ai)
+        // swap real and imaginary parts
+        auto swapped = eve::swap_adjacent(in_v);
+        // selectively negate
+        auto mul_with_i = eve::minus[even_mask](swapped);
+        // normalize
+        auto in_norm = mul_with_i * XYNorm;
+
+        eve::store(kx_v * in_norm, out_dx(i));
+        eve::store(ky_v[i] * in_norm, out_dy(i));
+      }
+    }
+
+    // tail
+    for (; kx < grid.KX; ++kx) {
+      for (int i = 0; i < KY_TILE; ++i) {
+        out.DX(kx, ky + i) = kx_(kx) * 1i * in(kx, ky + i) * XYNorm;
+        out.DY(kx, ky + i) = ky_(ky + i) * 1i * in(kx, ky + i) * XYNorm;
+      }
+    }
+  }
+}
 } // namespace ahr

cpp/lib/PrepareDerivatives.hpp

@@ -2,10 +2,12 @@
 
 #include "constants.hpp"
 #include "grid.hpp"
+#include <eve/wide.hpp>
 
 namespace ahr {
 
 class PrepareDerivatives {
+protected:
   using View = Grid::View;
   template <class T> using DxDy = Grid::DxDy<T>;
 
@@ -14,7 +16,7 @@ class PrepareDerivatives {
 
   void operator()(View::C_XY const &in, DxDy<View::C_XY> out) const;
 
-private:
+protected:
   Grid const &grid;
 
   /// Normalization factor for FFT
@@ -26,4 +28,22 @@ class PrepareDerivatives {
   }
   [[nodiscard]] Real kx_(Dim kx) const { return Real(kx); }
 };
+
+class PrepareDerivativesVector : protected PrepareDerivatives {
+public:
+  explicit PrepareDerivativesVector(Grid const &grid) : PrepareDerivatives(grid) {}
+
+  void operator()(View::C_XY const &in, DxDy<View::C_XY> out) const;
+
+protected:
+  using VReal = eve::wide<Real>;
+
+  // TODO vec_utils
+  using VIdx = eve::wide<long long>;
+  static auto constexpr R_WIDTH = VReal::size();
+  static auto constexpr C_WIDTH = VReal::size() / 2;
+
+  static constexpr auto KY_TILE = 4;
+};
+
 } // namespace ahr

cpp/test/CMakeLists.txt

@@ -23,6 +23,7 @@ add_viriato_test(exps.cpp)
 add_viriato_test(filter.cpp)
 add_viriato_test(grid.cpp)
 add_viriato_test(transformer.cpp)
+add_viriato_test(prepare-derivatives.cpp)
 
 add_viriato_test(naive-energies.cpp)
 add_viriato_test(naive-moments.cpp)

cpp/test/prepare-derivatives.cpp

@@ -0,0 +1,39 @@
+#include "PrepareDerivatives.hpp"
+
+#include "debug.hpp"
+#include "util.hpp"
+
+#include <gtest/gtest.h>
+
+namespace ahr {
+
+template <typename TestedPrepare> class TestPrepareDerivatives : public ::testing::Test {
+protected:
+  Grid grid{5, 32, 32};
+  PrepareDerivatives prepare{grid};
+  TestedPrepare prepare_t{grid};
+};
+
+using Types = ::testing::Types<PrepareDerivativesVector>;
+TYPED_TEST_SUITE(TestPrepareDerivatives, Types);
+
+TYPED_TEST(TestPrepareDerivatives, Prepare) {
+  auto const KX = this->grid.KX, KY = this->grid.KY;
+  auto buf = this->grid.cBufXY(), buf_t = this->grid.cBufXY();
+  Grid::DxDy<Grid::Buf::C_XY> bufD{KX, KY};
+  Grid::DxDy<Grid::Buf::C_XY> bufD_t{KX, KY};
+
+  // initialize buffers
+  this->grid.for_each_kxky([&](Dim kx, Dim ky) {
+    buf(kx, ky) = buf_t(kx, ky) =
+        1024.0 * Complex{std::sin(2 * pi * kx / KX), std::cos(2 * pi * ky / KY)};
+  });
+
+  this->prepare(buf, bufD);
+  this->prepare_t(buf_t, bufD_t);
+
+  ASSERT_THAT(bufD_t.DX.to_mdspan(), MdspanElementsAllClose(bufD.DX.to_mdspan(), 1e-16, 1e-15));
+  ASSERT_THAT(bufD_t.DY.to_mdspan(), MdspanElementsAllClose(bufD.DY.to_mdspan(), 1e-16, 1e-15));
+}
+
+} // namespace ahr