Skip to content
Snippets Groups Projects
dispatch.c 98.1 KiB
Newer Older
baptiste.coudray's avatar
baptiste.coudray committed
        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,
baptiste.coudray's avatar
baptiste.coudray committed
                  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);
baptiste.coudray's avatar
baptiste.coudray committed
        int send_count = (int) inner_envelope->front->south_east->count;
        MPI_Isend(inner_envelope->front->south_east->data, send_count, dc->datatype, dest_source, FRONT_SOUTH_EAST_TAG,
baptiste.coudray's avatar
baptiste.coudray committed
                  MPI_COMM_WORLD, &requests[i_request++]);

        /* Neighbour send north-west cell, which correspond to south-east cell */
        int dimensions[3] = {thickness, thickness, 1};
baptiste.coudray's avatar
baptiste.coudray committed
        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_east, dc->type, dimensions, start_y, start_x, 0, true);
        int recv_count = (int) outer_envelope->front->south_east->count;
        MPI_Irecv(outer_envelope->front->south_east->data, recv_count, dc->datatype, dest_source, FRONT_NORTH_WEST_TAG,
baptiste.coudray's avatar
baptiste.coudray committed
                  MPI_COMM_WORLD, &requests[i_request++]);
    }

    // South-West
    {
        int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
        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, dest_source_y, dest_source_x, 0);
baptiste.coudray's avatar
baptiste.coudray committed
        int send_count = (int) inner_envelope->front->south_west->count;
        MPI_Isend(inner_envelope->front->south_west->data, send_count, dc->datatype, dest_source, FRONT_SOUTH_WEST_TAG,
baptiste.coudray's avatar
baptiste.coudray committed
                  MPI_COMM_WORLD, &requests[i_request++]);
        /* Neighbour send north-east cell, which correspond to south-west cell */
        int dimensions[3] = {thickness, thickness, 1};
baptiste.coudray's avatar
baptiste.coudray committed
        int start_y = dc->chunks_info[dest_source].y;
baptiste.coudray's avatar
baptiste.coudray committed
        int start_x = dc->chunks_info[dest_source].x + dc->chunks_info[dest_source].dimensions[1] - dimensions[1];
        chunk_info_init(outer_envelope->front->south_west, dc->type, dimensions, start_y, start_x, 0, true);
        int recv_count = (int) outer_envelope->front->south_west->count;
        MPI_Irecv(outer_envelope->front->south_west->data, recv_count, dc->datatype, dest_source, FRONT_NORTH_EAST_TAG,
baptiste.coudray's avatar
baptiste.coudray committed
                  MPI_COMM_WORLD, &requests[i_request++]);
    }

    // North-West
    {
        int dest_source_y = (coordinate_y - 1) >= 0 ? coordinate_y - 1 : dc->network_dimensions[0] - 1;
        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, dest_source_y, dest_source_x, 0);
baptiste.coudray's avatar
baptiste.coudray committed
        int send_count = (int) inner_envelope->front->north_west->count;
        MPI_Isend(inner_envelope->front->north_west->data, send_count, dc->datatype, dest_source, FRONT_NORTH_WEST_TAG,
baptiste.coudray's avatar
baptiste.coudray committed
                  MPI_COMM_WORLD, &requests[i_request++]);
        /* Neighbour send south-east cell, which correspond to north-west cell */
        int dimensions[3] = {thickness, thickness, 1};
baptiste.coudray's avatar
baptiste.coudray committed
        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 + dc->chunks_info[dest_source].dimensions[1] - dimensions[1];
        chunk_info_init(outer_envelope->front->north_west, dc->type, dimensions, start_y, start_x, 0, true);
        int recv_count = (int) outer_envelope->front->north_west->count;
        MPI_Irecv(outer_envelope->front->north_west->data, recv_count, dc->datatype, dest_source, FRONT_SOUTH_EAST_TAG,
baptiste.coudray's avatar
baptiste.coudray committed
                  MPI_COMM_WORLD, &requests[i_request++]);
    }

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

static envelope_t *get_outer_envelope_3d(struct dispatch_context *dc, int thickness, envelope_t *inner_envelope) {
baptiste.coudray's avatar
baptiste.coudray committed
    int coordinate_y = dc->coordinates[0];
    int coordinate_x = dc->coordinates[1];
    int coordinate_z = dc->coordinates[2];
    MPI_Request requests[52] = {0};
    int i_request = 0;

    envelope_t *outer_envelope = envelope_new();
baptiste.coudray's avatar
baptiste.coudray committed

    // Back
    {
        // Surface
        {
            int dest_source_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, coordinate_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->back->surface->count;
            MPI_Isend(inner_envelope->back->surface->data, send_count, dc->datatype, dest_source_z, BACK_SURFACE_TAG,
                      dc->communicators[DEPTH_COMMUNICATOR], &requests[i_request++]);

            // Neighbour send front-surface, which correspond to back-surface
            int dimensions[3] = {
                    dc->chunks_info[dest_source].dimensions[0],
                    dc->chunks_info[dest_source].dimensions[1],
                    min(thickness, dc->chunks_info[dest_source].dimensions[2]),
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->back->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->back->surface->data, (int) outer_envelope->back->surface->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source_z, FRONT_SURFACE_TAG, dc->communicators[DEPTH_COMMUNICATOR], &requests[i_request++]);
        }

        // East
        {
            int dest_source_x = (coordinate_x + 1) % dc->network_dimensions[1];
            int dest_source_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->back->east->count;

            MPI_Isend(inner_envelope->back->east->data, send_count, dc->datatype, dest_source, BACK_EAST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send front-west, which correspond to back-east
            int dimensions[3] = {
                    dc->chunks_info[dest_source].dimensions[0],
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2]),
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->back->east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->back->east->data, (int) outer_envelope->back->east->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, FRONT_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // West
        {
            int dest_source_x = (coordinate_x - 1) >= 0 ? (coordinate_x - 1) : dc->network_dimensions[1] - 1;
            int dest_source_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->back->west->count;

            MPI_Isend(inner_envelope->back->west->data, send_count, dc->datatype, dest_source, BACK_WEST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send front-east, which correspond to back-west
            int dimensions[3] = {
                    dc->chunks_info[dest_source].dimensions[0],
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2]),
            };
            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];
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->back->west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->back->west->data, (int) outer_envelope->back->west->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, FRONT_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }
baptiste.coudray's avatar
baptiste.coudray committed

    // Bottom
    {
        // Surface
        {
baptiste.coudray's avatar
baptiste.coudray committed
            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, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->surface->count;

            MPI_Isend(inner_envelope->bottom->surface->data, send_count, dc->datatype, dest_source_y,
                      BOTTOM_SURFACE_TAG, dc->communicators[COLUMN_COMMUNICATOR], &requests[i_request++]);

            // Neighbour send top-surface, which correspond to bottom-surface
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    dc->chunks_info[dest_source].dimensions[1],
                    dc->chunks_info[dest_source].dimensions[2],
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->bottom->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->surface->data, (int) outer_envelope->bottom->surface->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source_y, TOP_SURFACE_TAG, dc->communicators[COLUMN_COMMUNICATOR], &requests[i_request++]);
        }

        // North-West
        {
            int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
baptiste.coudray's avatar
baptiste.coudray committed
            int dest_source_x = (coordinate_x - 1) >= 0 ? (coordinate_x - 1) : dc->network_dimensions[1] - 1;
            int dest_source_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->north_west->count;
            MPI_Isend(inner_envelope->bottom->north_west->data, send_count, dc->datatype, dest_source,
                      BOTTOM_NORTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-south-east, which correspond to bottom-north-west
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            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];
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->bottom->north_west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->north_west->data, (int) outer_envelope->bottom->north_west->count,
baptiste.coudray's avatar
baptiste.coudray committed
                      dc->datatype, dest_source, TOP_SOUTH_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // North-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_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->north_east->count;

            MPI_Isend(inner_envelope->bottom->north_east->data, send_count, dc->datatype, dest_source,
                      BOTTOM_NORTH_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-south-west, which correspond to bottom-north-east
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->bottom->north_east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->north_east->data, (int) outer_envelope->bottom->north_east->count,
baptiste.coudray's avatar
baptiste.coudray committed
                      dc->datatype, dest_source, TOP_SOUTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // South-West
        {
            int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
            int dest_source_x = (coordinate_x - 1) >= 0 ? (coordinate_x - 1) : dc->network_dimensions[1] - 1;
            int dest_source_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->south_west->count;

            MPI_Isend(inner_envelope->bottom->south_west->data, send_count, dc->datatype, dest_source,
                      BOTTOM_SOUTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-north-east, which correspond to bottom-south-west
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            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];
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->bottom->south_west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->south_west->data, (int) outer_envelope->bottom->south_west->count,
baptiste.coudray's avatar
baptiste.coudray committed
                      dc->datatype, dest_source, TOP_NORTH_EAST_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_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->south_east->count;

            MPI_Isend(inner_envelope->bottom->south_east->data, send_count, dc->datatype, dest_source,
                      BOTTOM_SOUTH_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-north-west, which correspond to bottom-south-east
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->bottom->south_east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->south_east->data, (int) outer_envelope->bottom->south_east->count,
                      dc->datatype, dest_source, TOP_NORTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
baptiste.coudray's avatar
baptiste.coudray committed
        }

        // 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, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->east->count;

            MPI_Isend(inner_envelope->bottom->east->data, send_count, dc->datatype, dest_source, BOTTOM_EAST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-west, which correspond to bottom-east
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    dc->chunks_info[dest_source].dimensions[2]
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->bottom->east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->east->data, (int) outer_envelope->bottom->east->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, TOP_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // West
        {
            int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
            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, dest_source_y, dest_source_x, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->west->count;

            MPI_Isend(inner_envelope->bottom->west->data, send_count, dc->datatype, dest_source, BOTTOM_WEST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-east, which correspond to bottom-west
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    dc->chunks_info[dest_source].dimensions[2]
            };
            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];
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->bottom->west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->west->data, (int) outer_envelope->bottom->west->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, TOP_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // North
        {
            int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
            int dest_source_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, coordinate_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->north->count;

            MPI_Isend(inner_envelope->bottom->north->data, send_count, dc->datatype, dest_source, BOTTOM_NORTH_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-south, which correspond to bottom-north
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    dc->chunks_info[dest_source].dimensions[1],
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->bottom->north, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->north->data, (int) outer_envelope->bottom->north->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, TOP_SOUTH_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // South
        {
            int dest_source_y = (coordinate_y + 1) % dc->network_dimensions[0];
            int dest_source_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, coordinate_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->bottom->south->count;

            MPI_Isend(inner_envelope->bottom->south->data, send_count, dc->datatype, dest_source, BOTTOM_SOUTH_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send top-north, which correspond to bottom-south
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    dc->chunks_info[dest_source].dimensions[1],
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->bottom->south, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->bottom->south->data, (int) outer_envelope->bottom->south->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, TOP_NORTH_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }
    }

    // Front
    {
        // Surface
        {
            int dest_source_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, coordinate_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->front->surface->count;

            MPI_Isend(inner_envelope->front->surface->data, send_count, dc->datatype, dest_source_z, FRONT_SURFACE_TAG,
                      dc->communicators[DEPTH_COMMUNICATOR], &requests[i_request++]);

            // Neighbour send back-surface, which correspond to front-surface
            int dimensions[3] = {
                    dc->chunks_info[dest_source].dimensions[0],
                    dc->chunks_info[dest_source].dimensions[1],
                    min(thickness, dc->chunks_info[dest_source].dimensions[2]),
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->front->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->front->surface->data, (int) outer_envelope->front->surface->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source_z, BACK_SURFACE_TAG, dc->communicators[DEPTH_COMMUNICATOR], &requests[i_request++]);
        }

        // East
        {
            int dest_source_x = (coordinate_x + 1) % dc->network_dimensions[1];
            int dest_source_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->front->east->count;

            MPI_Isend(inner_envelope->front->east->data, send_count, dc->datatype, dest_source, FRONT_EAST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send back-west, which correspond to front-east
            int dimensions[3] = {
                    dc->chunks_info[dest_source].dimensions[0],
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->front->east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->front->east->data, (int) outer_envelope->front->east->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BACK_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // West
        {
            int dest_source_x = (coordinate_x - 1) >= 0 ? (coordinate_x - 1) : dc->network_dimensions[1] - 1;
            int dest_source_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, coordinate_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->front->west->count;

            MPI_Isend(inner_envelope->front->west->data, send_count, dc->datatype, dest_source, FRONT_WEST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send back-east, which correspond to front-west
            int dimensions[3] = {
                    dc->chunks_info[dest_source].dimensions[0],
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2]),
            };
            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];
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->front->west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->front->west->data, (int) outer_envelope->front->west->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BACK_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }
    }

    // Left
    {
        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, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

        int send_count = (int) inner_envelope->left->surface->count;

        MPI_Isend(inner_envelope->left->surface->data, send_count, dc->datatype, dest_source_x, LEFT_SURFACE_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);

        // Neighbour send right-surface, which correspond to left-surface
        int dimensions[3] = {
                dc->chunks_info[dest_source].dimensions[0],
                min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                dc->chunks_info[dest_source].dimensions[2],
        };
        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];
        int start_z = dc->chunks_info[dest_source].z;
        chunk_info_init(outer_envelope->left->surface, dc->type, dimensions, start_y, start_x, start_z, true);
        MPI_Irecv(outer_envelope->left->surface->data, (int) outer_envelope->left->surface->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                  dest_source_x, RIGHT_SURFACE_TAG, dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);
    }

    // Right
    {
        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, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

        int send_count = (int) inner_envelope->right->surface->count;

        MPI_Isend(inner_envelope->right->surface->data, send_count, dc->datatype, dest_source_x, RIGHT_SURFACE_TAG,
                  dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);

        // Neighbour send left-surface, which correspond to right-surface
        int dimensions[3] = {
                dc->chunks_info[dest_source].dimensions[0],
                min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                dc->chunks_info[dest_source].dimensions[2],
        };
        int start_y = dc->chunks_info[dest_source].y;
        int start_x = dc->chunks_info[dest_source].x;
        int start_z = dc->chunks_info[dest_source].z;
        chunk_info_init(outer_envelope->right->surface, dc->type, dimensions, start_y, start_x, start_z, true);
        MPI_Irecv(outer_envelope->right->surface->data, (int) outer_envelope->right->surface->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                  dest_source_x, LEFT_SURFACE_TAG, dc->communicators[ROW_COMMUNICATOR], &requests[i_request++]);
    }

    // Top
    {
        // Surface
        {
baptiste.coudray's avatar
baptiste.coudray committed
            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, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->surface->count;

            MPI_Isend(inner_envelope->top->surface->data, send_count, dc->datatype, dest_source_y, TOP_SURFACE_TAG,
                      dc->communicators[COLUMN_COMMUNICATOR], &requests[i_request++]);

            // Neighbour send bottom-surface, which correspond to top-surface
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    dc->chunks_info[dest_source].dimensions[1],
                    dc->chunks_info[dest_source].dimensions[2],
            };
            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;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->top->surface, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->surface->data, (int) outer_envelope->top->surface->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source_y, BOTTOM_SURFACE_TAG, dc->communicators[COLUMN_COMMUNICATOR],
                      &requests[i_request++]);
        }

        // North-West
        {
            int dest_source_y = (coordinate_y - 1) >= 0 ? (coordinate_y - 1) : dc->network_dimensions[0] - 1;
            int dest_source_x = (coordinate_x - 1) >= 0 ? (coordinate_x - 1) : dc->network_dimensions[1] - 1;
            int dest_source_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->north_west->count;

            MPI_Isend(inner_envelope->top->north_west->data, send_count, dc->datatype, dest_source,
                      TOP_NORTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-south-east, which correspond to top-north-west
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            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 + dc->chunks_info[dest_source].dimensions[1] - dimensions[1];
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->top->north_west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->north_west->data, (int) outer_envelope->top->north_west->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BOTTOM_SOUTH_EAST_TAG, MPI_COMM_WORLD, &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_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->north_east->count;

            MPI_Isend(inner_envelope->top->north_east->data, send_count, dc->datatype, dest_source,
                      TOP_NORTH_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-south-west, which correspond to top-north-east
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            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;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->top->north_east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->north_east->data, (int) outer_envelope->top->north_east->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BOTTOM_SOUTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // South-West
        {
            int dest_source_y = (coordinate_y - 1) >= 0 ? (coordinate_y - 1) : dc->network_dimensions[0] - 1;
            int dest_source_x = (coordinate_x - 1) >= 0 ? (coordinate_x - 1) : dc->network_dimensions[1] - 1;
            int dest_source_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->south_west->count;

            MPI_Isend(inner_envelope->top->south_west->data, send_count, dc->datatype, dest_source,
                      TOP_SOUTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-north-east, which correspond to top-south-west
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            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];
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->top->south_west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->south_west->data, (int) outer_envelope->top->south_west->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BOTTOM_NORTH_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // South-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_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, dest_source_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->south_east->count;

            MPI_Isend(inner_envelope->top->south_east->data, send_count, dc->datatype, dest_source,
                      TOP_SOUTH_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-north-west, which correspond to top-south-east
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[2])
            };
            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];
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->top->south_east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->south_east->data, (int) outer_envelope->top->south_east->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BOTTOM_NORTH_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // 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, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->east->count;

            MPI_Isend(inner_envelope->top->east->data, send_count, dc->datatype, dest_source, TOP_EAST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-west, which correspond to top-east
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    dc->chunks_info[dest_source].dimensions[2]
            };
            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;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->top->east, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->east->data, (int) outer_envelope->top->east->count, dc->datatype,
                      dest_source,
baptiste.coudray's avatar
baptiste.coudray committed
                      BOTTOM_WEST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // West
        {
            int dest_source_y = (coordinate_y - 1) >= 0 ? (coordinate_y - 1) : dc->network_dimensions[0] - 1;
            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, dest_source_y, dest_source_x, coordinate_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->west->count;

            MPI_Isend(inner_envelope->top->west->data, send_count, dc->datatype, dest_source, TOP_WEST_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-east, which correspond to top-west
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    min(thickness, dc->chunks_info[dest_source].dimensions[1]),
                    dc->chunks_info[dest_source].dimensions[2]
            };
            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 + dc->chunks_info[dest_source].dimensions[1] - dimensions[1];
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->top->west, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->west->data, (int) outer_envelope->top->west->count, dc->datatype,
                      dest_source,
baptiste.coudray's avatar
baptiste.coudray committed
                      BOTTOM_EAST_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // North
        {
            int dest_source_y = (coordinate_y - 1) >= 0 ? (coordinate_y - 1) : dc->network_dimensions[0] - 1;
            int dest_source_z = (coordinate_z + 1) % dc->network_dimensions[2];
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, coordinate_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->north->count;

            MPI_Isend(inner_envelope->top->north->data, send_count, dc->datatype, dest_source, TOP_NORTH_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-south, which correspond to top-north
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    dc->chunks_info[dest_source].dimensions[1],
                    min(thickness, dc->chunks_info[dest_source].dimensions[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;
            int start_z = dc->chunks_info[dest_source].z;
            chunk_info_init(outer_envelope->top->north, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->north->data, (int) outer_envelope->top->north->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BOTTOM_SOUTH_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }

        // South
        {
            int dest_source_y = (coordinate_y - 1) >= 0 ? (coordinate_y - 1) : dc->network_dimensions[0] - 1;
            int dest_source_z = (coordinate_z - 1) >= 0 ? (coordinate_z - 1) : dc->network_dimensions[2] - 1;
            int dest_source = cart_rank_to_comm_rank(dc, dest_source_y, coordinate_x, dest_source_z);
baptiste.coudray's avatar
baptiste.coudray committed

            int send_count = (int) inner_envelope->top->south->count;

            MPI_Isend(inner_envelope->top->south->data, send_count, dc->datatype, dest_source, TOP_SOUTH_TAG,
                      MPI_COMM_WORLD, &requests[i_request++]);

            // Neighbour send bottom-north, which correspond to top-south
            int dimensions[3] = {
                    min(thickness, dc->chunks_info[dest_source].dimensions[0]),
                    dc->chunks_info[dest_source].dimensions[1],
                    min(thickness, dc->chunks_info[dest_source].dimensions[2]),
            };
            int start_y = dc->chunks_info[dest_source].y;
            int start_x = dc->chunks_info[dest_source].x;
            int start_z = dc->chunks_info[dest_source].z + dc->chunks_info[dest_source].dimensions[2] - dimensions[2];
            chunk_info_init(outer_envelope->top->south, dc->type, dimensions, start_y, start_x, start_z, true);
            MPI_Irecv(outer_envelope->top->south->data, (int) outer_envelope->top->south->count, dc->datatype,
baptiste.coudray's avatar
baptiste.coudray committed
                      dest_source, BOTTOM_NORTH_TAG, MPI_COMM_WORLD, &requests[i_request++]);
        }
    }

    MPI_Waitall(i_request, requests, MPI_STATUSES_IGNORE);
    return outer_envelope;
extern envelope_t *get_outer_envelope(struct dispatch_context *dc, struct futhark_context *fc, int thickness) {
    envelope_t *inner_envelope = get_inner_envelope(dc, fc, thickness);
    envelope_t *outer_envelope = NULL;
    switch (dc->n_dimensions) {
        case 1:
            outer_envelope = get_outer_envelope_1d(dc, thickness, inner_envelope);
            outer_envelope = get_outer_envelope_2d(dc, thickness, inner_envelope);
            outer_envelope = get_outer_envelope_3d(dc, thickness, inner_envelope);
            break;
        default:
            fprintf(stderr, "Invalid dimensions size.");
            MPI_Abort(MPI_COMM_WORLD, 1);
            break;
    }
    envelope_free(inner_envelope);
baptiste.coudray's avatar
baptiste.coudray committed
static uint8_t *
chunk_info_to_futhark_struct(chunk_info_t *ci, struct dispatch_context *dc, uint8_t *out, char *futhark_type) {
    /* Taken from Futhark generated code */
    *out++ = 'b';
    *out++ = 2;
    *out++ = dc->n_dimensions;
    memcpy(out, futhark_type, 4);
    out += 4;

    int64_t dimensions64[3];
    dimensions64[0] = (int64_t) ci->dimensions[0];
    dimensions64[1] = (int64_t) ci->dimensions[1];
    dimensions64[2] = (int64_t) ci->dimensions[2];

    if (dc->n_dimensions == 1) {
        dimensions64[1] *= dc->type;
baptiste.coudray's avatar
baptiste.coudray committed
        memcpy(out, &dimensions64[1], 1 * sizeof(int64_t));
    } else if (dc->n_dimensions == 2) {
        dimensions64[1] *= dc->type;
baptiste.coudray's avatar
baptiste.coudray committed
        memcpy(out, dimensions64, 2 * sizeof(int64_t));
    } else {
        dimensions64[2] *= dc->type;
baptiste.coudray's avatar
baptiste.coudray committed
        memcpy(out, dimensions64, 3 * sizeof(int64_t));
    }
    out += dc->n_dimensions * sizeof(int64_t);
baptiste.coudray's avatar
baptiste.coudray committed
    if (ci->data != NULL) {
        memcpy(out, ci->data, ci->count * dc->type);
    }
baptiste.coudray's avatar
baptiste.coudray committed
    out += ci->count * dc->type;
    return out;
}

static void *
futhark_outer_envelope_1d_new(struct dispatch_context *dc, struct futhark_context *fc, envelope_t *outer_envelope,
                              void *f(struct futhark_context *, const void *), char *futhark_type) {

    /* Taken from Futhark generated code */
    int64_t size_0 = 7 + 1 * sizeof(int64_t) + outer_envelope->front->east->count * dc->type;
    int64_t size_1 = 7 + 1 * sizeof(int64_t) + outer_envelope->front->west->count * dc->type;

    void *opaque_struct = calloc(size_0 + size_1, sizeof(uint8_t));
    uint8_t *opaque_struct8 = (uint8_t *) opaque_struct;

    // East
    opaque_struct8 = chunk_info_to_futhark_struct(outer_envelope->front->east, dc, opaque_struct8, futhark_type);

    // West
    chunk_info_to_futhark_struct(outer_envelope->front->west, dc, opaque_struct8, futhark_type);

    void *fut_opaque_struct = f(fc, opaque_struct);
    assert(fut_opaque_struct != NULL);
baptiste.coudray's avatar
baptiste.coudray committed
    futhark_context_sync(fc);

    free(opaque_struct);
    return fut_opaque_struct;
}

static void *
futhark_outer_envelope_2d_new(struct dispatch_context *dc, struct futhark_context *fc, envelope_t *outer_envelope,
                              void *f(struct futhark_context *, const void *), char *futhark_type) {
    size_t total_size = 0;

    for (int i = 0; i < NB_CHUNKS; ++i) {
        int j = FUTHARK_CHUNKS_ORDER[i];
        if (j != INDEX_CHUNK_SURFACE) {
            total_size += (7 + 2 * sizeof(int64_t) + outer_envelope->front->chunks[j].count * dc->type);
        }
    }
    void *opaque_struct = calloc(total_size, sizeof(uint8_t));
    uint8_t *opaque_struct8 = (uint8_t *) opaque_struct;
    for (int i = 0; i < NB_CHUNKS; ++i) {
        int j = FUTHARK_CHUNKS_ORDER[i];
        if (j != INDEX_CHUNK_SURFACE) {
            opaque_struct8 = chunk_info_to_futhark_struct(&outer_envelope->front->chunks[j], dc, opaque_struct8,
                                                          futhark_type);
        }
    }

    void *fut_opaque_struct = f(fc, opaque_struct);
    futhark_context_sync(fc);

    free(opaque_struct);
    return fut_opaque_struct;
}

static void *
futhark_outer_envelope_3d_new(struct dispatch_context *dc, struct futhark_context *fc, envelope_t *outer_envelope,
                              void *f(struct futhark_context *, const void *), char *futhark_type) {
    size_t total_size = 0;
    for (int i = 0; i < NB_SIDES; ++i) {
        for (int j = 0; j < NB_CHUNKS; ++j) {
            int k = FUTHARK_CHUNKS_ORDER[j];
            total_size += (7 + 3 * sizeof(int64_t) + outer_envelope->sides[i].chunks[k].count * dc->type);
        }
    }

    void *opaque_struct = calloc(total_size, sizeof(uint8_t));
    uint8_t *opaque_struct8 = (uint8_t *) opaque_struct;

    for (int i = 0; i < NB_SIDES; ++i) {
        for (int j = 0; j < NB_CHUNKS; ++j) {
            int k = FUTHARK_CHUNKS_ORDER[j];
            opaque_struct8 = chunk_info_to_futhark_struct(&outer_envelope->sides[i].chunks[k], dc, opaque_struct8,
                                                          futhark_type);
        }
    }

    void *fut_opaque_struct = f(fc, opaque_struct);
    futhark_context_sync(fc);

    free(opaque_struct);
    return fut_opaque_struct;
}

static void *
get_chunk_with_envelope_1d(struct dispatch_context *dc, struct futhark_context *fc, envelope_t *outer_envelope,
baptiste.coudray's avatar
baptiste.coudray committed
                           void *f(struct futhark_context *, const void *, int64_t)) {
    struct futhark_u8_1d *fut_chunk_xs = futhark_new_u8_1d(fc, dc->chunk_info->data,
                                                           dc->chunk_info->dimensions[1] * dc->type);
    struct futhark_u8_1d *fut_chunk_xs_with_envelope;
    struct futhark_opaque_envelope_1d_t *fut_outer_envelope =
            futhark_outer_envelope_1d_new(dc, fc, outer_envelope, futhark_restore_opaque_envelope_1d_t, "  i8");
    futhark_context_sync(fc);
    futhark_entry_augment_1d(fc, &fut_chunk_xs_with_envelope, fut_chunk_xs, fut_outer_envelope);
    futhark_context_sync(fc);

    int64_t dim = futhark_shape_u8_1d(fc, fut_chunk_xs_with_envelope)[0];
    uint8_t *data = calloc((size_t) dim, sizeof(uint8_t));
    futhark_values_u8_1d(fc, fut_chunk_xs_with_envelope, data);
    futhark_context_sync(fc);

    void *res = f(fc, data, dim / dc->type);
    futhark_context_sync(fc);

    free(data);
    futhark_free_u8_1d(fc, fut_chunk_xs);
    futhark_free_u8_1d(fc, fut_chunk_xs_with_envelope);
    futhark_free_opaque_envelope_1d_t(fc, fut_outer_envelope);

    return res;
}

static void *
get_chunk_with_envelope_2d(struct dispatch_context *dc, struct futhark_context *fc, envelope_t *outer_envelope,
baptiste.coudray's avatar
baptiste.coudray committed
                           void *f(struct futhark_context *, const void *, int64_t, int64_t)) {
    struct futhark_u8_2d *fut_chunk_elems = futhark_new_u8_2d(fc, dc->chunk_info->data,
                                                              dc->chunk_info->dimensions[0],
                                                              dc->chunk_info->dimensions[1] * dc->type);
    struct futhark_u8_2d *fut_chunk_elems_with_envelope;
    struct futhark_opaque_envelope_2d_t *fut_outer_envelope =
            futhark_outer_envelope_2d_new(dc, fc, outer_envelope, futhark_restore_opaque_envelope_2d_t, "  i8");
    futhark_context_sync(fc);
    futhark_entry_augment_2d(fc, &fut_chunk_elems_with_envelope, fut_chunk_elems, fut_outer_envelope);
    futhark_context_sync(fc);

    const int64_t *dims = futhark_shape_u8_2d(fc, fut_chunk_elems_with_envelope);
    uint8_t *data = calloc((size_t) dims[0] * (size_t) dims[1], sizeof(uint8_t));
    futhark_values_u8_2d(fc, fut_chunk_elems_with_envelope, data);
    futhark_context_sync(fc);

    void *res = f(fc, data, dims[0], dims[1] / dc->type);
    futhark_context_sync(fc);

    free(data);
    futhark_free_u8_2d(fc, fut_chunk_elems);
    futhark_free_u8_2d(fc, fut_chunk_elems_with_envelope);
    futhark_free_opaque_envelope_2d_t(fc, fut_outer_envelope);

    return res;
}

static void *
get_chunk_with_envelope_3d(struct dispatch_context *dc, struct futhark_context *fc, envelope_t *outer_envelope,
baptiste.coudray's avatar
baptiste.coudray committed
                           void *f(struct futhark_context *, const void *, int64_t, int64_t, int64_t)) {
    struct futhark_u8_3d *fut_chunk_cube = 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);
    struct futhark_u8_3d *fut_chunk_cube_with_envelope;
    struct futhark_opaque_envelope_3d_t *fut_outer_envelope =
            futhark_outer_envelope_3d_new(dc, fc, outer_envelope, futhark_restore_opaque_envelope_3d_t, "  i8");
    futhark_context_sync(fc);
    futhark_entry_augment_3d(fc, &fut_chunk_cube_with_envelope, fut_chunk_cube, fut_outer_envelope);
    futhark_context_sync(fc);

    const int64_t *dims = futhark_shape_u8_3d(fc, fut_chunk_cube_with_envelope);
    uint8_t *data = calloc((size_t) dims[0] * (size_t) dims[1] * (size_t) dims[2], sizeof(uint8_t));
    futhark_values_u8_3d(fc, fut_chunk_cube_with_envelope, data);
    futhark_context_sync(fc);

    void *res = f(fc, data, dims[0], dims[1], dims[2] / dc->type);
    futhark_context_sync(fc);

    free(data);
    futhark_free_u8_3d(fc, fut_chunk_cube);
    futhark_free_u8_3d(fc, fut_chunk_cube_with_envelope);
    futhark_free_opaque_envelope_3d_t(fc, fut_outer_envelope);

    return res;