// File: /workspace/github_c_makefile_star2_O1/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 0x257d
movl 0x4(%rbx), %ecx
testl %ecx, %ecx
jle 0x240f
xorl %edx, %edx
movslq (%r14,%rdx,4), %rcx
incl (%rax,%rcx,4)
incq %rdx
movslq 0x4(%rbx), %rcx
cmpq %rcx, %rdx
jl 0x23fc
testl %ecx, %ecx
jle 0x2466
movl (%rbx), %esi
xorl %edi, %edi
leal (%rdi,%rcx), %r8d
shrl %r8d
movl $0x0, %r9d
testl %esi, %esi
jle 0x2450
xorl %edx, %edx
xorl %r10d, %r10d
xorl %r9d, %r9d
movl (%rax,%rdx,4), %r11d
addl %r11d, %r10d
xorl %r14d, %r14d
cmpl %r8d, %r10d
cmovgl %r11d, %r10d
setg %r14b
addl %r14d, %r9d
incq %rdx
cmpq %rdx, %rsi
jne 0x2430
leal 0x1(%r8), %edx
cmpl %ebp, %r9d
cmovll %r8d, %ecx
cmovll %edi, %edx
movl %edx, %edi
cmpl %ecx, %edx
jl 0x2417
jmp 0x2468
xorl %edx, %edx
xorps %xmm0, %xmm0
movups %xmm0, 0x8(%rbx)
movslq (%rbx), %rcx
testq %rcx, %rcx
jle 0x24d8
xorl %esi, %esi
xorl %edi, %edi
xorl %r8d, %r8d
xorl %r9d, %r9d
movl (%rax,%rsi,4), %r10d
addl %r10d, %r9d
cmpl %edx, %r9d
jle 0x24b7
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 0x24cc
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, %rcx
jne 0x2481
popq %rbx
popq %r14
popq %rbp
retq
nopl (%rax)