Нужно распараллелить данные MPI в функции без потери данных

Необходимо распараллелить данные в функции kernel_adi, но у меня вообще ничего не получается. Сейчас данные сходятся, но, если буду распараллеливать функцию, то данные уже будут выводиться неправильные. Сам код

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <sys/time.h>
#include <mpi.h>
#include "adi.h"
 
double second();
 
DATA_TYPE X[N][N];
DATA_TYPE A[N][N];
DATA_TYPE B[N][N];
 
static void init_array(int n, DATA_TYPE X[N][N], DATA_TYPE A[N][N], DATA_TYPE B[N][N]) {
    for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++) {
            X[i][j] = ((DATA_TYPE) i*(j+1) + 1) / n;
            A[i][j] = ((DATA_TYPE) i*(j+2) + 2) / n;
            B[i][j] = ((DATA_TYPE) i*(j+3) + 3) / n;
        }
}
 
static void print_array(int n, DATA_TYPE X[N][N]) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            printf("%f ", X[i][j]);
        }
        printf("\n");
    }
}
 
static void kernel_adi(int tsteps, int n, DATA_TYPE X[N][N], DATA_TYPE A[N][N], DATA_TYPE B[N][N], int rank, int size) {
    int chunk_size = (n + size - 1) / size;
    int start_row = rank * chunk_size;
    int end_row = (start_row + chunk_size > n) ? n : (start_row + chunk_size);
    
    for (int t = 0; t < _PB_TSTEPS; t++) {
        for (int i1 = start_row; i1 < end_row; i1++) {
            for (int i2 = 1; i2 < _PB_N; i2++) {
                X[i1][i2] -= X[i1][i2-1] * A[i1][i2] / B[i1][i2-1];
                B[i1][i2] -= A[i1][i2] * A[i1][i2] / B[i1][i2-1];
            }
        }
        MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, X, chunk_size * N, MPI_DOUBLE, MPI_COMM_WORLD);
        MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, B, chunk_size * N, MPI_DOUBLE, MPI_COMM_WORLD);
        
        for (int i1 = start_row; i1 < end_row; i1++)
            X[i1][_PB_N-1] /= B[i1][_PB_N-1];
        MPI_Allreduce(MPI_IN_PLACE, X, N * N, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
 
        for (int i1 = start_row; i1 < end_row; i1++) {
            for (int i2 = 0; i2 < _PB_N-2; i2++) {
                X[i1][_PB_N-i2-2] = (X[i1][_PB_N-2-i2] - X[i1][_PB_N-2-i2-1] * A[i1][_PB_N-i2-3]) / B[i1][_PB_N-3-i2];
            }
        }
        MPI_Allreduce(MPI_IN_PLACE, X, N * N, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    }
}
 
int main(int argc, char** argv) {
    int n = N;
    int tsteps = TSTEPS;
    int rank, size;
    double time0, time1;
 
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
 
    if (rank == 0) {
        init_array(n, X, A, B);
    }
    MPI_Bcast(X, N*N, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(A, N*N, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(B, N*N, MPI_DOUBLE, 0, MPI_COMM_WORLD);
 
    MPI_Barrier(MPI_COMM_WORLD);
    time0 = MPI_Wtime();
    
    kernel_adi(tsteps, n, X, A, B, rank, size);
    
    MPI_Barrier(MPI_COMM_WORLD);
    time1 = MPI_Wtime();
 
    if (rank == 0) {
        printf("\nn=%d\n", n);
        printf("\ntime=%f\n", time1 - time0);
        print_array(n, X);
    }
    
    MPI_Finalize();
    return 0;
}

Функция, которую необходимо распараллелить

static void kernel_adi(int tsteps, int n, DATA_TYPE X[N][N], DATA_TYPE A[N][N], DATA_TYPE B[N][N], int rank, int size) {
    int chunk_size = (n + size - 1) / size;
    int start_row = rank * chunk_size;
    int end_row = (start_row + chunk_size > n) ? n : (start_row + chunk_size);
    
    for (int t = 0; t < _PB_TSTEPS; t++) {
        for (int i1 = start_row; i1 < end_row; i1++) {
            for (int i2 = 1; i2 < _PB_N; i2++) {
                X[i1][i2] -= X[i1][i2-1] * A[i1][i2] / B[i1][i2-1];
                B[i1][i2] -= A[i1][i2] * A[i1][i2] / B[i1][i2-1];
            }
        }
        MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, X, chunk_size * N, MPI_DOUBLE, MPI_COMM_WORLD);
        MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, B, chunk_size * N, MPI_DOUBLE, MPI_COMM_WORLD);
        
        for (int i1 = start_row; i1 < end_row; i1++)
            X[i1][_PB_N-1] /= B[i1][_PB_N-1];
        MPI_Allreduce(MPI_IN_PLACE, X, N * N, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
 
        for (int i1 = start_row; i1 < end_row; i1++) {
            for (int i2 = 0; i2 < _PB_N-2; i2++) {
                X[i1][_PB_N-i2-2] = (X[i1][_PB_N-2-i2] - X[i1][_PB_N-2-i2-1] * A[i1][_PB_N-i2-3]) / B[i1][_PB_N-3-i2];
            }
        }
        MPI_Allreduce(MPI_IN_PLACE, X, N * N, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    }
}