dispatch.c

/**
 * Author: Baptiste Coudray
 * School: HEPIA
 * Class: ITI-3
 * Year: 2020-2021
 */

#include <stdio.h>
#include <stdint.h>
#include <mpi.h>
#include <math.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include "dispatch.h"

#define min(a, b) (((a) <= (b)) ? (a) : (b))
#define max(a, b) (((a) >= (b)) ? (a) : (b))

#define INDEX_2D_TO_1D(y, x, nb_columns) ((x) + (y) * (nb_columns))
#define INDEX_3D_TO_1D(y, x, z, nb_columns, nb_depths) ((x) + (nb_columns) * ((y) + (nb_depths) * (z)))

#define ROW_COMMUNICATOR 1
#define COLUMN_COMMUNICATOR 2
#define DEPTH_COMMUNICATOR 3

#define CHUNK_DATA_TAG 8
#define ROOT_RANK 0

#define FRONT_EAST_TAG 0
#define FRONT_WEST_TAG 1
#define FRONT_NORTH_TAG 2
#define FRONT_SOUTH_TAG 3
#define FRONT_NORTH_EAST_TAG 4
#define FRONT_SOUTH_WEST_TAG 5
#define FRONT_SOUTH_EAST_TAG 6
#define FRONT_NORTH_WEST_TAG 7
#define FRONT_SURFACE_TAG 8

#define BACK_EAST_TAG 9
#define BACK_WEST_TAG 10
#define BACK_SURFACE_TAG 11

#define LEFT_SURFACE_TAG 12
#define RIGHT_SURFACE_TAG 13

#define TOP_EAST_TAG 14
#define TOP_WEST_TAG 15
#define TOP_NORTH_TAG 16
#define TOP_SOUTH_TAG 17
#define TOP_NORTH_EAST_TAG 18
#define TOP_SOUTH_WEST_TAG 19
#define TOP_SOUTH_EAST_TAG 20
#define TOP_NORTH_WEST_TAG 21
#define TOP_SURFACE_TAG 22

#define BOTTOM_EAST_TAG 23
#define BOTTOM_WEST_TAG 24
#define BOTTOM_NORTH_TAG 25
#define BOTTOM_SOUTH_TAG 26
#define BOTTOM_NORTH_EAST_TAG 27
#define BOTTOM_SOUTH_WEST_TAG 28
#define BOTTOM_SOUTH_EAST_TAG 29
#define BOTTOM_NORTH_WEST_TAG 30
#define BOTTOM_SURFACE_TAG 31

const int FUTHARK_CHUNKS_ORDER[NB_CHUNKS] = {INDEX_CHUNK_EAST, INDEX_CHUNK_NORTH, INDEX_CHUNK_NORTH_EAST,
                                             INDEX_CHUNK_NORTH_WEST, INDEX_CHUNK_SOUTH, INDEX_CHUNK_SOUTH_EAST,
                                             INDEX_CHUNK_SOUTH_WEST, INDEX_CHUNK_SURFACE, INDEX_CHUNK_WEST};


/* Taken from Futhark generated code */
struct futhark_opaque_envelope_1d_t {
    struct futhark_i8_1d *v0;
    struct futhark_i8_1d *v1;
};

struct futhark_opaque_envelope_2d_t {
    struct futhark_i8_2d *v0;
    struct futhark_i8_2d *v1;
    struct futhark_i8_2d *v2;
    struct futhark_i8_2d *v3;
    struct futhark_i8_2d *v4;
    struct futhark_i8_2d *v5;
    struct futhark_i8_2d *v6;
    struct futhark_i8_2d *v7;
};

struct futhark_opaque_envelope_3d_t {
    struct futhark_i8_3d *v0;
    struct futhark_i8_3d *v1;
    struct futhark_i8_3d *v2;
    struct futhark_i8_3d *v3;
    struct futhark_i8_3d *v4;
    struct futhark_i8_3d *v5;
    struct futhark_i8_3d *v6;
    struct futhark_i8_3d *v7;
    struct futhark_i8_3d *v8;
    struct futhark_i8_3d *v9;
    struct futhark_i8_3d *v10;
    struct futhark_i8_3d *v11;
    struct futhark_i8_3d *v12;
    struct futhark_i8_3d *v13;
    struct futhark_i8_3d *v14;
    struct futhark_i8_3d *v15;
    struct futhark_i8_3d *v16;
    struct futhark_i8_3d *v17;
    struct futhark_i8_3d *v18;
    struct futhark_i8_3d *v19;
    struct futhark_i8_3d *v20;
    struct futhark_i8_3d *v21;
    struct futhark_i8_3d *v22;
    struct futhark_i8_3d *v23;
    struct futhark_i8_3d *v24;
    struct futhark_i8_3d *v25;
    struct futhark_i8_3d *v26;
    struct futhark_i8_3d *v27;
    struct futhark_i8_3d *v28;
    struct futhark_i8_3d *v29;
    struct futhark_i8_3d *v30;
    struct futhark_i8_3d *v31;
    struct futhark_i8_3d *v32;
    struct futhark_i8_3d *v33;
    struct futhark_i8_3d *v34;
    struct futhark_i8_3d *v35;
    struct futhark_i8_3d *v36;
    struct futhark_i8_3d *v37;
    struct futhark_i8_3d *v38;
    struct futhark_i8_3d *v39;
    struct futhark_i8_3d *v40;
    struct futhark_i8_3d *v41;
    struct futhark_i8_3d *v42;
    struct futhark_i8_3d *v43;
    struct futhark_i8_3d *v44;
    struct futhark_i8_3d *v45;
    struct futhark_i8_3d *v46;
    struct futhark_i8_3d *v47;
    struct futhark_i8_3d *v48;
    struct futhark_i8_3d *v49;
    struct futhark_i8_3d *v50;
    struct futhark_i8_3d *v51;
    struct futhark_i8_3d *v52;
    struct futhark_i8_3d *v53;
};

struct dispatch_context {
    int my_rank;
    int world_size;
    int my_cart_rank;
    int coordinates[3];
    MPI_Comm communicators[4]; /* cart_comm, row_comm, column_comm, depth_comm */
    int network_dimensions[3];
    int data_dimensions[3];
    MPI_Datatype datatype;
    size_t count;
    int n_dimensions;
    chunk_info_t *chunk_info;
    int type;
    chunk_info_t *chunks_info;
    chunk_info_t active_domain;
};

extern void envelope_init_accessors(envelope_t *envelope) {
    envelope->back = &envelope->sides[INDEX_SIDE_BACK];
    envelope->bottom = &envelope->sides[INDEX_SIDE_BOTTOM];
    envelope->front = &envelope->sides[INDEX_SIDE_FRONT];
    envelope->left = &envelope->sides[INDEX_SIDE_LEFT];
    envelope->right = &envelope->sides[INDEX_SIDE_RIGHT];
    envelope->top = &envelope->sides[INDEX_SIDE_TOP];

    for (int i = 0; i < NB_SIDES; ++i) {
        envelope->sides[i].east = &envelope->sides[i].chunks[INDEX_CHUNK_EAST];
        envelope->sides[i].north = &envelope->sides[i].chunks[INDEX_CHUNK_NORTH];
        envelope->sides[i].north_east = &envelope->sides[i].chunks[INDEX_CHUNK_NORTH_EAST];
        envelope->sides[i].north_west = &envelope->sides[i].chunks[INDEX_CHUNK_NORTH_WEST];
        envelope->sides[i].south = &envelope->sides[i].chunks[INDEX_CHUNK_SOUTH];
        envelope->sides[i].south_east = &envelope->sides[i].chunks[INDEX_CHUNK_SOUTH_EAST];
        envelope->sides[i].south_west = &envelope->sides[i].chunks[INDEX_CHUNK_SOUTH_WEST];
        envelope->sides[i].surface = &envelope->sides[i].chunks[INDEX_CHUNK_SURFACE];
        envelope->sides[i].west = &envelope->sides[i].chunks[INDEX_CHUNK_WEST];
    }
}

static void get_world_size(struct dispatch_context *dc) {
    MPI_Comm_size(MPI_COMM_WORLD, &dc->world_size);
}

static void get_my_rank(struct dispatch_context *dc) {
    MPI_Comm_rank(MPI_COMM_WORLD, &dc->my_rank);
}

static void find_best_factors(int n, int factors[2]) {
    int result[2] = {0};
    int limit = (int) sqrt(n);
    bool first_pass = true;
    for (int i = 1; i <= limit; ++i) {
        if (n % i == 0) {
            int factor1 = i;
            int factor2 = n / i;
            int current_difference = abs(result[0] - result[1]);
            int new_difference = abs(factor1 - factor2);
            if (first_pass || current_difference > new_difference) {
                result[0] = factor1;
                result[1] = factor2;
            }
            first_pass = false;
        }
    }
    factors[0] = result[0];
    factors[1] = result[1];
}

static void find_network_dimensions(struct dispatch_context *dc) {
    /* 1D */
    if (dc->n_dimensions == 1) {
        dc->network_dimensions[0] = 1;
        dc->network_dimensions[1] = dc->world_size;
        dc->network_dimensions[2] = 1;
    } else {
        /* 2D */
        find_best_factors(dc->world_size, dc->network_dimensions);
        dc->network_dimensions[2] = 1;
        /* 3D */
        if (dc->n_dimensions == 3) {
            int factors[2] = {0};
            find_best_factors(dc->network_dimensions[1], factors);
            if (factors[0] < factors[1]) {
                dc->network_dimensions[1] = factors[1];
                dc->network_dimensions[2] = factors[0];
            } else {
                dc->network_dimensions[1] = factors[0];
                dc->network_dimensions[2] = factors[1];
            }
        }
    }
}

static void create_network_communicators(struct dispatch_context *dc) {
    int periods[3] = {true, true, true}; // Cyclic on row-column-depth

    MPI_Cart_create(MPI_COMM_WORLD, 3, dc->network_dimensions, periods, 1, &dc->communicators[0]);

    /* Create row communicator */
    int remain_dims[3] = {false, true, false};
    MPI_Cart_sub(dc->communicators[0], remain_dims, &dc->communicators[ROW_COMMUNICATOR]);

    /* Create column communicator */
    remain_dims[0] = true; // row
    remain_dims[1] = false; // column
    remain_dims[2] = false; // depth
    MPI_Cart_sub(dc->communicators[0], remain_dims, &dc->communicators[COLUMN_COMMUNICATOR]);

    /* Create depth communicator */
    remain_dims[0] = false; // row
    remain_dims[1] = false; // column
    remain_dims[2] = true; // depth
    MPI_Cart_sub(dc->communicators[0], remain_dims, &dc->communicators[DEPTH_COMMUNICATOR]);

    MPI_Comm_rank(dc->communicators[0], &dc->my_cart_rank);

    MPI_Cart_coords(dc->communicators[0], dc->my_cart_rank, 3, dc->coordinates);
}

static void divide_data(struct dispatch_context *dc) {
    dc->chunks_info = calloc((size_t) dc->world_size, sizeof(chunk_info_t));

    int nb_rows_per_process = dc->data_dimensions[0] / dc->network_dimensions[0];
    int remaining_rows = dc->data_dimensions[0] % dc->network_dimensions[0];

    int nb_columns_per_process = dc->data_dimensions[1] / dc->network_dimensions[1];
    int remaining_columns = dc->data_dimensions[1] % dc->network_dimensions[1];

    int nb_depths_per_process = dc->data_dimensions[2] / dc->network_dimensions[2];
    int remaining_depths = dc->data_dimensions[2] % dc->network_dimensions[2];

    printf("nb_rows = %d, rem = %d\n", nb_rows_per_process, remaining_rows);
    printf("nb_cols = %d, rem = %d\n", nb_columns_per_process, remaining_columns);
    printf("nb_depths = %d, rem = %d\n", nb_depths_per_process, remaining_depths);

    for (int i = 0, y = 0, x = 0, z = 0; i < dc->world_size; ++i) {
        int nb_rows = nb_rows_per_process;
        if (remaining_rows > 0) {
            ++nb_rows;
        }
        int nb_columns = nb_columns_per_process;
        if (remaining_columns > 0) {
            ++nb_columns;
            --remaining_columns;
        }
        int nb_depths = nb_depths_per_process;
        if (remaining_depths > 0) {
            ++nb_depths;
            --remaining_depths;
        }

        int dimensions[3] = {nb_rows, nb_columns, nb_depths};
        chunk_info_init(&dc->chunks_info[i], dc->type, dimensions, y, x, z, i == dc->my_rank);

        z += nb_depths;
        if (z >= dc->data_dimensions[2]) {
            z = 0;
            remaining_depths = dc->data_dimensions[2] % dc->network_dimensions[2];
            x += nb_columns;
            if (x >= dc->data_dimensions[1]) {
                x = 0;
                y += dc->chunks_info[max(i - 1, 0)].dimensions[0];
                remaining_columns = dc->data_dimensions[1] % dc->network_dimensions[1];
                --remaining_rows;
            }
        }
    }
    dc->chunk_info = &dc->chunks_info[dc->my_rank];
}

static int cart_rank_to_comm_rank(struct dispatch_context *dc, int y, int x, int z) {
    int rank;
    int coordinates[3] = {y, x, z};
    MPI_Cart_rank(dc->communicators[0], coordinates, &rank);
    return rank;
}

extern struct dispatch_context *dispatch_context_new(const int *dimensions, MPI_Datatype datatype, int n_dimensions) {
    struct dispatch_context *dc = calloc(1, sizeof(struct dispatch_context));
    assert(dc != NULL);
    get_world_size(dc);
    dc->n_dimensions = n_dimensions;
    dc->datatype = datatype;
    MPI_Type_size(dc->datatype, &dc->type);

    switch (n_dimensions) {
        case 1:
            dc->data_dimensions[0] = 1;
            dc->data_dimensions[1] = dimensions[0];
            dc->data_dimensions[2] = 1;
            break;
        case 2:
            dc->data_dimensions[0] = dimensions[0];
            dc->data_dimensions[1] = dimensions[1];
            dc->data_dimensions[2] = 1;
            break;
        case 3:
            dc->data_dimensions[0] = dimensions[0];
            dc->data_dimensions[1] = dimensions[1];
            dc->data_dimensions[2] = dimensions[2];
            break;
        default:
            fprintf(stderr, "Invalid dimensions size.");
            MPI_Abort(MPI_COMM_WORLD, 1);
            break;
    }
    dc->count = (size_t) dc->data_dimensions[0] * (size_t) dc->data_dimensions[1] * (size_t) dc->data_dimensions[2];

    find_network_dimensions(dc);
    create_network_communicators(dc);
    get_my_rank(dc);
    divide_data(dc);
    dc->active_domain = *dc->chunk_info;
    return dc;
}

extern void dispatch_context_print(struct dispatch_context *dc) {
    printf("[dispatch_context] my_rank = %d, world_size = %d, network_dimensions = [%d][%d][%d], n_dimensions = %d, data_dimensions = [%d][%d][%d], coordinates = (%d, %d, %d)\n",
           dc->my_rank, dc->world_size, dc->network_dimensions[0], dc->network_dimensions[1], dc->network_dimensions[2],
           dc->n_dimensions, dc->data_dimensions[0], dc->data_dimensions[1], dc->data_dimensions[2],
           dc->coordinates[0], dc->coordinates[1], dc->coordinates[2]);
}

static envelope_t get_inner_envelope_1d(struct dispatch_context *dc, struct futhark_context *fc, int thickness) {
    struct futhark_u8_1d *fut_chunk_data = futhark_new_u8_1d(fc, dc->chunk_info->data,
                                                             dc->chunk_info->dimensions[1] * dc->type);
    int thickness_x = min(thickness, dc->chunk_info->dimensions[1]);
    int dimensions[3] = {1, thickness_x, 1};
    struct futhark_opaque_envelope_1d_t *fut_inner_envelope;

    futhark_context_sync(fc);
    futhark_entry_get_envelope_1d(fc, &fut_inner_envelope, fut_chunk_data, dimensions[1] * dc->type);
    futhark_context_sync(fc);

    envelope_t inner_envelope = (envelope_t) {0};
    envelope_init_accessors(&inner_envelope);

    // East
    {
        int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - dimensions[1];
        chunk_info_init(inner_envelope.front->east, dc->type, dimensions, 0, start_x, 0, true);
        futhark_values_i8_1d(fc, fut_inner_envelope->v0, inner_envelope.front->east->data);
    }

    // West
    {
        int start_x = dc->chunk_info->x;
        chunk_info_init(inner_envelope.front->west, dc->type, dimensions, 0, start_x, 0, true);
        futhark_values_i8_1d(fc, fut_inner_envelope->v1, inner_envelope.front->west->data);
    }

    futhark_context_sync(fc);
    futhark_free_opaque_envelope_1d_t(fc, fut_inner_envelope);
    futhark_free_u8_1d(fc, fut_chunk_data);
    return inner_envelope;
}

static envelope_t get_inner_envelope_2d(struct dispatch_context *dc, struct futhark_context *fc, int thickness) {
    struct futhark_u8_2d *fut_chunk_data = futhark_new_u8_2d(fc, dc->chunk_info->data,
                                                             dc->chunk_info->dimensions[0],
                                                             dc->chunk_info->dimensions[1] * dc->type);

    int thickness_y = min(thickness, dc->chunk_info->dimensions[0]);
    int thickness_x = min(thickness, dc->chunk_info->dimensions[1]);
    struct futhark_opaque_envelope_2d_t *fut_inner_envelope;

    futhark_context_sync(fc);
    futhark_entry_get_envelope_2d(fc, &fut_inner_envelope, fut_chunk_data, thickness_y, thickness_x * dc->type);
    futhark_context_sync(fc);

    envelope_t inner_envelope = (envelope_t) {0};
    envelope_init_accessors(&inner_envelope);

    // East
    {
        int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, 1};
        int start_y = dc->chunk_info->y;
        int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - dimensions[1];
        chunk_info_init(inner_envelope.front->east, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v0, inner_envelope.front->east->data);
    }

    // North
    {
        int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], 1};
        int start_y = dc->chunk_info->y;
        int start_x = dc->chunk_info->x;
        chunk_info_init(inner_envelope.front->north, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v1, inner_envelope.front->north->data);
    }

    // North-East
    {
        int dimensions[3] = {thickness_y, thickness_x, 1};
        int start_y = dc->chunk_info->y;
        int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - dimensions[1];
        chunk_info_init(inner_envelope.front->north_east, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v2, inner_envelope.front->north_east->data);
    }

    // North-West
    {
        int dimensions[3] = {thickness_y, thickness_x, 1};
        int start_y = dc->chunk_info->y;
        int start_x = dc->chunk_info->x;
        chunk_info_init(inner_envelope.front->north_west, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v3, inner_envelope.front->north_west->data);
    }

    // South
    {
        int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], 1};
        int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - dimensions[0];
        int start_x = dc->chunk_info->x;
        chunk_info_init(inner_envelope.front->south, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v4, inner_envelope.front->south->data);
    }

    // South-East
    {
        int dimensions[3] = {thickness_y, thickness_x, 1};
        int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - dimensions[0];
        int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - dimensions[1];
        chunk_info_init(inner_envelope.front->south_east, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v5, inner_envelope.front->south_east->data);
    }

    // South-West
    {
        int dimensions[3] = {thickness_y, thickness_x, 1};
        int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - dimensions[0];
        int start_x = dc->chunk_info->x;
        chunk_info_init(inner_envelope.front->south_west, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v6, inner_envelope.front->south_west->data);
    }

    // West
    {
        int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, 1};
        int start_y = dc->chunk_info->y;
        int start_x = dc->chunk_info->x;
        chunk_info_init(inner_envelope.front->west, dc->type, dimensions, start_y, start_x, 0, true);
        futhark_values_i8_2d(fc, fut_inner_envelope->v7, inner_envelope.front->west->data);
    }

    futhark_context_sync(fc);
    futhark_free_u8_2d(fc, fut_chunk_data);
    futhark_free_opaque_envelope_2d_t(fc, fut_inner_envelope);
    return inner_envelope;
}

static envelope_t get_inner_envelope_3d(struct dispatch_context *dc, struct futhark_context *fc, int thickness) {
    struct futhark_u8_3d *fut_chunk_data = futhark_new_u8_3d(fc, dc->chunk_info->data,
                                                             dc->chunk_info->dimensions[0],
                                                             dc->chunk_info->dimensions[1],
                                                             dc->chunk_info->dimensions[2] * dc->type);


    int thickness_y = min(thickness, dc->chunk_info->dimensions[0]);
    int thickness_x = min(thickness, dc->chunk_info->dimensions[1]);
    int thickness_z = min(thickness, dc->chunk_info->dimensions[2]);
    struct futhark_opaque_envelope_3d_t *fut_inner_envelope;

    futhark_context_sync(fc);
    futhark_entry_get_envelope_3d(fc, &fut_inner_envelope, fut_chunk_data, thickness_y, thickness_x,
                                  thickness_z * dc->type);
    futhark_context_sync(fc);

    envelope_t inner_envelope = {0};
    envelope_init_accessors(&inner_envelope);

    // Back
    {
        // East
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - dimensions[1];
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - dimensions[2];
            chunk_info_init(inner_envelope.back->east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v0, inner_envelope.back->east->data);
        }

        // Surface
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], dc->chunk_info->dimensions[1], thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.back->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            const int64_t *dims = futhark_shape_i8_3d(fc, fut_inner_envelope->v7);
//            if(dc->my_rank == 0) printf("[%lld][%lld][%lld]\n", dims[0], dims[1], dims[2]);
            futhark_values_i8_3d(fc, fut_inner_envelope->v7, inner_envelope.back->surface->data);
        }

        // West
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.back->west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v8, inner_envelope.back->west->data);
        }
    }

    // Bottom
    {
        // East
        {
            int dimensions[3] = {thickness_y, thickness_x, dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.bottom->east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v9, inner_envelope.bottom->east->data);
        }

        // North
        {
            int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], thickness_z};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.bottom->north, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v10, inner_envelope.bottom->north->data);
        }

        // North-East
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.bottom->north_east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v11, inner_envelope.bottom->north_east->data);
        }

        // North-West
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.bottom->north_west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v12, inner_envelope.bottom->north_west->data);
        }

        // South
        {
            int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], thickness_z};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.bottom->south, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v13, inner_envelope.bottom->south->data);
        }

        // South-East
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.bottom->south_east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v14, inner_envelope.bottom->south_east->data);
        }

        // South-West
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.bottom->south_west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v15, inner_envelope.bottom->south_west->data);
        }

        // Surface
        {
            int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.bottom->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v16, inner_envelope.bottom->surface->data);
        }

        // West
        {
            int dimensions[3] = {thickness_y, thickness_x, dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y + dc->chunk_info->dimensions[0] - thickness_y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.bottom->west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v17, inner_envelope.bottom->west->data);
        }
    }

    // Front
    {
        // East
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.front->east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v18, inner_envelope.front->east->data);
        }

        // Surface
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], dc->chunk_info->dimensions[1], thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.front->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v25, inner_envelope.front->surface->data);
        }

        // West
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.front->west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v26, inner_envelope.front->west->data);
        }

    }

    // Left
    {
        // Surface
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.left->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v34, inner_envelope.left->surface->data);
        }
    }

    // Right
    {
        // Surface
        {
            int dimensions[3] = {dc->chunk_info->dimensions[0], thickness_x, dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.right->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v43, inner_envelope.right->surface->data);
        }
    }

    // Top
    {
        // East
        {
            int dimensions[3] = {thickness_y, thickness_x, dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.top->east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v45, inner_envelope.top->east->data);
        }

        // North
        {
            int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.top->north, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v46, inner_envelope.top->north->data);
        }

        // North-East
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.top->north_east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v47, inner_envelope.top->north_east->data);
        }

        // North-West
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z + dc->chunk_info->dimensions[2] - thickness_z;
            chunk_info_init(inner_envelope.top->north_west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v48, inner_envelope.top->north_west->data);
        }

        // South
        {
            int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.top->south, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v49, inner_envelope.top->south->data);
        }

        // South-East
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x + dc->chunk_info->dimensions[1] - thickness_x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.top->south_east, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v50, inner_envelope.top->south_east->data);
        }

        // South-West
        {
            int dimensions[3] = {thickness_y, thickness_x, thickness_z};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.top->south_west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v51, inner_envelope.top->south_west->data);
        }

        // Surface
        {
            int dimensions[3] = {thickness_y, dc->chunk_info->dimensions[1], dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.top->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v52, inner_envelope.top->surface->data);
        }

        // West
        {
            int dimensions[3] = {thickness_y, thickness_x, dc->chunk_info->dimensions[2]};
            int start_y = dc->chunk_info->y;
            int start_x = dc->chunk_info->x;
            int start_z = dc->chunk_info->z;
            chunk_info_init(inner_envelope.top->west, dc->type, dimensions, start_y, start_x, start_z, true);
            futhark_values_i8_3d(fc, fut_inner_envelope->v53, inner_envelope.top->west->data);
        }
    }

    futhark_context_sync(fc);

    futhark_free_u8_3d(fc, fut_chunk_data);
    futhark_free_opaque_envelope_3d_t(fc, fut_inner_envelope);
    return inner_envelope;
}

extern envelope_t get_inner_envelope(struct dispatch_context *dc, struct futhark_context *fc, int thickness) {
    envelope_t inner_envelope = {0};
    switch (dc->n_dimensions) {
        case 1:
            inner_envelope = get_inner_envelope_1d(dc, fc, thickness);
            break;
        case 2:
            inner_envelope = get_inner_envelope_2d(dc, fc, thickness);
            break;
        case 3:
            inner_envelope = get_inner_envelope_3d(dc, fc, thickness);
            break;
        default:
            fprintf(stderr, "Invalid dimensions size.");
            MPI_Abort(MPI_COMM_WORLD, 1);
            break;
    }
    return inner_envelope;
}

static envelope_t get_outer_envelope_1d(struct dispatch_context *dc, int thickness, envelope_t *inner_envelope) {
    int coordinate_x = dc->coordinates[1];
    MPI_Request requests[4];
    int i_request = 0;

    envelope_t outer_envelope = (envelope_t) {0};
    envelope_init_accessors(&outer_envelope);

    // East-part
    {
        int dest_source_x = (coordinate_x + 1) % dc->network_dimensions[1];
        int dest_source = dest_source_x;

        int send_count = (int) inner_envelope->front->east->count;
        MPI_Isend(inner_envelope->front->east->data, send_count, dc->datatype, dest_source_x, FRONT_EAST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);

        int dimensions[3] = {1, min(thickness, dc->chunks_info[dest_source].dimensions[1]), 1};
        int start_x = dc->chunks_info[dest_source].x;
        chunk_info_init(outer_envelope.front->east, dc->type, dimensions, 0, start_x, 0, true);

        int recv_count = (int) outer_envelope.front->east->count;
        MPI_Irecv(outer_envelope.front->east->data, recv_count, dc->datatype, dest_source_x, FRONT_WEST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);
    }

    // West-part
    {
        int dest_source_x = (coordinate_x - 1 >= 0) ? coordinate_x - 1 : dc->network_dimensions[1] - 1;
        int dest_source = dest_source_x;

        int send_count = (int) inner_envelope->front->west->count;
        MPI_Isend(inner_envelope->front->west->data, send_count, dc->datatype, dest_source_x, FRONT_WEST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);

        int dimensions[3] = {1, min(thickness, dc->chunks_info[dest_source].dimensions[1]), 1};
        int start_x = dc->chunks_info[dest_source].x + dc->chunks_info[dest_source].dimensions[1] - dimensions[1];
        chunk_info_init(outer_envelope.front->west, dc->type, dimensions, 0, start_x, 0, true);

        int recv_count = (int) outer_envelope.front->west->count;
        MPI_Irecv(outer_envelope.front->west->data, recv_count, dc->datatype, dest_source_x, FRONT_EAST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);
    }

    MPI_Waitall(i_request, requests, MPI_STATUSES_IGNORE);
    return outer_envelope;
}

static envelope_t get_outer_envelope_2d(struct dispatch_context *dc, int thickness, envelope_t *inner_envelope) {
    int coordinate_y = dc->coordinates[0];
    int coordinate_x = dc->coordinates[1];
    MPI_Request requests[16] = {0};
    int i_request = 0;

    envelope_t outer_envelope = (envelope_t) {0};
    envelope_init_accessors(&outer_envelope);

    // North
    {
        int dest_source_y = (coordinate_y - 1) >= 0 ? (coordinate_y - 1) : (dc->network_dimensions[0] - 1);
        int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, coordinate_x, 0);

        int send_count = (int) inner_envelope->front->north->count;
        MPI_Isend(inner_envelope->front->north->data, send_count, dc->datatype, dest_source_y, FRONT_NORTH_TAG,
                  dc->communicators[COLUMN_COMMUNICATOR], &requests[i_request++]);

        /* Neighbour send south row, which correspond to north envelope */
        int dimensions[3] = {
                min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                dc->chunks_info[dest_source].dimensions[1], 1
        };
        int start_y = dc->chunks_info[dest_source].y + dc->chunks_info[dest_source].dimensions[0] - dimensions[0];
        int start_x = dc->chunks_info[dest_source].x;
        chunk_info_init(outer_envelope.front->north, dc->type, dimensions, start_y, start_x, 0, true);

        int recv_count = (int) outer_envelope.front->north->count;
        MPI_Irecv(outer_envelope.front->north->data, recv_count, dc->datatype, dest_source_y, FRONT_SOUTH_TAG,
                  dc->communicators[COLUMN_COMMUNICATOR], &requests[i_request++]);
    }

    // East
    {
        int dest_source_x = (coordinate_x + 1) % dc->network_dimensions[1];
        int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, dest_source_x, 0);

        int send_count = (int) inner_envelope->front->east->count;
        MPI_Isend(inner_envelope->front->east->data, (int) send_count, dc->datatype, dest_source_x, FRONT_EAST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);

        /* Neighbour send west column, which correspond to east envelope */
        int dimensions[3] = {
                dc->chunks_info[dest_source].dimensions[0],
                min(thickness, dc->chunks_info[dest_source].dimensions[1]), 1
        };
        int start_y = dc->chunks_info[dest_source].y;
        int start_x = dc->chunks_info[dest_source].x;
        chunk_info_init(outer_envelope.front->east, dc->type, dimensions, start_y, start_x, 0, true);

        int recv_count = (int) outer_envelope.front->east->count;
        MPI_Irecv(outer_envelope.front->east->data, recv_count, dc->datatype, dest_source_x, FRONT_WEST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);
    }

    // South
    {
        int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
        int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, coordinate_x, 0);

        int send_count = (int) inner_envelope->front->south->count;
        MPI_Isend(inner_envelope->front->south->data, send_count, dc->datatype, dest_source_y, FRONT_SOUTH_TAG,
                  dc->communicators[COLUMN_COMMUNICATOR], &requests[i_request++]);

        /* Neighbour send north row, which correspond to south envelope */
        int dimensions[3] = {
                min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                dc->chunks_info[dest_source].dimensions[1], 1
        };
        int start_y = dc->chunks_info[dest_source].y;
        int start_x = dc->chunks_info[dest_source].x;
        chunk_info_init(outer_envelope.front->south, dc->type, dimensions, start_y, start_x, 0, true);

        int recv_count = (int) outer_envelope.front->south->count;
        MPI_Irecv(outer_envelope.front->south->data, recv_count, dc->datatype, dest_source_y, FRONT_NORTH_TAG,
                  dc->communicators[COLUMN_COMMUNICATOR], &requests[i_request++]);
    }

    // West
    {
        int dest_source_x = (coordinate_x - 1) >= 0 ? coordinate_x - 1 : dc->network_dimensions[1] - 1;
        int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, dest_source_x, 0);

        int send_count = (int) inner_envelope->front->west->count;
        MPI_Isend(inner_envelope->front->west->data, send_count, dc->datatype, dest_source_x, FRONT_WEST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);

        /* Neighbour send east column, which correspond to west envelope */
        int dimensions[3] = {
                dc->chunks_info[dest_source].dimensions[0],
                min(thickness, dc->chunks_info[dest_source].dimensions[1]), 1
        };
        int start_y = dc->chunks_info[dest_source].y;
        int start_x = dc->chunks_info[dest_source].x + dc->chunks_info[dest_source].dimensions[1] - dimensions[1];
        chunk_info_init(outer_envelope.front->west, dc->type, dimensions, start_y, start_x, 0, true);

        int recv_count = (int) outer_envelope.front->west->count;
        MPI_Irecv(outer_envelope.front->west->data, recv_count, dc->datatype, dest_source_x, FRONT_EAST_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);
    }

    // North-East
    {
        int dest_source_y = (coordinate_y - 1) >= 0 ? coordinate_y - 1 : dc->network_dimensions[0] - 1;
        int dest_source_x = (coordinate_x + 1) % dc->network_dimensions[1];
        int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, 0);

        int send_count = (int) inner_envelope->front->north_east->count;
        MPI_Isend(inner_envelope->front->north_east->data, send_count, dc->datatype, dest_source, FRONT_NORTH_EAST_TAG,
                  MPI_COMM_WORLD, &requests[i_request++]);

        /* Neighbour send south-west cell, which correspond to north-east cell */
        int dimensions[3] = {thickness, thickness, 1};
        int start_y = dc->chunks_info[dest_source].y + dc->chunks_info[dest_source].dimensions[0] + dimensions[0];
        int start_x = dc->chunks_info[dest_source].x;
        chunk_info_init(outer_envelope.front->north_east, dc->type, dimensions, start_y, start_x, 0, true);

        int recv_count = (int) outer_envelope.front->north_east->count;
        MPI_Irecv(outer_envelope.front->north_east->data, recv_count, dc->datatype, dest_source, FRONT_SOUTH_WEST_TAG,
                  MPI_COMM_WORLD, &requests[i_request++]);
    }

    // South-East
    {
        int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
        int dest_source_x = (coordinate_x + 1) % dc->network_dimensions[1];
        int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, 0);