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:
(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?
