Refactor standalone n-body codes to be more similar

examples/nbody_code_comp/nbody-AoS-baseline.cpp

@@ -13,43 +13,24 @@ constexpr int steps = 5, problemSize = 64 * 1024;
 
 struct Vec3 {
     FP x, y, z;
-
-    auto operator*=(Vec3 v) -> Vec3& {
-        x *= v.x;
-        y *= v.y;
-        z *= v.z;
-        return *this;
-    }
-
-    auto operator+=(Vec3 v) -> Vec3& {
-        x += v.x;
-        y += v.y;
-        z += v.z;
-        return *this;
-    }
-
-    friend auto operator-(Vec3 a, Vec3 b) -> Vec3 {
-        return Vec3{a.x - b.x, a.y - b.y, a.z - b.z};
-    }
-
-    friend auto operator*(Vec3 a, FP s) -> Vec3 {
-        return Vec3{a.x * s, a.y * s, a.z * s};
-    }
 };
-
 struct Particle {
     Vec3 pos, vel;
     FP mass;
 };
 
 inline void pPInteraction(Particle& pi, const Particle& pj) {
-    auto dist = pi.pos - pj.pos;
-    dist *= dist;
+    auto dist = Vec3{pi.pos.x - pj.pos.x, pi.pos.y - pj.pos.y, pi.pos.z - pj.pos.z};
+    dist.x *= dist.x;
+    dist.y *= dist.y;
+    dist.z *= dist.z;
     const auto distSqr = eps2 + dist.x + dist.y + dist.z;
     const auto distSixth = distSqr * distSqr * distSqr;
     const auto invDistCube = FP{1} / std::sqrt(distSixth);
     const auto sts = pj.mass * timestep * invDistCube;
-    pi.vel += dist * sts;
+    pi.vel.x += dist.x * sts;
+    pi.vel.y += dist.y * sts;
+    pi.vel.z += dist.z * sts;
 }
 
 void update(std::span<Particle> particles) {
@@ -64,8 +45,11 @@ void update(std::span<Particle> particles) {
 
 void move(std::span<Particle> particles) {
 #pragma GCC ivdep
-    for(int i = 0; i < problemSize; i++)
-        particles[i].pos += particles[i].vel * timestep;
+    for(int i = 0; i < problemSize; i++) {
+        particles[i].pos.x += particles[i].vel.x * timestep;
+        particles[i].pos.y += particles[i].vel.y * timestep;
+        particles[i].pos.z += particles[i].vel.z * timestep;
+    }
 }
 
 auto main() -> int {

examples/nbody_code_comp/nbody-AoSoA.cpp

@@ -14,12 +14,13 @@ constexpr int steps = 5, problemSize = 64 * 1024;
 constexpr auto Lanes = 16;
 constexpr auto blocks = problemSize / Lanes;
 
+struct alignas(Lanes * sizeof(FP)) Vec3Block {
+    FP x[Lanes];
+    FP y[Lanes];
+    FP z[Lanes];
+};
 struct alignas(Lanes * sizeof(FP)) ParticleBlock {
-    struct {
-        FP x[Lanes];
-        FP y[Lanes];
-        FP z[Lanes];
-    } pos, vel;
+    Vec3Block pos, vel;
     FP mass[Lanes];
 };
 

examples/nbody_code_comp/nbody-ported.cpp

@@ -12,8 +12,16 @@ constexpr FP timestep = 0.0001f, eps2 = 0.01f;
 constexpr int steps = 5, problemSize = 64 * 1024;
 
 struct Pos{}; struct Vel{}; struct X{}; struct Y{}; struct Z{}; struct Mass{};
-using V3 = llama::Record<llama::Field<X, FP>, llama::Field<Y, FP>, llama::Field<Z, FP>>;
-using Particle = llama::Record<llama::Field<Pos, V3>, llama::Field<Vel, V3>, llama::Field<Mass, FP>>;
+using Vec3 = llama::Record<
+    llama::Field<X, FP>,
+    llama::Field<Y, FP>,
+    llama::Field<Z, FP>
+>;
+using Particle = llama::Record<
+    llama::Field<Pos, Vec3>,
+    llama::Field<Vel, Vec3>,
+    llama::Field<Mass, FP>
+>;
 
 LLAMA_FN_HOST_ACC_INLINE void pPInteraction(auto&& pi, auto&& pj) {
     auto dist = pi(Pos{}) - pj(Pos{});