// File: /workspace/github_c_makefile_star2_O3/kkanellis[P]matvec-sparse/partition.c
// Function: partition_equal_nz_elements
void partition_equal_nz_elements(proc_info_t *proc_info, int nprocs, const int *row_idx) 
{
    /* calculate the number of non-zero elements for each row */
    int *row_elements = (int *)calloc_or_exit( proc_info[0].N, sizeof(int) );
    for (int i = 0; i < proc_info[0].NZ; i++) {
        row_elements[ row_idx[i] ]++;
    }

    /* Binary search for minimuzing the sum of workload (nz-elements)
     * per process. Each process receives consequtive rows */
    int l = 0, h = proc_info[0].NZ;
    while (l < h) {
        int nz = (l + h) / 2;
        
        /* calculate how many process are needed so that each
         * process receives up to nz elements */
        int sum = 0, procs_need = 0;
        for (int i = 0; i < proc_info[0].N; i++) {
            if (sum + row_elements[i] > nz ) {
                sum = row_elements[i];
                procs_need++;
            }
            else sum += row_elements[i];
        }

        /* update searching range */
        if (procs_need <= (nprocs - 1))
            h = nz;
        else
            l = nz + 1;
    }
    
    /* initialize process 0 info */
    proc_info[0] = (proc_info_t) {proc_info[0].N, proc_info[0].NZ, 0, 0, 0, 0};

    /* split rows to process - maximum workload is `l` non-zero elements */
    int sum = 0, total_sum = 0, k = 0;
    for (int i = 0; i < proc_info[0].N; i++) {
        if (sum + row_elements[i] > l) {
            sum = row_elements[i];
            k++;

            /* update process info */
            proc_info[k].row_count = 1;
            proc_info[k].row_start_idx = i;
            proc_info[k].nz_count = row_elements[i];
            proc_info[k].nz_start_idx = total_sum;
        }
        else {
            sum += row_elements[i];
            
            /* update process info */
            proc_info[k].row_count++;
            proc_info[k].nz_count += row_elements[i];
        }
        total_sum += row_elements[i];
    }
    
    for (int k = 0; k < nprocs; k++) {
        debug("Processor %d takes rows %3d till %3d [%3d] (entries %6d till %6d) [%6d]\n", 
                        k, 
                        proc_info[k].row_start_idx, 
                        proc_info[k].row_start_idx + proc_info[k].row_count - 1,
                        proc_info[k].row_count,
                        proc_info[k].nz_start_idx, 
                        proc_info[k].nz_start_idx + proc_info[k].nz_count - 1,
                        proc_info[k].nz_count
        );
    }
}

// Assembly:
pushq %rbp
pushq %r14
pushq %rbx
movq %rdx, %r14
movl %esi, %ebp
movq %rdi, %rbx
movslq (%rdi), %rdi
movl $0x4, %esi
callq 0x2594
cmpl $0x0, 0x4(%rbx)
jle 0x247a
xorl %ecx, %ecx
movslq (%r14,%rcx,4), %rdx
incl (%rax,%rdx,4)
incq %rcx
movslq 0x4(%rbx), %rsi
cmpq %rsi, %rcx
jl 0x2415
movl (%rbx), %edx
testl %esi, %esi
jle 0x247c
xorl %ecx, %ecx
leal (%rcx,%rsi), %edi
shrl %edi
movl $0x0, %r8d
testl %edx, %edx
jle 0x2468
xorl %r9d, %r9d
xorl %r10d, %r10d
xorl %r8d, %r8d
movl (%rax,%r9,4), %r11d
addl %r11d, %r10d
xorl %r14d, %r14d
cmpl %edi, %r10d
cmovgl %r11d, %r10d
setg %r14b
addl %r14d, %r8d
incq %r9
cmpq %r9, %rdx
jne 0x2448
cmpl %ebp, %r8d
cmovll %edi, %esi
leal 0x1(%rdi), %edi
cmovgel %edi, %ecx
cmpl %esi, %ecx
jl 0x2430
jmp 0x247e
movl (%rbx), %edx
xorl %ecx, %ecx
xorps %xmm0, %xmm0
movups %xmm0, 0x8(%rbx)
testl %edx, %edx
jle 0x24ec
movl %edx, %edx
xorl %esi, %esi
xorl %edi, %edi
xorl %r8d, %r8d
xorl %r9d, %r9d
movl (%rax,%rsi,4), %r10d
addl %r10d, %r9d
cmpl %ecx, %r9d
jle 0x24cb
incl %edi
movslq %edi, %r9
leaq (%r9,%r9,2), %r9
movl $0x1, 0x10(%rbx,%r9,8)
movl %esi, 0x14(%rbx,%r9,8)
movl (%rax,%rsi,4), %r11d
movl %r11d, 0x8(%rbx,%r9,8)
movl %r8d, 0xc(%rbx,%r9,8)
movl %r10d, %r9d
jmp 0x24e0
movslq %edi, %r10
leaq (%r10,%r10,2), %r10
incl 0x10(%rbx,%r10,8)
movl (%rax,%rsi,4), %r11d
addl %r11d, 0x8(%rbx,%r10,8)
addl (%rax,%rsi,4), %r8d
incq %rsi
cmpq %rsi, %rdx
jne 0x2495
popq %rbx
popq %r14
popq %rbp
retq
nopl (%rax)