// File: /workspace/github_cpp_makefile_star2_O3/WeihuaZhu[P]hpc_prog2/mpi_jacobi.cpp
// Function: distributed_jacobi(int, double*, double*, double*, ompi_communicator_t*, int, double)
void distributed_jacobi(const int n, double* local_A, double* local_b, double* local_x,
                MPI_Comm comm, int max_iter, double l2_termination)
{
    // obtain the necessary parameters required in thie method
    int rank;
    MPI_Comm_rank(comm, &rank);
    int p;
    MPI_Comm_size(comm, &p);
    int q = sqrt(p);
    int dims[2];
    int coords[2];
    int periods[2];
    MPI_Cart_get(comm, 2, dims, periods, coords);

    MPI_Status stat;
    int num_iter = 0; //num_iter counts the numebr of iterations in jacobi iteration
    int sign = 1; //sign defines whether to continue iteration in jacobi iteration
    MPI_Comm n_comm;
    double Ax_b_square; //stores the intermediate value of (Ax - b) ^ 2
    double norm; //norm stores the l2-norm value
    //initialization of the dimention size in each processor:
    int count_dim_0 = block_decompose_by_dim(n, comm, 0);
    int count_dim_1 = block_decompose_by_dim(n, comm, 1);

    //initialization of D, R
    double *local_D = (double*) malloc(sizeof(double) * count_dim_0);
    double *local_R = (double*) malloc(sizeof(double) * (count_dim_0 * count_dim_1));
    for (int i = 0; i < count_dim_0; i++) {
        for (int j = 0; j < count_dim_1; j++) {
            local_R[count_dim_1 * i + j] = local_A[count_dim_1 * i + j];
        }
    }
    if (coords[0] == coords[1]) {
        for (int i = 0; i < count_dim_0; i++) {
            local_D[i] = local_A[count_dim_1 * i + i];
            local_R[count_dim_1 * i + i] = 0;
        }
    }

    //initialization of x
    for (int i = 0; i < count_dim_0; i++) {
        local_x[i] = 0;
    }

    // diagonal processors send D to first column, A and R are in each processor, D is in the
    // first column
    if (coords[0] == coords[1]) {
        MPI_Send(local_D, count_dim_0, MPI_DOUBLE, coords[0] * q, 0, comm);
    }
    if (coords[1] == 0) { //coords[0] != 0 is not required, don't nned to avoid dead lock
        MPI_Recv(local_D, count_dim_0, MPI_DOUBLE, (coords[0] * (q + 1)), 0, comm, &stat);
    }

    //split the first colum as a communicator, and all other processors as another comminicator
    int color = 1;
    if (coords[1] == 0) {
        color = 0;
    }
    MPI_Comm_split(comm, color, rank, &n_comm);

    // local_Ax and local_Rx stores the Ax and Rx values
    double *local_Ax = (double*) malloc(sizeof(double) * count_dim_0);
    double *local_Rx = (double*) malloc(sizeof(double) * count_dim_0);

    //jacobi iteration
    while(sign == 1) {
        //calculate Ax and stores the solution in the first column
        distributed_matrix_vector_mult(n, local_A, local_x, local_Ax, comm);
        //compute the 2nd norm of Ax - b in the first column
        if (coords[1] == 0) {
            Ax_b_square = 0;
            for (int i = 0; i < count_dim_0; i++) {
                local_Ax[i] -= local_b[i];
                Ax_b_square += local_Ax[i] * local_Ax[i];
            }
        }
        // calculate the square of the second norm using reduce on first column processors,
        // and store the solution in the first processor
        MPI_Reduce(&Ax_b_square, &norm, 1, MPI_DOUBLE, MPI_SUM, 0, n_comm);

        //calculate norm and update sign & num_iter in the first processor
        if (rank == 0) {
            norm = sqrt(norm);
            if (norm < l2_termination || num_iter > max_iter) {
                sign = 0;
            }
        }
        num_iter++;

        //Broadcast sign to all the processors to avoid infinite loop
        MPI_Bcast(&sign, 1, MPI_INT, 0, comm);

        //update x in the first column
        if (sign == 1) {
            //calculate Rx, and store it in the first column
            distributed_matrix_vector_mult(n, local_R, local_x, local_Rx, comm);
            //calculate x = (D^-1)(b - Rx) in the first column
            if (coords[1] == 0) {
                for (int i = 0; i < count_dim_0; i++) {
                    local_Rx[i] = local_b[i] - local_Rx[i];
                    local_x[i] = local_Rx[i] / local_D[i];
                }
            }
        }
    }
    // free the useless memory
    MPI_Comm_free(&n_comm);
    free(local_Ax);
    free(local_Rx);
    free(local_D);
    free(local_R);
}

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0xb8, %rsp
movsd %xmm0, 0x78(%rsp)
movl %r9d, %ebp
movq %r8, %r14
movq %rcx, 0x38(%rsp)
movq %rdx, %r12
movq %rsi, 0x40(%rsp)
movl %edi, 0x14(%rsp)
leaq 0x34(%rsp), %rsi
movq %r8, %rdi
callq 0x6790
leaq 0x8c(%rsp), %rbx
movq %r14, %rdi
movq %rbx, %rsi
callq 0x65e0
xorps %xmm0, %xmm0
cvtsi2sdl (%rbx), %xmm0
xorpd %xmm1, %xmm1
ucomisd %xmm1, %xmm0
movl %ebp, 0x5c(%rsp)
jb 0x27c29
sqrtsd %xmm0, %xmm0
jmp 0x27c2e
callq 0x63a0
cvttsd2si %xmm0, %eax
movl %eax, 0x30(%rsp)
leaq 0xb0(%rsp), %rdx
leaq 0xa8(%rsp), %rcx
leaq 0xc(%rsp), %r8
movq %r14, %rdi
movl $0x2, %esi
callq 0x6420
movl $0x1, 0x8(%rsp)
leaq 0x90(%rsp), %rbx
leaq 0x50(%rsp), %rcx
leaq 0x48(%rsp), %r13
movq %r14, %rdi
movl $0x2, %esi
movq %rbx, %rdx
movq %r13, %r8
callq 0x6420
movl 0x14(%rsp), %ebp
movl %ebp, %eax
cltd
idivl (%rbx)
xorl %r15d, %r15d
cmpl (%r13), %edx
setg %r15b
addl %eax, %r15d
leaq 0x90(%rsp), %r13
leaq 0x50(%rsp), %rcx
leaq 0x48(%rsp), %rbx
movq %r14, 0x20(%rsp)
movq %r14, %rdi
movl $0x2, %esi
movq %r13, %rdx
movq %rbx, %r8
callq 0x6420
movl %ebp, %eax
cltd
idivl 0x4(%r13)
xorl %ebp, %ebp
cmpl 0x4(%rbx), %edx
setg %bpl
addl %eax, %ebp
movslq %r15d, %r14
leaq (,%r14,8), %rdi
movq %rdi, 0x70(%rsp)
callq 0x6720
movq %rax, %rbx
movl %ebp, %eax
imull %r15d, %eax
movslq %eax, %rdi
shlq $0x3, %rdi
callq 0x6720
movq %r15, 0x18(%rsp)
movl %r15d, %r15d
testl %r14d, %r14d
movq %rax, 0x60(%rsp)
jle 0x2803e
leal -0x1(%rbp), %ecx
leaq 0x8(,%rcx,8), %rcx
movq %rcx, 0x28(%rsp)
xorl %r13d, %r13d
movq %r15, %r14
movq 0x40(%rsp), %rdx
movq %r12, 0x80(%rsp)
testl %ebp, %ebp
jle 0x27d78
movl %r13d, %edx
movq 0x40(%rsp), %rcx
leaq (%rcx,%rdx,8), %rsi
movq %rbx, %r12
movq %rbp, %rbx
movq %rax, %rbp
leaq (%rax,%rdx,8), %rdi
movq 0x28(%rsp), %rdx
callq 0x63e0
movq %rbp, %rax
movq %rbx, %rbp
movq %r12, %rbx
movq 0x80(%rsp), %r12
movq 0x40(%rsp), %rdx
addl %ebp, %r13d
decq %r14
jne 0x27d3b
movl 0xc(%rsp), %ecx
movl 0x10(%rsp), %r13d
cmpl %r13d, %ecx
jne 0x27dba
cmpl $0x0, 0x18(%rsp)
jle 0x27dba
incl %ebp
xorl %esi, %esi
xorl %edi, %edi
movslq %esi, %rsi
movsd (%rdx,%rsi,8), %xmm0
movsd %xmm0, (%rbx,%rdi,8)
movq $0x0, (%rax,%rsi,8)
incq %rdi
addl %ebp, %esi
cmpq %rdi, %r15
jne 0x27d9b
movq 0x18(%rsp), %r14
testl %r14d, %r14d
jle 0x2805a
leal -0x1(%r14), %eax
leaq 0x8(,%rax,8), %rdx
movq 0x38(%rsp), %rdi
xorl %esi, %esi
movl %ecx, %ebp
callq 0x6280
movl %ebp, %ecx
cmpl %r13d, %ebp
jne 0x27e0d
imull 0x30(%rsp), %ecx
movq 0x111ab(%rip), %rdx     # 0x38fa0
movq %rbx, %rdi
movl %r14d, %esi
xorl %r8d, %r8d
movq 0x20(%rsp), %r9
callq 0x61c0
movl 0x10(%rsp), %r13d
movl $0x1, %esi
testl %r13d, %r13d
jne 0x27e53
movl 0x30(%rsp), %ecx
incl %ecx
imull 0xc(%rsp), %ecx
leaq 0x90(%rsp), %rax
movq %rax, (%rsp)
movq 0x1116b(%rip), %rdx     # 0x38fa0
movq %rbx, %rdi
movl %r14d, %esi
xorl %r8d, %r8d
movq 0x20(%rsp), %r9
callq 0x63d0
xorl %esi, %esi
cmpl $0x0, 0x10(%rsp)
setne %sil
movl 0x34(%rsp), %edx
leaq 0x50(%rsp), %rcx
movq 0x20(%rsp), %r14
movq %r14, %rdi
callq 0x68d0
movq 0x70(%rsp), %rbp
movq %rbp, %rdi
callq 0x6720
movq %rax, %r13
movq %rbp, %rdi
callq 0x6720
movq %rax, %rbp
movq $0x0, 0x28(%rsp)
movq 0x38(%rsp), %rdx
movl 0x14(%rsp), %edi
movq 0x40(%rsp), %rsi
movq %r13, %rcx
movq %r14, %r8
callq 0x27a01
cmpl $0x0, 0x10(%rsp)
jne 0x27eed
movq $0x0, 0x48(%rsp)
cmpl $0x0, 0x18(%rsp)
jle 0x27eed
xorpd %xmm0, %xmm0
xorl %eax, %eax
movsd (%r13,%rax,8), %xmm1
subsd (%r12,%rax,8), %xmm1
movsd %xmm1, (%r13,%rax,8)
mulsd %xmm1, %xmm1
addsd %xmm1, %xmm0
incq %rax
cmpq %rax, %r15
jne 0x27ec3
movsd %xmm0, 0x48(%rsp)
movq 0x50(%rsp), %rax
movq %rax, (%rsp)
leaq 0x48(%rsp), %rdi
leaq 0x68(%rsp), %rsi
movl $0x1, %edx
movq 0x11094(%rip), %rcx     # 0x38fa0
movq 0x11055(%rip), %r8      # 0x38f68
xorl %r9d, %r9d
callq 0x6710
cmpl $0x0, 0x34(%rsp)
jne 0x27f6b
movsd 0x68(%rsp), %xmm0
ucomisd 0x10f8(%rip), %xmm0     # 0x29028
jb 0x27f38
sqrtsd %xmm0, %xmm0
jmp 0x27f3d
callq 0x63a0
movsd %xmm0, 0x68(%rsp)
movsd 0x78(%rsp), %xmm1
ucomisd %xmm0, %xmm1
seta %al
movq 0x28(%rsp), %rcx
cmpl 0x5c(%rsp), %ecx
setg %cl
orb %al, %cl
cmpb $0x1, %cl
jne 0x27f6b
movl $0x0, 0x8(%rsp)
leaq 0x8(%rsp), %rdi
movl $0x1, %esi
movq 0x1103c(%rip), %rdx     # 0x38fb8
xorl %ecx, %ecx
movq 0x20(%rsp), %r14
movq %r14, %r8
callq 0x60d0
cmpl $0x1, 0x8(%rsp)
jne 0x28000
movl 0x14(%rsp), %edi
movq 0x60(%rsp), %rsi
movq 0x38(%rsp), %rdx
movq %rbp, %rcx
movq %r14, %r8
callq 0x27a01
movq 0x38(%rsp), %rdx
cmpl $0x0, 0x10(%rsp)
jne 0x27fe4
cmpl $0x0, 0x18(%rsp)
jle 0x27fe4
xorl %eax, %eax
movsd (%r12,%rax,8), %xmm0
subsd (%rbp,%rax,8), %xmm0
movsd %xmm0, (%rbp,%rax,8)
divsd (%rbx,%rax,8), %xmm0
movsd %xmm0, (%rdx,%rax,8)
incq %rax
cmpq %rax, %r15
jne 0x27fc0
movq 0x28(%rsp), %rax
incl %eax
movq %rax, 0x28(%rsp)
cmpl $0x1, 0x8(%rsp)
movq 0x20(%rsp), %r14
je 0x27e92
leaq 0x50(%rsp), %rdi
callq 0x68b0
movq %r13, %rdi
callq 0x6630
movq %rbp, %rdi
callq 0x6630
movq %rbx, %rdi
callq 0x6630
movq 0x60(%rsp), %rdi
callq 0x6630
addq $0xb8, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
movl 0xc(%rsp), %ecx
movl 0x10(%rsp), %r13d
cmpl %r13d, %ecx
movq 0x18(%rsp), %r14
je 0x27de9
jmp 0x27e0d
cmpl %r13d, %ecx
je 0x27de9
jmp 0x27e0d