// File: /workspace/github_c_makefile_star2_O2/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
pushq $0x4
popq %rsi
callq 0x243d
xorl %ecx, %ecx
movslq 0x4(%rbx), %rdx
cmpq %rdx, %rcx
jge 0x22b9
movslq (%r14,%rcx,4), %rdx
incl (%rax,%rdx,4)
incq %rcx
jmp 0x22a4
xorl %ecx, %ecx
cmpl %edx, %ecx
jge 0x2307
leal (%rcx,%rdx), %esi
shrl %esi
movl (%rbx), %edi
xorl %r8d, %r8d
testl %edi, %edi
cmovlel %r8d, %edi
xorl %r9d, %r9d
xorl %r10d, %r10d
cmpq %r8, %rdi
je 0x22f7
movl (%rax,%r8,4), %r11d
addl %r11d, %r10d
xorl %r14d, %r14d
cmpl %esi, %r10d
setg %r14b
cmovgl %r11d, %r10d
addl %r14d, %r9d
incq %r8
jmp 0x22d5
cmpl %ebp, %r9d
cmovll %esi, %edx
leal 0x1(%rsi), %esi
cmovll %ecx, %esi
movl %esi, %ecx
jmp 0x22bb
xorps %xmm0, %xmm0
movups %xmm0, 0x8(%rbx)
movl (%rbx), %edx
xorl %esi, %esi
testl %edx, %edx
cmovlel %esi, %edx
xorl %r9d, %r9d
xorl %edi, %edi
xorl %r8d, %r8d
cmpq %rsi, %rdx
je 0x2378
movl (%rax,%rsi,4), %r10d
addl %r10d, %r9d
cmpl %ecx, %r9d
jle 0x235b
incl %r8d
movslq %r8d, %r9
imulq $0x18, %r9, %r9
movl $0x1, 0x10(%rbx,%r9)
movl %esi, 0x14(%rbx,%r9)
movl (%rax,%rsi,4), %r11d
movl %r11d, 0x8(%rbx,%r9)
movl %edi, 0xc(%rbx,%r9)
movl %r10d, %r9d
jmp 0x2370
movslq %r8d, %r10
imulq $0x18, %r10, %r10
incl 0x10(%rbx,%r10)
movl (%rax,%rsi,4), %r11d
addl %r11d, 0x8(%rbx,%r10)
addl (%rax,%rsi,4), %edi
incq %rsi
jmp 0x231f
popq %rbx
popq %r14
popq %rbp
retq
nopl (%rax)