benchmarks game: zig

nbody-zig.zig

//! The Computer Language Benchmarks Game
//! https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
//!
//! This implementation is written in simple and idiomatic Zig.
//! The only optimization is using SIMD vector intrinsics.

const std = @import("std");

/////////////////////////////////////////////////////////////////////////
// important constants.  ALLCAPS constants are not idiomatic Zig, but it's
// easier for C practitioners to read.

const PI: f64 = 3.141592653589793;
const SOLAR_MASS: f64 = 4 * PI * PI;
const DPY: f64 = 365.24;

///////////////////////////////////////////////////////////////////////////
// useful types.  _t is not really idiomatic Zig notation for types but it
// more closely matches C convention, so it's easier for C practitioners to
// read.

/// SIMD float 64 type
const vec_t = @Vector(3, f64);

/// planets or the sun
const body_t = struct{
  pos: vec_t,
  vel: vec_t,
  mass: f64,
};

const planets = [_]body_t{
  .{ //jupiter
    .pos  = .{4.84143144246472090e+0,       -1.16032004402742839e+0,       -1.03622044471123109e-1},
    .vel  = .{1.66007664274403694e-3 * DPY,  7.69901118419740425e-3 * DPY, -6.90460016972063023e-5 * DPY},
    .mass = 9.54791938424326609e-4 * SOLAR_MASS,
  },
  .{ //saturn
    .pos = .{8.34336671824457987e+0,        4.12479856412430479e+0,      -4.03523417114321381e-1},
    .vel = .{-2.76742510726862411e-3 * DPY, 4.99852801234917238e-3 * DPY, 2.30417297573763929e-5 * DPY},
    .mass = 2.85885980666130812e-4 * SOLAR_MASS
  },
  .{ //uranus
    .pos = .{1.28943695621391310e+1,       -1.51111514016986312e+1,       -2.23307578892655734e-1},
    .vel = .{2.96460137564761618e-3 * DPY,  2.37847173959480950e-3 * DPY, -2.96589568540237556e-5 * DPY},
    .mass = 4.36624404335156298e-5 * SOLAR_MASS
  },
  .{ //neptune
    .pos = .{1.53796971148509165e+1,       -2.59193146099879641e+1,        1.79258772950371181e-1},
    .vel = .{2.68067772490389322e-3 * DPY,  1.62824170038242295e-3 * DPY, -9.51592254519715870e-5 * DPY},
    .mass = 5.15138902046611451e-5 * SOLAR_MASS
  }
};

/// calculates the initial solar velocity so that the entire system has zero
/// momentum and therefore a fixed center
fn sun_vel(comptime plnts: []const body_t) comptime vec_t {
  var p : vec_t = .{0.0, 0.0, 0.0};
  for (plnts) |plnt| {
    p = p - plnt.vel * @splat(3, plnt.mass);
  }
  return p / @splat(3, SOLAR_MASS);
}

const sun = [1]body_t{
  .{.pos = .{0.0, 0.0, 0.0},
    .vel = sun_vel(planets[0..]),
    .mass = SOLAR_MASS}
};

// build the runtime list of all bodies using comptime `++` operator.
var bodies: [planets.len + 1]body_t = planets ++ sun;

/// this type is used at comptime to sanely unroll pairs of heavenly
/// bodies.
const pair_t = struct{ l: u8, r: u8 };

// the following could probably be assigned at comptime:
const pairs = [_]pair_t{.{.l=0, .r=1}, .{.l=0, .r=2}, .{.l=0, .r=3},
                        .{.l=0, .r=4}, .{.l=1, .r=2}, .{.l=1, .r=3},
                        .{.l=1, .r=4}, .{.l=2, .r=3}, .{.l=2, .r=4},
                        .{.l=3, .r=4}};

/// basic SIMD dot product
fn dot(v1: vec_t, v2: vec_t) f64 {
  var prod = v1 * v2;
  var res: f64 = 0.0;

  comptime var i = 0;
  inline while (i < 3) : (i += 1) {
    res += prod[i];
  }
  return res;
}

/// positional difference between each pair of heavenly bodies
var delta_pos : [pairs.len]vec_t = undefined;

/// magnitude of gravitational attraction between pairs of heavenly bodies.
var mag : [pairs.len]f64 = undefined;

const builtin = @import("builtin");

/// one turn of the simulation.  Note that this function heavily causes
/// side effects in some global arrays.
fn advance() void {
  // equivalent to -ffastmath
  @setFloatMode(builtin.FloatMode.Optimized);
  inline for (pairs) |pair, idx| {
    delta_pos[idx] = bodies[pair.l].pos - bodies[pair.r].pos;
    var dsq = dot(delta_pos[idx], delta_pos[idx]);
    var rd =  1/@sqrt(dsq);

    // perform the iteration.
    mag[idx] = 0.01 * rd / dsq;
  }

  // update velocities using precalculated magnitudes.
  comptime var p_idx = 0;
  comptime var l_idx = 0;
  var vel_l: vec_t = undefined;
  inline while (l_idx < bodies.len - 1) : (l_idx += 1) {
    // keep track of the left-side velocity separate from the right, which
    // will be dynamically updated.
    vel_l = bodies[l_idx].vel;
    comptime var r_idx = l_idx + 1;
    inline while (r_idx < bodies.len) : (r_idx += 1) {
      vel_l             -= delta_pos[p_idx] * @splat(3, mag[p_idx] * bodies[r_idx].mass);
      bodies[r_idx].vel += delta_pos[p_idx] * @splat(3, mag[p_idx] * bodies[l_idx].mass);
      p_idx += 1;
    }
    bodies[l_idx].vel = vel_l;
  }

  // update body positions
  comptime var idx = 0;
  inline while (idx < bodies.len) : ( idx +=1 ) {
    bodies[idx].pos += bodies[idx].vel * @splat(3, @floatCast(f64, 0.01));
  }
}

//////////////////////////////////////////////////////////////////////////
// ENERGY CALCULATIONS

fn energy() f64 {
  // calculate KE of each body alone.
  var e : f64 = 0.0;
  for (bodies) |b| {
    e += 0.5 * b.mass * dot(b.vel, b.vel);
  }

  // calculate GPE of each body as a pair with the other.
  comptime var i = 0;
  inline for (pairs) |p| {
    var deltap = bodies[p.l].pos - bodies[p.r].pos;
    var dist = @sqrt(dot(deltap, deltap));
    e -= bodies[p.l].mass * bodies[p.r].mass / dist;
  }
  return e;
}

fn print_energy() void {
  std.debug.warn("{d:.9}\n", .{energy()});
}

fn strlen(str: [*:0]u8) u8 {
  var idx:u8 = 0;
  while (str[idx] != 0) : (idx += 1) {}
  return idx;
}

pub fn main() !void {
  if (std.os.argv.len == 2) {
    var iter_str: [*:0] u8 = std.os.argv[1];
    var iter_len = strlen(iter_str);
    var iters = try std.fmt.parseInt(i64, iter_str[0..iter_len], 10);
    print_energy();
    var idx : i64 = 0;
    while (idx < iters) : (idx += 1) { advance(); }
    print_energy();
  } else {
    std.debug.panic("wrong number of arguments!\n", .{});
  }
}