comm.impala

struct Comm {
    me: i32,                    // current process rank
    local_start: i32,           // offset where local particles start
    send_buffer: ArrayData,     // send buffer
    send_neighbors: Buffer,     // neighbor processes in send buffer
    send_rank_offsets: Buffer,  // offsets in send buffer
    send_rank_lengths: Buffer,  // lengths in send buffer
    send_offsets: ArrayData,    // serialized particle offsets
    send_pbc: ArrayData,        // serialized particle PBC flags
    copy_list: ArrayData,       // copy list
    send_capacity: i32,         // send capacity
    send_noffsets: i32,         // number of send offsets
    recv_buffer: ArrayData,     // receive buffer
    recv_neighbors: Buffer,     // receive neighbor processes buffer
    recv_rank_offsets: Buffer,  // offsets in recv_buffer 
    recv_rank_lengths: Buffer,  // lenghts in recv_buffer
    recv_capacity: i32,         // receive capacity
    recv_noffsets: i32,         // number of receive offsets
    max_neighs: i32,            // maximum number of neighbors
    neighs: i32                 // current number of neighbors
};

// Cartesian info
static mut xprev: i32;
static mut xnext: i32;
static mut yprev: i32;
static mut ynext: i32;
static mut zprev: i32;
static mut znext: i32;

// Communicate ghost particles? (useful for 6D stencil diagonal communication)
fn @communicate_ghost_particles() -> bool { select(use_walberla(), false, true) }

// Maximum number of neighbor ranks
fn @get_initial_maximum_neighbor_ranks() -> i64 { select(use_walberla(), 30, 6) as i64 }

// Initialize and finalize MPI
fn mpi_initialize() -> () { mpi().init(); }
fn mpi_finalize() -> () { mpi().finalize(); }

// Types for condition and communication functions
type CondFn = fn(Vector3D, fn(Vector3D, PBCFlags) -> ()) -> ();
type CommFn = fn(i32, i32, CondFn, CondFn) -> ();

// Communication pattern using walberla or 6-stencil neighbors
fn communication_ranks(grid: Grid, body: CommFn) -> () {
    let nbh = neighborhood;
    let world_aabb = grid.world_aabb;
    let spacing = grid.spacing;
    let no_pbc_flags = PBCFlags { x: 0 as i8, y: 0 as i8, z: 0 as i8 };

    if use_walberla() {
        let nranks = get_number_of_neighbor_ranks();
        let sizes = Vector3D {
            x: world_aabb.xmax - world_aabb.xmin,
            y: world_aabb.ymax - world_aabb.ymin,
            z: world_aabb.zmax - world_aabb.zmin
        };

        range(0, nranks as i32, |i| {
            let rank = get_neighborhood_rank(i);
            let naabbs = get_rank_number_of_aabbs(i);
            let offset = get_rank_offset(i);

            body(rank, rank,
                @|p, f| {
                    if is_within_aabb_radius(p, world_aabb, spacing) {
                        let mut i = 0;
                        while i < naabbs {
                            if distance_point_aabb(p, get_neighborhood_aabb(nbh, offset + i)) < spacing * spacing {
                                f(p, no_pbc_flags);
                                break()
                            }

                            i++;
                        }
                    } else {
                        let mut i = 0;
                        while i < naabbs {
                            distance_point_aabb_periodic(p, get_neighborhood_aabb(nbh, offset + i), sizes, @|dis, p_adj, pbc_flags| {
                                if dis < spacing * spacing {
                                    f(p_adj, pbc_flags);
                                    break()
                                }
                            });

                            i++;
                        }
                    }
                },
                @|p, f| {
                    let mut i = 0;
                    while i < naabbs {
                        let corrected_pos = pbc_correct(p, grid);
                        if is_within_domain(corrected_pos, get_neighborhood_aabb(nbh, offset + i)) {
                            f(corrected_pos, no_pbc_flags);
                            break()
                        }

                        i++;
                    }
                }
            );
        });
    } else {
        let dspacing = spacing * 2.0;
        let f_pbc = @|f: fn(Vector3D, PBCFlags) -> (), pos: Vector3D, pbc_x: i32, pbc_y: i32, pbc_z: i32| {
            let adj_pbc = get_pbc_flags_from_position(pos, pbc_x, pbc_y, pbc_z, grid);
            let mut adj_pos = pos;
            adj_pos.x += grid.xlength * adj_pbc.x as real_t;
            adj_pos.y += grid.ylength * adj_pbc.y as real_t;
            adj_pos.z += grid.zlength * adj_pbc.z as real_t;
            f(adj_pos, adj_pbc);
        };

        body(xnext, xprev, @|p, f| { if p.x > grid.aabb.xmax - dspacing { f_pbc(f, p, -1, 0, 0); }},
                           @|p, f| { if p.x > grid.aabb.xmax - spacing  { f(p, no_pbc_flags); }});
        body(xprev, xnext, @|p, f| { if p.x < grid.aabb.xmin + dspacing { f_pbc(f, p, 1, 0, 0); }},
                           @|p, f| { if p.x < grid.aabb.xmin + spacing  { f(p, no_pbc_flags); }});
        body(ynext, yprev, @|p, f| { if p.y > grid.aabb.ymax - dspacing { f_pbc(f, p, 0, -1, 0); }},
                           @|p, f| { if p.y > grid.aabb.ymax - spacing  { f(p, no_pbc_flags); }});
        body(yprev, ynext, @|p, f| { if p.y < grid.aabb.ymin + dspacing { f_pbc(f, p, 0, 1, 0); }},
                           @|p, f| { if p.y < grid.aabb.ymin + spacing  { f(p, no_pbc_flags); }});
        body(znext, zprev, @|p, f| { if p.z > grid.aabb.zmax - dspacing { f_pbc(f, p, 0, 0, -1); }},
                           @|p, f| { if p.z > grid.aabb.zmax - spacing  { f(p, no_pbc_flags); }});
        body(zprev, znext, @|p, f| { if p.z < grid.aabb.zmin + dspacing { f_pbc(f, p, 0, 0, 1); }},
                           @|p, f| { if p.z < grid.aabb.zmin + spacing  { f(p, no_pbc_flags); }});
    }
}

fn communication_local(me: i32, grid: Grid, body: CommFn) -> () {
    communication_ranks(grid, |send_rank, recv_rank, border_positions, exchange_positions| {
        if send_rank == me {
            body(send_rank, recv_rank, border_positions, exchange_positions);
        }
    });
}

fn communication_remote(me: i32, grid: Grid, body: CommFn) -> () {
    communication_ranks(grid, |send_rank, recv_rank, border_positions, exchange_positions| {
        if send_rank != me {
            body(send_rank, recv_rank, border_positions, exchange_positions);
        }
    });
}

fn communication_sep(me: i32, grid: Grid, body: CommFn) -> () {
    @@communication_remote(me, grid, body);
    @@communication_local(me, grid, body);
}

fn @pbc_correct(pos: Vector3D, grid: Grid) -> Vector3D {
    let world_aabb = grid.world_aabb;
    let mut corrected_pos = pos;
    if pos.x < world_aabb.xmin { corrected_pos.x += grid.xlength; }
    if pos.x > world_aabb.xmax { corrected_pos.x -= grid.xlength; }
    if pos.y < world_aabb.ymin { corrected_pos.y += grid.ylength; }
    if pos.y > world_aabb.ymax { corrected_pos.y -= grid.ylength; }
    if pos.z < world_aabb.zmin { corrected_pos.z += grid.zlength; }
    if pos.z > world_aabb.zmax { corrected_pos.z -= grid.zlength; }
    corrected_pos
}

// Initialize grid communication
fn alloc_comm(grid: Grid) -> Comm {
    let mpih = mpi();
    let max_neighs = get_initial_maximum_neighbor_ranks();
    let max_faces_dim = math.max(math.max(grid.nx * grid.ny, grid.nx * grid.nz), grid.ny * grid.nz) as i64;
    let send_capacity = (max_neighs * max_faces_dim) as i32 * 20;
    let recv_capacity = (max_neighs * max_faces_dim) as i32 * 20;
    let null_buf = Buffer {
        device: 0,
        data: 0 as &[i8],
        size: 0 as i64
    };

    Comm {
        me: get_process_rank(),
        local_start: 0,
        send_buffer: allocate_array(send_capacity, 7, sizeof[real_t](), true),
        send_neighbors: alloc_cpu(max_neighs * sizeof[i32]()),
        send_rank_offsets: alloc_cpu(max_neighs * sizeof[i32]()),
        send_rank_lengths: alloc_cpu(max_neighs * sizeof[i32]()),
        send_offsets: allocate_array(send_capacity, 1, sizeof[i32](), true),
        send_pbc: allocate_array(send_capacity, 3, sizeof[i8](), true),
        copy_list: allocate_array(send_capacity, 1, sizeof[i32](), true),
        send_capacity: send_capacity,
        send_noffsets: 0,
        recv_buffer: allocate_array(recv_capacity, 7, sizeof[real_t](), true),
        recv_neighbors: alloc_cpu(max_neighs * sizeof[i32]()),
        recv_rank_offsets: alloc_cpu(max_neighs * sizeof[i32]()),
        recv_rank_lengths: alloc_cpu(max_neighs * sizeof[i32]()),
        recv_capacity: recv_capacity,
        recv_noffsets: 0,
        max_neighs: max_neighs as i32,
        neighs: 0
    }
}

// Release communication buffers
fn release_comm(comm: Comm) -> () {
    release_array(comm.send_buffer);
    release_array(comm.send_offsets);
    release_array(comm.send_pbc);
    release_array(comm.recv_buffer);
    release_array(comm.copy_list);

    if comm.send_capacity > 0 {
        release(comm.send_neighbors);
        release(comm.send_rank_offsets);
        release(comm.send_rank_lengths);
    }

    if comm.recv_capacity > 0 {
        release(comm.recv_neighbors);
        release(comm.recv_rank_offsets);
        release(comm.recv_rank_lengths);
    }
}

fn resize_max_neighbors_capacity(comm: &mut Comm, max_neighs: i32) -> () {
    print_i32_with_rank("resize_max_neighbors_capacity()", max_neighs);
    let realloc_max_neigh_buf = @|buf: &mut Buffer| {
        let new_buf = alloc_cpu(max_neighs as i64 * sizeof[i32]());
        copy(*buf, new_buf);
        release(*buf);
        *buf = new_buf;
    };

    realloc_max_neigh_buf(&mut comm.send_neighbors);
    realloc_max_neigh_buf(&mut comm.send_rank_offsets);
    realloc_max_neigh_buf(&mut comm.send_rank_lengths);
    realloc_max_neigh_buf(&mut comm.recv_neighbors);
    realloc_max_neigh_buf(&mut comm.recv_rank_offsets);
    realloc_max_neigh_buf(&mut comm.recv_rank_lengths);
    comm.max_neighs = max_neighs;
}

fn resize_send_capacity(comm: &mut Comm, send_capacity: i32) -> () {
    print_i32_with_rank("resize_send_capacity()", send_capacity);
    reallocate_array(&mut comm.send_buffer, send_capacity, 7, sizeof[real_t](), true);
    reallocate_array(&mut comm.send_offsets, send_capacity, 1, sizeof[i32](), true);
    reallocate_array(&mut comm.send_pbc, send_capacity, 3, sizeof[i8](), true);
    reallocate_array(&mut comm.copy_list, send_capacity, 1, sizeof[i32](), true);
    comm.send_capacity = send_capacity;
}

fn resize_recv_capacity(comm: &mut Comm, recv_capacity: i32) -> () {
    print_i32_with_rank("resize_recv_capacity()", recv_capacity);
    reallocate_array(&mut comm.recv_buffer, recv_capacity, 7, sizeof[real_t](), true);
    comm.recv_capacity = recv_capacity;
}

fn get_node_config(xlength: real_t, ylength: real_t, zlength: real_t, destx: &mut i32, desty: &mut i32, destz: &mut i32) -> () {
    let mpih = mpi();
    let areax = xlength * ylength;
    let areay = xlength * zlength;
    let areaz = ylength * zlength;
    let mut bestsurf = 2.0 * (areax + areay + areaz) as f64;
    let mut world_size: i32;

    *destx = 1;
    *desty = 1;
    *destz = 1;

    mpih.comm_size(mpih.comms.world, &mut world_size);

    for i in range(1, world_size) {
        if world_size % i == 0 {
            let rem_yz = world_size / i;

            for j in range(1, rem_yz) {
                if rem_yz % j == 0 {
                    let k = rem_yz / j;
                    let surf = areax / i as f64 / j as f64 + areay / i as f64 / k as f64 + areaz / j as f64 / k as f64;

                    if surf < bestsurf {
                        *destx = i;
                        *desty = j;
                        *destz = k;
                        bestsurf = surf;
                    }
                }
            }
        }
    }
}

// Get bounding box for current node
fn get_node_bounding_box(aabb: AABB) -> AABB {
    let mpih = mpi();
    let mut gx: i32;
    let mut gy: i32;
    let mut gz: i32;
    let mut locx: i32;
    let mut locy: i32;
    let mut locz: i32;

    // Number of cells in each dimension
    let xtotallength = aabb.xmax - aabb.xmin;
    let ytotallength = aabb.ymax - aabb.ymin;
    let ztotallength = aabb.zmax - aabb.zmin;

    // Get configuration of nodes
    get_node_config(xtotallength, ytotallength, ztotallength, &mut gx, &mut gy, &mut gz);

    // Dimensions length for each rank
    let xlength = xtotallength / (gx as real_t);
    let ylength = ytotallength / (gy as real_t);
    let zlength = ztotallength / (gz as real_t);

    // 3D cartesian position of current rank
    mpih.cart(gx, gy, gz, &mut locx, &mut locy, &mut locz, &mut xprev, &mut xnext, &mut yprev, &mut ynext, &mut zprev, &mut znext);

    // Calculate boundaries using lengths in each dimension
    let xmin = aabb.xmin + xlength * locx as real_t;
    let xmax = xmin + xlength;
    let ymin = aabb.ymin + ylength * locy as real_t;
    let ymax = ymin + ylength;
    let zmin = aabb.zmin + zlength * locz as real_t;
    let zmax = zmin + zlength;

    AABB {
        xmin: xmin,
        xmax: xmax,
        ymin: ymin,
        ymax: ymax,
        zmin: zmin,
        zmax: zmax
    }
}

// Synchronize ghost layer cells with neighbor ranks
fn synchronize_ghost_layer(grid: Grid, comm: Comm) -> () {
    let mpih = mpi();
    let nlocal = comm.send_noffsets - comm.local_start;
    let send_data = get_array_real_ref(array_dev, comm.send_buffer);
    let recv_data = get_array_real_ref(array_dev, comm.recv_buffer);
    let particle = make_particle(grid, array_dev, ParticleDataLayout(), null_layout());

    device().loop_1d(false, comm.send_noffsets, |i| {
        let send_offsets = get_array_i32_ref(array_dev, comm.send_offsets);
        let send_pbc = get_array_i8_ref(array_dev, comm.send_pbc);
        let position = particle.get_position(send_offsets(i));
        send_data(i * 3 + 0) = position.x + grid.xlength * send_pbc(i * 3 + 0) as real_t;
        send_data(i * 3 + 1) = position.y + grid.ylength * send_pbc(i * 3 + 1) as real_t;
        send_data(i * 3 + 2) = position.z + grid.zlength * send_pbc(i * 3 + 2) as real_t;
    });

    if comm.local_start > 0 {
        transfer_array_to_host(comm.send_buffer);
    }

    range(0, comm.neighs, |nreq| {
        let recv_offset = bitcast[&[i32]](comm.recv_rank_offsets.data)(nreq) * 3;

        mpih.irecv(
            &get_array_real_ref(array_host, comm.recv_buffer)(recv_offset) as MPI_MutBuf,
            bitcast[&[i32]](comm.recv_rank_lengths.data)(nreq) * 3,
            mpih.double_t, bitcast[&[i32]](comm.recv_neighbors.data)(nreq), 0, mpih.comms.world, nreq);
    });

    range(0, comm.neighs, |nreq| {
        let send_offset = bitcast[&[i32]](comm.send_rank_offsets.data)(nreq) * 3;

        mpih.isend(
            &get_array_real_ref(array_host, comm.send_buffer)(send_offset) as MPI_Buf,
            bitcast[&[i32]](comm.send_rank_lengths.data)(nreq) * 3,
            mpih.double_t, bitcast[&[i32]](comm.send_neighbors.data)(nreq), 0, mpih.comms.world, nreq);
    });

    mpih.wait_all(comm.neighs);

    if comm.recv_noffsets > 0 {
        transfer_array_to_device(comm.recv_buffer);
    }

    device().loop_1d(false, comm.recv_noffsets, |i| {
        particle.set_position(grid.nparticles + i,
                              Vector3D { x: recv_data(i * 3 + 0), y: recv_data(i * 3 + 1), z: recv_data(i * 3 + 2) });
    });

    device().loop_1d(false, nlocal, |i| {
        let index = (comm.local_start + i) * 3;
        particle.set_position(grid.nparticles + comm.recv_noffsets + i,
                              Vector3D { x: send_data(index + 0), y: send_data(index + 1), z: send_data(index + 2), });
    });
}

// Exchange ghost layer particles with neighbor ranks (here the number of
// particles is also updated)
fn exchange_particles(grid: &mut Grid, comm: &mut Comm) -> () {
    let mut isend = 0;
    let mut neigh = 0;

    grid.nghost = 0;
    communication_remote(comm.me, *grid, |_, _, _, _| {
        if neigh >= comm.max_neighs {
            resize_max_neighbors_capacity(comm, neigh + 8);
        }

        bitcast[&mut[i32]](comm.send_rank_offsets.data)(neigh) = -1;
        bitcast[&mut[i32]](comm.recv_rank_offsets.data)(neigh) = -1;
        neigh++;
    });

    let send_rank_offsets = bitcast[&mut[i32]](comm.send_rank_offsets.data);
    let send_rank_lengths = bitcast[&mut[i32]](comm.send_rank_lengths.data);
    let exchange_and_unpack_if = @|cond: bool, nb: i32, comm_fn: fn(i32, Grid, CommFn) -> ()| {
        if cond {
            let (nirecv, _) = two_step_comm(comm_fn, *grid, comm, nb, 0, 0, true);
            let exchg_start = add_local_slots(nirecv, grid);
            unpack_remote_particles(*grid, *comm, exchg_start, 0, nirecv, true);
        }
    };

    neigh = 0;

    // Pack particles in each direction to exchange
    communication_remote(comm.me, *grid, |send_rank, recv_rank, border_positions, exchange_positions| {
        send_rank_offsets(neigh) = isend;

        let mut resize = 1;
        while resize > 0 {
            let const_grid = *grid;
            let const_comm = *comm;
            let const_isend = isend;
            let world_aabb = const_grid.world_aabb;

            set_counter(isend, const_grid);
            reset_resize(const_grid);

            particles_scalar(false, const_grid, |i, particle| {
                let send_flags = get_array_i8_ref(array_dev, const_grid.send_flags);
                let buffer = get_array_real_ref(array_dev, const_comm.send_buffer);
                let send_offsets = get_array_i32_ref(array_dev, const_comm.send_offsets);
                let particle_pos = particle.get_position(i);

                send_flags(i) = 0 as i8;
                exchange_positions(particle_pos, @|pos, _| {
                    let counter = add_counter(const_grid);
                    if counter >= const_comm.send_capacity {
                        grow_resize(counter, const_grid);
                    } else {
                        let vel = particle.get_velocity(i);
                        let index = counter * 7;
                        buffer(index + 0) = particle.get_mass(i);
                        buffer(index + 1) = pos.x;
                        buffer(index + 2) = pos.y;
                        buffer(index + 3) = pos.z;
                        buffer(index + 4) = vel.x;
                        buffer(index + 5) = vel.y;
                        buffer(index + 6) = vel.z;
                        send_offsets(counter - const_isend) = i;
                        send_flags(i) = 1 as i8;
                    }
                });
            });

            resize = get_resize(const_grid);

            if resize > 0 {
                resize_send_capacity(comm, resize * 2);
            }
        }

        transfer_array_to_host(comm.send_offsets);
        transfer_array_to_host(grid.send_flags);

        let counter = get_counter(*grid);
        let packed = counter - isend;
        let nparticles = grid.nparticles - packed;
        let send_flags_host = get_array_i8_ref(array_host, grid.send_flags);
        let send_offsets_host = get_array_i32_ref(array_host, comm.send_offsets);
        let copy_list_host = get_array_i32_ref(array_host, comm.copy_list);
        let mut send_pos = grid.nparticles - 1;

        range(0, packed, |i| {
            if send_offsets_host(i) < nparticles {
                while(send_flags_host(send_pos) == 1 as i8) {
                    send_pos--;
                }

                copy_list_host(i) = send_pos;
                send_pos--;
            } else {
                copy_list_host(i) = -1;
            }
        });

        transfer_array_to_device(comm.copy_list);

        let const_comm = *comm;
        let const_grid = *grid;
        let particle = make_particle(const_grid, array_dev, ParticleDataLayout(), null_layout());
        let copy_list = get_array_i32_ref(array_dev, const_comm.copy_list);

        device().loop_1d(false, packed, |i| {
            if copy_list(i) > 0 {
                let src = copy_list(i);
                let dst = get_array_i32_ref(array_dev, const_comm.send_offsets)(i);
                particle.set_mass(dst, particle.get_mass(src));
                particle.set_position(dst, particle.get_position(src));
                particle.set_velocity(dst, particle.get_velocity(src));
            }
        });

        grid.nparticles -= packed;
        isend = counter;
        send_rank_lengths(neigh) = isend - send_rank_offsets(neigh);

        // Communicate for each dimension
        if neigh % 2 != 0 {
            exchange_and_unpack_if(communicate_ghost_particles(), neigh - 1, @|me: i32, grid: Grid, f: CommFn| {
                f(recv_rank, send_rank, border_positions, exchange_positions);
                f(send_rank, recv_rank, border_positions, exchange_positions);
                isend = 0;
            });
        }

        neigh++;
    });

    exchange_and_unpack_if(!communicate_ghost_particles(), 0, communication_remote);
}

// Define particle borders to synchronize during next iterations
fn borders(grid: &mut Grid, comm: &mut Comm) -> () {
    let mut isend = 0;
    let mut irecv = 0;
    let mut ilocal = 0;
    let mut local_start = 0;
    let mut nremote = 0;
    let mut nlocal = 0;
    let mut neigh = 0;

    grid.nghost = 0;
    communication_ranks(*grid, |send_rank, _, _, _| {
        if neigh >= comm.max_neighs {
            resize_max_neighbors_capacity(comm, neigh + 8);
        }

        if send_rank != comm.me {
            nremote++;
        } else {
            nlocal++;
        }

        bitcast[&mut[i32]](comm.send_rank_offsets.data)(neigh) = -1;
        bitcast[&mut[i32]](comm.recv_rank_offsets.data)(neigh) = -1;
        neigh++;
    });

    let send_rank_offsets = bitcast[&mut[i32]](comm.send_rank_offsets.data);
    let send_rank_lengths = bitcast[&mut[i32]](comm.send_rank_lengths.data);
    let communicate_and_unpack_if = @|cond: bool, remote: bool, local: bool, nb: i32, comm_fn: fn(i32, Grid, CommFn) -> ()| {
        if cond {
            if remote {
                let oirecv = irecv;
                let (irecv_, local_start_) = two_step_comm(comm_fn, *grid, comm, nb, irecv, local_start, false);
                irecv = irecv_;
                local_start = local_start_;
                ilocal = local_start;
                add_ghost_slots(irecv - oirecv, grid);
                unpack_remote_particles(*grid, *comm, grid.nparticles, oirecv, irecv, false);
            }

            if local {
                add_ghost_slots(isend - ilocal, grid);
                unpack_local_particles(*grid, *comm, isend, irecv, ilocal, local_start);
                ilocal = isend;
            }
        }
    };

    neigh = 0;

    // Pack particles in each direction to send during next iterations
    communication_sep(comm.me, *grid, |send_rank, recv_rank, border_positions, exchange_positions| {
        send_rank_offsets(neigh) = isend;

        let mut resize = 1;
        while resize > 0 {
            let const_grid = *grid;
            let const_comm = *comm;
            let const_isend = isend;
            set_counter(isend, const_grid);
            reset_resize(const_grid);

            particles_scalar(true, const_grid, |i, particle| {
                let send_flags = get_array_i8_ref(array_dev, const_grid.send_flags);
                let buffer = get_array_real_ref(array_dev, const_comm.send_buffer);
                let send_offsets = get_array_i32_ref(array_dev, const_comm.send_offsets);
                let send_pbc = get_array_i8_ref(array_dev, const_comm.send_pbc);
                let particle_pos = particle.get_position(i);

                border_positions(particle_pos, @|pos, pbc| {
                    let counter = add_counter(const_grid);
                    if counter >= const_comm.send_capacity {
                        grow_resize(counter, const_grid);
                    } else {
                        let buf_index = counter * 4;
                        let pbc_index = counter * 3;
                        buffer(buf_index + 0) = particle.get_mass(i);
                        buffer(buf_index + 1) = pos.x;
                        buffer(buf_index + 2) = pos.y;
                        buffer(buf_index + 3) = pos.z;
                        send_pbc(pbc_index + 0) = pbc.x;
                        send_pbc(pbc_index + 1) = pbc.y;
                        send_pbc(pbc_index + 2) = pbc.z;
                        send_offsets(counter) = i;
                    }
                });
            });

            resize = get_resize(const_grid);
            if resize > 0 {
                resize_send_capacity(comm, resize * 2);
            }
        }

        isend = get_counter(*grid);
        send_rank_lengths(neigh) = isend - send_rank_offsets(neigh);

        // Communicate for each dimension
        if neigh % 2 != 0 {
            let remote = comm.me != send_rank;
            communicate_and_unpack_if(communicate_ghost_particles(), remote, !remote, neigh - 1, @|me: i32, grid: Grid, f: CommFn| {
                f(recv_rank, send_rank, border_positions, exchange_positions);
                f(send_rank, recv_rank, border_positions, exchange_positions);
            });
        }

        neigh++;
    });

    communicate_and_unpack_if(!communicate_ghost_particles(), nremote > 0, nlocal > 0, 0, communication_remote);
    comm.local_start = local_start;
    comm.send_noffsets = isend;
    comm.recv_noffsets = irecv;
    comm.neighs = nremote;
}

fn @two_step_comm(
    comm_fn: fn(i32, Grid, CommFn) -> (), grid: Grid, comm: &mut Comm,
    neigh: i32, irecv: i32, local_start: i32, @velocities: bool) -> (i32, i32) {

    let mpih = mpi();
    let send_rank_offsets = bitcast[&mut[i32]](comm.send_rank_offsets.data);
    let send_rank_lengths = bitcast[&mut[i32]](comm.send_rank_lengths.data);
    let recv_rank_offsets = bitcast[&mut[i32]](comm.recv_rank_offsets.data);
    let recv_rank_lengths = bitcast[&mut[i32]](comm.recv_rank_lengths.data);
    let elems = select(velocities, 7, 4);
    let mut nirecv = irecv;
    let mut nlocal_start = local_start;
    let mut nreq = 0;

    transfer_array_to_host(comm.send_buffer);
    comm_fn(comm.me, grid, |send_rank, recv_rank, _, _| {
        let i = neigh + nreq;
        mpih.irecv(&mut recv_rank_lengths(i) as MPI_MutBuf, 1, mpih.int_t, recv_rank, 0, mpih.comms.world, nreq);
        mpih.isend(&send_rank_lengths(i) as MPI_Buf, 1, mpih.int_t, send_rank, 0, mpih.comms.world, nreq);
        nreq++;
    });

    mpih.wait_all(nreq);
    range(0, nreq, |r| {
        let i = neigh + r;
        recv_rank_offsets(i) = nirecv;

        if nirecv + recv_rank_lengths(i) >= comm.recv_capacity {
            resize_recv_capacity(comm, (nirecv + recv_rank_lengths(i)) * 2);
        }

        nirecv += recv_rank_lengths(i);
        nlocal_start += send_rank_lengths(i);
    });

    nreq = 0;
    comm_fn(comm.me, grid, |send_rank, recv_rank, _, _| {
        let i = neigh + nreq;
        let send_offset = send_rank_offsets(i) * elems;
        let recv_offset = recv_rank_offsets(i) * elems;

        mpih.irecv(
            &get_array_real_ref(array_host, comm.recv_buffer)(recv_offset) as MPI_MutBuf,
            recv_rank_lengths(i) * elems, mpih.double_t, recv_rank, 0, mpih.comms.world, nreq);

        mpih.isend(
            &get_array_real_ref(array_host, comm.send_buffer)(send_offset) as MPI_Buf,
            send_rank_lengths(i) * elems, mpih.double_t, send_rank, 0, mpih.comms.world, nreq);

        bitcast[&mut[i32]](comm.send_neighbors.data)(i) = send_rank;
        bitcast[&mut[i32]](comm.recv_neighbors.data)(i) = recv_rank;
        nreq++;
    });

    mpih.wait_all(nreq);
    (nirecv, nlocal_start)
}

fn unpack_remote_particles(grid: Grid, comm: Comm, start: i32, oirecv: i32, nirecv: i32, @velocities: bool) -> () {
    let particle = make_particle(grid, array_dev, ParticleDataLayout(), null_layout());
    let recv_data = get_array_real_ref(array_dev, comm.recv_buffer);
    let nrecv = nirecv - oirecv;
    let elems = select(velocities, 7, 4);

    if nrecv > 0 {
        transfer_array_to_device(comm.recv_buffer);
    }

    device().loop_1d(false, nrecv, |i| {
        let index = (oirecv + i) * elems;
        let offset = start + oirecv + i;
        particle.set_mass(offset, recv_data(index));
        particle.set_position(offset, Vector3D { x: recv_data(index + 1), y: recv_data(index + 2), z: recv_data(index + 3) });
        if velocities {
            particle.set_velocity(offset, Vector3D { x: recv_data(index + 4), y: recv_data(index + 5), z: recv_data(index + 6) });
        }
    });
}

fn unpack_local_particles(grid: Grid, comm: Comm, isend: i32, irecv: i32, ilocal: i32, local_start: i32) -> () {
    let particle = make_particle(grid, array_dev, ParticleDataLayout(), null_layout());
    let send_data = get_array_real_ref(array_dev, comm.send_buffer);
    let recv_offset = grid.nparticles + irecv + ilocal - local_start;

    device().loop_1d(false, isend - ilocal, |i| {
        let index = (ilocal + i) * 4;
        let offset = recv_offset + i;
        particle.set_mass(offset, send_data(index));
        particle.set_position(offset, Vector3D { x: send_data(index + 1), y: send_data(index + 2), z: send_data(index + 3) });
    });
}

fn reduce_f64_sum(local_value: f64, global_value: &mut f64) -> () {
    let mpih = mpi();
    let mut local = local_value;
    mpih.allreduce(&mut local as MPI_Buf, global_value as MPI_MutBuf, 1, mpih.double_t, mpih.ops.sum, mpih.comms.world);
}

fn reduce_f64_max(local_value: f64, global_value: &mut f64) -> () {
    let mpih = mpi();
    let mut local = local_value;
    mpih.allreduce(&mut local as MPI_Buf, global_value as MPI_MutBuf, 1, mpih.double_t, mpih.ops.max, mpih.comms.world);
}

fn reduce_i32_sum(local_value: i32, global_value: &mut i32) -> () {
    let mpih = mpi();
    let mut local = local_value;
    mpih.allreduce(&mut local as MPI_Buf, global_value as MPI_MutBuf, 1, mpih.int_t, mpih.ops.sum, mpih.comms.world);
}

fn reduce_i64_sum(local_value: i64, global_value: &mut i64) -> () {
    let mpih = mpi();
    let mut local = local_value;
    mpih.allreduce(&mut local as MPI_Buf, global_value as MPI_MutBuf, 1, mpih.int64_t, mpih.ops.sum, mpih.comms.world);
}