Stories by Abhijit Mondal on Medium

Distributed Matrix Multiplication using MPI on Azure

Abhijit Mondal — Sat, 26 Oct 2024 09:31:51 GMT

*Note — This post assumes that you are already familiar with OpenMPI and MPI syntax in C/C++.

There could be several reasons as to why one would perform matrix multiplication in a distributed manner.

Size of input matrices are larger than available memory on a single node machine.
Although size of input matrices are small enough to fit in the available memory of a single node machine but the time complexity is O(n³) and for large n, this is prohibitively large.

Matrix multiplication finds a wide range of applications in the ML field and is quite heavily used in different ML libraries and algorithms.

Thus, we see that we have a lot of motivation to understand how to build a distributed algorithm for doing matrix multiplication.

Before any implementation let us try to see how to write a distributed algorithm for matrix multiplication.

There is no single best algorithm as different algorithms tries to solve different constraints and thus works differently.

Optimize for low memory on single node machines.
Optimize for the network bandwidth and topology of a cluster of nodes.
Optimize for faster run-time performance.

You will notice that most of the work on the 1st and 2nd pointers i.e. optimize for low memory and network topology were done quite a few years back when memory sizes were comparatively very small to today’s machines and network bandwidths were also low.

Some well established algorithms in these space are:

Cannon’s algorithm.
Fox’s algorithm.

Both these algorithms assume that the nodes on which computations are run have low memory and the cluster of nodes are arranged in a grid (or rather a Torus) like topology.

Let’s take a look at how Fox’s algorithm works internally:

In a non-distributed setting, given matrices A and B of dimensions NxM and MxP respectively, we can calculate C=A.B in the following way:

The algorithm is run in multiple stages. For a given stage S:

For each i from 0 to N-1, multiply scalar A[i,(i+S) %M] with the row vector B[(i+S)%M,:]. This will give a partial values for the row :

C[i, :] += A[i, (i+S) %M] . B[(i+S)%M, :]

Repeat for all stages S from 0 to M-1. Pseudocode is below:

Input:
A - NxM
B - MxP

Output:
C - NxP

For each stage j from [0, M)
  For each i from [0, N)
    Calculate C[i,:] += A[i,(i+stage)%N].B[(i+stage)%N,:]

return C

In a distributed setting, lets assume that we have KxK nodes arranged in a grid topology. The concept of a grid topology implies that it is faster to connect to T[i+1,j] or T[i-1,j] or T[i,j+1] or T[i,j-1] from T[i,j] where T[i,j] refers to the node at index (i,j) in the grid.

Thus, each node at (i, j) is responsible for computing the values of the sub-matrix :

Initially the B matrix is distributed to the nodes s.t. node at grid index (i,j) i.e. T[i,j] gets the sub-matrix :

In stage S, the scalar values A[j, (j+S) % M] for j from Ui to Ui+N/K are broadcasted to all nodes in the i-th row of the grid :

i.e. T[i,0], T[i,1], … T[i,K-1]

After each stage, the B matrix is ‘shifted up’ across all the nodes in the grid i.e. node T[i,j] sends the values B[Ui, Vj:Vj+P/K] to the node T[(i-1)%K, j] and the process is repeated for all M stages.

Once all the nodes have finished all their stages, we can merge all the sub-matrices C from each node.

Let’ take an example to demonstrate this. Assume we have the following 4x4 matrices A and B.

Also, assume that we have 4 nodes for distributed computing arranged in a 2x2 grid :

The matrix B is initially assigned to the nodes as follows:

For stage 0, the values of A are assigned as follows to the nodes :

In each node, we multiply the corresponding A and B values as shown in the above diagram and add them to output. For e.g. in T[0,0], we maintain the following:

C[0,0] += A[0,0]*B[0,0]

C[0,1] += A[0,0]*B[0,1]

C[1,0] += A[1,1]*B[1,0]

C[1,1] += A[1,1]*B[1,1]

In stage 1, values of A are shifted right by 1 column to the right and values of B are shifted by 1 row up. Both these are with wrap-around.

Note that the new values of B need not come from the master node as for a given node, the new row of B that it needs, is already present in the node located below it in the grid.

For node T[0,0] we have:

C[0,0] += A[0,1]*B[1,0]

C[0,1] += A[0,1]*B[1,1]

C[1,0] += A[1,2]*B[2,0]

C[1,1] += A[1,2]*B[2,1]

Similarly for the other nodes too.

In stage 2, again we shift A values by 1 column to right and B by 1 row up with wrap-around.

Lastly for stage 3, we do the same :

After all the stages, the products A[i,j]*B[j,k] in each cell are added up and the partial values of C[i,k] are calculated.

If you observe carefully, the A matrix is rotating in the horizontal direction whereas the B matrix in a vertical direction as if they are lying on a Torus.

The advantage of this algorithm is that each node is dealing with only O(N) memory when each matrix is of size NxN.

For a cluster with K nodes.

In each stage we are sending values of A to each node over the network and each node is sending a row of B to the node located just above itself. Thus, the number of network I/O is:

Total Network I/O = O(M.K), where M is the number of columns in A and is also equal to the number of stages.

Although we are using low memory footprint per node, but we are trading it off with higher number of network communications with smaller payloads.

Thus, this approach works well when we have low memory on each node and network bandwidths are low but its performance will be slow due to lots of network I/O.

Given today’s modern servers with high memory and high network bandwidths, a simpler distributed algorithm is as follows:

If we take the j-th column of A and j-th row of B, then their dot product will be a NxN matrix. If we sum these matrices for all j from 0 to M-1, then we will get the final matrix C.

We can distribute the columns j to j+M/K and rows j to j+M/K to a node. We are assuming that there are K nodes (without any topology). Then the node with rank=k, will compute the partial matrix C_k as follows :

This node will then send the partial matrices C_k to the master node with rank=0. The master node will then sum up all the individual elements of each partial matrix obtained from all nodes to get the resultant matrix C.

Input:
A - NxM
B - MxP
K = Num nodes

Output:
C - NxP

Let h = M/K

If rank == 0:
  Compute C = A[:, :M/K] . B[:M/K, :]
  For j in [M/K, M):
    Send A_sub = A[:, j:j+M/K] to node with rank=j/h
    Send B_sub = B[j:j+M/K, :] to node with rank=j/h

If rank != 0:
  Obtain A_sub and B_sub from node with rank=0
  Compute C = A_sub . B_sub
  Send C to node with rank=0

If rank = 0:
  For k in [0, K):
    Obtain C_k from node with rank=k
    C += C_k

return C

If there are K nodes, then the number of network I/O is O(K), because we are sending columns from A and rows from B to a node only once.

Finally, each node sends back its matrix to the master node i.e. O(K) number of communications.

But in each node we require O(N²) memory as compared to O(N) in the Fox’s algorithm. Thus we are optimizing for less number of network I/O in this algorithm but higher memory requirements.

One drawback with this approach is that all the nodes will be sending their matrix of size NxP to the master node with rank=0. Thus, the master node will be doing O(K) additions of O(NxP) matrices.

Time complexity for final step is O(K.N.P).

Alternative approach to consolidate the matrices is to merge in stages. Assuming K is a power of 2, for e.g. if K=8, the nodes will be :

0 1 2 3 4 5 6 7

In the 1st round, each node at even positions, will send its matrix to the node left of it i.e. 1 to 0, 3 to 2 and so on. The nodes at even position will obtain these matrices and sum with its own matrix.

In the next stage, we only consider the nodes:

0 2 4 6

Because these have the most up-to-date data. Follow the same process as above i.e. 6 sends its matrix to 4, 2 sends its data to 0. Nodes 0 and 4 will add the obtained matrices with its own matrices.

In the last stage, we will have the following:

0 4

Now 4 sends its data to node 0. Note that now node 0 contains all the data from nodes 0 to 7.

In this approach any node will have at-most log(K) matrices in its buffer from any other node as compared to K matrices in the earlier version. Thus node 0 will add up log(K) matrices instead of K matrices. Time complexity for final addition is O(log(K).N.P).

The C++ code using MPI to run this algorithm is as follows:

#include 
// include other header files as required

using namespace std;

// Function to generate a random normal matrix
void generate(double *inp, int n, int m){
    std::random_device rd;
    std::mt19937 engine(rd());
    std::normal_distribution dist(0.0, 1.0);

    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            inp[i*m+j] = dist(engine);
        }
    }
}

// Function to transpose a matrix
double *transpose(const double *a, const int n, const int m) {
    double *b = new double[n*m];

    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            b[j*n+i] = a[i*m+j];
        }
    }

    return b;
}

// Function called by all nodes except master node
void dot_mpi(const int n, const int m, const int p, const int rank, const int size) {
    int h = m/size;

    double *a_r = new double[h*n];
    double *b_r = new double[h*p];
    double *merged = new double[h*(n+p)];

    // Receive columns of A and rows of B from master node
    MPI_Recv(merged, h*(n+p), MPI_DOUBLE, 0, rank, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

    std::copy(merged, merged+h*n, a_r);
    std::copy(merged+h*n, merged+h*n+h*p, b_r);

    // Calculate own product
    double *out = new double[n*p];
    for (int i = 0; i < n*p; i++) out[i] = 0.0;

    for (int i = 0; i < h; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < p; k++) {
                out[j*p+k] += a_r[i*n+j]*b_r[i*p+k];
            }
        }
    }

    // Send the matrices in stages to the master node
    for (int stage = 1; stage < size; stage *= 2) {
        if ((rank % stage == 0) && (rank-stage >= 0) && ((rank - stage) % (stage*2) == 0)) {
            MPI_Send(out, n*p, MPI_DOUBLE, rank-stage, rank, MPI_COMM_WORLD);
        }
        else if ((rank % stage == 0) && (rank+stage < size) && (rank % (stage*2) == 0)) {
            double *out_r = new double[n*p];
            MPI_Recv(out_r, n*p, MPI_DOUBLE, rank+stage, rank+stage, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            for (int j = 0; j < n*p; j++) {
                out[j] += out_r[j];
            }
        }
    }
}

// Function called for master node
double *dot_mpi_root(const double *a, const double *b, const int n, const int m, const int p, const int size) {
    double *out = new double[n*p];
    for (int i = 0; i < n*p; i++) out[i] = 0.0;

    // Instead of selecting columns of A, transpose A and select rows of A
    // This allows us to efficiently select elements of A to send.
    double *aT = transpose(a, n, m);
    int h = m/size;

    MPI_Request request = MPI_REQUEST_NULL;

    for (int i = h; i < m; i += h) {
        double *merged = new double[h*(n+p)];
        std::copy(aT+i*n, aT+(i+h)*n, merged);
        std::copy(b+i*p, b+(i+h)*p, merged+h*n);

        // Merge columns of A and rows of B before sending to a node
        // Reduces the number of network I/O
        MPI_Isend(merged, h*(n+p), MPI_DOUBLE, i/h, i/h, MPI_COMM_WORLD, &request);
    }

    // Do own contribution to the product
    for (int i1 = 0; i1 < h; i1++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < p; k++) {
                out[j*p+k] += aT[i1*n+j]*b[i1*p+k];
            }
        }
    }

    // Receive the computed matrices from other nodes and sum them
    for (int stage = 1; stage < size; stage *= 2) {
        double *out_r = new double[n*p];
        MPI_Recv(out_r, n*p, MPI_DOUBLE, stage, stage, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        for (int j = 0; j < n*p; j++) {
            out[j] += out_r[j];
        }
    }

    return out;
}

double *dot(const double *a, const double *b, const int n, const int m, const int p) {
    double *out = new double[n*p];

    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            for (int k = 0; k < p; k++) {
                out[i*p+k] += a[i*m+j]*b[j*p+k];
            }
        }
    }

    return out;
}

int main(int argc, char *argv[]) {
    int n = atoi(argv[1]);
    int m = atoi(argv[2]);
    int p = atoi(argv[3]);

    MPI_Init(&argc, &argv);

    int size;
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    double *a, *b;

    if (rank == 0) {
        // Form the input matrices only in the master node, other nodes
        // will have smaller memory footprints
        a = new double[n*m];
        b = new double[m*p];

        generate(a, n, m);
        generate(b, m, p);
    }

    double *out;

    auto start = std::chrono::high_resolution_clock::now();
    if (rank == 0) out = dot_mpi_root(a, b, n, m, p, size);
    else dot_mpi(n, m, p, rank, size);
    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast(stop - start);

    if (rank == 0) {
        std::cout << duration.count() << std::endl;

        // Compare the duration for distributed computation with single node
        // computation.
        start = std::chrono::high_resolution_clock::now();
        double *out = dot(a, b, n, m, p);
        stop = std::chrono::high_resolution_clock::now();
        duration = std::chrono::duration_cast(stop - start);

        std::cout << duration.count() << std::endl;
    }

    MPI_Finalize();
    return 0;
}

So far we have described how to write a distributed matrix multiplication algorithm using C++ and MPI. If you wish you can also run the above code on your local machine instead of cloud based server.

To compile the above program (using -O3 flag for optimized binary file):

mpicxx -O3 -o matrix_multiplication matrix_multiplication.cpp

Now in order to run and test your code on your local machine, run the following command. Here I am assuming that there are 4 cores on my machine.

mpirun -n 4 ./matrix_multiplication 500 500 500

The program utilizes all the 4 cores while running the code. It will not run if you specify more number of processes than the number of cores.

But why run on local machine when you can run on a cluster of nodes where it will be actually useful to run since in a local machine you are limited by the amount of memory available.

To setup a simple HPC cluster on Azure with 16 standard_d4s_v5 nodes, follow these steps (assume you have az cli installed and you are using Ubuntu 22.04):

The steps are slightly modified from the ones posted here: Quick MPI Cluster Setup on Azure (glennklockwood.com).

Note that this is not the most favorable approach for running a HPC system on cloud but it will help you develop the intuition behind everything that is happening in behind.

Step 1:

Setup the cloud-init script. Create a file named cloud-init.txt in your home folder with the following content:

#cloud-config
package_upgrade: true
packages:
  - clustershell
  - openmpi-bin
  - libopenmpi-dev
  - mpich

clustershell tool is used to run the same linux command on multiple nodes. This is a very handy tool as we will see.
mpich is required for using MPI with C++. For C, mpich is not required.

Step 2:

Create a resource group in your subscription:

# if having multiple subscriptions
az account set --subscription 
az group create --name mpi_test_rg --location eastus

Step 3:

Create a proximity placement group (ppg). Every VM we put in here will be within the same low-latency network bubble.

az ppg create --name mpi_ppg \
              --resource-group mpi_test_rg \
              --intent-vm-sizes standard_d4s_v5

Step 4:

Create 16 Ubuntu 22.04 nodes in this proximity placement group:

az vm create --name mpicluster \
             --resource-group mpi_test_rg \
             --image Ubuntu2204 \
             --ppg mpi_ppg \
             --generate-ssh-keys \
             --size standard_d4s_v5 \
             --accelerated-networking true \
             --custom-data cloud-init.txt \
             --count 16

--generate-ssh-keys puts your local ssh key ~/.ssh/id_rsa.pub into the authorized_keys file in all the nodes.
--custom-data cloud-init.txt passes in our cloud-init.txt file to the provisioning process. This installs the libraries mentioned in the cloud-init.txt file on all nodes.

Step 5:

Get the public and private IP addresses of your nodes:

az vm list-ip-addresses --resource-group mpi_test_rg -o table

Note down the public and private IP address of the node with the hostname mpicluster0. This is the master node.

Step 6:

Copy the private key from your local machine to the master node. Replace the IP with the public IP address of master node.

This is a highly insecure step but is required for you to connect to all nodes from your master node.

scp ~/.ssh/id_rsa 20.25.28.201:.ssh/
ssh 20.25.28.201

The next steps assume that we are running the commands on the master node i.e. mpicluster0

Step 7:

To get passwordless ssh up and working (so we can noninteractively run commands on other nodes through SSH), first add all the compute nodes’ host keys to your SSH known_hosts file.

The ssh-keyscan command connects to other hosts and retrieves its host keys, which we then store in ~/.ssh/known_hosts on our main mpicluster0 node.

for i in {0..15};do ssh-keyscan mpicluster${i};done > ~/.ssh/known_hosts

Step 8:

clush command runs the same command on all nodes mentioned using -w. Here it copies the private key and known_hosts file on master node to all other nodes.

This allows any node to connect to any other node without IP addresses.

clush -w mpicluster[0-15] -c ~/.ssh/id_rsa ~/.ssh/known_hosts

Step 9:

Install NFS server and client on all the nodes. Although NFS server is required only on master node. This allows us to mount and share filesystem on master node with all other nodes.

clush -w mpicluster[0-15] sudo apt -y install nfs-common nfs-kernel-server

Step 10:

Create a folder on all nodes. This folder will mount another folder on the master node via NFS.

clush -w mpicluster[0-15] sudo mkdir /scratch

Step 11:

On the master node, create a folder /shared which will be mounted and shared by all other nodes. Also add the current user as owner of this folder.

sudo mkdir /shared
sudo chown user:user /shared

Step 12:

Update the /etc/exports file in the master node with the possible IP address range of all the other nodes sharing the folder.

Here we are using 10.0.0.0/16. But the actual CIDR notation will depend on the private IP addresses from Step 5.

sudo bash -c 'echo "/shared 10.0.0.0/16(rw,no_root_squash,no_subtree_check)" >> /etc/exports'
sudo exportfs -a

Step 13:

Run mount command on all nodes to mount the /shared directory on master node.

The IP address 10.0.0.7 is the private IP address of the master node i.e. mpicluster0. Update this according to your output from Step 5.

clush -w mpicluster[0-15] sudo mount -t nfs -o vers=3,rsize=1048576,wsize=1048576,nolock,proto=tcp,nconnect=8 10.0.0.7:/shared /scratch

Step 14:

On the master node create a file named matrix_multiplication.cpp and copy the code above for distributed matrix multiplication.

Then compile the code as below. The compiled binary will be located inside the /shared folder so that all nodes will have access to the compiled binary for running the code.

mpicxx -O3 -o /shared/matrix_multiplication matrix_multiplication.cpp

Step 15:

Finally we run the compiled code on all our 16 nodes as below. Since we are using 16 nodes, we are choosing the matrix sizes multiples of 16. The matrices here are of sizes 1600x1600.

mpirun --host mpicluster0,mpicluster1,mpicluster2,mpicluster3,mpicluster4,mpicluster5,mpicluster6,mpicluster7,mpicluster8,mpicluster9,mpicluster10,mpicluster11,mpicluster12,mpicluster13,mpicluster14,mpicluster15 --mca btl tcp,vader,self --wdir /scratch ./matrix_multiplication 1600 1600 1600

In our code, we are also comparing the runtime of a non-distributed matrix multiplication on the master node with the distributed one.

For smaller matrix sizes, obviously the non-distributed algorithm works much faster as compared to the distributed one because in distributed setting we have lots of network I/O which degrades performance.

But for larger matrices such as 1600x1600, the non-distributed algorithm is 4 times slower than the distributed one. It takes around 1.11 seconds with 1600x1600 matrices in non-distributed while it takes only 0.3 seconds on the distributed cluster.

As we can see that the effect of network I/O is much less as compared to the time complexity of matrix multiplication as the sizes of matrices become larger.

If you are interested in the Fox’s algorithm implementation using MPI, the code is uploaded in github: hpc_mpi/matrix_multiplication.cpp at main · funktor/hpc_mpi (github.com)

References

Distributed Matrix Multiplication using MPI on Azure was originally published in Gaussian Machine on Medium, where people are continuing the conversation by highlighting and responding to this story.

Implementing SVD algorithm in Rust

Abhijit Mondal — Tue, 08 Oct 2024 19:07:47 GMT

Singular Value Decomposition also popularly known as SVD is one of the fundamental algorithms used in ML for doing low rank decomposition of large matrices. Some notable applications of SVD are as follows:

Doing dimensionality reduction of large and sparse matrices used in document search and retrieval systems such as search engines. SVD helps to reduce the time complexity for search.
In recommender systems, SVD can be used to find latent vectors for users and items. The latent vectors can then be used to find similar users and items.
Image compression by reducing the rank of the original image matrix.
Finding pseudo-inverse of non-square matrices. This is useful for many regression problems where the number of features (or variables) may not be equal to the number of training instances.

In this post I’m not going to deep dive into the applications, as I just want to focus of how to implement SVD algorithm from scratch. If you are interested in understanding more about the applications, check out my slides here:

SVD_New — Google Slides

I am assuming that you are familiar with what SVD is and what are the different components of SVD. For refresher, a matrix A of dimensions NxM can be decomposed into 3 matrices U, S and V as follows:

U is NxN matrix, S is NxM matrix and V is MxM matrix.

Columns of matrix U are the eigenvectors of A.A^T and columns of matrix V are the eigenvectors of the matrix A^T.A

The matrix S is a diagonal matrix where the diagonal elements are the square root of the eigenvalues of the matrix A.A^T (or A^T.A).

Thus, we now have sufficient information to calculate the matrices U, S and V. We need to find the eigenvectors of the matrices A.A^T and A^T.A and the eigenvalues for any one of those matrices. That brings us to the question of how to find eigenvalues and eigenvectors of a matrix.

If you are a Pythonista, you might find them quite easily using the Numpy library as follows:

import numpy as np

n, m = 100, 100
a = np.random.normal(0,1,(n,m))

eigenvalues, eigenvectors = np.linalg.eig(a)

print(eigenvalues)
print(eigenvectors)

Note that the eigenvalues and eigenvectors may be complex numbers. This is because when we are finding the eigenvalues, we are actually solving the following equation:

det(A-yI) = 0

y are the eigenvalues of the matrix A. det() is the determinant of matrix. Thus, we will obtain a polynomial in y. Not all roots of y will be real numbers.

But note that if A is a symmetric matrix, then all the eigenvalues are real numbers and also the eigenvectors are all real. The matrices A.A^T and A^T.A are symmetric matrices and thus the eigenvalues and eigenvectors are real.

linear algebra — Prove that the eigenvalues of a real symmetric matrix are real — Mathematics Stack Exchange

Thus, it is pretty easy to calculate the matrices U and V using the above.

n, m = 100, 100
a = np.random.normal(0,1,(n,m))

eigenvalues_u, u = np.linalg.eig(a @ a.T)
eigenvalues_v, v = np.linalg.eig(a.T @ a)

s = np.diag(np.sqrt(eigenvalues_u))

np.testing.assert_allclose(u @ s @ v.T, a)

Did you notice anything after running the above code ? You will most likely get assertion error. But why ?

This is because of sign errors of the eigenvectors.

Note that if you revert the sign of an eigenvector, it will still be a valid eignvector of the matrix. This is because in the eigenvector equation:

(A-yI).x = 0

Multiplying both sides by -1 has no effect.

If we compute U and V independently, we will land into this trouble where multiplying U, S and V^T does not give back the original matrix A. To overcome this problem, we can compute V first and then compute U:

S inverse is 1.0 over the original values of S for all non-zero diagonal elements, for zero diagonal elements, we can keep them as 0.

n, m = 100, 100
a = np.random.normal(0,1,(n,m))

eigenvalues_v, v = np.linalg.eig(a.T @ a)
s = np.diag(np.sqrt(eigenvalues_v))
s_inv = np.where(s == 0, 0, 1.0/s)
u = a @ v @ s_inv

np.testing.assert_allclose(u @ s @ v.T, a)

But we shall see that the above approach is not a stable approach when the values of matrix A.A^T or A^T.A are smaller than the machine precision. For e.g. in the above code, if we change the values of matrix A to be sampled from 0 to 1e-200, we will notice that the approach fails with assertion error.

n, m = 100, 100
a = np.random.normal(0,1e-200,(n,m))

eigenvalues_v, v = np.linalg.eig(a.T @ a)
s = np.diag(np.sqrt(eigenvalues_v))
s_inv = np.where(s == 0, 0, 1.0/s)
u = a @ v @ s_inv

np.testing.assert_allclose(u @ s @ v.T, a)




AssertionError                            Traceback (most recent call last)
Cell In[87], line 9
      6 s_inv = np.where(s == 0, 0, 1.0/s)
      7 u = a @ v @ s_inv
----> 9 np.testing.assert_allclose(u @ s @ v.T, a)

    [... skipping hidden 1 frame]

File C:\Program Files\WindowsApps\PythonSoftwareFoundation.Python.3.10_3.10.3056.0_x64__qbz5n2kfra8p0\lib\contextlib.py:79, in ContextDecorator.__call__..inner(*args, **kwds)
     76 @wraps(func)
     77 def inner(*args, **kwds):
     78     with self._recreate_cm():
---> 79         return func(*args, **kwds)

File ~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\numpy\testing\_private\utils.py:862, in assert_array_compare(comparison, x, y, err_msg, verbose, header, precision, equal_nan, equal_inf, strict)
    858         err_msg += '\n' + '\n'.join(remarks)
    859         msg = build_err_msg([ox, oy], err_msg,
    860                             verbose=verbose, header=header,
    861                             names=('x', 'y'), precision=precision)
--> 862         raise AssertionError(msg)
    863 except ValueError:
    864     import traceback

AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0

Mismatched elements: 25 / 25 (100%)
Max absolute difference: 1.90398337e-200
Max relative difference: 1.
 x: array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],...
 y: array([[-1.834095e-200, -2.743976e-201, -1.228526e-200,  1.444880e-200,
         1.197758e-200],
       [-5.629092e-202, -6.866588e-201, -1.684314e-201,  8.075045e-201,...

Instead, we will look into one of the earliest stable algorithm for SVD calculations known as the Golub-Kahan-Reisch SVD algorithm.

In the rest of this post we will understand this algorithm and implement it using Rust. Note that this algorithm is not the fastest out there and many algorithms such as the divide and conquer approach are used in standard packages.

The idea behind this algorithm is to multiply the matrix A by some orthogonal matrices L1, L2, … Lk on the left and some orthogonal matrices R1, R2, ..Rk on the right to form a diagonal matrix say X, i.e.

Then we can write A as follows:

The above follows from the fact that the matrices L and R are orthogonal i.e. L.L^T = I and L^T.L=I and similarly for R.

If we denote the following:

and

Then we can rewrite A as follows:

This form looks familiar and it is infact the SVD of matrix A because X is a diagonal matrix and U and V are orthogonal. Now the only question remaining is how to find these orthogonal matrices L and R ?

Let’s deviate a bit from here into another related algorithm.

Remember QR decomposition from linear algebra class ?

QR decomposition of a NxM matrix A, produces 2 matrices Q and R where Q is a NxN orthogonal matrix and R is a NxM upper diagonal matrix.

The dimensionality of Q and R can be varied depending on the use case. For e.g. if M < N, we can instead use NxM matrix for Q (first M columns) and MxM matrix for R (top M rows).

Note that this will not change the results because is N > M and R is an upper diagonal matrix, then it is filled with all 0s in the lower (N-M)xM matrix which does not add any value.

There are 2 very popular algorithms to perform QR decomposition:

Householder reflection
Givens’ rotation

In Householder reflection method, firstly for column 0, all values in rows 1 to n-1 (0-indexed) are zeroed out, then for column 1, rows 2 to n-1 are zeroed out and so on, till we have zeroed out all sub-diagonal elements.

The algorithm is as follows:

Let G = A - NxM matrix
Let H = I - NxN identity matrix

For each column i from 0 to min(N-1, M-1), repeat steps 1 to 6:
  1. Let x = G[i:,i]
  2. Let l2 = L2_norm(x)
  3. Update x[0] = x[0]-sgn(x[0])*l2
  4. Let u = x/L2_norm(x)
  5. Update G[i:,i:] = G[i:,i:] - 2*(u @ (u.T @ G[i:,i:]))
  6. Update H[i:,:] = H[i:,:]- 2*(u @ (u.T @ H[i:,:]))

Let K = min(N, M)
return Q=H.T[:,:K], R=G[:K,:]

Time complexity of the above algorithm is O(NM²) ignoring the constant terms and lower order terms. Thus, if M > N, then it is beneficial to perform QR decomposition on the transpose of A.

More elaborate explanation can be found in these slides:

Applied Numerical Linear Algebra. Lecture 9 (chalmers.se)

It can be shown that once a column has been zeroed out below its main diagonal, the algorithm for subsequent columns do not affect the already zeroed out columns.

Implementing the above algorithm using Rust:

// returns sign of a real number
fn sgn(x:f64) -> f64 {
    if x < 0.0 {
        return -1.0;
    }
    return 1.0;
}

// generates an nxn identity matrix
fn identity(n:usize) -> Vec {
    let mut q = vec![0.0;n*n];
    for i in 0..n {
        q[i*(n+1)] = 1.0;
    }

    return q;
}

// transpose a matrix
fn transpose(a:&[f64], n:usize, m:usize) -> Vec {
    let mut b = vec![0.0;n*m];

    for i in 0..n {
        for j in 0..m {
            b[j*n+i] = a[i*m+j];
        }
    }

    return b;
}

// householder reflection algorithm
fn householder_reflection_left_multiply(a:&[f64], n:usize, m:usize) -> (Vec, Vec) {
    let mut q_lt = identity(n);
    let mut r = a.to_vec();
    let w = min(n, m);

    for i in 0..w {
        let n1 = n-i;

        let mut nm = 0.0;
        let mut u = vec![0.0;n1];

        for i1 in i..n {
            // compute the L2 norm
            nm += r[i1*m+i]*r[i1*m+i];
            u[i1-i] = r[i1*m+i];
        }

        u[0] -= sgn(r[i*m+i])*nm.sqrt();
        nm = (nm-r[i*m+i]*r[i*m+i]+u[0]*u[0]).sqrt();

        if z > 0.0 {
            for i1 in 0..n1 {
                u[i1] = u[i1]/nm;
            }

            // r1 stores the values of u.T @ R for each column
            let mut r1 = vec![0.0;m-i];
            for i1 in i..n {
                for j1 in i..m {
                    r1[j1-i] += u[i1-i]*r[i1*m+j1];
                }
            }

            // update the matrix R
            for i1 in i..n {
                for j1 in i..m {
                    r[i1*m+j1] -= 2.0*u[i1-i]*r1[j1-i];
                }
            }

            // q1 stores the values of u.T @ Q for each column
            let mut q1 = vec![0.0;n];
            for i1 in i..n {
                for j1 in 0..n {
                    q1[j1] += u[i1-i]*q_lt[i1*n+j1];
                }
            }

            // update the matrix Q
            for i1 in i..n {
                for j1 in 0..n {
                    q_lt[i1*n+j1] -= 2.0*u[i1-i]*q1[j1];
                }
            }
        }
    }

    return (q_lt, r);
}

pub fn householder_reflection_qr(a:&[f64], n:usize, m:usize) -> (Vec, Vec, usize, usize, usize) {
    // if n < m, transpose and solve qr
    if n < m {
        let aT = transpose(&a, n, m);
        let (q, r) = householder_reflection_left_multiply(&aT, m, n);
        return (transpose(&r, m, n), q);
    }
    else {
        let (q, r) = householder_reflection_left_multiply(&a, n, m);
        return (transpose(&q, n, n), r);
    }
}

fn main() {
    let n = 500;
    let m = 500;

    //generate matrix of nxm shape from normal distribution
    let mut rng = thread_rng();
    let normal:Normal = Normal::new(0.0, 1.0).ok().unwrap();
    let mut a:Vec = vec![0.0;n*m];

    for i in 0..n*m {
        a[i] = normal.sample(&mut rng);
    }

    let (q, r) = householder_reflection_qr(&a, n, m);
}

In Givens Rotation, instead of targetting entire column at a time, individual cell values at (i, j) are targetted. For e.g. if we have to zero out A[i,j], then we can left multiply the 2xM matrix A[i-1:i+1, :] by a 2x2 matrix H, where:

H.A[i-1:i+1, :] only affects rows i-1 and i. Thus, we have to zero out the values in a column in bottom up fashion. The algorithm is as follows:

Let G = A - NxM matrix
Let H = I - NxN identity matrix

For each column j from 0 to M-1:
  For each row i from N-1 to i+1:
    1. Let r = sqrt(G[i-1,j]**2 + G[i,j]**2)
    2. Let c = G[i-1,j]/r
    3. Let s = -G[i,j]/r
    4. Let Q = [[c, -s], [s, c]]
    5. Update G[i-1:i+1, j:] = Q @ G[i-1:i+1, j:]
    6. Update H[i-1:i+1, j:] = Q @ H[i-1:i+1, j:]

Let K = min(N, M)
return Q=H.T[:,:K], R=G[:K,:]

The time complexity of the above algorithm is O(NM²) which is similar to the Householder reflection method.

One advantage of Givens Rotation over Householder Reflection method is that in Givens Rotation, some cells can be zeroed out in parallel.

Since the Givens Rotation for (i, j) affects only rows i-1 and i, thus once we have zeroed out A[n-1,0] and A[n-2,0], then we can start to zero out columns 0 and 1 in parallel because column 0 will be working from row n-3 upwards and column 1 from row n-1 upwards.

Similar logic applies for all subsequent columns too.

Slide from Parallel Numerical Algorithms — Chapter 11 — QR Factorization (illinois.edu)

In the above diagram, cells with same values can be zeroed in parallel. For e.g. cells (3,0), (5,1) and (7,2) all having seq. num 5 can be run in parallel.

Here is an implementation for Givens Rotation using Rust (without the parallel implementation):

fn givens_left_rotation(a:&[f64], _n:usize, m:usize, i:usize, j:usize) -> (f64, f64) {
    let x = a[(i-1)*m+j];
    let y = a[i*m+j];
    let r = (x*x + y*y).sqrt();

    return (x/r, -y/r);
}

fn givens_left_rotation_multiply(a:&mut [f64], _n:usize, m:usize, c:f64, s:f64, i:usize, _j:usize) {
    for j1 in 0..m {
        let p = a[(i-1)*m+j1];
        let q = a[i*m+j1];
        a[(i-1)*m+j1] = c*p - s*q;
        a[i*m+j1] = s*p + c*q;
    }
}

pub fn givens_rotation_qr(a:&[f64], n:usize, m:usize,) -> (Vec, Vec, usize, usize, usize) {
    // if m > sqrt(n), transpose and solve qr
    if n < m {
        let mut r = transpose(&a, n, m);
        let mut q = identity(m);

        for j in 0..n {
            for i in (j+1..m).rev() {
                let b = givens_left_rotation(&r, m, n, i, j);
                givens_left_rotation_multiply(&mut r, m, n, b.0, b.1, i, j);
                givens_left_rotation_multiply(&mut q, m, m, b.0, b.1, i, j);
            }
        }

        return (transpose(&r, m, n), q);
    }
    else {
        let mut r = a.to_vec();
        let mut q = identity(n);

        for j in 0..m {
            for i in (j+1..n).rev() {
                let b = givens_left_rotation(&r, n, m, i, j);
                givens_left_rotation_multiply(&mut r, n, m, b.0, b.1, i, j);
                givens_left_rotation_multiply(&mut q, n, n, b.0, b.1, i, j);
            }
        }

        return (transpose(&q, n, n), r);
    }
}

The best part is that both Householder reflectors matrices and Givens Rotation matrices are orthogonal matrices.

Now, coming back to our SVD algorithm.

The first step of the algorithm is the ‘bidiagonalization’ of the input matrix A. The main diagonal and upper diagonal have non-zero values. The following matrix is a “upper-bidiagonal” matrix.

Bidiagonalization has quite a few advantages:

It is “close” in shape to a diagonal matrix, which we ultimately want.
The matrix is very sparse and we can efficiently perform different operations such as matrix multiplications quite cheaply.
If we denote the bidiagonal matrix with B, then the matrix B^T.B has a tridiagonal form, which is again sparse. We would be interested in finding the eigenvalues and eigenvectors of B^T.B.

Bidiagonalization of a matrix can be done by using similar methods as in QR decomposition. Assuming that we choose the Householder reflection method, then first, we zero out all the elements below the main diagonal similar to the QR step seen earlier.

Then, we zero out all elements above the main diagonal by performing the Householder reflection on each row instead of each column. Note that it is not possible to zero out the diagonal above the main diagonal as then it will change the already zeroed out elements below the main diagonal.

Bidiagonalization using Rust:

fn householder_reflection_bidiagonalization(a:&[f64], n:usize, m:usize) -> (Vec, Vec, Vec) {
    let mut q_lt = identity(n);
    let mut q_rt = identity(m);

    let mut r = a.to_vec();

    for i in 0..min(n, m) {
        let n1 = n-i;

        let mut nm = 0.0;
        let mut u = vec![0.0;n1];

        for i1 in i..n {
            nm += r[i1*m+i]*r[i1*m+i];
            u[i1-i] = r[i1*m+i];
        }

        u[0] -= sgn(r[i*m+i])*nm.sqrt();
        let z = (nm-r[i*m+i]*r[i*m+i]+u[0]*u[0]).sqrt();

        if z > 0.0 {
            for i1 in 0..n1 {
                u[i1] = u[i1]/z;
            }

            let mut r1 = vec![0.0;m-i];
            for i1 in i..n {
                for j1 in i..m {
                    r1[j1-i] += u[i1-i]*r[i1*m+j1];
                }
            }

            for i1 in i..n {
                for j1 in i..m {
                    r[i1*m+j1] -= 2.0*u[i1-i]*r1[j1-i];
                }
            }

            let mut q1 = vec![0.0;n];
            for i1 in i..n {
                for j1 in 0..n {
                    q1[j1] += u[i1-i]*q_lt[i1*n+j1];
                }
            }

            for i1 in i..n {
                for j1 in 0..n {
                    q_lt[i1*n+j1] -= 2.0*u[i1-i]*q1[j1];
                }
            }
        }

        if m-i-1 > 0 {
            let n1 = m-i-1;

            let mut nm = 0.0;
            let mut u = vec![0.0;n1];

            for j1 in i+1..m {
                nm += r[i*m+j1]*r[i*m+j1];
                u[j1-i-1] = r[i*m+j1];
            }

            u[0] -= sgn(r[i*m+i+1])*nm.sqrt();
            let z = (nm-r[i*m+i+1]*r[i*m+i+1]+u[0]*u[0]).sqrt();

            if z > 0.0 {
                for i1 in 0..n1 {
                    u[i1] = u[i1]/z;
                }
                
                let mut r1 = vec![0.0;n-i];
                for i1 in i..n {
                    for j1 in i+1..m {
                        r1[i1-i] += u[j1-i-1]*r[i1*m+j1];
                    }
                }

                for i1 in i..n {
                    for j1 in i+1..m {
                        r[i1*m+j1] -= 2.0*u[j1-i-1]*r1[i1-i];
                    }
                }
                
                let mut q1 = vec![0.0;m];
                for i1 in 0..m {
                    for j1 in i+1..m {
                        q1[i1] += u[j1-i-1]*q_rt[i1*m+j1];
                    }
                }

                for i1 in 0..m {
                    for j1 in i+1..m {
                        q_rt[i1*m+j1] -= 2.0*u[j1-i-1]*q1[i1];
                    }
                }
            }
        }
    }

    return (q_lt, r, q_rt);
}

Thus, after bidiagonalization step, we will obtain a bidiagonal matrix B:

But note that, we still have not diagonalized the matrix A and have only bidiagonalized it. The matrices Q_L and Q_R are all orthogonal and so are their products.

Thus, finding SVD of A is reduced to finding SVD of B. If the SVD of B is represented as follows:

Then, the SVD of A can be written as:

The next natural step, would be to find the SVD of B, which again we can do so by finding the eigenvalues and eigenvectors of B^T.B as seen above.

But note that we were motivated to use this method for SVD calculation in the first place, was because calculations of A^T.A or B^T.B can have very small values smaller than the machine precision, thus giving unstable results when the eigenvalues are small.

Thus, we should avoid going the route of calculating B^T.B. Then, what are the alternative ways to diagonalize the matrix B ? The idea is to use B^T.B without explicitly forming the matrix B^T.B.

Let T = B^T.B

If B is a bidiagonal matrix, then T is a tridiagonal symmetric matrix as follows:

Assuming there are Givens Rotation matrices G to be multiplied to the right of T and G^T to the left of T, that will eventually zero out the non-diagonal elements of T i.e.

Note that in QR decomposition once we zero out a cell (i, j) using Givens Left Rotation, we do not have to be bothered about the other columns that are affected in rows i-1 and i. But here, it is not so straightforward.

Remember that a Givens Left Rotation matrix for zeroing out (i, j) will affect only rows i-1 and i, whereas a Givens Right Rotation matrix will affect only columns j-1 and j.

Thus, if we zero out (0, 1) in T using G1 (Right Rotation), then it will affect only the columns 0 and 1 and (2,0) will be non-zero. Rest all zeros in columns 0 and 1 will be unaffected (if both columns have 0 in same row, both 0s are unaffected). Thus when we are zeroing out one cell, we are introducing a non-zero in another cell.

Similarly, if we apply G1^T to T.G1, then it is application of Givens Left Rotation and will affect rows 0 and 1. Thus, it will introduce a new non-zero values at (0,1) and (0,2). Columns where both rows 0 and 1 are 0 are unaffected.

Note, that application of G1^T to T.G1, will not zero out (1,0) because (1,0) value was changed after T.G1.

Repeat the process to zero out (0,2) and (2,0) by multiplying with G2 (right rotation) and G2^T (left rotation). This will introduce new non-zero values at (1,3) and (3,1).

Continue to repeat this process, until the new non-zero values are all chased out from T and we again obtain the tridiagonal form.

This entire process of “chasing out” a new non-zero element from T and restoring the tridiagonal form of T using Givens Left and Right rotation matrices is known as the Golub Kahan step.

In the below figure, the ‘x’ symbolizes a non-zero value. See how the 2 additional ‘x’’s are chased out to regain the original tridiagonal form.

After doing one round of Golub Kahan step as described above, the off-diagonal elements inches closer to 0 but may not yet be zero. After repeated application of the above steps, some of the off-diagonal elements starts to be closer to 0. We can assume that whenever the absolute value of an off-diagonal element is smaller than some threshold, consider it as 0.

To fasten the process of convergence to 0, Wilkinson Shift strategy is applied, which is basically using (T-uI) inplace of original T, where u is the eigenvalue of the bottom right 2x2 sub-matrix of T, that is closer to the last diagonal element.

Using (T-uI) inplace of T in the above steps will improve convergence rate for the off-diagonal elements going to 0.

But so far we have been doing left and right multiplications with Givens matrices on T but remember that we do not want to calculate T as because this will lead to error in calculations when the eigenvalues of T are smaller than the machine precision.

In the equation:

If we replace T with B^T.B, we get:

Thus, instead of applying G1, G2, … Gk to T, we can apply them directly to B. But note that applying only the right rotation matrices G1, G2, … to B will introduce new non-zero elements in B in the lower diagonal.

These new non-zero elements introduced, can be chased out by Givens Left Rotation matrices.

For e.g. let there be Givens Left rotation matrices Q1, Q2… s.t.

Thus, left multiplying by matrices Q to keep the bidiagonal matrix B in the bidiagonal form does not alter the product as Q are orthogonal matrices.

Similar to the tridiagonal matrix, where multiplying left and right by Givens matrices chased away the new non-zero elements, we can see the similar thing happening with B also:

Note that, the above steps of multiplying B on the left and right by Givens rotation matrices to gradually remove the off-diagonal elements as well as chase away unwanted non-zero elements introduced during the process is applicable to only a bidiagonal matrix form.

Once some off-diagonal elements or some diagonal elements become 0, we have to divide the current matrix s.t. each part is bidiagonal and separately process each bidiagonal matrix.

For e.g. if a non-diagonal element say B[i,i+1] becomes 0 after sometime, then we separately process B[:i+1,:i+1] and B[i+1:,i+1:].

In the below example i=2.

In the event that a diagonal element becomes 0 (implying at-least one eigenvalue of T is 0), then we can handle this situation as below:

For e.g. if B[i,i] = 0 and B[i,i+1] != 0, then applying Givens left rotation on rows i and i+1 to B to zero out B[i,i+1], will introduce a new non-zero at B[i,i+2]. Again applying Givens left rotation on rows i and i+1 to zero out B[i,i+2] will introduce new non-zero at B[i,i+3] and so on till B[i,M-1].

After that, the final Givens Rotation matrix will chase out the unwanted non-zero value from the matrix.

In the below example B[2,2] = 0

Once we obtain multiple bidiagonal sub-matrices, we can process the Golub Kahan steps in parallel for each. The above Golub Kahan steps works fine even if the last diagonal of a bidiagonal matrix is 0.

Since we are multiplying the bidiagonal matrix left and right with Givens matrices after doing shift by the eigenvalue of the bottom right 2x2 sub-matrix closer to the last diagonal, the last super diagonal i.e. B[N-2,N-1] will be the first to go to zero. Then, after we process the remaining (N-2)x(N-2) matrix, the next off diagonal to go to zero would be B[N-3,N-2] and so on.

Thus, the off diagonal elements in B starts to go to zero starting from bottom right. Starting from i=N-2, for each i, we can check if B[i,i+1]= 0, if it is 0 then continue going to the top by decreasing i else stop and process the sub-matrix B[:i+1,:i+1].

Return after the matrix B have become a diagonal matrix.

The algorithm is as follows :

Instead of processing multiple bidiagonal matrices after splits in parallel, here we are processing them in bottom up fashion. We can use multithreading to process them in parallel but would not expect lot of performance gains as these are CPU bound calculations.

Let A - NxM matrix

1. B = P @ A @ Q where B is bidiagonal matrix (Householder Reflection)
2. Loop 3-8
  3. Set B[i,i+1] = 0 if abs(B[i,i+1]) < eps*(abs(B[i,i) + abs(B[i+1,i+1))
  4. Find q s.t. B[i,i+1]=0 for all i in (q+1, r-1), if not found, set q = 0
  5. Break if q = 0
  6. Find p s.t. B[i,i+1]!=0 for all i in (p, q), if not found, set p = 0
  7. If for any i in (p, q), B[i,i] = 0, remove the off-diagonal at B[i,i+1], 
     using Givens Left Rotation as described above and update the matrices 
     P and Q in step 1.
  8. If no i exists in (p, q), s.t. B[i,i] = 0, then run Golub Kahan step 
     for the sub-matrix (p,p) to (q,q) and update the matrices 
     P and Q in step 1.

return (P, B, Q)

Let U^T be the product of all orthogonal matrices applied so far to the left of A including the Givens left rotation matrices to B and the householder reflection matrices for bidiagonalization from A to B and V product of all orthogonal matrices applied so far to the right of A. Then,

Thus, we have found the SVD for A. S is a diagonal matrix.

Here is the Rust implementation for the above algorithm:

use std::simd::prelude::*;
use rand_distr::{Distribution, Normal};
use rand::thread_rng;
use std::cmp::min;
use std::time::SystemTime;

// Get sub-matrix (r_start, c_start) and (r_end, c_end)
fn sub_mat(a:&[f64], _n:usize, m:usize, r_start:usize, r_end:usize, c_start:usize, c_end:usize) -> Vec {
    let mut b:Vec = vec![0.0;(r_end-r_start+1)*(c_end-c_start+1)];
    let u = c_end-c_start+1;
    let mut k = 0;

    for i in r_start..(r_end+1) {
        b[k*u..(k+1)*u].copy_from_slice(&a[i*m+c_start..i*m+c_end+1]);
        k += 1;
    }

    return b;
}

// Bidiagonalization using Householder transformation
fn householder_reflection_bidiagonalization(a:&[f64], n:usize, m:usize) -> (Vec, Vec, Vec) {
    let mut q_lt = identity(n);
    let mut q_rt = identity(m);

    let mut r = a.to_vec();

    for i in 0..min(n, m) {
        let n1 = n-i;

        let mut nm = 0.0;
        let mut u = vec![0.0;n1];

        for i1 in i..n {
            nm += r[i1*m+i]*r[i1*m+i];
            u[i1-i] = r[i1*m+i];
        }

        u[0] -= sgn(r[i*m+i])*nm.sqrt();
        let z = (nm-r[i*m+i]*r[i*m+i]+u[0]*u[0]).sqrt();

        if z > 0.0 {
            for i1 in 0..n1 {
                u[i1] = u[i1]/z;
            }

            let mut r1 = vec![0.0;m-i];
            for i1 in i..n {
                for j1 in i..m {
                    r1[j1-i] += u[i1-i]*r[i1*m+j1];
                }
            }

            for i1 in i..n {
                for j1 in i..m {
                    r[i1*m+j1] -= 2.0*u[i1-i]*r1[j1-i];
                }
            }

            let mut q1 = vec![0.0;n];
            for i1 in i..n {
                for j1 in 0..n {
                    q1[j1] += u[i1-i]*q_lt[i1*n+j1];
                }
            }

            for i1 in i..n {
                for j1 in 0..n {
                    q_lt[i1*n+j1] -= 2.0*u[i1-i]*q1[j1];
                }
            }
        }

        if m-i-1 > 0 {
            let n1 = m-i-1;

            let mut nm = 0.0;
            let mut u = vec![0.0;n1];

            for j1 in i+1..m {
                nm += r[i*m+j1]*r[i*m+j1];
                u[j1-i-1] = r[i*m+j1];
            }

            u[0] -= sgn(r[i*m+i+1])*nm.sqrt();
            let z = (nm-r[i*m+i+1]*r[i*m+i+1]+u[0]*u[0]).sqrt();

            if z > 0.0 {
                for i1 in 0..n1 {
                    u[i1] = u[i1]/z;
                }
                
                let mut r1 = vec![0.0;n-i];
                for i1 in i..n {
                    for j1 in i+1..m {
                        r1[i1-i] += u[j1-i-1]*r[i1*m+j1];
                    }
                }

                for i1 in i..n {
                    for j1 in i+1..m {
                        r[i1*m+j1] -= 2.0*u[j1-i-1]*r1[i1-i];
                    }
                }
                
                let mut q1 = vec![0.0;m];
                for i1 in 0..m {
                    for j1 in i+1..m {
                        q1[i1] += u[j1-i-1]*q_rt[i1*m+j1];
                    }
                }

                for i1 in 0..m {
                    for j1 in i+1..m {
                        q_rt[i1*m+j1] -= 2.0*u[j1-i-1]*q1[i1];
                    }
                }
            }
        }
    }

    return (q_lt, r, q_rt);
}

// Get the eigenvalue of the bottom right 2x2 sub-matrix which is closer
// to the last diagonal element in B
fn eigenvalue_bidiagonal(a:&[f64], n:usize, m:usize, i1:usize, i2:usize, j1:usize, j2:usize) -> f64{
    let h = min(i2, j2)+1;

    let mut d1 = 0.0;
    let mut d2 = 0.0;
    let mut d3 = 0.0;
    let mut d4 = 0.0;

    if h >= 3 {
        d1 = a[(h-3)*m+h-2];
    }

    if h >= 2 {
        d2 = a[(h-2)*m+h-2];
        d3 = a[(h-2)*m+h-1];
    }

    if h >= 1 {
        d4 = a[(h-1)*m+h-1];
    }

    
    let a1 = d2*d2 + d1*d1;
    let a2 = d2*d3;
    let a3 = d2*d3;
    let a4 = d4*d4 + d3*d3;

    let u = 1.0;
    let b = -(a1+a4);
    let c = a1*a4-a2*a3;

    let v1 = (-b + (b*b-4.0*u*c).sqrt())/(2.0*u);
    let v2 = (-b - (b*b-4.0*u*c).sqrt())/(2.0*u);

    if (v1-a4).abs() < (v2-a4).abs() {
        return v1;
    }
        
    return v2;
}

// Get (c, s) value for Givens right rotation matrix
fn givens_right_rotation(a:&[f64], _n:usize, m:usize, i:usize, j:usize, flip:bool) -> (f64, f64) {
    let x = a[i*m+j-1];
    let y = a[i*m+j];
    let r = (x*x + y*y).sqrt();

    if flip {
        return (y/r, -x/r);
    }

    return (x/r, -y/r);
}

// Multiple a by Givens right rotation matrix
fn givens_right_rotation_multiply(a:&mut [f64], n:usize, m:usize, c:f64, s:f64, _i:usize, j:usize, r1:usize, r2:usize, c1:usize, c2:usize) {
    for i1 in r1..r2+1 {
        let p = a[i1*m+j-1];
        let q = a[i1*m+j];
        a[i1*m+j-1] = c*p - s*q;
        a[i1*m+j] = s*p + c*q;
    }
}

// Get (c, s) value for Givens left rotation matrix
fn givens_left_rotation(a:&[f64], _n:usize, m:usize, i:usize, j:usize, flip:bool) -> (f64, f64) {
    let x = a[(i-1)*m+j];
    let y = a[i*m+j];
    let r = (x*x + y*y).sqrt();

    if flip {
        return (y/r, -x/r);
    }

    return (x/r, -y/r);
}

// Multiple a by Givens left rotation matrix
fn givens_left_rotation_multiply(a:&mut [f64], _n:usize, m:usize, c:f64, s:f64, i:usize, _j:usize, r1:usize, r2:usize, c1:usize, c2:usize) {
    for j1 in c1..c2+1 {
        let p = a[(i-1)*m+j1];
        let q = a[i*m+j1];
        a[(i-1)*m+j1] = c*p - s*q;
        a[i*m+j1] = s*p + c*q;
    }
}

// Golub Kahan step on biadiagonal matrix
fn golub_kahan(a:&mut [f64], l:&mut [f64], r:&mut [f64], n:usize, m:usize, z:usize, i:usize, j:usize) {
    let mu = eigenvalue_bidiagonal(&a, z, z, i, j, i, j);
    
    let u = a[i*z+i];
    let v = a[i*z+i+1];
    
    a[i*z+i] = u*u-mu;
    a[i*z+i+1] = u*v;
    
    for k in i..j {
        let mut x;
        let mut y;

        if k > i {
            x = k-1;
            y = k+1; 
        }
        else {
            x = i;
            y = i+1; 
        }

        let b = givens_right_rotation(&a, z, z, x, y, false);

        if k == i {
            a[i*z+i] = u;
            a[i*z+i+1] = v;
        }

        givens_right_rotation_multiply(a, z, z, b.0, b.1, x, y, i, j, i, j);
        givens_right_rotation_multiply(r, m, z, b.0, b.1, x, y, 0, m-1, 0, z-1);

        if k > i {
            x = k+1;
            y = k;
        }
        else {
            x = i+1;
            y = i;
        }
            
        let b = givens_left_rotation(&a, z, z, x, y, false);

        givens_left_rotation_multiply(a, z, z, b.0, b.1, x, y, i, j, i, j);
        givens_left_rotation_multiply(l, z, n, b.0, b.1, x, y, 0, z-1, 0, n-1);
    }
}

// Golub Reisch SVD algorithm
pub fn golub_reisch_svd(a:&[f64], mut n:usize, mut m:usize) -> (Vec, Vec, Vec) {
    let mut a1 = a.to_vec();
    let mut do_transpose = false;

    if n < m {
        do_transpose = true;
        a1 = transpose(&a1, n, m);
        let g = n;
        n = m;
        m = g;
    }
    
    let hr = householder_reflection_bidiagonalization(&a1, n, m);
    let r = min(n, m);

    let mut u = sub_mat(&hr.0, n, n, 0, r-1, 0, n-1);
    a1 = sub_mat(&hr.1, n, m, 0, r-1, 0, r-1);
    let mut v = sub_mat(&hr.2, m, m, 0, m-1, 0, r-1);
    
    let eps = 1.0e-7;
    
    loop {
        for i in 0..r-1 {
            if a1[i*r+i+1].abs() < eps*(a1[i*r+i].abs() + a1[(i+1)*r+i+1].abs()) {
                a1[i*r+i+1] = 0.0;
            }
        }

        let mut q = 0;
        for i in (0..r-1).rev() {
            if a1[i*r+i+1].abs() > 0.0 {
                q = i+1;
                break;
            }
        }

        if q == 0 {
            break;
        }

        let mut p = 0;
        for i in (0..q).rev() {
            if a1[i*r+i+1].abs() == 0.0 {
                p = i+1;
                break;
            }
        }

        let mut flag: bool = false;
        
        for i in p..q {
            if a1[i*r+i].abs() == 0.0 {
                flag = true;
                for j in i+1..r {
                    let b = givens_left_rotation(&a1, r, r, i+1, j, true);
                    
                    givens_left_rotation_multiply(&mut a1, r, r, b.0, b.1, i+1, j, 0, r-1, 0, r-1);
                    givens_left_rotation_multiply(&mut u, r, n, b.0, b.1, i+1, j, 0, r-1, 0, n-1);
                }
            }
        }

        if !flag && p < q {
            golub_kahan(&mut a1, &mut u, &mut v, n, m, r, p, q);
        }
    }

    if do_transpose {
        return (v, transpose(&a1, r, r), u);
    }
        
    return (transpose(&u, r, n), a1, transpose(&v, m, r));
}

// Multiply matrices inp1 and inp2 using SIMD instructions
pub fn matrix_multiply_simd(inp1:&[f64], inp2:&[f64], n:usize, m:usize, p:usize) -> Vec {
    const LANES:usize = 64;
    let mut out:Vec = vec![0.0;n*p];

    for i in 0..n {
        for k in 0..m {
            let a:Simd = Simd::splat(inp1[i*m+k]);
            for j in (0..p).step_by(LANES) {
                if k*p+j+LANES > (k+1)*p {
                    let mut r:usize = i*p+j;
                    for h in k*p+j..(k+1)*p {
                        out[r] += inp1[i*m+k]*inp2[h];
                        r += 1;
                    }
                }
                else {
                    let x:Simd = Simd::from_slice(&inp2[k*p+j..k*p+j+LANES]);
                    let z:Simd = a*x;
                    let mut c:Simd = Simd::from_slice(&out[i*p+j..i*p+j+LANES]);
                    c += z;
                    Simd::copy_to_slice(c, &mut out[i*p+j..i*p+j+LANES]);
                }
            }
        }
    }
    return out;
}

fn main() {
    let n = 500;
    let m = 500;

    let mut rng = thread_rng();
    let normal:Normal = Normal::new(0.0, 1.0).ok().unwrap();
    
    let mut a:Vec = vec![0.0;n*m];

    for i in 0..n*m {
        a[i] = normal.sample(&mut rng);
    }

    let start_time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_millis();
    let b = golub_reisch_svd(&a, n, m);
    let end_time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_millis();
    println!("{:?}", end_time-start_time);
    
    let r = min(n, m);
    let mut c = matrix_multiply_simd(&b.0, &b.1, n, r, r);
    c = matrix_multiply_simd(&c, &b.2, n, r, m);

    for i in 0..n*m {
        assert!((c[i]-a[i]).abs() < 1e-5, "Some issue in SVD !!!");
    }
}

The Rust codes are not highly optimized although after build and run with — release flags and compilation with C++ “target-cpu=native”, we saw some pretty good performance for relatively smaller matrices of the order of N=500 and M=500, which are comparable to the “np.linalg.svd” algorithm, which is implemented using LAPACK’s highly optimized FORTRAN routines.

RUSTFLAGS="-C target-cpu=native" cargo run --release

Possible optimizations we could do (which may or may not improve performance are):

Using multithreading for parallel operations.
Using SIMD instructions for block operations.

But the above code is not yet ready to handle very small values of the order of 1e-100. One of the reasons being that if a value x ~ O(1e-100), then x² ~ O(1e-200). With such small values, sqrt() method in Rust is going to 0 and thus causing errors in calculations.

One way is to handle this is to use Taylor Series approximation of sqrt() method when the value is smaller than some threshold.

For e.g. in the Givens Rotation matrixes, calculating the values of c and s could be done as follows (assuming x > y) :

We can terminate the series at the point where the absolute value of the i-th term becomes less than say 1% of the total sum.

The following Rust code can be used to calculate the value of r in the above equations.

fn hypot(x:f64, y:f64) -> f64 {
    let mut i = 1.0;
    let mut s = 1.0;
    let mut r = 0.5;
    let w;
    let u;
    if x.abs() > y.abs() {
        w = (y/x).abs();
        u = x.abs();
    }
    else {
        w = (x/y).abs();
        u = y.abs();
    }
    let mut z = w*w;

    loop {
        let p = (r/i)*z;
        if p.abs()/s < 0.01 {
            break;
        }
        s += p;
        r *= 0.5-i;
        i *= i+1.0;
        z *= w*w;
    }

    return s.abs()*u;
}

Similarly, we should take care of all possible operations wherever we are computing x² terms and probably avoid them.

Given that the time complexity of the above algorithm is O(NM²), we can try to optimize the run-time performance of the SVD algorithm by using a lower value of M (i.e. number of columns)

But, how do we choose a lower M ? By projecting the full matrix of M columns into a lower dimension.

Let P be a matrix of dimensions MxL where L < M. If we multiply A.P, then we will get a new matrix A’ of dimensions NxL. Now if we do SVD on A, the time complexity of the algorithm would be O(NL²). If L << M, it will improve the run time performance a lot.

A’ = A.P

If we do SVD of A, then we will have:

To get the SVD of A, we require P to be an orthogonal matrix. How do we choose an orthogonal matrix, is the question here ?

Instead of choosing an orthogonal matrix, which is a non-trivial operation, we can first find the QR decomposition of A’=A.P i.e.

So, Q is an orthogonal matrix. Instead of multiplying A.P we do Q^T.A i.e.

Now since both Q and U’ are orthogonal matrices, their product is also an orthogonal matrix, thus the matrix U corresponding to A is Q.U’ where U’ is the matrix corresponding to SVD of A’.

Since A.P is of dimensions NxL, where N > L, Q is also an orthogonal matrix of dimensions NxL (see above) and R is an upper diagonal matrix of dimensions LxL.

A’=Q^T.A is a matrix of dimensions LxM. We can perform SVD on A’ by first transposing A’ to a MxL matrix where L << M, and then again transposing the final matrices U, S and V as follows:

The initial matrix P of shape MxL can be chosen from a random normal distribution. This process is popularly known as the Randomzied SVD or Truncated SVD algorithm and is used widely with low rank matrices such as in text classification problems or recommendation problems etc.

The results obtained after Truncated SVD may or may not be close to the true values of the actual SVD because this is a probabilistic algorithm. On increasing the value of L, the results would be closer to the true SVD of A.

The task is to find a good balance between speed and accuracy by choosing an optimal value of L.

pub fn randomized_svd(a:&[f64], n:usize, m:usize, l:usize) -> (Vec, Vec, Vec){
    let mut rng = thread_rng();
    let normal:Normal = Normal::new(0.0, 1.0).ok().unwrap();
    
    let mut p:Vec = vec![0.0;m*l];

    for i in 0..m*l {
        p[i] = normal.sample(&mut rng);
    }

    let b = matrix_multiply_simd(&a, &p, n, m, l);
    let (mut q, _r, n1, m1, _p1) = givens_rotation_qr(&b, n, l);
    let w= min(m1, l);
    q = sub_mat(&q, n1, m1, 0, n1-1, 0, w-1);

    let c = &matrix_multiply_simd(&transpose(&q, n1, w), &a, w, n1, m);
    let (mut u, s, v) = golub_reisch_svd(&c, w, m);
    u = matrix_multiply_simd(&q, &u, n1, w, w);
    return (u, s, v);
}

Link to the entire code in Github: rust_machine_learning/src/svd.rs at main · funktor/rust_machine_learning (github.com)

References

Implementing SVD algorithm in Rust was originally published in Gaussian Machine on Medium, where people are continuing the conversation by highlighting and responding to this story.

Anatomy of a High Level System Design

Abhijit Mondal — Fri, 30 Aug 2024 03:22:32 GMT

System design interviews can be daunting if a candidate is unprepared. Even if a candidate is prepared well enough, during an interview they might crumble under pressure while structuring their system design and such a situation can lead to following debacles:

Candidate presents a very unstructured design (although they mention all important points but do not connect them in a logical flow)
Candidate misunderstands a problem statement and gives a system design of a similar but different system. Candidate does not realizes this until 15–20 minutes into the interview after interviewer probes the design.
Candidate blacks out in the middle when they realize that their approach is leading to nowhere.

After giving several system design interviews in both software engineering and machine learning and also taking several interviews at some big tech companies, here are some pointers that candidates should keep handy during an interview.

Since most interviews nowadays are online, candidates can list down all these pointers in a tab or notepad (most interviewers or companies do not consider this to be cheating).

Requirement Gathering

This should be the first step towards your solution. Ask the interviewer questions relevant to the problem. Some of these could be:

Who are the users of the system, are they internal customers, enterprise customers, or general public ?
Where are the users located, are they geographically distributed or in a specific region ?
How many concurrent users are going to use this service ?
How many requests per second the service is expected to handle ?
Is this a read heavy or write heavy system ?
Does the system needs to be strongly consistent or can be eventually consistent ?
If system deals with file uploads/downloads, how large files or data is uploaded or downloaded on an average ?

… and so on. Most of the questions will be related to the system you are building. But the idea of this exercise is to do :

Capacity planning i.e. number of CPU cores needed, amount of memory required, throughput, disk space, caching requirements etc.
Selecting the appropriate database type i.e. SQL vs. NoSQL vs. Graph DB
Choosing an appropriate database schema i.e. normalized vs. denormalized.
Choosing the optimum number of database partitions.
Whether to go for a single leader or multiple leaders configuration.
Should we go for geo-redundancy or regional redundancy of the service ?
Will the service require a distributed queueing system to queue the incoming requests ?
What kind of database indexing engine is appropriate i.e. B-Tree vs. LSM Tree and so on.

Most of the decisions and tradeoffs are driven by the requirements of the system. All the requirement gathering can be done at the start or as in most practical system design, requirement gathering is a continuous process.

High Level Flow Diagram — Read and Write Paths

This step requires drawing the different high level components required for a working prototype in production. The different component comprises of but not limited to are:

Web server(s) to accept the incoming requests
Load balancer(s) in case multiple web servers are serving requests.
A queueing system like Kafka to queue the incoming requests in the event that there are too many concurrent requests.
A rate limiter to handle too many concurrent requests.
The application server(s) where the main service is running. Again this could be behind some load balancer.
The database server(s) running the DB service (again possibly behind some load balancer). The database service usually will have multiple partitions and each partition will have multiple read replicas.
Caching service to serve faster read requests.
Some cloud based system to collect logs and run monitoring and observability of the different components.

… and so on. Note that these are just some common high level components that most distributed systems would use.

Some “mandatory” boxes in system design

API Design

Many system design interviews will not focus on API Design but there are some very good companies that focus on API design. This round is specifically to understand how much the candidate understands about different API nuances for example:

Define the CRUD (Create, Read, Update, Delete) operations for the system.
Define apis for the CRUD operations.
How to handle pagination during reads through the APIs ?
How to handle SQL injection and sanitize user inputs in GET and POST queries ?
How to serialize and deserialize the POST data into the HTTP request ?
How to deal with error case scenarios and send appropriate HTTP response codes ?

…and so on.

Database Service

This is probably the most “discussed” part in system design interviews in big tech companies as it highlights lots of understanding of distributed system design. The different areas where the interviewers focusses on are as follows:

Does the service need a distributed database or we can work with a distributed blob storage instead of a database ?
Choosing the right database technology — SQL vs. NoSQL vs. Graph Databases vs. In Memory databases and the different tradeoffs associated with it.
Whether a row-oriented filesystem is appropriate or a column oriented filesystem depending on query patterns ?
Designing the table schemas based on the different read and write patterns.
Defining different indexing strategies and comparing different data structures such as B-Tree vs LSM Tree.
How to partition the database, the partitioning columns, number of partitions etc. ?
How to read and write data across different partitions using Consistent Hashing ?
How to choose the number of replicas for each database partition ?
How to select the leader and readers based on the consistency requirements, i.e. strong vs eventual consistency ?
How to sync multiple leaders ?
What happens when a replica server crashes and a new replica joins ?
What happens if a leader crashes and a new leader joins ? How leader election algorithm works ?

… and so on.

This section will probably be the longest if you are interviewing at FAANG compaies as there are lots of different strategies and tradeoffs one can discuss here. Good part is that this is also the most popular section among candidates and one can find very good resources over the internet for all of the above.

Caching Service

If you are designing your system for high throughput, caching is a very critical component. The design of a caching service is quite similar to a database service as often caching is implemented using in-memory databases with limited memory. Some areas where interviewer would want you to focus are:

Do you need a cache and why ?
What data are you going to cache and why ?
What data strcutures are you going to use to cache the data ?
How are you going to update the cache ?
What kind of cache updation strategy will you use such as LRU vs LFU for your problem ?
How many cache servers would be required ?

… and so on.

Quality and Reliability

Since reliability of a system is very critical for a product, this section allows the interviewers to evaluate candidates based on how do they think of different edge cases, failure scenarios and so on. Some pointers that the interviewer will look out for:

How will you debug if your application crashes ?
How will you setup logging ?
What different metrics will you monitor from the VMs for e.g. CPU usage, memory usage etc. ?
How will you autoscale you system based on the monitoring metrics ?
Will you be setting up geo-redundancy or regional redundancy for your service ?
How will you deal with issues like missing data or data duplication in your database or cache ?
How are you going to setup distributed locking for common resources like filesystem ?
How will you debug situations when write queries fails vs. when read queries fails ?

… and so on.

Security

While most companies do not expect candidates to handle security in their system design especially at or below mid-level but for principal and above levels, handling security in you design is of utmost importance. Here are some pointers that interviewers will look out for in your design:

How to handle and sanitize user inputs before querying the database ?
How to handle SQL injections ?
How does authentication and authorization works for each of your component ?
How are you going to setup private networks for your service and how to do active development over private networks ?
How are you going to deal with data with various security levels such as general vs public vs confidential etc. ?

… and so on.

Incorporating quality and security into your system design will fetch bonus points during the interview.

Implementing L-BFGS optimization from scratch

Abhijit Mondal — Thu, 04 Jul 2024 17:51:17 GMT

source: quanta magazine

In this series of posts we are going to look at how we can use L-BFGS-B optimization in dynamic pricing algorithms. The posts are inspired by our team’s work at Microsoft in building a distributed and scalable spot virtual machine dynamic pricing engine.

The posts are broken down as follows:

Understanding and implementing L-BFGS optimization from scratch.
Understanding and implementing automatic differentiation from scratch.
Understanding and implementing L-BFGS-B for bounded constraint problems.
Improving the run-time performance and memory utilization in the L-BFGS-B optimization algorithm.
Using L-BFGS-B for dynamic pricing problems with different types of constraints.
Optimizing our L-BFGS-B implementation to run on GPU.
Making the L-BFGS-B optimization run in a distributed cluster for large number of parameters.

In this post, we are going to understand the basics of L-BFGS optimization algorithm and how it compares with some other 1st order and 2nd order optimization algorithms.

Assuming we are solving an unconstrained optimization problem of a single variable x : minimize f(x).

Then, we can write f(x) about a point x0 using the Taylor Series as follows:

f’(x0) is the 1st order derivative of f w.r.t. x evaluated at x=x0, and similarly f’’(x0) is the 2nd order derivative of f w.r.t. x evaluated at x=x0 and so on.

Note that the above formula holds true for any arbitrary complex function f(x).

1st order approximation of f(x) implies that we truncate the above formula for Taylor Series at the 2nd term i.e.

Similarly, 2nd order approximation of f(x) implies that we truncate the formula for Taylor Series at the 3rd term i.e.

For any x, which is very close to x0 i.e. x-x0 -> 0, then the higher order terms such as (x-x0)³, (x-x0)⁴ and so on becomes very small and thus the 2nd order approximation is very close to the actual value of f(x).

Remember that to find the minimum value of a function f(x) (assuming f is continuous and both f’(x) and f’’(x) exists), we take the 1st derivative of f(x) w.r.t. x and set them to 0, i.e.

Once we solve the the above equation and find the corresponding x_opt at which f’(x_opt) = 0, we can then evaulate the 2nd order derivative at x_opt:

If the 2nd order derivative is greater than 0, then x_opt is a minima. If the 2nd order derivative is less than 0, then x_opt is a maxima. Note that if f(x) is a convex function, then there is only one global minima, but if f(x) is non-convex, then f(x) can have multiple local minima apart from a global minima.

The condition f’(x_opt) = 0 and f’’(x_opt) > 0 holds true for local minima. Thus, with non-convex functions, x_opt can be a local minima.

Solving the equation f’(x) = 0, when f(x) is some complex function may not be always feasible or can be too complicated. Thus, we cannot always use f(x) in its original form to find x_opt.

In such scenarios, the Taylor Series approximations comes handy.

For the 1st order approximation scenario:

Denoting h = x-x0, we have:

We want to find an optimum h, s.t. f(x0+h) is minimum among all possible values of h. Note that we cannot use the method of solving df/dh = 0 and then checking d²f/dh² > 0 because in the above formula, the 2nd derivative d²f/dh² do not exist.

Alternatively, we can find h s.t. f(x0+h) < f(x0), i.e.

Thus, f’(x0)*h must be less than 0.

If we choose h = -a*f’(x0) where a > 0, then we have f’(x0)*h = -a*(f’(x0)² which is < 0 because both a and (f’(x0)² > 0.

Thus, if we choose h = -a*f’(x0), or the next point from x0 as x1 = x0-a*f’(x0), then we will find f(x1) which is smaller than f(x0). This gives us an idea that if we move to x1 then we can lower the value from f(x0) and can bring us closer to the minima.

If we use the following update strategy for x_n in the n-th iteration, then:

This is what usually happens in the gradient descent algorithm, where ‘a’ is the learning rate.

Note that we have to carefully choose the constant ‘a’.

If a is too high, then the next parameter x_n+1 might overshoot the minima and thus it might happen that either the iterations oscillates and do not converge at all or it converges but very slowly.

If ‘a’ is too low, then although the iterations will converge to a local minima, but it might be very slow.

Another approach one can take is instead of using a constant value for ‘a’, one can update ‘a’ in each iteration. One can use heuristics to come up with satisfactory values for ‘a’ such as the sufficient decrease condition:

It says that after an update to x0, f(x1) should be “sufficiently” reduced w.r.t. f(x0) so that we do not converge towards the minima too slow. Another way to look at this is that if we plot the curve of f w.r.t. a, then the curve lies below the line f(x0) — c1*a*(f’(x0))², where the slope of the line is -c1*(f’(x0))².

The other is the curvature condition:

It says that, if we use the same descent direction -f’(x0) at both x0 and x1, then at x1, the angle between the descent direction and the slope is larger than the angle between the decent direction and the slope at x0. It implies that at x1, we are more closer to the minima than at x0.

In general, these conditions are known as the Wolfe Conditions and are used in the Line Search algorithm we will see later in this post.

If f(x) is a convex function, then it is guarenteed to converge to a global minima when the values of ‘a’ are small enough. For non-convex functions, most often we will land up in a local minima with some initial starting point x0. One needs to run the iterations with different initial points to be able to reach a “better” minima or a global minima.

Often the 1st order approach we took above can be slow to reach a local minima. To overcome this, one usually uses a 2nd order approximation to the Taylor Series to compute the descent directions.

In this case, we can use the approach where we set df/dh = 0 and check if d²f/dh² > 0 because d²f/dh² exists.

df/dh = 0 implies that :

If we take the 2nd derivative of f w.r.t. h, we get: d²f/dh² = f’’(x0). Thus in order for f(x0 + h) to be a minima, we also need to satisfy f’’(x0) > 0.

Using the 2nd order approximation, the update to x_n are as follows:

To find the optimum value of ‘a’ at each iteration we can perform the above mentioned line search.

Instead of a function of a single parameter x, we can similarly solve the above optimization for functions taking a vector of parameters X. The 1st derivative is replaced by the gradient vector and the 2nd derivative by the Hessian matrix.

For the 2nd order approximation, we can rewrite the equations in vector and matrix terms as follows:

X0 is a vector of size m, h is another vector of size m, f is the optimization function that takes in a vector of size m and returns a scalar.

g is the gradient vector and g^T is the transpose of the vector.

H is the hessian matrix and is of size m x m:

The update equation looks like:

The corresponding condition f’’(x0) > 0 for a vector of parameters X0, is that the hessian matrix H is positive definite matrix i.e. for any non-zero vector x, we have the following:

If the above condition is satisfied for the hessian matrix H, then we can move in the direction -B_n.g_n to reduce the value of the optimization function.

When the number of parameters i.e. the size of the vector X, m is very large ~ 1⁰⁶ or so, storing and computing the hessian matrix or its inverse is not scalable. To store a matrix of 1⁰¹² floating point values each of 64 bits, would require approximately 8TB of RAM.

Also, computing the inverse of a m x m matrix has a time complexity of O(m³).

In the BFGS approach, instead of computing the inverse of the hessian, its values are approximated using the below formula.

where we have the following:

Thus, we see that the inverse hessian matrix at iteration n+1 i.e. B_n+1 is only dependent on the hessian at iteration n, the difference between the gradients at iterations n and n-1 and the difference between the parameters X at iterations n and n-1.

Now, we do not need to compute the inverse hessian matrix. The initial inverse hessian i.e. B0 can be taken as the identity matrix. Another nice property of the above updates are that if B0 is positive definite, then B_n is positive definite for all n > 0.

But even with this modifications, we still have to keep the matrix B_n of size m x m in memory.

In the Limited Memory BFGS approach, we do not need to store the B_n matrices except for B0 which is an identity matrix of size m. Instead of B_n+1 using the recursion from the n-th iteration, if we recompute all of the previous values at each iteration n i.e.

Let:

Then we can write by expanding B_n in the update equation as follows:

Similarly, expanding the B_n-1 term completely, we get the following:

where

Note that for each iteration n, we have to compute the entire product from iteration 0 to n-1. Moreover, we need to keep all the values y_n and s_n in memory to compute the above terms.

If the number of iterations required to converge are very high, then the amount of memory required as well as the time complexity of this approach will very well exceed the vanilla BFGS strategy where we just needed to keep the inverse hessian matrix from last iteration in memory.

In L-BFGS, instead of using all the previous iterations, only the latest k iterations’ values are considered and the rest are discarded i.e.

Similarly, we use only the latest k values for Wn.

Note that this is an approximation and usually for k between 5 and 10, the inverse hessian computed with latest k values are very much close to the true inverse hessian.

Let’s estimate the time and space complexity with this approach.

Instead of computing B_n and then multiplying with the gradient vector g_n, we can calculate B_n.g_n together. This is because when we are computing B_n beforehand, we need to store a matrix of size m x m in memory which we saw earlier would take up around 8 TB of RAM.

Instead if we compute the following:

Computing the dot product yT_n.g_n, where yT_n is of dimensions 1 x m and g_n is of dimensions m x 1, would result in a scalar quantity. Thus uT_n.g_n is a vector of dimension m x 1 and only requires keeping O(m) values in memory.

Next we take the dot product of the above quantity with uT_n-1.

Thus, again we are only storing O(m) data in memory. Continue the above multiplications till we multiply with the matrix B0. Since B0 is an identity matrix, the multiplication of B0 with the above chain of multiplications results in a vector of size m x 1.

The analysis is same for W_n.g_n. Thus, we see that the space complexity of doing the dot product of B_n.g_n is O(m). To compute the terms u_n and w_n, we need to keep the last k pairs of y_n and s_n in memory i.e. a space complexity of O(km).

Thus the overall space complexity of the above approach is O(k*m), where k is the maximum history we are using.

Time complexity is O(k²m) because when we are taking the dot product of W_n with g_n. There are k terms in W_n and for each term, taking the dot product with g_n has time complexity of O(km), infact it is O(jm) for the j-th term in W_n.

But since k = O(10), the overall time complexity per iteration can still be considered to be O(m).

Once we have obtained the dot product B_n.g_n, then we can go ahead and update the parameter vector X_n+1 as follows:

Only thing remaining is finding the optimum value for the constant ‘a’. There are several ways to calculate the value of ‘a’.

One such strategy is as follows:

Let us define: F(a)=f(X_n)-c1*a*B_n.g_n

For any 2 step sizes, a_low and a_high, if we take some step in between them, let’s say a_mid = (a_low + a_high)/2.

If F(a_mid) > f(X_n+1), then a_mid is a ‘good’ step and we can search for larger step sizes in the interval (a_mid, a_high]. If we get any ‘good’ step in the updated interval then we take that else we use the current ‘good’ step.

Similarly, if F(a_mid) ≤ f(X_n+1), then a_mid is not a ‘good’ step. In this case, we can search in the low interval [a_low, a_mid). But it may happen that there are ‘good’ steps in (a_mid, a_high] but not in [a_low, a_mid).

If we do not find any ‘good’ step in [a_low, a_mid), then we return some default small value for ‘a’ as the step size.

A ‘good’ step is defined based on the Wolfe conditions shown below:

Let the maximum value of a = 1.0 and a_default = 1e-7, then a_low = 0.0 and a_high = 1.0

Find a_mid = (a_low + a_high)/2
Calculate X_n+1 and g_n+1 using a_mid in the above formula.
If the sufficient decrease condition is satisfied then a_mid is a “good” candidate for ‘a’ and we try to search for larger values of ‘a’ s.t. the conditions still holds good by setting a_low = a_mid + 1e-7 and repeat from step 1.
In step 3 if for a_mid, the curvature condition is also satisfied, then we return a_mid as the optimum value for a.
If the sufficient decrease condition do not satisfy, then reduce the value of ‘a’ by setting a_high = a_mid — 1e-7 and repeat from step 1.

Once we have found the optimum value of X_n, we can stop iterating when:

Let’s get down to some coding.

Python implementation of the above L-BFGS algorithm applied on the 2D Rosenbrock function:

import numpy as np
from collections import deque

def fn(x):
    # Compute the optimization function
    return (1.0-x[0,0])**2 + 100*(x[1,0]-x[0,0]**2)**2

def gfn(x):
    # Gradient of fn(x)
    return np.array([[-2*(1.0-x[0,0]) - 400*x[0,0]*(x[1,0]-x[0,0]**2)], 
                    [200*(x[1,0]-x[0,0]**2)]])

def line_search(x, g, d, a0=1e-5, e=1e-6):
    # Perform line (binary) search to find optimal value of a
    left, right = 0.0, 1.0
    a, c1, c2 = a0, 1e-4, 0.9
    
    while left < right:
        m = (left + right)/2
        
        x_new = x - m*d
        g_new = gfn(x_new)

        # sufficient decrease condition
        if fn(x_new)-fn(x) < -c1*m*np.dot(g.T, d):
            # curvature condition
            if abs(np.dot(d.T, g_new)) < c2*abs(np.dot(d.T, g)):
                return m
                
            a = m
            left = m + e
        else:
            right = m - e
            
    return a
    
def optimize(x0, k=10, epsilon=1e-7, max_iter=1e5):
    # l-bfgs optimization algorithm
    m = len(x0)

    x = x0
    g = gfn(x)
    
    # Store the last k values of y and s vectors
    prev_values = deque([(x, g)])

    while max_iter > 0:
        print(f"Current objective value = {fn(x)}")
        max_iter -= 1
        
        prev_x = x
        prev_g = g

        q = g
        for i in range(len(prev_values)-1, -1, -1):
            y, s = prev_values[i]
            r = 1.0/np.dot(y.T, s)
            q = q - r*np.dot(y.T, q)*s

        for i in range(len(prev_values)):
            y, s = prev_values[i]
            r = 1.0/np.dot(y.T, s)
            q = q - r*np.dot(s.T, q)*y

        p = 0
        for i in range(len(prev_values)-1, -1, -1):
            w = g
            
            for j in range(len(prev_values)-1, -1, -1):
                if j > i:
                    y, s = prev_values[j]
                    r = 1.0/np.dot(y.T, s)
                    w = w - r*np.dot(y.T, w)*s
                else:
                    break
                    
            y, s = prev_values[i]
            r = 1.0/np.dot(y.T, s)
            w = r*np.dot(s.T, w)*s

            for j in range(len(prev_values)):
                if j < len(prev_values)-1-i:
                    y, s = prev_values[j]
                    r = 1.0/np.dot(y.T, s)
                    w = w - r*np.dot(s.T, w)*y
                else:
                    break
            p += w

        # Descent direction
        d = p+q

        # Perform line search
        a = line_search(x, g, d)
        
        x = x - a*d
        g = gfn(x)

        y = x-prev_x
        s = g-prev_g
        
        prev_values.append((y, s))

        # Maintain only latest k values
        if len(prev_values) > k:
            prev_values.popleft()

        # Exit condition
        if np.abs(y).all() < epsilon:
            break
        
    return x

if __name__ == '__main__':
    x0 = np.array([[0.003], [0.006]])
    x = optimize(x0, k=10, epsilon=1e-7)
    print(f"Final solution = {x}")

In the above code, we had to explicitly specify the gradient function because the gradient is specific to the objective function and for each new function we have to update the gfn(x) gradient function. Also if the objective function is complicated, then finding the derivative calculations can be erroneous and cumbersome.

In the next post we will see how we can use the automatic differentiation technique to automatically find the gradient without having to update it for every new function.

References:

Implementing L-BFGS optimization from scratch was originally published in Gaussian Machine on Medium, where people are continuing the conversation by highlighting and responding to this story.

Automatic Differentiation from scratch

Abhijit Mondal — Thu, 04 Jul 2024 17:34:21 GMT

source: quanta magazine

In the previous part of this series, we understood L-BFGS optimization algorithm and implemented the same from scratch in Python.

Implement L-BFGS optimization from scratch | by Abhijit Mondal | Jul, 2024 | Medium

One of the issues in the implementation we found was that, we needed to explicitly define the gradient for the objective function. While this works for simple objective functions, but for complicated functions calculating gradient can be tedious and erroneous.

For dynamic pricing, the objective function can be quite complicated with millions of parameters interacting with one another. In this post, we take a detour from L-BFGS and show an automatic method for computing the gradients without explicitly calculating the derivatives.

Recall, that from the 1st order Taylor series approximation, we can write the following:

The above approximation holds good when h is very small and is close to 0.

Then, we can compute the derivative f’ at some point x, using the above formula as:

Thus, we can use the above strategy to find the derivative of the objective function f at some point x. But the above has few problems such as:

Holds good only in the limit h->0. On most computers it can cause truncation errors. For e.g. if h is very small, then f(x+h) — f(x) would be very small and the computer will make it 0. Thus f’(x) would be 0.
If the function f(x) is complicated and require lots of CPU cycles, we are computing f 2 times for derivative. Now, if x is a vector of parameters of size n, then computing the gradient for each parameter would require 2n calculations.
Even if the error from the above approach is small, the errors can propagate and can blow up. For e.g. in neural networks with 1000s of layers, the error will propagate to all weights and all layers eventually, leading to significant errors in calculations.

To overcome the above issues, researchers have come up with the “automatic differentiation” approach. There are 2 different types of automatic differentiation strategies:

Forward mode and Reverse mode.

The forward mode is most intuitive and works as follows:

Given a function, which is a composition of multiple functions, for e.g. f(g1(h1(x,y), h2(y)), g2(h3(x,y), h4(x))). Then we can find the partial derivatives of f w.r.t. x and y as follows:

If we draw the computation graph for the above function:

Then we can obtain the partial derivatives of f w.r.t. x and y by backtracking from f in the above diagram to the corresponding nodes x and y. Note that we have to backtrack across all paths that lead to either x or y.

The backtracking paths are shown in red in the below diagram for partial derivative w.r.t. x.

For each edge from a to b in red, if we take the derivative da/db and take the product of all the derivatives for each edge in a given path and then sum over all paths, we will obtain the partial derivative formulae described above.

Thus, either we can use recursion or find the topological sorting of the computational graph to calculate the partial derivatives.

For e.g. using recursion we can write, G(f,x) = (df/dg1).G(g1,x) + (df/dg2).G(g2,x) where G(f,x) is the partial derivative of f w.r.t. x. or in general:

Another way to represent the above partial derivatives is through Jacobian matrices. Let us represent the following Jacobian matrices:

Then we can obtain the partial derivatives df/dx and df/dy simply by multiplying the above Jacobians.

Let’s look at the time complexity.

Let us assume that Jh is a matrix of size m1 x n (where n is the number of inputs and here n=2 i.e. x and y)

Let the dimension of Jg be m2 x m1 (The number of columns is m1 because g is a function of h and h has m1 rows).

Similarly, Jf will have m2 columns, but has only 1 row hence dimension is 1 x m2

Assuming that we are doing forward mode automatic differentiation, then the Jacobian Jh will be calculated first (due to recursion in the computational graph), then Jg and then Jf and thus the multiplication happens from right to left. Thus the time complexity is

O(m2*m1*n + 1*m2*n) = O(n³) assuming m1 = m2 = O(n).

Instead if we have done the matrix multiplications from left to right, the time complexity would have been:

O(1*m2*m1 + 1*m1*n) = O(n²)

This is only possible if we compute the Jacobian Jf first and then Jg and lastly Jh, which requires us to compute the automatic differentiation in the reverse order which is known as the Backward or Reverse mode of automatic differentiation.

An efficient way to calculate the Jacobians in the reverse mode is to first evaluate each function for the given inputs and then cache the function results and use them for differentiation. For e.g. in the above example given x and y as inputs calculate the following in the order:

h1, h2, h3 and h4
g1 and g2
f

Let us call it the forward pass.

Once we cache the results of the above function calls, then we can directly use them in the partial derivatives. For e.g. df/dg1 is simply the derivative of f w.r.t. g1 evaluated at g1 calculated in the forward pass.

Something similar happens in Neural Networks where you can assume that the (x, y) is the input layer, (h1, h2, h3, h4) is the 1st hidden layer, (g1, g2) is the next hidden layer and f is the output unit. Usually we make a forward pass first with x and y and then do backpropagation in the reverse order.

Thus the reverse mode of automatic differentiation is essentially the backpropagation algorithm.

Firstly, lets start by implementing the forward mode of automatic differention. Then we will gradually go and implement the reverse mode.

But the question is how to form the computational graph given any objective function. Let’s take an example:

Here we have only a single input variable x. Here is a way to break up that complicated looking function into intermediate functions as follows:

The computational graph will look as follows:

Note that after breaking up the function f into the above computational graph, we have obtained much simpler functions such as sin, cos, ln, exp, x^k etc. apart from addition and multiplication.

Computing the derivatives of the simpler functions is straightforward and moreover we can explictly define the derivatives for such functions given that there will be a handful of basic operations:

sin, cos, tan, exp, pow, log should cover most useful objective functions.

Here is a simple Python implementation of automatic differentiation for a single variable from scratch:

from functools import wraps
import math

# Automatic differentiation of exp(1 + x^4*(cos(log(x)) + sin(log(x))))

class Node:
    def __init__(self):
        self.function = None
        self.d_function = None

    def __add__(self, other):
        return add(self, other)

    def __sub__(self, other):
        return sub(self, other)

    def __mul__(self, other):
        return mul(self, other)

    def __truediv__(self, other):
        return div(self, other)

    def __call__(self, *args):
        return self.function(*args)

class Variable:
    def __init__(self):
        pass
        
    def __new__(cls):
        cls.node = Node()
        cls.node.function = lambda z: z
        cls.node.d_function = lambda : Constant(1.0)
        return cls.node

class Constant:
    def __init__(self, val):
        pass
        
    def __new__(cls, val):
        cls.node = Node()
        cls.node.function = lambda z: val
        cls.node.d_function = lambda : Constant(0.0)
        return cls.node

def my_decorator(f):
    @wraps(f)
    def wrapper(*args, **kwds):
        node = Node()
        new_args = []

        for x in args:
            if isinstance(x, Node):
                new_args += [x]
            else:
                raise Exception("Invalid input parameter !!!")

        evaluate, derive = f(*new_args, **kwds)
        
        node.function = evaluate
        node.d_function = derive
        
        return node
        
    return wrapper 

@my_decorator
def cos(x):
    def evaluate(p):
        return math.cos(x.function(p))
    def derive():
        return Constant(-1.0)*sin(x)*grad(x)
    return evaluate, derive

@my_decorator
def sin(x):
    def evaluate(p):
        return math.sin(x.function(p))
    def derive():
        return cos(x)*grad(x)
    return evaluate, derive

@my_decorator
def log(x):
    def evaluate(p):
        return math.log(x.function(p))
    def derive():
        return (Constant(1.0)/x)*grad(x)
    return evaluate, derive

@my_decorator
def exp(x):
    def evaluate(p):
        return math.exp(x.function(p))
    def derive():
        return exp(x)*grad(x)
    return evaluate, derive

@my_decorator
def pow(x, y):
    def evaluate(p):
        return math.pow(x.function(p), y.function(p))
    def derive():
        return y*pow(x, y-Constant(1))*grad(x)
    return evaluate, derive

@my_decorator
def add(x, y):
    def evaluate(p):
        return x.function(p) + y.function(p)
    def derive():
        return grad(x)+grad(y)
    return evaluate, derive

@my_decorator
def sub(x, y):
    def evaluate(p):
        return x.function(p) - y.function(p)
    def derive():
        return grad(x)-grad(y)
    return evaluate, derive

@my_decorator
def mul(x, y):
    def evaluate(p):
        return x.function(p) * y.function(p)
    def derive():
        return x*grad(y) + y*grad(x)
    return evaluate, derive

@my_decorator
def div(x, y):
    def evaluate(p):
        return x.function(p) / y.function(p)
    def derive():
        return Constant(-1.0)*x*pow(grad(y), Constant(-2.0)) + pow(y, Constant(-1.0))*grad(x)
    return evaluate, derive

@my_decorator
def grad(f):
    def evaluate(p):
        return f.d_function().function(p)
    def derive():
        return f.d_function().d_function()
    return evaluate, derive

def standard_fn(x):
    f1 = math.log(x)
    f2 = math.cos(f1)
    f3 = math.sin(f1)
    f4 = f2 + f3
    f5 = math.pow(x, 4.0)
    f6 = f5*f4
    f7 = math.exp(1.0+f6)
    return f7

def grad_numerical(f, x, h=1e-6):
    return (f(x+h)-f(x))/h

def grad_ad(x):
    y = Variable()
    f = exp(Constant(1.0) + pow(y, Constant(4.0))*(cos(log(y)) + sin(log(y))))
    return grad(f)(x)

if __name__ == '__main__':
    y = Variable()
    z = pow(y, Constant(5))
    assert z(2) == 32
    assert grad(z)(3) == 405
    assert grad(grad(z))(4) == 1280
    assert grad(grad(grad(z)))(5) == 1500

    x = 0.679
    ad = grad_ad(x)
    print(f"Autograd = {ad}")
    nm = grad_numerical(standard_fn, x)
    print(f"Numerical grad = {nm}")
    ex = math.exp(1.0 + (x**4)*(math.cos(math.log(x)) + math.sin(math.log(x))))*(4*(x**3)*(math.cos(math.log(x)) + math.sin(math.log(x))) + ((x**4)*(-math.sin(math.log(x))/x + math.cos(math.log(x))/x)))
    print(f"Exact grad = {ex}")

The idea is as follows:

We are wrapping the conventional python math functions inside a Node object. Each Node object has a “function” method and a “d_function” method. The “function” method evaluates the corresponding math function at some value p. “d_function” method returns a Node object corresponding to the derivative function.

Taking an example from above:

The corresponding representation of the above function in our implementation would be:

x = Variable()
sin(exp(x)+pow(x, Constant(3.0))) + \
  pow(x, Constant(2.0))*cos(Constant(1.0) + log(exp(x) + pow(x, Constant(3.0))))

Calling one of these functions sin(f(x)), cos(f(x)), pow(f(x), g(x)), log(f(x)), will first call the ‘my_decorator’ decorator which creates a Node object and checks if the input arguments are Node objects or not. Each argument must be a Node object, whether its another function, a variable or a constant.

After checking that all arguments are Node objects, execute the code inside the called method and return the “function” and “d_function” methods and attach them to the new Node object.

The above implementation for automatic differentiation for the given function may not be most optimal because note that the term (e^x + x³) occurs twice in the function f(x).

Instead of creating 2 different Node objects corresponding to (e^x + x³), we could have created a single Node object and re-used it like in the computational graph where x3 is used by both x4 and x6. Moreover we could cache the results of the “function” call at each Node.

Now let’s update the code to handle a vector of parameters and compute partial derivatives w.r.t. one of the parameters at a time. Also, adding caching of function results.

from functools import wraps
import math, numpy as np

class Node:
    def __init__(self):
        self.function = None
        self.d_function = None
        self.id = None
        self.cached_function = {}

    def __add__(self, other):
        return add(self, other)

    def __sub__(self, other):
        return sub(self, other)

    def __mul__(self, other):
        return mul(self, other)

    def __truediv__(self, other):
        return div(self, other)

    def __call__(self, *args):
        return self.function(*args)

class Variable:
    def __init__(self):
        pass
        
    def __new__(cls, id):
        cls.node = Node()
        cls.node.id = id
        cls.node.function = lambda z: z[id,0] if id < len(z) else 0
        cls.node.d_function = lambda j: Constant(1.0) if j == id else Constant(0.0)
        return cls.node

class Constant:
    def __init__(self, val):
        pass
        
    def __new__(cls, val):
        cls.node = Node()
        cls.node.function = lambda z: val
        cls.node.d_function = lambda j: Constant(0.0)
        return cls.node

def my_decorator(f):
    @wraps(f)
    def wrapper(*args, **kwds):
        node = Node()
        new_args = []

        for i in range(len(args)):
            if i == len(args)-1 and f.__name__ == 'grad':
                new_args += [args[i]]
            else:
                if isinstance(args[i], Node):
                    new_args += [args[i]]
                else:
                    raise Exception("Invalid input parameter !!!")

        evaluate, derive = f(*new_args, **kwds)
        
        node.function = evaluate
        node.d_function = derive
        
        return node
        
    return wrapper 

def get_cached_function(x, p):
    q = tuple(p.reshape((1, p.shape[0]))[0].tolist())
        
    if q not in x.cached_function:
        x.cached_function[q] = x.function(p) 
    return x.cached_function[q]

@my_decorator
def cos(x):
    def evaluate(p):
        z = get_cached_function(x, p)
        return math.cos(z)
    def derive(j):
        return Constant(-1.0)*sin(x)*grad(x,j)
        
    return evaluate, derive

@my_decorator
def sin(x):
    def evaluate(p):
        z = get_cached_function(x, p)
        return math.sin(z)
    def derive(j):
        return cos(x)*grad(x,j)
        
    return evaluate, derive

@my_decorator
def log(x):
    def evaluate(p):
        z = get_cached_function(x, p)
        return math.log(z)
    def derive(j):
        return (Constant(1.0)/x)*grad(x,j)
        
    return evaluate, derive

@my_decorator
def exp(x):
    def evaluate(p):
        z = get_cached_function(x, p)
        return math.exp(z)
    def derive(j):
        return exp(x)*grad(x,j)
        
    return evaluate, derive

@my_decorator
def pow(x, y):
    def evaluate(p):
        z = get_cached_function(x, p)
        w = get_cached_function(y, p)
        return math.pow(z, w)
    def derive(j):
        return y*pow(x, y-Constant(1.0))*grad(x,j)
        
    return evaluate, derive

@my_decorator
def add(x, y):
    def evaluate(p):
        z = get_cached_function(x, p)
        w = get_cached_function(y, p)
        return z + w
    def derive(j):
        return grad(x,j)+grad(y,j)
        
    return evaluate, derive

@my_decorator
def sub(x, y):
    def evaluate(p):
        z = get_cached_function(x, p)
        w = get_cached_function(y, p)
        return z - w
    def derive(j):
        return grad(x,j)-grad(y,j)
        
    return evaluate, derive

@my_decorator
def mul(x, y):
    def evaluate(p):
        z = get_cached_function(x, p)
        w = get_cached_function(y, p)
        return z * w
    def derive(j):
        return x*grad(y,j) + y*grad(x,j)
        
    return evaluate, derive

@my_decorator
def div(x, y):
    def evaluate(p):
        z = get_cached_function(x, p)
        w = get_cached_function(y, p)
        return z / w
    def derive(j):
        return Constant(-1.0)*x*pow(grad(y,j), Constant(-2.0)) + pow(y, Constant(-1.0))*grad(x,j)
        
    return evaluate, derive

@my_decorator
def grad(f,j):
    def evaluate(p):
        return f.d_function(j).function(p)
    def derive(k):
        return f.d_function(j).d_function(k)
    return evaluate, derive

if __name__ == '__main__':
    x = np.array([[Variable(0)], [Variable(1)]])
    f = cos(pow(x[0,0],Constant(3.0)) + pow(x[1,0],Constant(2.0)) + Constant(2.0)*x[0,0]*x[1,0])
    inp = np.array([[0.543], [0.678]])

    # forward pass to cache the function values
    f(inp)
    assert grad(grad(f, 1), 0)(inp) == -3.1197960007398184

To first compute the partial derivative w.r.t. y and then w.r.t. x:

at the point [0.543, 0.678].

When you compare the results of the gradients computed using automatic differentiation vs. numerical methods as described earlier, we can see that the result from automatic differentiation is much closer to the true value and thus the error is almost 0.

The updated L-BFGS code (introduced in the previous post of this series) using automatic differentiation instead of explicit gfn() method described earlier.

import numpy as np
from collections import deque

def p_grad(f, n):
    def eval(x0):
        out = []
        for i in range(n):
            out += [[grad(f, i)(x0)]]
    
        return np.array(out)
    return eval

def fn(x):
    # Compute the optimization function
    return pow(Constant(1.0)-x[0,0], Constant(2.0)) + Constant(100.0)*pow((x[1,0]-pow(x[0,0], Constant(2.0))), Constant(2.0))

def line_search(x, g, d, fnode, gnode, a0=1e-5, e=1e-6):
    # Perform line (binary) search to find optimal value of a
    left, right = 0.0, 1.0
    a, c1, c2 = a0, 1e-4, 0.9
    
    while left < right:
        m = (left + right)/2
        
        x_new = x - m*d
        f_new = fnode(x_new)
        g_new = gnode(x_new)

        # sufficient decrease condition
        if f_new-fnode(x) < -c1*m*np.dot(g.T, d):
            # curvature condition
            if abs(np.dot(d.T, g_new)) < c2*abs(np.dot(d.T, g)):
                return m
                
            a = m
            left = m + e
        else:
            right = m - e
            
    return a
    
def optimize(x0, k=10, epsilon=1e-7, max_iter=1e5):
    # l-bfgs optimization algorithm
    m = len(x0)
    v = np.array([[Variable(i)] for i in range(m)])

    fnode = fn(v)
    gnode = p_grad(fnode, m)
    
    x = x0
    f = fnode(x)
    g = gnode(x)
    
    # Store the last k values of y and s vectors
    prev_values = deque([(x, g)])

    while max_iter > 0:
        print(f"Current objective value = {f}")
        max_iter -= 1
        
        prev_x = x
        prev_g = g

        q = g
        for i in range(len(prev_values)-1, -1, -1):
            y, s = prev_values[i]
            r = 1.0/np.dot(y.T, s)
            q = q - r*np.dot(y.T, q)*s

        for i in range(len(prev_values)):
            y, s = prev_values[i]
            r = 1.0/np.dot(y.T, s)
            q = q - r*np.dot(s.T, q)*y

        p = 0
        for i in range(len(prev_values)-1, -1, -1):
            w = g
            
            for j in range(len(prev_values)-1, -1, -1):
                if j > i:
                    y, s = prev_values[j]
                    r = 1.0/np.dot(y.T, s)
                    w = w - r*np.dot(y.T, w)*s
                else:
                    break
                    
            y, s = prev_values[i]
            r = 1.0/np.dot(y.T, s)
            w = r*np.dot(s.T, w)*s

            for j in range(len(prev_values)):
                if j < len(prev_values)-1-i:
                    y, s = prev_values[j]
                    r = 1.0/np.dot(y.T, s)
                    w = w - r*np.dot(s.T, w)*y
                else:
                    break
            p += w

        # Descent direction
        d = p+q

        # Perform line search
        a = line_search(x, g, d, fnode, gnode)
        
        x = x - a*d
        f = fnode(x)
        g = gnode(x)

        y = x-prev_x
        s = g-prev_g
        
        prev_values.append((y, s))

        # Maintain only latest k values
        if len(prev_values) > k:
            prev_values.popleft()

        # Exit condition
        if np.abs(y).all() < epsilon:
            break
        
    return x

if __name__ == '__main__':
    x0 = np.array([[0.003], [0.006]])
    x = optimize(x0, k=10, epsilon=1e-7)
    print(x)

Two popular libraries for automatic differentiation are autograd and JAX. I would encourage the reader to try out JAX for the above L-BFGS implementation.

In the next post we will dive deeper into L-BFGS algorithm for problems involving constraints and understand L-BFGS-B i.e. L-BFGS Bounded, algorithm.

References:

Automatic Differentiation from scratch was originally published in Gaussian Machine on Medium, where people are continuing the conversation by highlighting and responding to this story.

A hate story of Rust and Non Linear Data Structures

Abhijit Mondal — Mon, 19 Feb 2024 04:30:48 GMT

source: quanta magazine

Quite recently Rust has gained a lot of popularity as a high level programming language that is fast and memory safe. The high performance of Rust primarily stems from the fact that it is a compiled language similar to C/C++.

I have explored the memory management and memory safety aspect of Rust in my previous post. Do give it a read.

The time I had mixed feelings about Rust … | by Abhijit Mondal | Feb, 2024 | Medium

In this post, I wanted to show and explain how I struggled to implement some non-linear data structure operations using Rust especially Trees.

Here in this post I will show some possible ways by which we can implement AVL Tree (insertion operation only) and see what kind of challenges we are going to face as a newbie Rust adopter.

AVL Tree is a self balanced binary search tree that allows O(logN) worst case time complexity for search, insert and delete operations.

Although a standard binary search tree also has O(logN) time complexity but it is the average and best case complexity. In the worst case, a skewed BST has time complexity of O(N) for all these operations.

AVL Tree uses self balancing techniques after each insertion and deletion operation in order to ensure that the height difference between left and right subtree for any node is atmost 1. This ensures that the height of AVL Tree is O(logN) for N data points. In this post I will only show insertion operations on an AVL Tree.

Before we can explain why Rust is challenging for an AVL Tree, we need to explain how AVL Tree achieves height re-balancing for these conditions after each insertion.

For the 1st case, there can be 2 scenarios after a new node is inserted (note that the AVL Tree before insertion is height balanced).

Case 1: Insert node 10

In the diagram node 80 violates the height balancing property. Will assume that node 80 is the 1st node in the path from the new node 10 to root node that violates height balancing property.

Case 2: Insert node 60

Case 1 is what is called the left left rotation and we can rebalance the tree as follows. Here we are assuming that the root node of sub-tree T1 is 40.

Note that if instead of node 40, value of this node was 5, then node 10 would have gone to sub-tree T5 instead of T4. But even in that case the above rotation is valid. The main thing is that node 10 is less than node 50 for left left rotation.

Case 2 is what is called left right rotation because node 60 goes to left of node 80 and then to right of node 50. We will assume that the root node of sub-tree T2 is node 55.

This left right rotation is a 2 step process, where first we do right rotation on node 50, then we do left rotation on node 55.

Similar logic is applicable for the 2nd violation condition, when the right sub-tree has height more than 1 than the left sub-tree. In such a case, we do “right right” rotation if the new node has value greater than the right child and do “right left” rotation if the new node has value smaller than the right child. (I will not go into the details of that approach here).

Let’s dive into the approaches here:

Box

This is in continuation of our previous post, where we had used Box smart pointers to implement a BST node. We will continue to use it here and explore some of its advantages and disadvantages.

Let’s start by defining our AVLTree struct and the Node struct:

pub trait MyTrait: PartialOrd + Debug + Copy {}

impl MyTrait for isize {}
impl MyTrait for usize {}

impl MyTrait for f32 {}
impl MyTrait for f64 {}

impl MyTrait for &str {}

struct AVLTree where T: PartialOrd {
    root_node: Option>>,
}

impl AVLTree {
    fn new() -> Self {
        Self {
            root_node: None,
        }
    }
}

#[derive(Clone)]
struct Node {
    val: T,
    lt_node: Option>>,
    rt_node: Option>>,
    height: isize
}

impl Node {
    fn new(&val:&T) -> Self {
        Self {
            val,
            lt_node: None,
            rt_node: None,
            height: 1,
        }
    }
}

The Node of our AVLTree is assumed to hold the following data types in Rust: isize, usize, f32, f64 and &str. We can always extend this list but for demonstration purposes, we will keep it small.

Next we define some helper methods that we will use for height calculations. These methods are required during re-balancing stage.

impl AVLTree {
    fn get_lt_rt_heights(&self, root_node:&Option>>) -> (isize, isize) {
        // returns tuple of left and right subtree heights
 
        let mut lheight:isize = 0;
        let mut rheight:isize = 0;

        match root_node {
            Some(root) => {
                let lt_node = &root.lt_node;
                let rt_node = &root.rt_node;

                match lt_node {
                    Some(x) => {
                        lheight = x.height;
                    }
                    None => {}
                }

                match rt_node {
                    Some(x) => {
                        rheight = x.height;
                    }
                    None => {}
                }
            }
            None => {}
        }

        return (lheight, rheight);
    }
}

impl AVLTree {
    fn height_diff(&self, root_node:&Option>>) -> isize {
        // returns difference b/w left subtree height and right subtree height
        let lr_heights = self.get_lt_rt_heights(&root_node);
        return lr_heights.0 - lr_heights.1;
    }
}

impl AVLTree {
    fn set_height(&mut self, root_node:&mut Option>>) {
        // calculate the height of a tree node
        let lr_heights = self.get_lt_rt_heights(&root_node);

        match root_node {
            Some(root) => {
                root.height = 1 + max(lr_heights.0, lr_heights.1);
            }
            None => {}
        }
    }
}

Note that in the 1st two methods above, we are passing root_node as an immutable reference and in the 3rd method i.e. set_height as a mutable reference.

If we do not pass as a reference i.e. borrow root_node, then root_node will move into the function and its scope will end with the function. The reason for passing as mutable reference is that we are updating the root_node struct properties in-place i.e. height of root_node.

Everytime we define a function that takes as input the root_node, we need to either pass it as a reference (if we are doing in-place modifications to the heap object) or move it inside the function but create a clone object in heap and return the clone object.

Something similar to below:

impl AVLTree {
    fn set_height(&mut self, root_node:Option>>)->Option>> {
        // moving root_node into set_height method
        let mut root_clone = root_node.clone();
        let lr_heights = self.get_lt_rt_heights(&root_clone);

        match root_clone {
            Some(ref mut root) => {
                root.height = 1 + max(lr_heights.0, lr_heights.1);
            }
            None => {}
        }

        return root_clone;
    }
}

But cloning a heap object has memory implications as with Box smart pointers, Box pointer owns the data thus when we clone the object in memory, the value is also copied.

But in the above code, Rust compiler detects that the original root_node that moved into the function was never used again, thus it deallocates the heap memory assigned to original root_node. Thus only a single copy of root_node exists at any point in time in the above code.

But it is not always as simple as this. One can track the peak memory usage using “heaptrack” utility.

Compile and run the program avl_tree.rs with heaptrack:

RUSTFLAGS=-g cargo build --release
heaptrack target/release/avl_tree

If you have heaptrack_gui installed, you can see the graph of memory usage overtime. The peak heap memory usage gives the maximum heap memory used by the program and peak RSS memory is all types of memory used by the program including stack, heap etc.

Now let’s look at the left rotation function. This function takes as input node X and updates the sub-tree rooted at X as follows:

The code is as follows:

impl AVLTree {
    fn left_rotate(&mut self, root_node:&mut Option>>) {
        match root_node {
            Some(node) => {
                let lt_node = &mut node.lt_node;
                let lt_node_clone = &mut lt_node.clone();

                match lt_node {
                    Some(x) => {
                        let lt_rt_node = &x.rt_node;
                        node.lt_node = lt_rt_node.clone();
                    }
                    None => {}
                }

                self.set_height(root_node);

                match lt_node_clone {
                    Some(x) => {
                        x.rt_node = root_node.clone();
                    }
                    None => {}
                }

                *root_node = lt_node_clone.clone();
                self.set_height(root_node);
            }
            None => {}
        }
    }
}

You would notice that we are doing a bunch clone() operations on the Node objects. Cloning is required because we are only borrowing the Node objects and we cannot move out by dereferencing inside the function because the root_node we are passing to the function is borrowed.

Simply put, if we are borrowing an object inside a function and the object is used after function scope ends, we cannot move out (or dereference or change ownership) of that object or any other object referenced by the object, inside the function.

Note that once we connect X to subtree T2, we need to recalculate the height of X. Height of subtree T2 is not affected. Next we recalculate height of Y. Also we are re-assigning the value inside root_node to Y.

Similarly, we can define the right rotation method:

impl AVLTree {
    fn right_rotate(&mut self, root_node:&mut Option>>) {
        match root_node {
            Some(node) => {
                let rt_node = &mut node.rt_node;
                let rt_node_clone = &mut rt_node.clone();

                match rt_node {
                    Some(x) => {
                        let rt_lt_node = &x.lt_node;
                        node.rt_node = rt_lt_node.clone();
                    }
                    None => {}
                }

                self.set_height(root_node);

                match rt_node_clone {
                    Some(x) => {
                        x.lt_node = root_node.clone();
                    }
                    None => {}
                }

                *root_node = rt_node_clone.clone();
                self.set_height(root_node);
            }
            None => {}
        }
    }
}

This is reverse of the left rotation.

Finally we define the insert method:

impl AVLTree {
    fn insert(&mut self, root_node:&mut Option>>, &val:&T) {
        let mut hdiff:isize = 0;

        match root_node {
            Some(node) => {

                if val <= node.val {
                    let lt_node = &mut node.lt_node;
                    match lt_node {
                        Some(_x) => {
                            self.insert(lt_node, &val);
                        }
                        None => {
                            let new_node:Node = Node::new(&val);
                            node.lt_node = Some(Box::new(new_node));
                        }
                    }
                }

                else {
                    let rt_node = &mut node.rt_node;
                    match rt_node {
                        Some(_x) => {
                            self.insert(rt_node, &val);
                        }
                        None => {
                            let new_node:Node = Node::new(&val);
                            node.rt_node = Some(Box::new(new_node));
                        }
                    }
                }

                self.set_height(root_node);
                hdiff = self.height_diff(root_node);
            }

            None => {
                let new_node:Node = Node::new(&val);
                *root_node = Some(Box::new(new_node));
            }
        }

        if hdiff > 1 {
            match root_node {
                Some(node) => {
                    let lt_node = &mut node.lt_node;
                    match lt_node {
                        Some(x) => {
                            if val <= x.val {
                                self.left_rotate(root_node);
                            }
                            else if val > x.val {
                                self.right_rotate(lt_node);
                                self.left_rotate(root_node);
                            }
                        }
                        None => {}
                    }
                }
                None => {}
            }
        }

        else if hdiff < -1 {
            match root_node {
                Some(node) => {
                    let rt_node = &mut node.rt_node;
                    match rt_node {
                        Some(x) => {
                            if val <= x.val {
                                self.left_rotate(rt_node);
                                self.right_rotate(root_node);
                            }
                            else if val > x.val {
                                self.right_rotate(root_node);
                            }
                        }
                        None => {}
                    }
                }
                None => {}
            }
        }
    }
}

Lets also define the main method, which creates a vector of 1 million integers from 1 to 1 million, shuffles them and inserts one by one into the tree.

fn main() {
    let mut avltree:AVLTree = AVLTree::new();
    let mut root:Option>> = None;

    let mut vec: Vec = (1..1000000).collect();
    vec.shuffle(&mut thread_rng());

    for v in vec.iter() {
        avltree.insert(&mut root, &v);
    }
}

We can track how much heap memory and total memory used by the program using the heaptrack program.

We can see that the maximum heap memory used at any point in time is around 62 MB. This will vary since the inputs are shuffled.

If we look at the heaptrack output for lines within the program that consumes that largest heap memory, we can see that all the clone() operations within the left and right rotations take up the largest heap memory.

If you observe carefully, for the left rotation, cloning node Y takes roughly twice the memory of cloning X simply because number of nodes in Y is roughly twice that in X.

Similarly, cloning X takes roughly twice memory than T2. For each unit increase in height, the memory consumption increases by 2x.

HashMap

In the next implementation, instead of using Box smart pointers for left and right subtrees of a Node, we are going to assign ids to Node and maintain a separate id to Node HashMap.

The usefulness of this approach is that when we are cloning the Node objects, instead of the (almost) entire tree being deep copied in heap memory, we are just copying individual Node objects with id fields.

struct AVLTree where T: PartialOrd {
    auto_inc_id: usize,
    id_to_node_map: HashMap>,
}

impl AVLTree {
    fn new() -> Self {
        Self {
            auto_inc_id: 0,
            id_to_node_map: HashMap::new(),
        }
    }
}

#[derive(Clone)]
struct Node {
    val: T,
    lt_node_id: Option,
    rt_node_id: Option,
    height: isize
}

impl Node {
    fn new(&val:&T) -> Self {
        Self {
            val,
            lt_node_id: None,
            rt_node_id: None,
            height: 1,
        }
    }
}

Note that we are using HashMap instead of a Vec in Rust because if we are implementing deletion, then it is efficient to delete from a HashMap rather than a Vector.

The id field is of type usize.

I am not going to show all the codes for HashMap implementation here but would recommend going through it in the below repo.

rust-algorithms/src/bin/avl_tree_hmap.rs at main · funktor/rust-algorithms (github.com)

An interesting observation from the heaptrack output is that the peak heap memory with HashMap is quite high as compared to Box approach above. Roughly 2–3 times the memory usage.

On deep diving we can see that HashMap is the culprit here as it requires additional data structure to hold the key-value pairs and the hashes.

As you can see that whenever a key-value pair is inserted into the HashMap, there is a possibility that the Hash Table will be resized depending on the load factor and often when they are resized, there are lots of wasted memory in the Hash Table.

Rc

In the next approach, instead of using Box smart pointers, I had used another type of smart pointers Rc and RefCell. Some important differences between Box and Rc are as follows:

In Box smart pointers the data behind the smart pointer has a single owner. Thus when you clone the data, another copy of that data is made in the heap memory as we saw above for left and right rotations.
In Rc, the data behind the smart pointer can have multiple owners. But Rc without RefCell allows only immutable access i.e. even if there are multiple owners of the data, they cannot mutate the values. Rc represents Reference Counting. Thus whenever a clone is made, the reference count of the object is incremented by 1, but not copied in heap memory.
Using RefCell along with Rc allows multiple owners to mutate the data behind the smart pointer. But unlike Box, where the ownership and borrow is checked at compile time, in Rc the borrow is checked at runtime. Thus you will not see “much” compiler complains but if it is the case that multiple owners are trying to mutate the data at runtime, the rust program will panic.

Here is the data structure definition:

struct AVLTree where T: PartialOrd {
    root_node: Option>>>,
}

impl AVLTree {
    fn new() -> Self {
        Self {
            root_node: None,
        }
    }
}

#[derive(Clone)]
struct Node {
    val: T,
    lt_node: Option>>>,
    rt_node: Option>>>,
    height: isize
}

impl Node {
    fn new(&val:&T) -> Self {
        Self {
            val,
            lt_node: None,
            rt_node: None,
            height: 1,
        }
    }
}

The methods to compute the heights of left and right subtrees and set height of self.

impl AVLTree {
    fn get_lt_rt_heights(&self, root_borrow:&RefMut<'_, Node>) -> (isize, isize) {
        let mut lheight:isize = 0;
        let mut rheight:isize = 0;

        let lt_node = &root_borrow.lt_node;
        let rt_node = &root_borrow.rt_node;

        match lt_node {
            Some(x) => {
                lheight = x.borrow().height;
            }
            None => {}
        }

        match rt_node {
            Some(x) => {
                rheight = x.borrow().height;
            }
            None => {}
        }

        return (lheight, rheight);
    }
}

impl AVLTree {
    fn height_diff(&self, root_borrow:&RefMut<'_, Node>) -> isize {
        let lr_heights = self.get_lt_rt_heights(&root_borrow);
        return lr_heights.0 - lr_heights.1;
    }
}

impl AVLTree {
    fn set_height(&self, root_borrow:&mut RefMut<'_, Node>) {
        let lr_heights = self.get_lt_rt_heights(&root_borrow);
        root_borrow.height = 1 + max(lr_heights.0, lr_heights.1);
    }
}

Note that instead of passing a mutable reference to the root_node unlike in the Box method, we are passing a mutable reference to the borrowed object.

This is because when you borrow an object mutably, then you cannot borrow that object again mutably or immutably and with Rc, anytime you have to modify the struct, you need to call borrow_mut() on the object.

Thus it is only possible to call borrow_mut() once on the object which in our example will not work.

The other pattern is get the RefMut object by calling borrow_mut() once and pass around the RefMut object.

Similarly, the left rotation and right rotation methods are as follows:

impl AVLTree {
    fn left_rotate(&mut self, root_borrow:&mut RefMut<'_, Node>)->Option>>> {
        let root_clone = root_borrow.clone();
        let lt_node = &root_clone.lt_node;

        match lt_node {
            Some(x) => {
                let lt_rt_node = &x.borrow().rt_node;
                match lt_rt_node {
                    Some(y) => {
                        root_borrow.lt_node = Some(y.clone());
                    }
                    None => {
                        root_borrow.lt_node = None;
                    }
                }
            }
            None => {}
        }

        self.set_height(root_borrow);

        match lt_node {
            Some(ref x) => {
                let xborrow = &mut x.borrow_mut();
                xborrow.rt_node = Some(Rc::new(RefCell::new(root_borrow.clone())));
                self.set_height(xborrow);
            }
            None => {}
        }

        return lt_node.clone();

    }
}

impl AVLTree {
    fn right_rotate(&mut self, root_borrow:&mut RefMut<'_, Node>)->Option>>> {
        let root_clone = root_borrow.clone();
        let rt_node = &root_clone.rt_node;

        match rt_node {
            Some(x) => {
                let rt_lt_node = &x.borrow().lt_node;
                match rt_lt_node {
                    Some(y) => {
                        root_borrow.rt_node = Some(y.clone());
                    }
                    None => {
                        root_borrow.rt_node = None;
                    }
                }
            }
            None => {}
        }

        self.set_height(root_borrow);

        match rt_node {
            Some(ref x) => {
                let xborrow = &mut x.borrow_mut();
                xborrow.lt_node = Some(Rc::new(RefCell::new(root_borrow.clone())));
                self.set_height(xborrow);
            }
            None => {}
        }

        return rt_node.clone();

    }
}

And finally the insert method:

impl AVLTree {
    fn insert(&mut self, root_node:&mut Option>>>, &val:&T)->Option>>> {        
        let mut output:Option>>> = root_node.clone();

        match root_node {
            Some(ref node) => {
                let mut node_borrow = node.borrow_mut();

                if val <= node_borrow.val {
                    let lt_node = &mut node_borrow.lt_node;
                    let new_lt:Option>>>;

                    match lt_node {
                        Some(_x) => {
                            new_lt = self.insert(lt_node, &val);
                        }
                        None => {
                            let new_node:Node = Node::new(&val);
                            new_lt = Some(Rc::new(RefCell::new(new_node)));
                        }
                    }

                    node_borrow.lt_node = new_lt;
                }
                else {
                    let rt_node = &mut node_borrow.rt_node;
                    let new_rt:Option>>>;

                    match rt_node {
                        Some(_x) => {
                            new_rt = self.insert(rt_node, &val);
                        }
                        None => {
                            let new_node:Node = Node::new(&val);
                            new_rt = Some(Rc::new(RefCell::new(new_node)));
                        }
                    }
                    node_borrow.rt_node = new_rt;
                }
            }
            None => {
                let new_node:Node = Node::new(&val);
                output = Some(Rc::new(RefCell::new(new_node)));
            }
        }

        let mut hdiff:isize = 0;

        match root_node {
            Some(ref node) => {
                let mut node_borrow = node.borrow_mut();

                self.set_height(&mut node_borrow);
                hdiff = self.height_diff(&node_borrow);
            }
            None => {}
        }

        match root_node {
            Some(ref node) => {
                let mut node_borrow = node.borrow_mut();
                
                if hdiff > 1 {
                    let lt_node = &mut node_borrow.lt_node;

                    match lt_node {
                        Some(ref x) => {
                            if val <= x.borrow().val {
                                output = self.left_rotate(&mut node_borrow);
                            }
                            else if val > x.borrow().val {
                                let new_lt = self.right_rotate(&mut x.borrow_mut());
                                node_borrow.lt_node = new_lt;
                                output = self.left_rotate(&mut node_borrow);
                            }
                        }
                        None => {}
                    }
                }
        
                else if hdiff < -1 {
                    let rt_node = &mut node_borrow.rt_node;

                    match rt_node {
                        Some(ref x) => {
                            if val <= x.borrow().val {
                                let new_rt = self.left_rotate(&mut x.borrow_mut());
                                node_borrow.rt_node = new_rt;
                                output = self.right_rotate(&mut node_borrow);
                            }
                            else if val > x.borrow().val {
                                output = self.right_rotate(&mut node_borrow);
                            }
                        }
                        None => {}
                    }
                }
            }
            None => {}
        }

        return output.clone();
    }
}

Note that we call borrow_mut() once before using the returned RefMut object in the left and right rotation function calls.

Now lets look at the heap memory usage and how the usage per line looks like using heaptrack.

Peak heap memory usage is somewhere around 64 MB.

The line in the left and right rotation functions, which has the most heap memory usages are:

xborrow.rt_node = Some(Rc::new(RefCell::new(root_borrow.clone())));

xborrow.lt_node = Some(Rc::new(RefCell::new(root_borrow.clone())));

Also the most memory usage lines in the insert function are as follows:

new_lt = Some(Rc::new(RefCell::new(new_node)));

new_rt = Some(Rc::new(RefCell::new(new_node)));

We can clearly see the pattern here. Whenever we are creating a new Rc object in the heap, the memory usage is high. This is because when we are calling root_borrow.clone() in the left and right rotation functions, we are creating copies in heap memory.

Note that calling .clone() on existing Rc objects will not create copies but calling clone on the RefMut object will.

Personally, I feel Box method is much “cleaner” and “safer” way simply because I get to know all the memory issues at compile time instead of getting “panciky” at runtime where the error messages are not always very helpful.

There could be better way to implement these functions which I have not yet explored or not aware of. But the idea I wanted to demonstrate here is that there are multiple different ways to implement a non-linear data structure in Rust each with its own set of challenges. There are lots of struggle involved to get it right.

But the good part is that we can analyze the memory usages with heaptrack tool to understand how each method behaves at runtime and thus one can find ways to optimize them.

References

The time I had mixed feelings about Rust …

Abhijit Mondal — Sat, 10 Feb 2024 13:18:46 GMT

source: quanta magazine

The time I had mixed feelings about Rust …

Due to Rust’s recent rise in popularity among software engineers around the world, I got interested in learning about Rust and how it would help me and my team write better production grade codes.

Note: I come from a background with extensive experience in Python and some bit of C/C++ and Java.

After reading some materials on Rust and its internal workings, here are few points that clearly stood out as to how and why Rust is different.

The most important aspect of Rust is how it manages memory. Instead of deploying a garbage collector at runtime it checks for single ownership and borrowing at compile time.
Unlike Python, Rust is a compiled programming language similar to C/C++.
Its not the responsibility of the OS or the Rust runtime to clean up unused heap memory or optimize memory usage but the onus is on the developer to take all the pain of writing production grade codes.

Let me elaborate a bit on point number 1.

There are 2 ways by which an OS allocates functions and variables in the computer’s memory — Stack and Heap.

CS 225 | Stack and Heap Memory (illinois.edu)

Functions (arguments etc.) and variables with primitive data types are stored on a stack data structure in the memory. The sizes of these objects are predefined prior to program run.

For e.g. in C++, “int” takes up 4 bytes, “long” 8 bytes and so on. The objects in the stack are deallocated using a LIFO strategy.

stack vs. heap allocation

Hence stack deallocation do not require explicit garbage collection or memory management as when the top object in the stack goes out of scope, it is deallocated and the memory is freed.

Objects having dynamic sizes i.e. the size can increase or decrease during program run such as a vector in C++, list in Python etc. are allocated in the Heap instead of Stack

(Reference to the actual object stored in Heap might be stored in the Stack).

Objects in the Heap are not automatically deallocated and require either explicit way such as using “delete” keyword in C/C++ or “del” in Python, or the garbage collector will clean them up after program completion.

Microsoft Word — CSEP521 Project — Garbage Collection Algorithms.doc (washington.edu)

Java garbage collection: What is it and how does it work? | New Relic

Garbage collection algorithms — Holly Cummins

C/C++ do not use garbage collection to clean up the Heap memory. Python and JAVA on the other hand uses garbage collection algorithms to clean up Heap memory.

Thus C/C++ can have memory leakage issues if an object that was allocated in the Heap is not explicitly deallocated.

Let’s think about this program in CPython:

import ctypes

# The list is stored in heap space
# The references x and y will be stored in stack
x = [1,2,3,4]
y = x

# x and y points to the same memory location in the heap
x_id = id(x)
y_id = id(y)

# They will print same id
print(f"x_id={x_id}, y_id={y_id}")

# changing either x or y will be reflected in both
y[2] = 5
x[0] = 7

print(f"x={x}, y={y}")

# get the refernce count of the address in memory holding the list
# It will be 2 (because x and y are both pointing to it)
print(ctypes.c_long.from_address(x_id).value)

# remove reference y
del y

# reference count will now be 1 
# (only x is pointing to the same memory location in heap)
print(ctypes.c_long.from_address(x_id).value)

After the above program ends, the cpython garbage collector will clean up the remaining reference i.e. reference x to the heap memory location and also free the heap memory.

What if we remove both the references x and y in the code:

# remove last reference x
del x

# check the contents of the memory address x_id
print(ctypes.cast(x_id, ctypes.py_object).value)

If the reference count of an object in heap space reaches 0, the object is deallocated and memory is freed.

But there could be certain scenarios where there are cyclic references and in that case reference count will never reach 0. Only option to the free the memory is by using a Garbage Collector.

An example of cyclic reference count in CPython.

class Node:
  def __init__(self):
    self.next = None

x = Node()
y = Node()

x.next = y
y.next = x

x_id = id(x)
y_id = id(y)

print(f"x_id={x_id}, y_id={y_id}")

# reference counts will be 2 for both x and y
print(ctypes.c_long.from_address(x_id).value)
print(ctypes.c_long.from_address(y_id).value)

Some disadvantages with this kind of memory manangement are as follows:

What if someone intended to do a cloning of the list instead of referencing ?
Any changes made to a one reference also updates all the other references.
If there are many objects with cyclic references as seen above, the heap space are not freed at runtime even if some of these objects are not used anymore.
Due to multiple references to the same object in heap space, the objects are not freed until number of references reduces to 0. Let’s take a simple example as below.

def sum_of_max_per_list(input_list:list[list]):
    input_list_id = id(input_list)
    n = len(input_list)
    
    outputs = [0]*n
    for i in range(n):
        outputs[i] = max(input_list[i])

    # reference count is 3
    # one for initization outside this method, one for function argument
    # and one for function stack
    print(ctypes.c_long.from_address(input_list_id).value)

    # some memory intensive workloads
    s = sum(outputs)
    
    return []

input_list = [[10,1,12,4],[5,87,13,65],[10,11,12,13,14],[46,1,32,90,87]]
s = max_per_list(input_list)

Note that after we find the maximum element per list, there is no further requirement of “input_list” variable. If afterwards we are doing some memory intensive task with the output, it could lead to OOM errors in Python because the memory pointed to by input_list is not freed.

A better approach would be:

def max_per_list(input_list:list[list]):
    n = len(input_list)
    
    outputs = [0]*n
    for i in range(n):
        outputs[i] = max(input_list[i])
    
    return outputs

input_list = [[10,1,12,4],[5,87,13,65],[10,11,12,13,14],[46,1,32,90,87]]
input_list_id = id(input_list)

outputs = max_per_list(input_list)

# memory corresponding to input_list is freed
del input_list

print(ctypes.c_long.from_address(input_list_id).value)
print(ctypes.cast(input_list_id, ctypes.py_object).value)

s = sum(outputs)

Something similar happens with C++, but unlike Python, C++ does not have garbage collection. Thus heap space occupied by unwanted objects must be freed manually by the user.

#include 

int main() {
    // pointer x is stored in stack but the actual array is stored in heap
    int* x = new int[10];
    // do some initialization
    for (int i = 0; i < 10; i++) x[i] = 2*i;
    
    // pointer y is stored in stack
    // points to same memory location as x
    int* y = x;
    
    y[2] = 5;
    x[0] = 7;
    
    for (int i = 0; i < 10; i++) std::cout << x[i] << " ";
    std::cout << std::endl;
    
    for (int i = 0; i < 10; i++) std::cout << y[i] << " ";
    std::cout << std::endl;
    
    // deleting y will also delete the pointer to x and the heap space will
    // be freed.
    delete [] y;
    
    return 0;
}

In the above C++ code, x and y are pointers in stack pointing to same object in the heap memory. When the pointer y is deleted, the memory is freed. This is unlike Python, where the reference counter is decremented by 1.

In C/C++ for every malloc we need to call free() and for every new we need to call delete else we would be having memory leaks, since there is no garbage collector to take care of unwanted objects lying around in heap.

In Rust, we would write the sum of max per list program as follows:

use std::isize::MIN;

fn sum_of_max_per_list(input_list:Vec>)->isize {
    let n:usize = input_list.len();
    let mut outputs:Vec = vec![MIN;n];
    
    for (i, _) in input_list.iter().enumerate() {
        for (_j, y) in input_list[i].iter().enumerate() {
            if *y >= outputs[i] {
                outputs[i] = *y;
            }
        }
    }
    
    let mut s:isize = 0;
    for x in outputs.iter() {
        s += *x;
    }
    
    return s;
}
fn main() {
    let input_list:Vec> = vec![
        vec![10, 1, 12, 4],
        vec![5, 87, 13, 65],
        vec![10, 11, 12, 13, 14],
        vec![46, 1, 32, 90, 87],
    ];
    
    let s:isize = sum_of_max_per_list(input_list);
    println!("{:?}", s);
}

There are few things to notice in the above code when we compare the same code in Python.

When we are passing the input_list as function argument, the ownership of the input_list changes and it is moved into the function. The implication of this is that if we wished to use input_list variable after calling the “sum_of_max_per_list” method, we would get a compilation error.

fn main() {
    let input_list:Vec> = vec![
        vec![10, 1, 12, 4],
        vec![5, 87, 13, 65],
        vec![10, 11, 12, 13, 14],
        vec![46, 1, 32, 90, 87],
    ];
    
    let s:isize = sum_of_max_per_list(input_list);
    
    println!("{:?}", s);
    println!("{:?}", input_list);
}

Compilation error:

error[E0382]: borrow of moved value: `input_list`
  --> src/main.rs:33:22
   |
23 |     let input_list:Vec> = vec![
   |         ---------- move occurs because `input_list` has type `Vec>`, which does not implement the `Copy` trait
...
30 |     let s:isize = sum_of_max_per_list(input_list);
   |                                       ---------- value moved here
...
33 |     println!("{:?}", input_list);
   |                      ^^^^^^^^^^ value borrowed here after move
   |
note: consider changing this parameter type in function `sum_of_max_per_list` to borrow instead if owning the value isn't necessary
  --> src/main.rs:3:35
   |
3  | fn sum_of_max_per_list(input_list:Vec>)->isize {
   |    -------------------            ^^^^^^^^^^^^^^^ this parameter takes ownership of the value
   |    |
   |    in this function
   = note: this error originates in the macro `$crate::format_args_nl` which comes from the expansion of the macro `println` (in Nightly builds, run with -Z macro-backtrace for more info)
help: consider cloning the value if the performance cost is acceptable
   |
30 |     let s:isize = sum_of_max_per_list(input_list.clone());
   |                                                 ++++++++

Thus, when we return from the sum_of_max_per_list method, the variable input_list goes out of scope and the heap memory is freed.

In Rust, there can be at-most one owner of any object in heap memory. If the owner goes out of scope, the object is deleted and memory is freed. Thus we do not need additional reference counting or garbage collection.

What if we needed to use the input_list after function call ?

In that case, the function can only borrow the input_list object but the owner is still the input_list variable.

The updated Rust code with borrowing:

fn sum_of_max_per_list(input_list:&Vec>)->isize {
    let n:usize = input_list.len();
    let mut outputs:Vec = vec![MIN;n];
    
    for (i, _) in input_list.iter().enumerate() {
        for (_j, y) in input_list[i].iter().enumerate() {
            if *y >= outputs[i] {
                outputs[i] = *y;
            }
        }
    }
    
    let mut s:isize = 0;
    for x in outputs.iter() {
        s += *x;
    }
    
    return s;
}
fn main() {
    let input_list:Vec> = vec![
        vec![10, 1, 12, 4],
        vec![5, 87, 13, 65],
        vec![10, 11, 12, 13, 14],
        vec![46, 1, 32, 90, 87],
    ];
    
    let s:isize = sum_of_max_per_list(&input_list);
    
    println!("{:?}", s);
    println!("{:?}", input_list);
}

When the function is borrowing the object, by default it is an immutable borrow i.e. the input_list cannot be changed inside the function. If we needed to update the input_list, we would have to pass a mutable reference

use std::isize::MIN;

fn sum_of_max_per_list(input_list:&mut Vec>)->isize {
    let n:usize = input_list.len();
    let mut outputs:Vec = vec![MIN;n];
    
    for (i, x) in input_list.iter_mut().enumerate() {
        for (_j, y) in x.iter_mut().enumerate() {
            if *y >= outputs[i] {
                outputs[i] = *y;
            }
            *y = -1;
        }
    }
    
    let mut s:isize = 0;
    for x in outputs.iter() {
        s += *x;
    }
    
    return s;
}
fn main() {
    let mut input_list:Vec> = vec![
        vec![10, 1, 12, 4],
        vec![5, 87, 13, 65],
        vec![10, 11, 12, 13, 14],
        vec![46, 1, 32, 90, 87],
    ];
    
    let s:isize = sum_of_max_per_list(&mut input_list);
    
    println!("{:?}", s);
    println!("{:?}", input_list);
}

But even if we are passing an object as a mutable reference to a function, we cannot move out i.e. change ownership of the object within the function scope. The below will throw compilation error.

fn sum_of_max_per_list(input_list:&mut Vec>)->isize {
    let n:usize = input_list.len();
    let mut outputs:Vec = vec![MIN;n];

    // will throw compilation error here
    let obj = *input_list;
    
    for (i, x) in input_list.iter_mut().enumerate() {
        for (_j, y) in x.iter_mut().enumerate() {
            if *y >= outputs[i] {
                outputs[i] = *y;
            }
            *y = -1;
        }
    }
    
    let mut s:isize = 0;
    for x in outputs.iter() {
        s += *x;
    }
    
    return s;
}

Compilation error:

error[E0507]: cannot move out of `*input_list` which is behind a mutable reference
 --> src/main.rs:7:15
  |
7 |     let obj = *input_list;
  |               ^^^^^^^^^^^ move occurs because `*input_list` has type `Vec>`, which does not implement the `Copy` trait
  |
help: consider removing the dereference here
  |
7 -     let obj = *input_list;
7 +     let obj = input_list;
  |

If we have to assign the value of the object to another variable in the above code, we need to clone or copy the original object (duplicating in the heap memory).

Here is how to do it:

let obj = input_list.clone();

Also be careful of the fact that for same object we cannot mix and match concurrent mutable and immutable references.

fn sum_of_max_per_list(input_list:&mut Vec>)->isize {
    let n:usize = input_list.len();
    let mut outputs:Vec = vec![MIN;n];
    
    // borrow immutably
    let first_element:&isize = &input_list[0][0];
    
    // borrow mutably
    for (i, x) in input_list.iter_mut().enumerate() {
        for (_j, y) in x.iter_mut().enumerate() {
            if *y >= outputs[i] {
                outputs[i] = *y;
            }
            *y = -1;
        }
    }
    
    let mut s:isize = 0;
    for x in outputs.iter() {
        s += *x;
    }
    
    // immutable borrow used here
    println!("{:?}", *first_element);
    
    return s;
}

So what we are doing above is, first we borrow input_list as immutable i.e. this value should not be changed to get the 1st element at index (0,0).

Then we borrow mutably, implying that the 1st element might be changed during this process but later on we are using the immutable borrow (to print the 1st element) hence compiler will throw an error.

error[E0502]: cannot borrow `*input_list` as mutable because it is also borrowed as immutable
  --> src/main.rs:9:19
   |
7  |     let first_element:&isize = &input_list[0][0];
   |                                 ---------- immutable borrow occurs here
8  |     
9  |     for (i, x) in input_list.iter_mut().enumerate() {
   |                   ^^^^^^^^^^ mutable borrow occurs here
...
23 |     println!("{:?}", *first_element);
   |                      -------------- immutable borrow later used here

If we did not use the immutable borrow later on (line no. 23) we will not get the error.

These are some of the common things that a Rust developer must be aware of:

Understanding when an object is moved. If an object is moved i.e. ownership changes, then any statement using the object later on (by non-owner) will cause compilation error.
In case an object is borrowed by a function, one cannot move out or change ownership within the function scope. Although can clone or copy the object.
In case an object is borrowed immutably by a function, one cannot change the value of the object within the function scope. Need to pass the reference as mutable.
If an object is borrowed mutably then we cannot simulataneously borrow the object as immutable or another mutably. This follows from the fact that if it is intended to update an object then until and unless the updates are completed, one cannot make another update simultaneously. Also if it is intended to read the object, then one cannot simulataneously update the object till the read is completed.

If you think carefully all of the above makes sense if we are not using reference counting or garbage collection for heap memory cleanup.

Now let’s look at a complete Rust example to implement a simple Binary Search Tree.

Here is the BST setup:

use std::fmt::Debug;
pub trait MyTrait: PartialOrd + Debug + Copy {}

struct BST where T: MyTrait {
    root_node: Option>>,
}

impl BST {
    fn new() -> Self {
        Self {
            root_node: None,
        }
    }
}

#[derive(Clone)]
struct Node where T: MyTrait {
    val: T,
    lt_node: Option>>,
    rt_node: Option>>,
}

impl Node {
    fn new(&val:&T) -> Self {
        Self {
            val,
            lt_node: None,
            rt_node: None,
        }
    }
}

Since we are going to implement a BST for a generic type i.e. the values stored in BST nodes is “generic” i.e. could be integer, float, string, other complex types etc. But we will limit the types to integer, float and string.

To define a generic type and its properties, we create a new trait MyTrait.

The properties of MyTrait are PartialOrd i.e. we can compare two elements of type T using less than or greater than operators. This is required since during BST insertion, we will compare an incoming value with the value in current node in BST using < or > inequalities and depending on that we will go to either left or right subtree.

Note: In order to use ≤ and ≥ instead of < or >, we have to use the property Ord instead of PartialOrd.

Debug property allows us to print the BST node values and Copy property allows us to copy or clone the node values. For e.g. when we delete a node, we copy the value of its successor node into itself.

“struct” in Rust is similar to C/C++ struct. The method “new” is like the class constructor.

One interesting thing to note is that when we define the Node struct, in C/C++, we define the left and the right nodes as pointers to the same Node struct and similarly in Python:

struct Node {
    int val;
    Node* lt_node;
    Node* rt_node;
};

But in Rust, we have something like:

struct Node {
    val: i32,
    lt_node: Option>,
    rt_node: Option>,
}

When Rust creates the struct in the heap memory, it has to know how much memory it has to allocate in the heap. If we just we use Option, then it becomes recursive type and it becomes impossible to infer the size in heap memory to allocate for Node object.

If we use something similar to C/C++ i.e.

struct Node {
    val: i32,
    lt_node: Option,
    rt_node: Option,
}

Then we get compilation error:

error[E0072]: recursive type `Node` has infinite size
 --> src/lib.rs:1:1
  |
1 | struct Node {
  | ^^^^^^^^^^^
2 |     val: i32,
3 |     lt_node: Option,
  |                     ---- recursive without indirection
  |
help: insert some indirection (e.g., a `Box`, `Rc`, or `&`) to break the cycle
  |
3 |     lt_node: Option>,
  |                     ++++    +

Box is a smart pointer in Rust. It can be used to break the infinite cycle.

The size of a Node object is the sum of the i32 value, and the sizes of the two Box pointers corresponding to lt_node and rt_node. I will not go into details of smart pointers in this post.

Next we define the insert method for BST struct as follows:

impl BST {
    fn insert(&mut self, root_node:&mut Option>>, &val:&T) {
        let new_node = Node::new(&val);

        match root_node {
            Some(node) => {
                // root_node is not None
                if val <= node.val {
                    if node.lt_node.is_none() {
                        node.lt_node = Some(Box::new(new_node));
                    }
                    else {
                        self.insert(&mut node.lt_node, &val);
                    }
                }
                else {
                    if node.rt_node.is_none() {
                        node.rt_node = Some(Box::new(new_node));
                    }
                    else {
                        self.insert(&mut node.rt_node, &val);
                    }
                }
            }
            None => {
                // root node is None
                *root_node = Some(Box::new(new_node));
            }
        }
    }
}

Check if the root_node is not None, if root_node is None, create a new root node and assign to root node. If root_node is not None, we add the new node by calling the insert method recursively.

Note that we are passing the root_node as a mutable reference.

It needs to be mutable because the tree is updated by adding a new node at the leaf. The mutability is propagated upwards, towards the root node.

Now let’s look at the delete method for BST:

impl BST {
    // delete a node with value val
    fn delete(&mut self, root_node:&mut Option>>, &val:&T) {  
        match root_node {
            Some(node) => {
                // if root_node is not None
                if val == node.val {
                    // found the node to delete
                    if (node.lt_node.is_none()) && (node.rt_node.is_none()) {
                        *root_node = None;
                    }
                    else if node.lt_node.is_none() {
                        match &node.rt_node {
                            Some(x) => {
                                node.val = x.val;
                                node.rt_node = None;
                            }
                            None => {}
                        }
                    }
                    else if node.rt_node.is_none() {
                        match &node.lt_node {
                            Some(x) => {
                                node.val = x.val;
                                node.lt_node = None;
                            }
                            None => {}
                        }
                    }
                    else {
                        let mut new_node = &node.rt_node;
                        let mut has_successor:bool = false;

                        // find successor node
                        while !new_node.is_none() {
                            match new_node {
                                Some(x) => {
                                    has_successor = true;
                                    node.val = x.val;
                                    new_node = &x.lt_node;
                                }
                                None => {}
                            }
                        }

                        if has_successor {
                            self.delete(&mut node.rt_node, &node.val);
                        }
                    }
                }

                else if val < node.val {
                    // go left to delete
                    self.delete(&mut node.lt_node, &val);
                }

                else {
                    // go right to delete
                    self.delete(&mut node.rt_node, &val);
                }
            }
            None => {}
        }
    }
}

The deletion algorithm follows standard deletion in BST. If the node to be deleted is a leaf node, we simply delete it. Else if the node to be deleted has either a left or right child, we replace the node with its left or right child.

If node to be deleted has both children, then we find the successor node of this node. Copy the node value of successor to itself and delete the successor node. Note that the successor node is either a leaf node or has only right child.

The final piece of the puzzle is defining the main function and how we create the BST and call the insert and delete methods on it.

impl BST {
    // print tree as shown below
    //  10
    //  -- 5
    //  ---- 2
    //  ---- 9
    //  -- 20
    //  ---- 15
    //  ---- 30
    fn print_tree(&self, root_node:&Option>>, level:usize) {  
        match root_node {
            Some(node) => {
                println!("{} {:?}", "-".repeat(2*level), node.val);
                self.print_tree(&node.lt_node, level+1);
                self.print_tree(&node.rt_node, level+1);
            }
            None => {}
        }
    }
}

fn main() {
    // BST with values as unsigned integers
    let mut bst:BST = BST::new();
    let mut root:Option>> = None;

    bst.insert(&mut root, &10);
    bst.insert(&mut root, &20);
    bst.insert(&mut root, &5);
    bst.insert(&mut root, &9);
    bst.insert(&mut root, &2);
    bst.insert(&mut root, &15);
    bst.insert(&mut root, &30);

    bst.print_tree(&root, 0);

    println!();
    bst.delete(&mut root, &10);
    bst.print_tree(&root, 0);
}

Although it looks like a clean implementation and it is infact true that if there are no compilation errors in Rust, the code is correct, but it takes a significant amount of practice to write compilation error free codes.

Most of the compiler errors stems from the earlier discussion points.

After implementing few algorithms and data structures in Rust, I believe that Rust codes are truly production grade and as a tech lead, I would want myself and my team members to use Rust extensively in all data intensive production workloads.

References

HyperLogLog, Probabilities and Dynamic Programming …

Abhijit Mondal — Tue, 02 Jan 2024 13:07:46 GMT

source: quanta magazine

HyperLogLog, Probabilities and Dynamic Programming …

In this post we will be discussing the HyperLogLog algorithm for cardinality estimation.

Given a stream of input data points for e.g. the input could be IP addresses, user ids on a social networking site, geolocations such as latitude and longitude pairs for GPS enabled device and so on.

In order to find the number of unique visitors to a website in the last 30 days assuming that there are 1 million requests per second, there will be 86.4 billion user ids (non-unique) logged into the system each day.

Multiply that by 30 days to get 2.6 trillion user ids in a month.

Assuming that each user id is of size 32 bytes, we will need 76 TB of data points to analyze and find the number of unique user ids from this data.

Ofcourse, one can store all the 76TB on disk and scan the data to find the number of unique user ids but each such disk scan is an expensive operation. Based on the latency numbers chart, it takes somewhere around 10-20ms to read 1 MB of data from disk, thus reading 76 TB it will take somewhere around 15-20 hours to read all of that data.

Latency Numbers Every Programmer Should Know (github.com)

If we store all that data in memory, we need to have at-least 4864 RAMs each of 16GB and since memory is expensive than storing on disk, this option seems infeasible at the moment.

Alternatively, we can use a ‘set’ data structure to store the stream of user ids. The set could be implemented using either a hash table or a balanced binary search tree such as AVL Tree or Skip List.

The space complexity is O(n) where n is the number of unique user ids. If n ~ O(10⁹), we would still requires somewhere around 30GB of RAM for storing the set.

Can we live with some error in the counts of the number of unique visitors but use somewhere around 1 KB memory only?

This can be a desired use-case when the exact count do not matter but approximate is enough such as counting the number of unique visitors to a website or the number of virtual machines created on a cloud platform etc.

The idea behind HyperLogLog is pretty intuitive.

Assuming we have a ‘good’ hash function that generates hashes with an uniform distribution in the range [0, 2³² -1].

If the maximum number of trailing zeros in the binary representation of the hash is K, then the approximate cardinality is 2^K.

For e.g. if we use the Murmur3 family of hash functions, then the 32 bit hash of the string “Hello world” is:

hsh = mmh3.hash("Hello world", 42, signed=False) & ((1 << 32)-1)
print(bin(hsh))

0b10010000010010111101110011111100

We can see that the number of trailing zeros is 2.

Let X be the event that we saw K trailing zeros in the binary representation of the hash.

The number of possible hashes with trailing K bits as 0 is 2^(N-K) because we can place 0 or 1 in each of the remaining N-K places.

Thus the probability that a hash ends with K zeros is :

To calculate the expected number of hashes before we see K trailing zeros, let the expected number of hashes be E. Then

If we see an event with K trailing zeros, then we add the term 1*p, else if we do not see this event, then we repeat again, thus if we do not see K trailing zeros, then the number of trials is 1+E.

Solving for E, we get:

Thus the expected number of hashes after which we see K trailing zeros is 2^K.

The standard deviation is given by:

Note that the above analysis do not assume that the maximum number of trailing zeros seen is K, it just assumes that we see an event with K trailing zeros.

In HLL, instead of considering the event that we see a hash with K trailing zeros we consider the event that the maximum number of trailing zeros seen is K.

The benefit of the later strategy in streaming data is that the maximum number of trailing zeros is a non-decreasing function which correlates with the fact that cardinality is always non-decreasing.

The former strategy is not useful because if at time T we see a hash with 10 trailing zeros and at time T+1, a hash with 5 trailing zeros then at time T+1, the expected cardinality is lower than at time T which contradicts with our intuition.

Let’s calculate the probability that the maximum number of trailing zeros seen so far is K. Let the number of inputs seen so far be n.

Let r = (2^N), N is the total number of bits.

Let P be the number of ways by which we can have n N-bit numbers such that the number of trailing zeros ranges from 0 to K for all n numbers.

Similarly, let Q be the number of ways by which we can have n N-bit numbers such that the number of trailing zeros ranges from 0 to K-1 for all n numbers.

Then the number of ways by which we can have n N-bit numbers such that at-least one of n has K trailing zeros = P-Q

Probability that the maximum number of trailing zeros seen is K:

The expected cardinality in HLL with a single hash function is given to be :

The variance of the cardinality is given as:

The standard deviation and standard error is given as:

For N=32 and n = 2³², the standard error is high as 51610.

One can experimentally verify this is true by running the following Python simulation:

import random, mmh3, math, numpy as np, string

def get_trailing_zeros(n, N):
    if n == 0:
        j = N
    else:
        j = 0
        while n > 0:
            if (n & 1) == 1:
                break

            n = n >> 1
            j += 1

    return j


N = 32
n = 1 << N
a = []

for j in range(1000):
    g = -1
    for i in range(n):
        data = ''.join(\
                random.choices(string.ascii_lowercase + string.digits, 
                                   k=random.randint(1, 20)))
        
        # generate hash between 0 and 2^N-1
        hsh = mmh3.hash(data, 42, signed=False) & ((1<
        # number of trailing zeros in binary of hsh
        b = get_trailing_zeros(hsh, N)
        g = max(g, b)

    # calculate cardinality
    a += [1<
# standard deviation
print(np.std(a))

Assuming that we have seen n inputs so far in the data stream and the leftmost position of the first one in the binary representations of the hashes is K i.e. the maximum number of trailing zeros is K-1.

......100
.......10
....10000
........1
..1000000

Given this data, can we estimate the expected cardinality of the stream ?

Let X be the total number of ways of obtaining a maximum of K-1 trailing zeros with n input hashes. Then X is calculated as:

Using the same reasoning as seen earlier.

Let a be the number of ways of obtaining K-1 trailing zeros and 1 at the K-th place in a N-bit hash .

Then:

Let b be the number of ways of obtaining less than K-1 trailing zeros in an N-bit hash.

Then:

To generate a sequence of n input hashes such that at-least one of the input hashes have K-1 trailing zeros and the remaining have less than K-1 trailing zeros.

Let the cardinality with n hashes be h i.e. out of n hashes seen, h are unique hashes.

I can choose j hashes from ‘a’ hashes (‘a’ is defined above) and h-j hashes from ‘b’ hashes varying j from 1 to h. After filling the h hashes out of n hashes, I can fill the remaining n-h hashes using repetition from the h unique hash values.

Let H(n, j, h-j) be the number of ways of arranging n input hashes with h unique hashes and j hashes from set ‘a’ i.e. equal to K-1 leading zeros and h-j hashes from set ‘b’ i.e. less than K-1 leading zeros.

Then we can use the following recursion relation to solve for H:

The 1st component refers to the case where we have formed n-1 hashes with j hashes from set ‘a’ and h-j from set ‘b’. For the remaining n-th hash, I can pick any one of the h hashes selected so far.

The 2nd component refers to the case where we have formed n-1 hashes with j-1 hashes from set ‘a’ and h-j from set ‘b’. For the remaining n-th hash, I have to pick one of the a-(j-1) hashes from set ‘a’ because I need to have j hashes from set ‘a’.

The 3rd component refers to the case where we have formed n-1 hashes with j hashes from set ‘a’ and h-j-1 from set ‘b’. For the remaining n-th hash, I have to pick one of the b-(h-j-1) hashes from set ‘b’ because I need to have h-j hashes from set ‘b’.

Here is the Python code to find H using memoization technique:

def H(n, x, y, dp):
    if x < 0 or y < 0:
        dp[(n, x, y)] = 0
        return dp[(n, x, y)]
        
    if n < x+y:
        dp[(n, x, y)] = 0
        return dp[(n, x, y)]

    if n == 1:
        if x+y == 1:
            if x == 1:
                dp[(n, x, y)] = a
            else:
                dp[(n, x, y)] = b
        else:
            dp[(n, x, y)] = 0
            
        return dp[(n, x, y)]

    u = H(n-1, x, y, dp) if (n-1, x, y) not in dp \
                else dp[(n-1, x, y)]

    v = H(n-1, x-1, y, dp) if (n-1, x-1, y) not in dp \
                else dp[(n-1, x-1, y)]

    w = H(n-1, x, y-1, dp) if (n-1, x, y-1) not in dp \
                else dp[(n-1, x, y-1)]

    dp[(n, x, y)] = u*(x+y) + (a-x+1)*v + (b-y+1)*w
        
    return dp[(n, x, y)]

Compute H for all different cardinalities h ranging from 1 to n.

To find the expected cardinality, multiply each cardinality h with H and divide by X (calculated above) and take their sum.

n = 300
N = 10
k = 7

q = (1 << N)
x = int((q*(1-(1/2)**k))**n - (q*(1-(1/2)**(k-1)))**n)

a = (1 << (N-k))
b = int(q*(1-(1/2)**(k-1)))

dp = {}
            
e = 0
for h in range(1, n+1):
    y = 0
    for j in range(1, h+1):
        y += H(n, j, h-j, dp) if (n, j, h-j) not in dp else dp[(n, j, h-j)]

    e += h*(y/x)

print(e)

The above code is for finding the expected cardinality when we have seen 300 input data points each of 10 bits and the maximum number of trailing zeros seen is 7.

The expected cardinality is somehwere around 260.

We can verify this is true by running a simulation.

import random, mmh3, math, string

v, c = 0, 0
n = 300
N = 10
k = 7

for j in range(10000):
    nums = random.choices(range(10**9), k=n)
    g = -1
    d = set()
    for i in range(len(nums)):
        data = ''.join(\
                random.choices(string.ascii_lowercase + string.digits, 
                                   k=random.randint(1, 20)))

        hsh = mmh3.hash(data, 42, signed=False) & ((1<        b = get_trailing_zeros(hsh, N)
        g = max(g, b)
        d.add(hsh)

    if g == k:
        v += len(d)
        c += 1

print(v/c)

We will see that the expected cardinality is also somewhere around 260.

This is one way to find the cardinality without having to store all the input data but the issue is that :

We are using O(n³) memory to store the cached values dp.
Running time complexity is O(n³).

Both these drawbacks suggest that we cannot use this analysis to calculate the expected cardinality.

Coming back to our HLL algorithm, instead of using a single hash function, what if we use multiple hash functions. Will that bring the standard error down?

Using m hash functions, let K1, K2, … Km be the maximum number of trailing zeros corresponding to each hash function.

Then the estimated cardinalities from each hash are respectively 2^K1, 2^K2 … and so on. To get the final cardinality we can take the geometric mean of the individual cardinalities i.e.

The expected cardinality is the sum over all possible combinations of K1, K2 … upto Km given to be:

Number of possible combinations of K1, K2, … Km is:

If we run the below Python code, we can see that the standard deviation decreases on increasing m i.e. the number of hash functions.

import math

# Function to generate all possible combinations of k1, k2, ..., km
def hn(N, m, curr, out):
    if m == 0:
        out += [curr]

    else:
        for k in range(N+1):
            hn(N, m-1, curr + [k], out)

N = 10
n = 1 << N

for m in range(1, 10):
    s1, s2 = 0, 0
    
    out = []
    hn(N, m, [], out)
    
    # For each combination of k1, k2, ..., km
    for x in out:
        p = 1

        # Mean of the k1, k2, ..., km
        q = 2**(1.0 + sum(x)/m)
        
        for y in x:
            p *= ((1-(1/2)**(y+1))**n - (1-(1/2)**y)**n)
    
        s1 += p*q**2
        s2 += p*q

    print(m)

    # print standard error
    print(math.sqrt(s1-s2**2)/math.sqrt(n))

The plot of the standard error vs. number of hash functions m, looks like one below:

standard error vs. number of hash functions (m)

Due to the combinatorial nature of the number of hash functions on the different combinations Ki, calculating for higher values of m becomes exponentially difficult, hence we have shown till m=6.

Here is a python code to demonstrate an estimation of cardinality using 1000 independent hash functions:

import random, mmh3, math, numpy as np, string

N = 32
m = 1000
k = 5123

# Generate m seeds for murmur3 hash corresponds to m hash functions
seeds = [random.randint(0, 10**5) for _ in range(m)]

# Generate k random strings
data = list(set([''.join(\
            random.choices(string.ascii_lowercase + string.digits, 
              k=random.randint(1, 20))) for _ in range(k)]))

# True cardinality
print(len(data))

g = [-1]*m

for i in range(10*k):
    h = random.sample(data, 1)[0]
    
    for j in range(m):
        # generate hash between 0 and 2^N-1
        hsh = mmh3.hash(h, seeds[j], signed=False) & ((1<
        # number of leading zeros in binary of hsh
        b = get_trailing_zeros(hsh, N)
        g[j] = max(g[j], b)

# Estimated cardinality
w = sum(g)/m
a = 0.79402*2**w

print(a)

In the above code, we are multiplying estimated cardinality by 0.79402. This constant is found out in the paper by the original authors.

One of the drawbacks with the above approach is that we have to use m different hash functions and if m is large i.e. m > 1000, the time complexity is quite high i.e. O(m*w) for each input string of length w present in the stream.

Another approach to simulate m different hash functions, is that we take the first p bits from the hash as a bucket. Thus with p bits, we will have 2^p number of buckets. Estimate the cardinality using the remaining N-p bits.

Similar to m hash functions, we take the geometric mean of the cardinality in each bucket but the final cardinality has to be multiplied by m.

Because the input stream is divide into m sub-streams unlike m hash functions which are applied to all inputs, here each bucket corresponds to approximately n/m inputs.

The updated code for using 1st p bits as a bucket identifier:

import random, mmh3, math, numpy as np, string

N = 32
p = 10
m = 1 << p
k = 51236

# Generate k random strings
data = list(set([''.join(\
            random.choices(string.ascii_lowercase + string.digits, 
              k=random.randint(1, 20))) for _ in range(k)]))

# True cardinality
print(len(data))

g = [-1]*m

for i in range(10*k):
    h = random.sample(data, 1)[0]
    hsh = mmh3.hash(h, 42, signed=False) & ((1<
    # First p bits is for bucket index r
    # Remaining N-p bits are for finding the cardinality
    u = N-p

    r = hsh >> u
    q = hsh & ((1<
    # number of leading zeros in binary of quotient q above
    b = get_trailing_zeros(q, u)
    g[r] = max(g[r], b)

# Estimated cardinality
w = sum(g)/m
a = 0.79402*m*2**w

print(a)

For each input in the stream, the time complexity is now reduced to O(w) from O(m*w).

The updated expected cardinality when we are using m buckets instead of m hashes can be written as follows:

Note that now we are using the additional terms because out of n inputs, n1 can go into 1st bucket, n2 into 2nd bucket and so on. The maximum number of trailing zeros in bucket j is Kj (same as before).

The code for computing standard errors is updated as follows:

# Compute nCr values for r from 0 to n using dynamic programming
def ncr(n):
    cache = [[0]*(n+1) for i in range(n+1)]
    
    for i in range(n+1):
        for j in range(i+1):
            if j == 0 or i == j:
                cache[i][j] = 1
            else:
                cache[i][j] = cache[i-1][j] + cache[i-1][j-1]
    
    return cache

# Compute the term inside the 2nd summation in the expectation formula
# dp = memoization variable
def gn(arr, i, n, m, u, dp):
    if m == 0:
        return 0 if n > 0 else 1
        
    s = 0
    for j in range(n+1):
        p = (1-(1/2)**(arr[i]+1))**j - (1-(1/2)**arr[i])**j

        a = gn(arr, i+1, n-j, m-1, u, dp) if (n-j, m-1) not in dp \
            else dp[(n-j, m-1)]

        s += p*u[n][j]*a

    dp[(n, m)] = s
    return s


N = 5
n = 1 << N
u = ncr(n)

for m in range(1, 10):
    # Get all combinations of k1, k2, ... km similar to previous
    out = []
    hn(N, m, [], out)
    
    s1, s2 = 0, 0
    
    # Compute the 1st summation in the above expectation formula
    for x in out:
        p = gn(x, 0, n, m, u, {})/m**n
        q = 2**(1.0 + sum(x)/m)
        
        s1 += p*q**2
        s2 += p*q
    
    print(m)
    print(math.sqrt(s1-s2**2)/math.sqrt(n))

Due to the combinatorial nature of the expectation, it is computationally expensive to compute the standard errors for values of N beyond 6 and m beyond 5.

The final version of HLL has few other modifications such as :

Using harmonic mean instead of geometric mean
Handling cases for small and large cardinalities separately.

The final code looks like as follows:

import mmh3, random, math
from scipy.stats import hmean
import numpy as np

class HyperLogLog:
    def __init__(self, p=10):
        self.N = 32
        self.p = p
        self.m = 1<        self.buckets = [-1]*self.m
        
        self.alpha = 0.7213/(1+1.079/self.m)

        if self.m == 16:
            self.alpha = 0.673
        elif self.m == 32:
            self.alpha = 0.697
        elif self.m == 64:
            self.alpha = 0.709
    
    def add(self, data):
        h = mmh3.hash(data, 42, signed=False) & ((1<
        # Exclude the 1st p bits in h
        v = self.N-self.p
        
        # Quotient using remaining 22 bits
        q = h & ((1<
        # Remainder is first p=10 bits
        r = h >> v

        # Number of trailing zeros of quotient q
        j = get_trailing_zeros(q, v)

        self.buckets[r] = max(self.buckets[r], j+1)
    
    def get_size(self):
        u = 1 << self.N
        v = [2**x for x in self.buckets]
        n = self.alpha*self.m*hmean(v)
        
        # Case for small cardinalities
        if n <= 2.5*self.m:
            z = sum([self.buckets[i] == 0 for i in range(self.m)])
            if z != 0:
                return self.m*math.log(self.m/z)
        
        # Case for very large cardinalities
        elif n > (1/30)*u:
            return -u*math.log(1-(n/u))
        
        return n

There are 1024 buckets and in each bucket we are storing the maximum number of trailing zeros K which ranges from 0 to 22 (10 bits out of 32 bits for bucket). To store K for range of values from 0 to 22, we need 5 bits.

Thus each of the 1024 buckets is storing 5 bits, thus total = 1024*5 bits = 0.625 KB. The error rate from the above code varies around 1–3%.

Compare this to using 30GB of memory for accurate cardinality estimation.

Using harmonic mean, the standard error values are lower as compared to using geometric mean.

standard error with harmonic mean

Read more about the theory behind the hyperloglog and error bounds in the original paper here.

HyperLogLog, Probabilities and Dynamic Programming … was originally published in Gaussian Machine on Medium, where people are continuing the conversation by highlighting and responding to this story.

Demystifying CPU Caches with Examples

Abhijit Mondal — Wed, 22 Nov 2023 08:44:01 GMT

source: quanta magazine

This post takes inspiration from one of the famous blogs on the different CPU cache effects:

Gallery of Processor Cache Effects (igoro.com)

Some notable factors that can affect the performance of a computer program utilizing the CPU cache are as follows:

Impact of Cache Lines
Cache Associativity
False Cache Line Sharing

Before getting into each of the above cache performance parameters, let us go over some important pointers regarding CPU caches.

Reading from and writing to main memory (RAM) is expensive as compared to CPU caches.
According to the CS61 course website, the access times for RAM (main memory) is 62.9 ns whereas it is 12.8 ns for L3 cache, 3.3 ns for L2 cache and 1.1 ns for L1 cache.

During each read query, it is first checked in CPU’s L1 cache, if it is found then returned, else check L2 cache and so on till L3 cache or L4 cache (if present).
If not found in any CPU caches, fetch from main memory and add to L1, L2 and L3 cache.
When a byte is read from main memory, not only the byte is returned but a block of bytes is returned known as the cache line.
Generally the cache line is 64 bytes. Any read for a byte in the same cache line is read from either L1, L2 or L3 cache because the cache line of 64 bytes is already cached in the CPU.
For e.g. if the cache lines are aligned s.t. 0–63 bytes is one cache line, 64–127 is next one and so on. Then if byte 10 is read, all bytes from 0 to 63 is transferred from RAM to CPU and cached. Next if request for byte 57 is made, it is already present in CPU cache since 0 <= 57 <= 63, hence read from cache.
A cache line read from main memory can occupy one of N cache slots i.e. if an existing cache line already occupies one of the N slots, then the incoming cache line can only occupy one of remaining N-1 slots.
If all the N slots are occupied then we have to evict one of the N slots to make room for the new cache line. LRU policy is used to evict cache slots.

Cache replacement policies — wikipedia

Finding L1, L2 and L3 cache sizes

In Linux, we can issue the following commands to find the L1, L2 and L3 cache sizes:

getconf -a | grep CACHE

The output:

LEVEL1_ICACHE_SIZE                 32768
LEVEL1_ICACHE_ASSOC
LEVEL1_ICACHE_LINESIZE             64
LEVEL1_DCACHE_SIZE                 49152
LEVEL1_DCACHE_ASSOC                12
LEVEL1_DCACHE_LINESIZE             64
LEVEL2_CACHE_SIZE                  1310720
LEVEL2_CACHE_ASSOC                 20
LEVEL2_CACHE_LINESIZE              64
LEVEL3_CACHE_SIZE                  12582912
LEVEL3_CACHE_ASSOC                 12
LEVEL3_CACHE_LINESIZE              64
LEVEL4_CACHE_SIZE                  0
LEVEL4_CACHE_ASSOC
LEVEL4_CACHE_LINESIZE

Size of L1 instruction cache (cache for storing the program) = 32768 bytes or 32 KB.

Size of L1 data cache (cache for storing the data) = 49152 bytes or 48 KB.

Note that the L1 cache exists in each core of the CPU i.e. if a CPU is 4-core, then each core has 32 KB of L1 instruction cache and 48KB of L1 data cache. Total size of L1 data cache is 48*4 KB = 192 KB.

Also L2 cache is present in each core of the CPU. Each core of CPU has 1.25 MB of cache, thus total size of L2 cache for a 4-core CPU is 1.25*4 = 5 MB.

But L3 cache is shared across all cores of the CPU. Total size of L3 cache is 12 MB.

Each byte present in L3 cache is also present in L2 and L1 cache.

Cache line size i.e. number of bytes prefetched for each byte read from main memory is 64 bytes.

L1 data cache associativity is 12 i.e. each cache line can occupy one of 12 slots in the L1 cache. Similarly for L2 cache, each cache line can occupy one of 20 slots and for L3 cache can occupy one of 12 cache slots.

Impact of Cache Lines

Assuming you have an integer array of size N, in the first case read M integers contiguously starting at index 0 to index M-1.

Since the cache line size is 64 bytes or 16 integers (each integer is 4 bytes), reading M integers contiguously requires accessing approximately M/16 cache lines.

More importantly each cache line is accessed from main memory at-most once.

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

    const int length = 512 * 1024 * 1024;
    const int cache_line_size = 16;
    const int m = length/cache_line_size;

    int *arr = new int[length];

    // contiguous access
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < m; i++) 
        arr[i]++;
    auto stop = std::chrono::high_resolution_clock::now();

    auto duration = 
    std::chrono::duration_cast(stop - start);
    
    std::cout << (float) duration.count() << std::endl;

    free(arr);
    return 0;
}

Compiling the above program with the default g++ compiler (v11.3.0) options:

g++ cache_line.cpp -o cache_line
./cache_line

Time taken in this case is 296 ms.

But if M integers are read with a jump of step=16 i.e. 0, 16, 32, … 16*(M-1), then although the number of elements scanned is still M but the number of cache lines accessed is 16*M/16 = M.

Thus in this case we are accessing 16 times more cache lines.

int *arr = new int[length];
    
auto start = std::chrono::high_resolution_clock::now();
for (int i = 0; i < m*cache_line_size; i+=cache_line_size) 
    arr[i]++;
auto stop = std::chrono::high_resolution_clock::now();

auto duration = 
std::chrono::duration_cast(stop - start);

std::cout << (float) duration.count() << std::endl;

Time taken in this stepped access case is 1871 ms.

This clearly shows that although we are iterating over same number of indices M, but the amount of work done in fetching the additional cache lines in the 2nd case is significantly high as compared to the 1st case.

In the 1st case when we fetch index 0, all indices from 0 to 15 are also fetched and cached in L1, L2 and L3 cache, hence updating the indices from 1 to 15 is done within the CPU cache which is not the case when we are doing stepped access.

In a 3rd scenario, we are iterating with a step size of 32, but we are updating integers at index i and i+16 i.e. we are still accessing same number of cache lines.

int *arr = new int[length];
    
auto start = std::chrono::high_resolution_clock::now();
for (int i = 0; i < m*cache_line_size; i+=2*cache_line_size) {
    arr[i]++;
    arr[i+cache_line_size]++;
}
auto stop = std::chrono::high_resolution_clock::now();

auto duration = 
std::chrono::duration_cast(stop - start);

std::cout << (float) duration.count() << std::endl;

Time taken in this case is around 1500 ms.

The purpose of the 3rd scenario is to show that although we are iterating over half the number of indexes in the array, but the running time should remain almost same since we are accessing the same number of cache lines.

Key Takeaway : Performance is dependent not only on how many elements we are looping over but also how many cache lines we are accessing from main memory.

False Cache Line Sharing

Let’s say there is an array of size N and using 4 threads A, B, C and D we repeatedly update the integers at indices i, j, k and l.

Now if i, j, k and l belong to the same cache line for e.g. i=0, j=1, k=2 and l=3, then whenever a thread updates an integer, the entire cache line is invalidated and the entire cache line is again fetched from main memory.

void update(int *inp, int j) {
    for (int i = 0; i < 1000000; i++) inp[j]++;
}

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

    const int length = 1024;
    int *arr = new int[length];

    long long b = 1000;
    float a = 0.0;

    // updates happen in different cache lines
    for (int r = 0; r < b; r++) {
        auto start = std::chrono::high_resolution_clock::now();

        std::thread thread1(update, arr, 0);
        std::thread thread2(update, arr, 1);
        std::thread thread3(update, arr, 2);
        std::thread thread4(update, arr, 3);

        thread1.join();
        thread2.join();
        thread3.join();
        thread4.join();

        auto stop = std::chrono::high_resolution_clock::now();

        auto duration = 
        std::chrono::duration_cast
                    (stop - start);

        a += (float) duration.count();
    }
        
    std::cout << float(a)/float(b) << std::endl << std::endl;

    return 0;
}

Average time taken when updates takes place in same cache line is around 9.0 ms.

Each thread runs in a different CPU core, and as we saw earlier that each core has its own L1 and L2 cache.

The advantage of having per core L1 and L2 cache is that when working in a non-threaded environment, reads and writes are very fast as the core where the program runs can directly use its own L1 or L2 cache.

But disadvantage is that with multiple threads, where each core updates its own L1 or L2 cache independently, the values needs to be in sync across all cores. Syncing takes time.

For e.g. if thread A updates the integer at index i=0, then the thread B cannot update the integer at index j=1 in its own L1/L2 cache concurrently because the entire cache line is invalidated and until the update is synced across all cores, the next update cannot happen in the same cache line.

But if i, j, k and l belong to different cache lines for e.g. i=0, j=16, k=32 and l=48, then an update made by a thread does not invalidate the other cache lines and hence threads updating the other cache lines will update the integers in its own core’s cache concurrently without having to sync across all cores.

float a = 0.0;
int *arr = new int[length];

// updates happen in same cache line
for (int r = 0; r < b; r++) {
    auto start = std::chrono::high_resolution_clock::now();

    std::thread thread1(update, arr, 0);
    std::thread thread2(update, arr, 16);
    std::thread thread3(update, arr, 32);
    std::thread thread4(update, arr, 48);

    thread1.join();
    thread2.join();
    thread3.join();
    thread4.join();

    auto stop = std::chrono::high_resolution_clock::now();

    auto duration = 
    std::chrono::duration_cast
              (stop - start);

    a += (float) duration.count();
}
    
std::cout << float(a)/float(b) << std::endl << std::endl;

Average time taken when updates takes place in different cache lines is around 3.3 ms.

Thus it takes almost 3 times more to update integers in the same cache line as compared to different cache lines when using multiple threads.

Key Takeaway : Care must be taken while using multi threading with CPU bound tasks such as above as the additional costs due to context switching, creation and deletion of threads and false cache line sharing can significantly reduce the gains.

Cache Associativity

If we look at only the highest cache i.e. L3 cache in my system, it has 12 MB cache size and 12-way associativity. It implies that each cache line can be placed in one of 12 slots in the L3 cache.

Number of cache lines that can be placed in L3 cache = 12 MB/64 bytes = 196608.

This is the total number of slots available in the L3 cache.

Since L3 is 12-ways associative, number of different 12-slots in the L3 cache = 196608/12 = 16384.

For brevity, let us call a “12-slot” as a “group”.

But how to determine which cache lines goes into which group ?

Since 0 to 16383 can be represented using 14 bits, thus if we index each cache line as 0, 1, 2 … and so on then all cache lines whose lower 14 bits are the same, can be put in the same group in L3 cache. If there are more than 12 cache lines with 14 lower bits same then there will be evictions in L3 cache.

Thus all integers separated by 16384*16 = 262144 positions in the array would occupy the same group in L3 cache.

If there are more than 262144*12 = 3145728 integers present in the array, there will be evictions and if we try to read the evicted values, they will have to be fetched again from main memory.

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

    const int length = 512*1024*1024;
    const int m = 262144;
    int *arr = new int[length];

    auto start = std::chrono::high_resolution_clock::now();
    for (int r = 0; r < 10000; r++) {
        for (int i = 0; i < length; i+=m) 
            arr[i]++;
    }
    auto stop = std::chrono::high_resolution_clock::now();

    auto duration = 
    std::chrono::duration_cast(stop - start);
    
    std::cout << (float) duration.count() << std::endl << std::endl;
    
    free(arr);
    return 0;
}

If we run the above program with values of m ranging from 262100 to 262200, we can see that the time taken for m=262144 is almost 3x of that for any other value of m in that range.

This is because for all m except for m=262144 gets distributed across multiple groups whereas for m=262144, all integers maps to the same group and hence the number of evictions in that group is very high as the number of slots is only 12 per group.

In-fact there are 512*1024*1024/262144 = 2048 different elements that maps to the same group (cache line index have same lower 14 bits).

For values of m that are multiples of 262144, they will similarly be placed in the same group.

For values of m lower than 262144 but powers of 2 such as 131072, 65536, 16384 … and so on, the integers will be distributed across multiple groups.

For e.g. with m=65536, the integers have the lower 12 bits same in cache line index and thus they will be distributed across 4 different groups (14–12=2 bits), similarly for m=16384, have lower 11 bits same in cache line index and hence distributed across 8 different groups (14–11=3 bits).

Although the integers are distributed across groups but the number of cache lines mapping to the same group remains the same since we are accessing more number of cache lines with m=256 as compared to m=262144.

Thus for values of m which are powers of 2, there is an upward spike in running time as compared to its neighboring values.

Key Takeaway : When accessing elements in an array in a stepped fashion, try not to use a step size that is a power of 2.

References

Fast Fourier Transform Optimizations

Abhijit Mondal — Wed, 08 Nov 2023 16:18:10 GMT

source: quanta magazine

In the last two posts we understood the basics of fourier transform, how to speed up the DFT calculations with FFT, then looked at how we can use the FFT algorithm to solve different problems such as time series analysis, pattern matching in strings, convolution etc.

Part 1

An excursion into Fast Fourier Transform — The Basics and Time Series Analysis | by Abhijit Mondal | Nov, 2023 | Medium | Medium

Part 2

An excursion into Fast Fourier Transform — Pattern Matching and Other Problems | by Abhijit Mondal | Nov, 2023 | Medium | Medium

In this post we are going to look at ways by which we can further optimize and speed up the FFT algorithm. There could be possible hundreds of ways by which we can gain speed-up but not all strategies can give significant improvements. Here we will look at techniques which can give significant improvements (at-least 1.5x to 2x speed-ups).

Bit Reversal

Recall that in the FFT algorithm, for a given input sequence, we extract the elements at the even indexes and the odd indexes, do FFT on them separately and merge them. This is done recursively.

I will post the original implementation here again:

def fft(x):
    n = len(x)
    
    if n == 1:
        return x

    # recursively do FFT on even indexes and odd indexes
    even_fft = fft(x[::2])
    odd_fft = fft(x[1::2])

    # multiplicative factor
    w = np.exp(-2j*np.pi/n)
    m = int(n/2)

    h = 1
    result = [0]*2*m
    
    # merge the even and odd ffts
    for i in range(m):
        result[i] = even_fft[i] + h*odd_fft[i]
        result[i+m] = even_fft[i] - h*odd_fft[i]
        h *= w
        
    return result

Although the time complexity is O(NlogN) for an input sequence of length N, but as you know that recursion has overhead. Each recursion call, adds a new entry to the function call stack and at the end we also have to roll up the function call stack.

Also, for each recursion call, we are passing 2 new arrays of size N/2 each. If we draw the recursion tree, we can see that at each level we are allocating O(N) memory and there are O(logN) levels, thus additional memory requirement is O(NlogN).

Can we implement the FFT algorithm without recursion?

One way we can replace recursion is with stacks. This is similar to doing Depth First Traversal on a tree using stacks.

But in here too, we have to store the intermediate results of each FFT operation and also the maximum size of the stack is O(logN). Total additional memory requirement is still O(NlogN).

Instead of a top down approach, we aim to do a bottom-up approach which in this case is possible because we can reproduce the sequence at the leaf nodes of the recursion tree without going thorugh all paths from the root.

Final sequence with N=8

Input: 0 1 2 3 4 5 6 7

          [0,1,2,3,4,5,6,7]
                 /\
                /  \
               /    \
              /      \
    [0,2,4,6]          [1,3,5,7]
       /\                 /\
      /  \               /  \
[0,4]     [2,6]     [1,5]    [3,7] 
 /\        /\        /\        /\
0  4      2  6      1  5      3  7


Final sequence : 0 4 2 6 1 5 3 7

For a given N=8, the input sequence is [0,1,2,3,4,5,6,7]

The final sequence which we evaluate at the last level of the recursion tree is [0,4,2,6,1,5,3,7]

The interesting property of the output sequence is that the binary representation of each element of the output seqence is reverse of the corresponding element in the input sequence.

For e.g. the element at index 1 in input is 1 with binary 001 (3 bits), whereas the element at index 1 in output is 4 with binary 100 which is reverse of 001.

Observe that the bits are reversed for all the elements.

Input         Output

0   000       0   000    
1   001       4   100
2   010       2   010
3   011       6   110
4   100       1   001
5   101       5   101
6   110       3   011
7   111       7   111

Thus we can obtain the output sequence directly from the input sequence without going through all the paths in the tree.

Then we can just run a bottom-up algorithm to calculate the FFT.

Here is the python function for FFT updated to take as input, permutation of the input x according to the final sequence:

def fft_bit_reversal(x, seq):
    n = len(seq)

    # seq : permutation of 0 to n-1 using bit reversal technique
    x = [x[i] for i in seq]

    k = 2
    
    # at each level, the even FFT lies from index i+j to i+j+u-1
    # and odd FFT lies from index i+j+u to i+j+k-1
    # merge the 2 FFTs into a single FFT from i+j to i+j+k-1
    while k <= n:
        w = np.exp(-2j*np.pi/k)
        u = int(k/2)
        
        for i in range(0, n, k):
            h = 1
            for j in range(u):
                a, b = x[i+j], x[i+j+u]
                
                x[i+j] = a + h*b
                x[i+j+u] = a - h*b
                h *= w
                
        k *= 2

    return x

Note that the time complexity of the algorithm is still O(NlogN).

But we have not yet shown the code to calculate the final sequence using bit reversal. I wanted to deliberately keep that part outside of the fft code because the sequence permutation using bit reversal is independent of the input sequence x and is only dependent on N i.e. the size of data.

For many applications, the size N is almost fixed for many inputs like in time series analysis, the number of time intervals remains constant. But for dynamic N, we can pre-compute a lookup table of final sequences.

For e.g.

lookup_table[8] = [0, 4, 2, 6, 1, 5, 3, 7]
lookup_table[16] = [0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15]

and so on.

Now we come to calculation of the lookup tables using bit reversal.

There are multiple ways we can do the calculations. In python, we can simply calculate the binary representation of integer using bin() and then reverse it and recalculate new integer using eval().

def get_bit_reversed_seq(n, m):
    # m is the number of bits
    seq = []
    
    for x in range(n):
        y = bin(x)[2:]
        y = '0'*(m-len(y)) + y
        y_rev = y[::-1]
        seq += [eval('0b' + y_rev)]

    return seq

In the above ‘m’ is the number of bits for e.g. if n=16, then m=4 because all integers 0 to 15 can be represented using 4 bits: 0000 to 1111

Time complexity of the above algorithm is O(NlogN) because there are N inputs and each input has O(logN) bits. But additional space complexity is O(N). Can we come up with an in-place algorithm?

To reverse a sequence, one such algorithm is to swap the adjacent elements first, then in pairs of 2, then in 4, then 8 and so on. For e.g. given a sequence x0, x1, … x15, this is how we are gonna reverse it.


x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 x10 x11 x12 x13 x14 x15

Swap adjacent pairs i.e. xi and xi+1 where i % 2 = 0

x1 x0 x3 x2 x5 x4 x7 x6 x9 x8 x11 x10 x13 x12 x15 x14

Then swap in batches of 2: x1 x0 with x3 x2 and so on.

x3 x2 x1 x0 x7 x6 x5 x4 x11 x10 x9 x8 x15 x14 x13 x12

Then swap in batches of 4: x1 x0 x3 x2 with x5 x4 x7 x6 and so on.

x7 x6 x5 x4 x3 x2 x1 x0 x15 x14 x13 x12 x11 x10 x9 x8

Then swap in batches of 8:

x15 x14 x13 x12 x11 x10 x9 x8 x7 x6 x5 x4 x3 x2 x1 x0

Now we have reversed the sequence

But note that we are dealing with integer sequences and so we have to first convert them to binary representation, do bit reversal in place and then calculate new integer representation all in-place.

Instead we can use bit manipulation to achieve the same without having to convert integer to binary and then binary to integer.

The idea goes something like this (assuming that the number of bits is power of 2):

Input x
Let us say that number of bits in x is 4

x = (x & 0101) << 1 | (x & 1010) >> 1
x = (x & 0011) << 2 | (x & 1100) >> 2

The final value of x will have its bits reversed.

In the 1st step we are doing bitwise AND operation to extract the even and odd bits. Then left shift the odd bits and right shift the even bits by 1 position and then do bitwise OR. This step is similar to swapping the adjacent elements in a sequence.

In the next step, similarly we do bitwise AND but this time with pairs of 2 bits at a time. Similarly left shift and right shift by 2 bits and do bitwise OR. This step is similar to swapping pairs of 2 elements in the sequence.

The above approach is not generic for any number of bits and is shown only for 4 bits. Also it is only applicable for number of bits which is a power of 2. i.e. 1, 2, 4, 8 … and so on.

Assuming that we need at-most 32 bit sequence, we can pre-compute the sequences quite easily. For number of bits which are not a power of 2, add 0s to the right of binary representation of x to make number of bits equal to the next higher power of 2.

Here is the full python code to pre-compute the lookup tables:

def generate_swapping_constants(n):
    # This function returns the constants used with bitwise AND
    # i.e. for extracting alternate elements, then pairs of 2, then
    # batches of 4 and so on.
    swap_constants = []
    
    i = 0
    while (1 << i) <= int(n/2):
        nbits = (1 << i)
        
        p = '0'*nbits + '1'*nbits
        q = '1'*nbits + '0'*nbits

        h = len(p)

        d, r = int(n/h), n % h
        
        a = '0b' + p*d + p[:r]
        b = '0b' + q*d + q[:r]

        swap_constants += [(nbits, eval(a), eval(b))]
        i += 1

    return swap_constants

def get_bit_reversed_seq(n, m, swap_constants, nxt_power):
    seq = []
    
    for x in range(n):
        x *= (1 << (nxt_power-m))
        
        for q, a, b in swap_constants:
            x = ((x & a) << q) | ((x & b) >> q)
            
        seq += [x]

    return seq


max_bits = 32
swap_table = []

nbits = 1
while nbits <= max_bits:
    swap_constants = generate_swapping_constants(nbits)
    swap_table += [(nbits, swap_constants)]
    nbits *= 2

sequences = {}
i = 0
for m in range(1, 25):
    n = 1 << m

    while i < len(swap_table) and m > swap_table[i][0]:
        i += 1

    sequences[m] = \
        get_bit_reversed_seq(n, m, 
            swap_table[i][1], swap_table[i][0])

The ‘sequences’ is python dict with key equal to the number of bits=m and value equal to the bit reversed sequence for N=2^m.

Thus for an input sequence x, we can call the fft_bit_reversal() method as follows:

def num_bits(n):
    c = 0
    while n > 0:
        n = n >> 1
        c += 1

    return c

def fft_optimized(x):
    n = len(x)

    # n is power of 2, but sequence is from 0 to n-1
    # number of bits required is 1 minus the number of bits in n
    m = num_bits(n)-1

    return fft_bit_reversal(x, sequences[m])

Doing a benchmark on the unoptimized FFT vs. the optimized FFT, with input sizes ranging from N=2⁸ to 2²⁰, we found that the optimized FFT is at-least 100% faster than the unoptimized FFT i.e. takes half the time to run.

Next we look into another optimization.

Bailey’s FFT algorithm

Till now we have been looking at Cooley-Tukey algorithm for FFT. In Bailey’s algorithm, we can optimize the calculations for parallel and distributed systems.

For this algorithm, assume that the size N is a perfect square power of 2 e.g. 4, 16, 64 etc.

1. Layout the N elements in the input sequence in a 2D matrix of 
   rows=sqrt(N) and cols=sqrt(N).

2. Transpose the matrix.

3. Perform FFT indepndently on each row of the matrix. This step can be
   distributed or done in parallel threads/processes.

4. Multiply each element at index (p, q) from step 3 with e^(-2*pi*i*p*q/N)

5. Again transpose the matrix from step 4.

6. Perform FFT indepndently on each row of the matrix. This step can be
   distributed or done in parallel threads/processes.

7. Transpose the matrix from step 6.

The final matrix from step 7 is laid out in row-major order and represents the result of FFT on whole input sequence.

Here is a python implementation for Bailey’s algorithm:

def transpose(inp, n):
    # transpose matrix in place
    for i in range(n):
        for j in range(i+1, n):
            temp = inp[j*n+i]
            inp[j*n+i] = inp[i*n+j]
            inp[i*n+j] = temp


def fft_bit_reversal_bailey(x, seq):
    # seq : permutation of 0 to sqrt(n)-1 calculated using bit reversal
    n = len(x)
    n_r = int(math.sqrt(n))

    transpose(x, n_r)

    # calculate fft on each of n_r rows using 4 processes.
    # due to python's GIL, using multiple threads will be of no use here.
    futures = []
    
    with ProcessPoolExecutor(max_workers=4) as executor:
        for i in range(n_r):
            start, end = i*n_r, (i+1)*n_r
            futures += [executor.submit(fft_bit_reversal2, 
                                            x[start:end], seq, i)]

        for future in as_completed(futures):
            res, i = future.result()
            start, end = i*n_r, (i+1)*n_r
            x[start:end] = res

    # multiply each element at (p, q) by factor np.exp(-2j*np.pi*p*q/n)
    w = np.exp(-2j*np.pi/n)
    h1 = 1
    
    for i in range(n_r):
        h2 = 1
        for j in range(n_r):
            x[i*n_r + j] *= h2
            h2 *= h1
        h1 *= w

    transpose(x, n_r)

    # calculate fft on each of n_r rows using 4 processes.
    # due to python's GIL, using multiple threads will be of no use here.
    futures = []
    
    with ProcessPoolExecutor(max_workers=4) as executor:
        for i in range(n_r):
            start, end = i*n_r, (i+1)*n_r
            futures += [executor.submit(fft_bit_reversal2, 
                                        x[start:end], seq, i)]

        for future in as_completed(futures):
            res, i = future.result()
            start, end = i*n_r, (i+1)*n_r
            x[start:end] = res

    # final transpose
    transpose(x, n_r)
    
    return x

To perform FFT on each row of the matrix, we are calling a modified version of the ‘fft_bit_reversal’ routine because we are tracking for which row index the FFT is done which is used to collect the futures from concurrent.futures ProcessPoolExecutor module.

def fft_bit_reversal2(x, seq, index):
    n = len(x)
    x = [x[i] for i in seq]
            
    k = 2
    
    while k <= n:
        w = np.exp(-2j*np.pi/k)
        u = int(k/2)
        
        for i in range(0, n, k):
            h = 1
            for j in range(u):
                a, b = x[i+j], x[i+j+u]
                
                x[i+j] = a + h*b
                x[i+j+u] = a - h*b
                h *= w
                
        k *= 2

    return x, index

Time complexity of Bailey’s algorithm is still O(NlogN).

Another important optimization that Bailey’s algorithm does is that since the size of each FFT is sqrt(N) instead of N, thus if sqrt(N) fits entirely in the CPU cache, then number of I/O operations using RAM will be greatly reduced and thus we can gain significant speed-ups.

The performance of the Bailey’s algorithm as compared to standard FFT algorithm (bit reversal) is not significant due to multiple factors:

We are using multi-processing to run the per row FFT in parallel. Forking and cleaning up multiple processes is expensive. We cannot use multi-threading as due to Python’s GIL there will be no observable improvements for CPU bound tasks.
There are multiple transpose operations involved.
Matrix transpose operations for large matrices are generally not cache efficient as there will be large number of cache misses.

The multi-process Bailey’s version starts to outperform bit reversal algorithm for input sizes of greater than 2²⁰.

In order to improve the cache performance of matrix transpose operation, we have seen that the following C++ version of a cache-oblivious matrix transpose generally performs 4–5x better than standard matrix transpose for input sizes beyond 1024 by 1024.

// standard matrix transpose
void transpose(int *inp_arr, int *out_arr, int n, int m) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            out_arr[j*n+i] = inp_arr[i*m+j];
        }
    }
}

// cache oblivious matrix transpose
void transpose_opt(int *inp_arr, int *out_arr, 
        int i_start, int j_start, int i_end, int j_end, int n, int m) {

    int n1 = i_end - i_start + 1;
    int m1 = j_end - j_start + 1;

    // small enough matrix size fits into 64 bytes cache line
    if (n1*m1 <= 16) {
        for (int i = i_start; i <= i_end; i++) {
            for (int j = j_start; j <= j_end; j++) {
                out_arr[j*n+i] = inp_arr[i*m+j];
            }
        }
    }
    else {
        // number of rows greater than number of columns then
        // split horizontally.
        if (n1 >= m1) {
            int mid = (i_start+i_end)/2;

            transpose_opt(inp_arr, out_arr, 
                          i_start, j_start, 
                          mid, j_end, 
                          n, m);

            transpose_opt(inp_arr, out_arr, 
                          mid+1, j_start, 
                          i_end, j_end, 
                          n, m);

        }

        // number of columns greater than number of rows then
        // split vertically.
        else {
            int mid = (j_start+j_end)/2;

            transpose_opt(inp_arr, out_arr, 
                          i_start, j_start, 
                          i_end, mid, 
                          n, m);

            transpose_opt(inp_arr, out_arr, 
                          i_start, mid+1, 
                          i_end, j_end, 
                          n, m);
        }
    }
}

Rewriting the FFT algorithm in C++ leveraging our cache oblivious matrix transpose operations:

// in-place bit reversal FFT in C++
void fft_bit_reversal(std::complex *inp, 
                std::vector permutation, int n) {

    // permute the input sequence according to bit reversal
    for (int i = 0; i < n; i++) {
        if (i < seq[i]) {
            std::complex temp = inp[i];
            inp[i] = inp[permutation[i]];
            inp[permutation[i]] = temp;
        }
    }

    int k = 2;
    const std::complex i_complex(0.0,1.0);

    while (k <= n) {
        double k_inv = 1.0/float(k);
        int u = k/2;

        std::complex w = std::exp(-2.0*i_complex*M_PI*k_inv);

        for (int i = 0; i < n; i += k) {
            std::complex h(1.0, 0.0);

            for (int j = 0; j < u; j++) {
                std::complex a = inp[i+j];
                std::complex b = inp[i+j+u];

                inp[i+j] = a + h*b;
                inp[i+j+u] = a - h*b;

                h *= w;
            }
        }

        k *= 2;
    }
}

// bailey's algorithm
void fft_bailey(std::complex *inp, 
                  std::vector seq, int n) {

    int n_r = int(sqrt(n));

    double n_inv = 1.0/float(n);
    const std::complex i_complex(0.0,1.0);
    std::complex w = std::exp(-2.0*i_complex*M_PI*n_inv);

    // 1st transpose
    std::complex *out = new std::complex[n];
    transpose_opt(inp, out, 0, 0, n_r-1, n_r-1, n_r, n_r);

    // do per row FFT using std::threads
    std::vector my_threads;

    for (int i = 0; i < n_r; i++) {
        std::thread newThread(fft_bit_reversal, out+i*n_r, seq, n_r);
        my_threads.push_back(move(newThread));
    }

    for (int i = 0; i < my_threads.size(); i++) {
        my_threads[i].join();
    }

    // multiply each element by constant
    std::complex h1(1.0, 0.0);
    for (int i = 0; i < n_r; i++) {
        std::complex h2(1.0, 0.0);
        for (int j = 0; j < n_r; j++) {
            out[i*n_r + j] *= h2;
            h2 *= h1;
        }
        h1 *= w;
    }

    // 2nd transpose
    std::complex *out2 = new std::complex[n];
    transpose_opt(out, out2, 0, 0, n_r-1, n_r-1, n_r, n_r);

    // do per row FFT using std::threads
    my_threads.clear();

    for (int i = 0; i < n_r; i++) {
        std::thread newThread(fft_bit_reversal, out2+i*n_r, seq, n_r);
        my_threads.push_back(move(newThread));
    }

    for (int i = 0; i < my_threads.size(); i++) {
        my_threads[i].join();
    }

    // final transpose
    transpose_opt(out2, inp, 0, 0, n_r-1, n_r-1, n_r, n_r);
}

Some cool references

Fast Fourier Transform Optimizations was originally published in Gaussian Machine on Medium, where people are continuing the conversation by highlighting and responding to this story.