Multiplying two- or three-dimensional arrays with broadcasting

Background

Disclaimer: Skip if you are not interested in the background. The code is far far simpler than this. The problem is more about programming than math. Here is the definition of multiplication with broadcasting in human readable language.

I am trying to optimize the multiplication of large n-dimensional arrays $A$ and $B$ with shape $(I_0,I_1,I_2, ..., I_{n-1})$ and $(J_0,J_1,J_2,..., J_{n-1})$ with broadcasting, as in matlab and numpy. $A$ and $B$ are stored in row-major order in memory.

The two arrays can be multiplied with broadcasting if for each $i$, at least one of the following three statements is true: Case S1: $I_i == J_i$, Case S2: $I_{i} == 1$, Case S3: $J_{i} == 1$.

From the shape, we can define stride. Stride is the numbers of linear index required to step through to take one step along each dimension of the n-dimensional array.

Linear index is the index of the 1D array in memory that stores the n-dimensional array.

The stride for $A$, which is $(P_1, P_2, P_3,...,P_n)$ has the following formula: $P_i = I_{i+1}I_{i+2}I_{i+3}...I_{n}$ except $P_{n} = 1$. Similarly strides for $B$ is $(Q_1, Q_2, Q_3, ... Q_n)$.

The general form of the nested loop is:

int cal_size(int* shape, int n){
int size = 1;
for(int i = 0; i < n; ++i) size*=shape[i];
}
int* cal_stride(int* shape, int n){
int size = cal_size(shape, n);
int* stride = new int[n];
stride[0] = size/shape[0];
for(int i = 0; i < n; ++i){
stride[i] = stride[i-1]/shape[i];
}
}

int n;
//the number of dimensions (given)
int I[] = {I0,I1,I2,...,I_last};
//shape of n-dimensional array A (given)
int J[] = {J0,J1,J2,...,J_last};
//shape of n-dimensional array B (given)
int A[cal_size(I, n)];
//1D array in memory that stores
//the n-dimensional array A in row major order
int B[cal_size(J, n)];
//1D array in memory that stores
//the n-dimensional array B in row major order

int* P = cal_stride(I, n);
int* Q = cal_stride(J, n);
int i[n];
int j[n];
//The nested loop
int Ai, Bi;
for(int i[0] = 0; i[0] < I[0]; ++i[0]){         //Case S1
for(int i[1] = 0; i[1] < J[1]; ++i[1]){     //Case S2
for(int i[2] = 0; i[2] < I[2]; ++i[2]){ //Case S3
...
Ai =  i[0]*P[0] //for case S1 or S3
+i[1]*0  //for case S2
+i[2]*P[2] //for case S1 or S3
...
+i[n-1]*P[n-1]; //for case S1 or S3

//The above line converts the multi-index
//(i[0],i[1],i[2],...,i[n-1]) over the n-dimensional
//array to the index of the 1D array
//in memory.

Bi =  i[0]*Q[0] //for case S1 or S2
+i[1]*Q[1] //for case S2 or S2
+i[2]*0  //for case S3
...
+i[n-1]*Q[n-1]; //for case S1 or S2

A[Ai] *= B[Bi];
}}}}}...}


(i think i can deal with this general part with meta-programming. For now, I am just concerned with optimization.)

Question

Is it possible to rewrite the following loop to give better performance?

Can memoization help?

I am not interested in unrolling because the actual array has 10,000+ elements.

#include <iostream>
using namespace std;
//version 1
int main(void){
int A[24];
int B[12];
int iA,iB;
//version 1.
//(only one version is used. but i show all versions i can think of.)
for(int i = 0; i < 2; ++i){
for(int j = 0; j < 3; ++j){
for(int k = 0; k < 4; ++k){
iA = i*12 + j *4 + k*1;
iB = 4*j + k*1;
A[iA]*=B[iB];
cout << iA <<"," <<iB <<endl;
}
}
}
//version 2.
for(int i = 0; i < 2; ++i){
for(int j = 0; j < 12; ++j){
iA = i*12 + j;
iB = j;
A[iA]*=B[iB];
cout << iA <<"," <<iB <<endl;
}
}

//version 3
for(int iA = 0; iA < 2*3*4; iA+=12){
for(int iB = 0; iB < 3*4; ++iA, ++iB){
A[iA]*=B[iB];
cout << iA <<"," <<iB <<endl;
}
}

}


Assembly instruction generated by gcc compiler:

gcc godbolt of version 1

gcc godbolt of version 3

(version 1 and 3 give same assembly code after fixing the bug)

Wait a minute... in the actual case, the shapes are not known during compliation time. Is version 1 and 3 still the same in this case?

• Can you please reformulate your question in a mathematically exact way. Why can A and B have different dimensionality with different number of elements. Do you only multiply the elements or do you want a full matrix multiplication? – miscco Sep 23 '16 at 20:02
• It is simpler than full matrix multiplication. I seriously doubt most programmer would want the formal mathematical definition. But it is here. Hopefully this won't make everyone face palm and start to throw error. – rxu Sep 23 '16 at 20:48
• hopefully this can work with openmp without too much trouble. – rxu Sep 24 '16 at 0:32

Bug

Version 3 does half as many multiplications than the other versions. The problem is in the loops:

  for(int iA = 0; iA < 2*3*4; iA+=12){
for(int iB = 0; iB < 3*4; ++iA, ++iB){


As you can see, you have both iA+=12 in the outer loop and ++iA in the inner loop. For this version, if you are incrementing iA in the inner loop, you shouldn't also increment it in the outer loop. The correct way would be:

  for(int iA = 0; iA < 2*3*4; ){
for(int iB = 0; iB < 3*4; ++iA, ++iB){