// File: /workspace/github_c_makefile_star2_O3/kkanellis[P]matvec-sparse/matvec_mpi_bcast.c
// Function: mat_vec_mult_parallel
double* mat_vec_mult_parallel(int rank, int nprocs, proc_info_t *all_proc_info,
                            int *i_idx, int *j_idx, double *values, double *x)
{
    proc_info_t proc_info;  /* info about submatrix; different per process */
    double *res;            /* result of multiplication res = A*x */

    /***** MPI MASTER (root) process only ******/
    int *count, *offset;    /* auxilliary arrays used for Scatterv/Gatherv */

    /* scatter to processors all info that will be needed */
    MPI_Scatter(all_proc_info, 1, proc_info_type, &proc_info,
                    1, proc_info_type, MASTER, MPI_COMM_WORLD);

    /* allocate memory for vectors and submatrixes */
    double *y = (double *)calloc_or_exit( proc_info.row_count, sizeof(double) );
    if (rank != MASTER) {
        x = (double *)malloc_or_exit( proc_info.N * sizeof(double) );
        i_idx = (int *)malloc_or_exit( proc_info.nz_count * sizeof(int) );
        j_idx = (int *)malloc_or_exit( proc_info.nz_count * sizeof(int) );
        values = (double *)malloc_or_exit( proc_info.nz_count * sizeof(double) );
    }
    else {
        res = (double *)malloc_or_exit( proc_info.N * sizeof(double) );
    }

    //debug("[%d] %d %d %d\n", rank, proc_info.N, proc_info.NZ, proc_info.nz_count);
    
    /* broadcast x vector */
    /* TODO: change late to Point-2-Point communication */
    MPI_Bcast(x, proc_info.N, MPI_DOUBLE, MASTER, MPI_COMM_WORLD);

    /* scatter matrix elements to processes */
    if (rank == MASTER) {
        count = (int *)malloc_or_exit( nprocs * sizeof(int) );
        offset = (int *)malloc_or_exit( nprocs * sizeof(int) );
        for (int p = 0; p < nprocs; p++) {
            count[p] = all_proc_info[p].nz_count;
            offset[p] = all_proc_info[p].nz_start_idx;
        }
    }
    MPI_Scatterv(i_idx, count, offset, MPI_INT, i_idx, 
                    proc_info.nz_count, MPI_INT, MASTER, MPI_COMM_WORLD);
    MPI_Scatterv(j_idx, count, offset, MPI_INT, j_idx, 
                    proc_info.nz_count, MPI_INT, MASTER, MPI_COMM_WORLD);
    MPI_Scatterv(values, count, offset, MPI_DOUBLE, values, 
                    proc_info.nz_count, MPI_DOUBLE, MASTER, MPI_COMM_WORLD);

    /* validation check */
    /*for (int i = 0; i < proc_info.nz_count; i++) {
        debug("[%d] %d %d %lf\n", rank, i_idx[i], j_idx[i], values[i]);
    }*/

    /* multiplication kernel */
    for (int k = 0 ; k < proc_info.nz_count; k++) {
        y[ i_idx[k] - proc_info.row_start_idx ] += values[k] * x[ j_idx[k] ];
    }

    /* gather y elements from processes and save it to res */
    if (rank == MASTER) {
        res = (double *)malloc_or_exit( proc_info.N * sizeof(double) );
        for (int p = 0; p < nprocs; p++) {
            count[p] = all_proc_info[p].row_count;
            offset[p] = all_proc_info[p].row_start_idx;
        }
    }
    MPI_Gatherv(y, proc_info.row_count, MPI_DOUBLE, res, count, 
                offset, MPI_DOUBLE, MASTER, MPI_COMM_WORLD);

    /* return final result */
    return res;
}

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0x48, %rsp
movq %r9, %r12
movq %r8, %rbp
movq %rcx, %r13
movq %rdx, %rbx
movl %esi, 0xc(%rsp)
movl %edi, %r15d
movq 0x2aa0(%rip), %rdx      # 0x5130
leaq 0x1c(%rsp), %r14
movq %rbx, %rdi
movl $0x1, %esi
movq %r14, %rcx
movl $0x1, %r8d
movq %rdx, %r9
pushq 0x2931(%rip)            # 0x4fe0
pushq $0x0
callq 0x1170
addq $0x10, %rsp
movslq 0x10(%r14), %rdi
movl $0x8, %esi
callq 0x2644
movq %rax, %r14
leaq 0x1c(%rsp), %rax
movslq (%rax), %rdi
shlq $0x3, %rdi
callq 0x2632
movl %r15d, 0x34(%rsp)
testl %r15d, %r15d
movq %rax, 0x40(%rsp)
je 0x2742
movslq 0x24(%rsp), %rdi
shlq $0x2, %rdi
movq %rax, %r15
callq 0x2632
movq %rax, %r13
movslq 0x24(%rsp), %rdi
shlq $0x2, %rdi
callq 0x2632
movq %rax, %rbp
movslq 0x24(%rsp), %rdi
shlq $0x3, %rdi
callq 0x2632
movq %rax, %r12
movl 0x1c(%rsp), %esi
movq 0x28ac(%rip), %rdx      # 0x4fd8
movq 0x28ad(%rip), %r8       # 0x4fe0
movq %r15, %rdi
xorl %ecx, %ecx
callq 0x1140
jmp 0x27c6
movq 0x80(%rsp), %r15
movl 0x1c(%rsp), %esi
movq 0x2883(%rip), %rdx      # 0x4fd8
movq 0x2884(%rip), %r8       # 0x4fe0
movq %r15, %rdi
xorl %ecx, %ecx
callq 0x1140
movslq 0xc(%rsp), %rax
movq %rax, 0x10(%rsp)
leaq (,%rax,4), %rdi
movq %rdi, 0x38(%rsp)
callq 0x2632
movq %rax, (%rsp)
movq 0x38(%rsp), %rdi
callq 0x2632
cmpl $0x0, 0x10(%rsp)
jle 0x27c2
movl 0xc(%rsp), %r8d
leaq 0xc(%rbx), %rcx
xorl %edx, %edx
movq (%rsp), %rdi
movl -0x4(%rcx), %esi
movl %esi, (%rdi,%rdx,4)
movl (%rcx), %esi
movl %esi, (%rax,%rdx,4)
incq %rdx
addq $0x18, %rcx
cmpq %rdx, %r8
jne 0x27a6
movq %rdi, %rdx
jmp 0x27c6
movq (%rsp), %rdx
movl 0x24(%rsp), %r9d
subq $0x8, %rsp
xorl %r10d, %r10d
movq 0x27d7(%rip), %rcx      # 0x4fb0
movq %r13, %rdi
movq %rdx, %rsi
movq %rdx, 0x8(%rsp)
movq %rax, %rdx
movq %rax, 0x18(%rsp)
movq %r13, %r8
movq 0x27ea(%rip), %rax      # 0x4fe0
pushq %rax
pushq %r10
pushq %rcx
callq 0x10c0
addq $0x20, %rsp
movl 0x24(%rsp), %r9d
subq $0x8, %rsp
movq %rbp, %rdi
movq 0x8(%rsp), %rsi
movq 0x18(%rsp), %rdx
movq 0x2790(%rip), %rcx      # 0x4fb0
movq %rbp, %r8
pushq 0x27b7(%rip)            # 0x4fe0
xorl %eax, %eax
pushq %rax
pushq %rcx
callq 0x10c0
addq $0x20, %rsp
movl 0x24(%rsp), %r9d
subq $0x8, %rsp
movq 0x2792(%rip), %rcx      # 0x4fd8
movq %r12, %rdi
movq 0x8(%rsp), %rsi
movq 0x18(%rsp), %rdx
movq %r12, %r8
pushq 0x2784(%rip)            # 0x4fe0
xorl %eax, %eax
pushq %rax
pushq %rcx
callq 0x10c0
addq $0x20, %rsp
movslq 0x24(%rsp), %rax
testq %rax, %rax
jle 0x28a7
movslq 0x30(%rsp), %rcx
xorl %edx, %edx
movsd (%r12,%rdx,8), %xmm0
movslq (%rbp,%rdx,4), %rsi
movslq (%r13,%rdx,4), %rdi
subq %rcx, %rdi
mulsd (%r15,%rsi,8), %xmm0
addsd (%r14,%rdi,8), %xmm0
movsd %xmm0, (%r14,%rdi,8)
incq %rdx
cmpq %rdx, %rax
jne 0x287a
cmpl $0x0, 0x34(%rsp)
movq 0x40(%rsp), %r15
jne 0x28f4
movslq 0x1c(%rsp), %rdi
shlq $0x3, %rdi
callq 0x2632
movq %rax, %r15
movl 0xc(%rsp), %eax
testl %eax, %eax
jle 0x28f4
movl %eax, %eax
addq $0x14, %rbx
xorl %ecx, %ecx
movq (%rsp), %rsi
movq 0x10(%rsp), %rdi
movl -0x4(%rbx), %edx
movl %edx, (%rsi,%rcx,4)
movl (%rbx), %edx
movl %edx, (%rdi,%rcx,4)
incq %rcx
addq $0x18, %rbx
cmpq %rcx, %rax
jne 0x28dd
movl 0x2c(%rsp), %esi
subq $0x8, %rsp
movq %r14, %rdi
movq 0x26d2(%rip), %rdx      # 0x4fd8
movq %r15, %rcx
movq 0x8(%rsp), %r8
movq 0x18(%rsp), %r9
pushq 0x26c7(%rip)            # 0x4fe0
pushq $0x0
pushq %rdx
callq 0x1200
addq $0x20, %rsp
movq %r15, %rax
addq $0x48, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq