Parallelize code which is doing bit wise operation

Question

I have this code to make a 100k*200k matrix occupy less space by packing 8 elements per byte of this AU matrix into A to reduce the memory consumption. This code takes forever to run as you can expect and i am planning on increasing the number of rows to 200k as well. I am running the code on a pretty powerful instance (CPU and GPU)and can scale it so can anyone help parallelize this code so that it is quicker.

import numpy as np
colm = int(2000000/8)
rows = 1000000
cols = int(colm*8)
AU = np.random.randint(2,size=(rows, cols),dtype=np.int8)
start_time = time.time()

A = np.empty((rows,colm), dtype=np.uint8)
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        A[i,j] = 0
        for k in range(8):
            if AU[i,(j*8)+k] == 1:
                A[i,j] = A[i,j] | (1<<(7-k))

Answer 1

Warning: You try to allocate a huge amount of memory: about 2 TB of memory that you probably do not have.

Assuming you have enough memory or you can reduce the size of the dataset, you can write a much much faster implementation using the Numba JIT. Moreover, you can parallelize the code and replace the slow conditional with a branchless implementation to significantly speed up the computation since AU is filled with binary values. Finally, you can unroll the inner loop working on k to make the code even faster. Here is the resulting implementation:

import numpy as np
import numba as nb
colm = int(2000000/8)
rows = 1000000
cols = int(colm*8)
AU = np.random.randint(2,size=(rows, cols),dtype=np.int8)
A = np.empty((rows,colm), dtype=np.uint8)

@nb.njit('void(uint8[:,:],int8[:,:])', parallel=True)
def compute(A, AU):
    for i in nb.prange(A.shape[0]):
        for j in range(A.shape[1]):
            offset = j * 8
            res = AU[i,offset] << 7
            res |= AU[i,offset+1] << 6
            res |= AU[i,offset+2] << 5
            res |= AU[i,offset+3] << 4
            res |= AU[i,offset+4] << 3
            res |= AU[i,offset+5] << 2
            res |= AU[i,offset+6] << 1
            res |= AU[i,offset+7]
            A[i,j] = res

compute(A, AU)

On my machine, this code is 37851 times faster than the original implementation on a smaller dataset (with colm=int(20000/8) and rows=10000). The original implementation took 6min3s while the optimized one took 9.6ms.

This code is memory bound on my machine. With the current inputs, this code is close to be optimal as it spends most of its time reading the AU input matrix. A good additional optimization would be to "compress" the AU matrix to a smaller one (if possible).

Answer 2

Python solution:

import numpy as np
import time

def compute(A, AU):
    A[:,:] = 0
    # Put every 8 columns in AU into A
    for i in range(A.shape[1]):
        A[:, i//8] = np.bitwise_or(A[:, i//8], np.left_shift(AU[:, i], i % 8))

colm = int(20000/8)
rows = 10000
cols = int(colm*8)
AU = np.random.randint(2,size=(rows, cols),dtype=np.int8)
start_time = time.time()

A = np.empty((rows,colm), dtype=np.uint8)

start_time = time.time()

compute(A, AU)
    
end_time = time.time()
print(end_time - start_time)

Packs bits in 1/2 second

Same code in C:

#include <stdio.h>
#include <stdlib.h>
#include <time.h>

int main(int argc, char* argv[]) {

    int colm = 200000/8;
    int rows = 10000;
    int cols = colm*8;
    unsigned char *A = (unsigned char *)malloc(rows * colm * sizeof(unsigned char)); 
    unsigned char *AU = (unsigned char *)malloc(rows * cols * sizeof(unsigned char)); 
    int i, j;
    clock_t begin;
    clock_t end;
    double time_spent;

    begin = clock();
        
    // Create AU
    for (i = 0; i < rows; i++)
        for (j = 0; j < cols; j++)
            *(AU + i*cols + j) = (unsigned char) (rand() & 0x01);  
            
    end = clock();
    time_spent = (double)(end - begin) / CLOCKS_PER_SEC;
    printf("%lf seconds to create AU\n", time_spent);
            
    begin = clock();
    
    // Create a zeroed out A
    for (i = 0; i < rows; i++)
        for (j = 0; j < colm; j++)
            *(A + i*colm + j) = (unsigned char) 0;  

    end = clock();
    time_spent = (double)(end - begin) / CLOCKS_PER_SEC;
    printf("%lf seconds to create A\n", time_spent);

    begin = clock();
            
    // Pack into bits
    for (i = 0; i < rows; i++)
        for (j = 0; j < colm; j++) {
            int au_idx = i*cols + j*8;
            for (int k=0; k<8; k++)
                *(A + i*colm + j) |= *(AU + au_idx + k) << k;
            }
            
    end = clock();
    time_spent = (double)(end - begin) / CLOCKS_PER_SEC;
    printf("%lf seconds to pack\n", time_spent);
            

    free(A); 
    free(AU);
    return 0;
}

Tested with colm=200,000. Bit packing takes 0.27 seconds against 0.64 seconds for the optimized Python version provided by Jérôme Richard. Calls to rand() are expensive and greatly increase overall runtime. In terms of memory, the C version peaks at 2GB against Python's 4.2GB. Further code optimization and parallelization would certainly reduce runtime.

Julia version:

using Random
colm = 200000÷8
rows = 30000
cols = colm*8

AU = zeros(UInt8, (rows, cols))

rand!(AU)
AU .&= 0x01

A = zeros(UInt8, (rows, colm))

function compute(A, AU)
    for i in 1:size(A)[2]
        start_col = (i-1) << 3
        @views A[:, i] .=  AU[:, start_col + 1] .| 
                   (AU[:, start_col + 2] .<< 1) .|
                   (AU[:, start_col + 3] .<< 2) .|
                   (AU[:, start_col + 4] .<< 3) .|
                   (AU[:, start_col + 5] .<< 4) .|
                   (AU[:, start_col + 6] .<< 5) .|
                   (AU[:, start_col + 7] .<< 6) .|
                   (AU[:, start_col + 8] .<< 7)        
    end
end

@time compute(A, AU)

Julia scales well in terms of performance. Results with colm=25,000 and rows=30,000:

Language  Total Run Time (secs)   Bit Packing Time (secs)  Peak Memory (GB)
Python    22.1                    3.0                      6
Julia     11.7                    1.2                      6                          

Answer 3

After reading your post initially yesterday I actually planned to write a routine myself with the just-in-time compiler numba but Jérôme was faster than me and delivered you an excellent solution. But I have an alternative to offer: Why reinvent the wheel when there already exists a numpy function which does exactly the same: numpy.packbits.

import numpy as np
A = np.packbits(AU,axis=-1)

does the job. From my tests it seems to be quite a bit slower than Jérôme's version but anyways way faster than your initial version.