Advice on Nbabel benchmark with Mojo?

[paugier]

I published in 2021 a letter in Nature Astro about "Reducing the ecological impact of computing through education and Python compilers" (https://www.nature.com/articles/s41550-021-01342-y, article available here).

This study is based on a simple benchmark on the Nbody problem (code here). I tried Mojo on this problem with 2 very simple implementations. I think nbabel/mojo/bench_no_part.mojo is buggy but I didn't spend too much time on this. The performance is good but not so good especially compared to https://github.com/paugier/nbabel/blob/main/julia/nbabel5_serial.jl (in Julia) which is significantly faster (something like 30%).

I don't know how to profile Mojo code but for other implementations in other languages, most of the time is spent in the accelerate function which is currently implemented in nbabel/mojo/bench.mojo as

alias Vec4floats = SIMD[DType.float64, 4]
alias vec4zeros = Vec4floats(0)
alias VecParticles = InlinedFixedVector[Particle, 4]


@register_passable("trivial")
struct Particle:
    var position: Vec4floats
    var velocity: Vec4floats
    var acceleration: Vec4floats
    var acceleration1: Vec4floats
    var mass: Float64

    ...

fn accelerate(inout particles: VecParticles) -> NoneType:
    for idx in range(len(particles)):
        var particle = particles[idx]
        particle.acceleration1 = particle.acceleration
        particle.acceleration = Vec4floats(0)
        particles[idx] = particle

    let nb_particules = len(particles)
    for i0 in range(nb_particules - 1):
        var p0 = particles[i0]
        for i1 in range(i0 + 1, nb_particules):
            var p1 = particles[i1]
            let delta = p0.position - p1.position
            let distance_cube = norm_cube(delta)
            p0.acceleration -= p1.mass / distance_cube * delta
            p1.acceleration += p0.mass / distance_cube * delta

            particles[i0] = p0
            particles[i1] = p1

I was wondering if this code is reasonable (first time I touch Mojo) and if I could have some advice to improve it (and potentially get better performance) ?

[paugier]

See https://discord.com/channels/1087530497313357884/1179347946978951179