Bitonic Sort (10/30)

#include <bits/stdc++.h>
#include <mpi.h>

using namespace std;

void bitonic_sort_mpi(float* chunk, int chunk_size, int rank, int size) {

    if(rank&1) sort(chunk, chunk+chunk_size,greater<float>());
    else sort(chunk,chunk+chunk_size);

    for (int step = 2; step <= size; step <<= 1) {
        for (int j = step / 2; j > 0; j /= 2) {

            int partner = rank ^ j;

            MPI_Status status;
            float* recv_chunk = new float[chunk_size];
            MPI_Sendrecv(chunk, chunk_size, MPI_FLOAT, partner, 0, recv_chunk, chunk_size, MPI_FLOAT, partner, 0, MPI_COMM_WORLD, &status);
            
            bool dir = ((rank & step) == 0);

            if (rank < partner) {
                if (dir) {
                    for (int i = 0; i < chunk_size; i++) {
                        if (chunk[i] > recv_chunk[i]) swap(chunk[i], recv_chunk[i]);
                    }
                } else {
                    for (int i = 0; i < chunk_size; i++) {
                        if (chunk[i] < recv_chunk[i]) swap(chunk[i], recv_chunk[i]);
                    }
                }
            } else {
                if (dir) {
                    for (int i = 0; i < chunk_size; i++) {
                        if (chunk[i] < recv_chunk[i]) swap(chunk[i], recv_chunk[i]);
                    }
                } else {
                    for (int i = 0; i < chunk_size; i++) {
                        if (chunk[i] > recv_chunk[i]) swap(chunk[i], recv_chunk[i]);
                    }
                }
            }

            // Re-sort after swapping
            if(!dir) sort(chunk, chunk+chunk_size,greater<float>());
            else sort(chunk, chunk+chunk_size);

            delete[] recv_chunk;
            
        }
    }
}

int main(int argc, char* argv[]) {
    MPI_Init(&argc, &argv);

    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    int N = std::atoi(argv[1]);
    int padding_N = 1 << (int)std::ceil(std::log2(N));
    char* input_filename = argv[2];
    char* output_filename = argv[3];

    int chunksize = padding_N / size;
    float* chunk = new float[chunksize];
    int total_offset = chunksize * rank;

    MPI_File input_file, output_file;
    MPI_File_open(MPI_COMM_WORLD, input_filename, MPI_MODE_RDONLY, MPI_INFO_NULL, &input_file);

    if (total_offset + chunksize < N) {
        MPI_File_read_at(input_file, rank * chunksize * sizeof(float), chunk, chunksize, MPI_FLOAT, MPI_STATUS_IGNORE);
    } else {
        if (total_offset >= N) {
            std::fill(chunk, chunk + chunksize, FLT_MAX);
        } else {
            MPI_File_read_at(input_file, sizeof(float) * total_offset, chunk, N - total_offset, MPI_FLOAT, MPI_STATUS_IGNORE);
            std::fill(chunk + (N - total_offset), chunk + chunksize, FLT_MAX);
        }
    }

    MPI_File_close(&input_file);

    bitonic_sort_mpi(chunk, chunksize, rank, size);

    MPI_File_open(MPI_COMM_WORLD, output_filename, MPI_MODE_CREATE | MPI_MODE_WRONLY, MPI_INFO_NULL, &output_file);

    MPI_File_write_at(output_file, rank * chunksize * sizeof(float), chunk, std::min(chunksize, N - total_offset), MPI_FLOAT, MPI_STATUS_IGNORE);

    MPI_File_close(&output_file);

    delete[] chunk;
    MPI_Finalize();

    return 0;
}