ipc_debug/src/ipc_debug.cpp

#include <new>
#include <numeric>
#include <unistd.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <cassert>
#include <errno.h>
#include <iostream>
#include <vector>
#include <mpi.h>
#include <map>
#include <iostream>
#include <fstream>
#include <string>

#include <ipc_debug.h>

bool debug_ipc_is_root = false;
FILE *debug_ipc_file = nullptr;

int rank, cluster_size;
std::map<uint64_t, int> start_indices_map; // maps indices to MPI ranks
std::vector<int> start_indices_list; // maps ranks to start indices
std::vector<int> recv_counts;

std::vector<char> buffer;
std::vector<double> float_buffer;

constexpr size_t INITIAL_BUFFER_SIZE = 5 * 1024 * 1024; // 5 MiB

/** 
 * Additionally to the data necessary to compare between the two processes,
 * the client also sends a single OK_FLAG to signal that it can continue
 * execution. If the assertion on the root process does not pass, it will enter
 * an endless loop, thus causing the write to block on client side as well.
 *
 * This is useful to make sure the client stops execution as well if the
 * assertion does not pass.
 */
const char OK_FLAG = 0x42;

void print_backtrace();

void debug_ipc_init() {
    int mpi_initialized;
    MPI_Initialized(&mpi_initialized);
    if (mpi_initialized) {
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        MPI_Comm_size(MPI_COMM_WORLD, &cluster_size);
    }

    if (rank == 0) {
        const char *root_env = std::getenv("IPC_DEBUG_ROOT");
        const char *file_env = std::getenv("IPC_DEBUG_FILE");
        
        if (file_env == nullptr) {
            return;
        }

        if (root_env != nullptr && 0 == strcmp(root_env, "1")) {
            debug_ipc_is_root = true;
            printf("[IPCDBG] I am host\n");
        } else {
            printf("[IPCDBG] I am client\n");
        }

        if (debug_ipc_is_root) {
            debug_ipc_file = fopen(file_env, "rb");
        } else {
            debug_ipc_file = fopen(file_env, "wb");
        }

        if (debug_ipc_file == nullptr) {
            printf("[IPCDBG] Error, could not open named pipe: %s", strerror(errno));
            exit(-1);
        }

        setvbuf(debug_ipc_file, nullptr, _IONBF, 0);


        buffer.resize(INITIAL_BUFFER_SIZE);
    }

}

template<typename T>
void debug_ipc_assert_equal(T value) {
    if (debug_ipc_file == nullptr) {
        return;
    }

    const size_t expected_size = sizeof(T);
    if (debug_ipc_is_root) {
        T other_value;
        size_t read = fread(&other_value, expected_size, 1, debug_ipc_file);

        if (read != 1) {
            printf("[IPCDBG] Could not read enough bytes. Error: %s\n", strerror(errno));
            exit(-1);
        }

        if (other_value != value) {
            std::cout << "[IPCDBG] Assertion failed!"
                      << " Root has " << value << " but client has " << other_value;
            print_backtrace();
            std::cout << "Entering endless loop, attach debugger to PID " << getpid();
            fflush(stdout);
            while (1) {
                sleep(1);
            }
        } else {
#ifdef TRACE
            std::cout << "[IPCDBG] Assertion passed, value = " << value << std::endl;
#endif
        }

        // Read ok flag so that client can continue
        char status; fread(&status, 1, 1, debug_ipc_file); assert(status == OK_FLAG);
    } else {
        size_t written = fwrite(&value, expected_size, 1, debug_ipc_file);

        if (written != 1) {
            printf("[IPCDBG] Could not write enough bytes. Error: %s\n", strerror(errno));
            exit(-1);
        }

        fflush(debug_ipc_file);
        fwrite(&OK_FLAG, 1, 1, debug_ipc_file);
        fflush(debug_ipc_file);

    }
}

template<typename T>
void debug_ipc_assert_equal_vector(std::vector<T> value) {
    if (debug_ipc_file == nullptr) {
        return;
    }

    debug_ipc_assert_equal(value.size());
    const size_t array_byte_length = sizeof(T) * value.size();

    if (debug_ipc_is_root) {
        buffer.resize(array_byte_length);
        size_t read = fread(buffer.data(), 1, array_byte_length, debug_ipc_file);

        if (read != array_byte_length) {
            printf("[IPCDBG] Could not read enough bytes. Error: %s\n", strerror(errno));
            exit(-1);
        }

        assert(reinterpret_cast<uint64_t>(buffer.data()) % 8 == 0); // Make sure the array is properly aligned
        T *local_array = value.data();
        T *other_array = reinterpret_cast<T *>(buffer.data());

        for (size_t i = 0; i < value.size(); ++i) {
            if (local_array[i] != other_array[i]) {
                std::cout << "[IPCDBG] Assertion failed in vector  at index " << i 
                    << ". Root has " << local_array[i] << " but client has " << other_array[i] << std::endl;
                print_backtrace();
                printf("Entering endless loop, attach debugger to PID %i \n", getpid());
                fflush(stdout);
                while (1) {
                    sleep(1);
                }
            }
        }

        // Read ok flag so that client can continue
        char status; fread(&status, 1, 1, debug_ipc_file); assert(status == OK_FLAG);
    } else {
        fwrite(value.data(), 1, array_byte_length, debug_ipc_file);
        fflush(debug_ipc_file);
        size_t result = fwrite(&OK_FLAG, 1, 1, debug_ipc_file);
        printf("Result: %lu\n", result);
    }
}

// Explicit template instantiation
template void debug_ipc_assert_equal<double>(double);
template void debug_ipc_assert_equal<float>(float);
template void debug_ipc_assert_equal<uint32_t>(uint32_t);
template void debug_ipc_assert_equal<int32_t>(int32_t);
template void debug_ipc_assert_equal<uint64_t>(uint64_t);
template void debug_ipc_assert_equal<int64_t>(int64_t);
template void debug_ipc_assert_equal<bool>(bool);
template void debug_ipc_assert_equal<std::string>(std::string);

void debug_ipc_assert_equal_double(double value) {
    debug_ipc_assert_equal<double>(value);
}
void debug_ipc_assert_equal_uint(uint32_t value) {
    debug_ipc_assert_equal<uint32_t>(value);
}
void debug_ipc_assert_equal_int(int32_t value) {
    debug_ipc_assert_equal<int32_t>(value);
}
void debug_ipc_assert_equal_int64(int64_t value) {
    debug_ipc_assert_equal<int64_t>(value);
}

void debug_ipc_assert_equal_array(void *value, size_t size) {
    static_assert(sizeof(char) == 1);
    if (debug_ipc_file == nullptr) {
        return;
    }
    debug_ipc_assert_equal(size); // Make sure arrays are the same size
    

    char *array = reinterpret_cast<char *>(value);
    if (debug_ipc_is_root) {
        buffer.resize(size);
        char *other_array = buffer.data();

        size_t read = fread(other_array, 1, size, debug_ipc_file);

        if (read != size) {
            printf("[IPCDBG] Could not read enough bytes. Error: %s\n", strerror(errno));
            exit(-1);
        }

        for (size_t i = 0; i < size; i++) {
            if (array[i] != other_array[i]) {
                printf("[IPCDBG] Assertion failed in byte %lu!\n", i);
                print_backtrace();
                printf("Entering endless loop, attach debugger to PID %i \n", getpid());
                while (1) {
                    sleep(1);
                }
            }
        }

        // Read ok flag so that client can continue
        char status; fread(&status, 1, 1, debug_ipc_file); assert(status == OK_FLAG);
    } else {
        size_t written = fwrite(value, 1, size, debug_ipc_file);

        if (written != size) {
            printf("[IPCDBG] Could not write enough bytes. Error: %s\n", strerror(errno));
            exit(-1);
        }

        fflush(debug_ipc_file);
        fwrite(&OK_FLAG, 1, 1, debug_ipc_file);
    }
}

void debug_ipc_assert_source_location(const char *source_file, const long int line_number) {
    std::string fname = source_file;

    debug_ipc_assert_equal_array(fname.data(), fname.size());
    debug_ipc_assert_equal(line_number);
}

#include <backward.hpp>

void print_backtrace()
{
    backward::StackTrace st;
    st.load_here(32);
    backward::Printer p;
    p.print(st);
}

void debug_ipc_mpi_set_data_distribution(int start_index) {
    if (rank == 0) {
        start_indices_map.clear();
        start_indices_list.resize(cluster_size + 1);
    }

    int my_start_index = start_index;
    MPI_Gather(&my_start_index, 1, MPI_INT,
            start_indices_list.data(), 1, MPI_INT,
            0, MPI_COMM_WORLD);

    if (rank == 0) {
        for (int i = 0; i < cluster_size; i++) {
            start_indices_map[start_indices_list[i]] = i;
        }
    }
}




void debug_ipc_assert_equal_mpi_double_array(double *array, size_t array_length) {
    std::vector<int> recv_counts(cluster_size);
    std::vector<int> displacements(cluster_size);
    int local_array_length = static_cast<int>(array_length);

    // Gather local lengths
    MPI_Gather(&local_array_length, 1, MPI_INT,
            recv_counts.data(), 1, MPI_INT,
            0, MPI_COMM_WORLD);

    if (rank == 0) {
        // Setup displacements and receive counts
        int global_index = 0;
        // Iterate over start indices from small to large
        for (const auto& [start_index, m_rank] : start_indices_map) {
            displacements[m_rank] = global_index;
            global_index += recv_counts[m_rank];
        }
        float_buffer.resize(global_index); // Make sure buffer is appropriately sized.
    }

    MPI_Gatherv(array, array_length, MPI_DOUBLE,
            float_buffer.data(), // recvbuf
            recv_counts.data(),  // recv counts
            displacements.data(), // displacements
            MPI_DOUBLE, 0, MPI_COMM_WORLD);

    if (rank == 0) {
        debug_ipc_assert_equal_vector(float_buffer);
    }
}