// File: /workspace/github_c_makefile_star2_O0/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
movq %rsp, %rbp
subq $0x70, %rsp
movq %rdi, -0x8(%rbp)
movl %esi, -0xc(%rbp)
movq %rdx, -0x18(%rbp)
movq -0x8(%rbp), %rax
movslq (%rax), %rdi
movl $0x4, %esi
callq 0x2ee0
movq %rax, -0x20(%rbp)
movl $0x0, -0x24(%rbp)
movl -0x24(%rbp), %eax
movq -0x8(%rbp), %rcx
cmpl 0x4(%rcx), %eax
jge 0x2b9f
movq -0x20(%rbp), %rax
movq -0x18(%rbp), %rcx
movslq -0x24(%rbp), %rdx
movslq (%rcx,%rdx,4), %rcx
movl (%rax,%rcx,4), %edx
addl $0x1, %edx
movl %edx, (%rax,%rcx,4)
movl -0x24(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x24(%rbp)
jmp 0x2b6f
movl $0x0, -0x28(%rbp)
movq -0x8(%rbp), %rax
movl 0x4(%rax), %eax
movl %eax, -0x2c(%rbp)
movl -0x28(%rbp), %eax
cmpl -0x2c(%rbp), %eax
jge 0x2c5a
movl -0x28(%rbp), %eax
addl -0x2c(%rbp), %eax
movl $0x2, %ecx
cltd
idivl %ecx
movl %eax, -0x30(%rbp)
movl $0x0, -0x34(%rbp)
movl $0x0, -0x38(%rbp)
movl $0x0, -0x3c(%rbp)
movl -0x3c(%rbp), %eax
movq -0x8(%rbp), %rcx
cmpl (%rcx), %eax
jge 0x2c37
movl -0x34(%rbp), %eax
movq -0x20(%rbp), %rcx
movslq -0x3c(%rbp), %rdx
addl (%rcx,%rdx,4), %eax
cmpl -0x30(%rbp), %eax
jle 0x2c19
movq -0x20(%rbp), %rax
movslq -0x3c(%rbp), %rcx
movl (%rax,%rcx,4), %eax
movl %eax, -0x34(%rbp)
movl -0x38(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x38(%rbp)
jmp 0x2c2a
movq -0x20(%rbp), %rax
movslq -0x3c(%rbp), %rcx
movl (%rax,%rcx,4), %eax
addl -0x34(%rbp), %eax
movl %eax, -0x34(%rbp)
jmp 0x2c2c
movl -0x3c(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x3c(%rbp)
jmp 0x2be2
movl -0x38(%rbp), %eax
movl -0xc(%rbp), %ecx
subl $0x1, %ecx
cmpl %ecx, %eax
jg 0x2c4c
movl -0x30(%rbp), %eax
movl %eax, -0x2c(%rbp)
jmp 0x2c55
movl -0x30(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x28(%rbp)
jmp 0x2bb0
movq -0x8(%rbp), %rax
movq -0x8(%rbp), %rcx
movl (%rcx), %ecx
movl %ecx, -0x54(%rbp)
movq -0x8(%rbp), %rcx
movl 0x4(%rcx), %ecx
movl %ecx, -0x50(%rbp)
movl $0x0, -0x4c(%rbp)
movl $0x0, -0x48(%rbp)
movl $0x0, -0x44(%rbp)
movl $0x0, -0x40(%rbp)
movq -0x54(%rbp), %rcx
movq %rcx, (%rax)
movq -0x4c(%rbp), %rcx
movq %rcx, 0x8(%rax)
movq -0x44(%rbp), %rcx
movq %rcx, 0x10(%rax)
movl $0x0, -0x58(%rbp)
movl $0x0, -0x5c(%rbp)
movl $0x0, -0x60(%rbp)
movl $0x0, -0x64(%rbp)
movl -0x64(%rbp), %eax
movq -0x8(%rbp), %rcx
cmpl (%rcx), %eax
jge 0x2dc0
movl -0x58(%rbp), %eax
movq -0x20(%rbp), %rcx
movslq -0x64(%rbp), %rdx
addl (%rcx,%rdx,4), %eax
cmpl -0x28(%rbp), %eax
jle 0x2d58
movq -0x20(%rbp), %rax
movslq -0x64(%rbp), %rcx
movl (%rax,%rcx,4), %eax
movl %eax, -0x58(%rbp)
movl -0x60(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x60(%rbp)
movq -0x8(%rbp), %rax
movslq -0x60(%rbp), %rcx
imulq $0x18, %rcx, %rcx
addq %rcx, %rax
movl $0x1, 0x10(%rax)
movl -0x64(%rbp), %ecx
movq -0x8(%rbp), %rax
movslq -0x60(%rbp), %rdx
imulq $0x18, %rdx, %rdx
addq %rdx, %rax
movl %ecx, 0x14(%rax)
movq -0x20(%rbp), %rax
movslq -0x64(%rbp), %rcx
movl (%rax,%rcx,4), %ecx
movq -0x8(%rbp), %rax
movslq -0x60(%rbp), %rdx
imulq $0x18, %rdx, %rdx
addq %rdx, %rax
movl %ecx, 0x8(%rax)
movl -0x5c(%rbp), %ecx
movq -0x8(%rbp), %rax
movslq -0x60(%rbp), %rdx
imulq $0x18, %rdx, %rdx
addq %rdx, %rax
movl %ecx, 0xc(%rax)
jmp 0x2da1
movq -0x20(%rbp), %rax
movslq -0x64(%rbp), %rcx
movl (%rax,%rcx,4), %eax
addl -0x58(%rbp), %eax
movl %eax, -0x58(%rbp)
movq -0x8(%rbp), %rax
movslq -0x60(%rbp), %rcx
imulq $0x18, %rcx, %rcx
addq %rcx, %rax
movl 0x10(%rax), %ecx
addl $0x1, %ecx
movl %ecx, 0x10(%rax)
movq -0x20(%rbp), %rax
movslq -0x64(%rbp), %rcx
movl (%rax,%rcx,4), %ecx
movq -0x8(%rbp), %rax
movslq -0x60(%rbp), %rdx
imulq $0x18, %rdx, %rdx
addq %rdx, %rax
addl 0x8(%rax), %ecx
movl %ecx, 0x8(%rax)
movq -0x20(%rbp), %rax
movslq -0x64(%rbp), %rcx
movl (%rax,%rcx,4), %eax
addl -0x5c(%rbp), %eax
movl %eax, -0x5c(%rbp)
movl -0x64(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x64(%rbp)
jmp 0x2cc0
movl $0x0, -0x68(%rbp)
movl -0x68(%rbp), %eax
cmpl -0xc(%rbp), %eax
jge 0x2de0
jmp 0x2dd1
jmp 0x2dd3
jmp 0x2dd5
movl -0x68(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x68(%rbp)
jmp 0x2dc7
addq $0x70, %rsp
popq %rbp
retq
nopw %cs:(%rax,%rax)