How to compile properly for performance tuning

Question

I have a question about how to compile functions properly for performance tuning.

Currently, I am trying to solve an eigenvalue problem of a large matrix. Initially I just used the conventional method to define a function using square brackets (ex. f[x_]=x). Now, I want to solve eigenvalues of much larger matrix, and it turned out my code is too slow to solve this problem. I was trying to optimize my code and I got advice that Compile is what I am supposed to look up.

Now, this is what I have been trying to do:

First, this is the base code:

msa = 10^3;
msb = 10^6;
mu0 = 10^-6;
h0 = 1/mu0;
la = 10^-16;
lb = 10^-16;

aa = la*mu0*(msa^2)/2;
ab = lb*mu0*(msb^2)/2;
alpa = 0.2;
alpb = 0.1;
f = 0.1;
a = 10*10^-8; 
rcyl = Sqrt[(a^2)*f/Pi];


mmax = 49;
kmax = 3;

gx[n_] := 2*n*Pi/a; 
gy[m_] := 2*m*Pi/a; 
gxy = Table[{gx[n], gy[m]}, {n, -kmax, kmax}, {m, -kmax, kmax}];
g = ArrayReshape[gxy, {mmax, 2}];


msg = ParallelTable[
   If[(g[[i, 1]] - g[[j, 1]]) != 0 || (g[[i, 2]] - g[[j, 2]]) != 
      0, (msa - msb)*(2 f)*
     BesselJ[1, (Sqrt[((g[[i, 1]] - g[[j, 1]])*(g[[i, 1]] - 
               g[[j, 1]])) + ((g[[i, 2]] - g[[j, 2]])*(g[[i, 2]] - 
               g[[j, 2]]))]*
         rcyl)]/(Sqrt[((g[[i, 1]] - g[[j, 1]])*(g[[i, 1]] - 
              g[[j, 1]])) + ((g[[i, 2]] - g[[j, 2]])*(g[[i, 2]] - 
              g[[j, 2]]))]*rcyl), (msa*f) + (msb*(1 - f))], {i, 1, 
    mmax}, {j, 1, mmax}]; 
q = ParallelTable[
   If[(g[[i, 1]] - g[[j, 1]]) != 0 || (g[[i, 2]] - g[[j, 2]]) != 
      0, ((2*aa/(h0*mu0*msa^2)) - (2*ab/(h0*mu0*msb^2)))*(2 f)*
     BesselJ[1, (Sqrt[((g[[i, 1]] - g[[j, 1]])*(g[[i, 1]] - 
               g[[j, 1]])) + ((g[[i, 2]] - g[[j, 2]])*(g[[i, 2]] - 
               g[[j, 2]]))]*
         rcyl)]/(Sqrt[((g[[i, 1]] - g[[j, 1]])*(g[[i, 1]] - 
              g[[j, 1]])) + ((g[[i, 2]] - g[[j, 2]])*(g[[i, 2]] - 
              g[[j, 2]]))]*rcyl), ((2*aa*f/(h0*mu0*msa^2)) + (2*
        ab*(1 - f)/(h0*mu0*msb^2)))], {i, 1, mmax}, {j, 1, mmax}];

Next, I want to define a function called sumfunc as following:

sumfunc[i_, j_, kx_, ky_] := 
 Sum[msg[[i, l]]*
   q[[l, j]]*(({kx, ky} + g[[j]]).({kx, ky} + g[[l]]) - (g[[i]] - 
        g[[j]]).(g[[i]] - g[[l]])), {l, mmax}]

By using compile, this is what I got:

sumcompile = Compile[{{i, _Integer}, {j, _Integer}, {kx, _Real}, {ky, _Real}}, Sum[
   msg[[i, l]]*
    q[[l, j]]*(({kx, ky} + g[[j]]).({kx, ky} + g[[l]]) - (g[[i]] - 
         g[[j]]).(g[[i]] - g[[l]])), {l, 49}]]

I tested speed of both methods and this is the result:

sumcompile[1, 1, 1, 1] // AbsoluteTiming
{0.0154722, 6.39621}

for compiled function and

sumfunc[1, 1, 1, 1] // AbsoluteTiming
{0.00106203, 6.39621}

for the original function.

In addition, I also tried to define a matrix function as following:

bxycompile = Compile[{{kx, _Real}, {ky, _Real}}, 
  Table[If[(g[[j]] + {kx, ky}) != {0, 0}, 
    KroneckerDelta[i, j] + 
     msg[[i, j]]*((g[[j, 2]] + ky)*(g[[j, 2]] + 
           ky))/(h0*((g[[j]] + {kx, ky}).(g[[j]] + {kx, ky}))) + 
     sumcompile[i, j, kx, ky], 
    KroneckerDelta[i, j] + msg[[i, j]]/(2*h0) + 
     sum[i, j, kx, ky]], {i, 1, 49}, {j, 1, 49}], 
  Parallelization -> True]

for the compiled matrix function and

bxytest[kx_, ky_] := 
  ParallelTable[If[(g[[j]] + {kx, ky}) != {0, 0}, 
    KroneckerDelta[i, j] + 
     msg[[i, j]]*((g[[j, 2]] + ky)*(g[[j, 2]] + 
           ky))/(h0*((g[[j]] + {kx, ky}).(g[[j]] + {kx, ky}))) + 
     sum[i, j, kx, ky], 
    KroneckerDelta[i, j] + msg[[i, j]]/(2*h0) + 
     sumfunc[i, j, kx, ky]], {i, 1, mmax}, {j, 1, mmax}];

for the original matrix function

For bxycompile, it took about 38 secs while bxytest took only 10 secs.

This is not what I expected, since speed of the conventional method is much faster than Compile. Definitely, I am missing something here, but I am not sure where I am supposed start from.

I would like to get some advice how to compile these functions properly to make the code faster.

Thank you in advance.

Answer 1

The main problem of your compiled function is that you do not feed it with the relevant tensors: There are msg, q, and g appearing in the body of Compile but they do not appear as arguments of the function. Hence, the notorious MainEvaluate will be called very often from the resulting CompiledFunction.

Another, less severe bottleneck is the needless construction of the vectors {kx, ky}.

As a further general advice, I suggest to start with computations in machine precision as early as possible (unless you do not encounter catastrophic loss of precision).

Here is how I would write things down:

msa = 1. 10^3;
msb = 1. 10^6;
mu0 = 1. 10^-6;
h0 = 1./mu0;
la = 1. 10^-16;
lb = 1. 10^-16;

aa = la*mu0*(msa^2)/2;
ab = lb*mu0*(msb^2)/2;
alpa = 0.2;
alpb = 0.1;
f = 0.1;
a = 1.*10^-7;
rcyl = Sqrt[(a^2)*f/Pi];
kmax = 3;
mmax = (2 kmax + 1)^2;
{g1, g2} = 
  Transpose@
   Tuples[Subdivide[-2. Pi/a kmax, 2. Pi/a kmax, 2 kmax], 2];
g1diff = Outer[Subtract, g1, g1];
g2diff = Outer[Subtract, g2, g2];
g1total = Total[g1];
g2total = Total[g2];

gdist = rcyl Sqrt[Plus[g1diff^2, g2diff^2]];
id = IdentityMatrix[Length[gdist], SparseArray, 
   WorkingPrecision -> MachinePrecision];
besselvals = Divide[BesselJ[1, gdist], (gdist + id)];
msg = ((msa - msb) (2 f)) besselvals + ((msa f) + (msb (1 - f))) id;
q = Plus[
   (((2 aa/(h0 mu0 msa^2)) - (2 ab/(h0*mu0*
             msb^2))) (2 f)) besselvals,
   ((2 aa f/(h0 mu0 msa^2)) + (2 ab (1 - f)/(h0 mu0 msb^2))) id
   ];

sumcompile = Compile[{
    {msg, _Real, 2}, {q, _Real, 2}, {g1, _Real, 1}, {g2, _Real, 1},
    {i, _Integer}, {j, _Integer}, {kx, _Real}, {ky, _Real}},
   Dot[
    msg[[i]] q[[All, j]],
    Plus[
     (kx + Compile`GetElement[g1, j]) (kx + g1),
     (ky + Compile`GetElement[g2, j]) (ky + g2),
     (Compile`GetElement[g1, i] - Compile`GetElement[g1, j]) g1 - 
      Compile`GetElement[g1, i],
     (Compile`GetElement[g2, i] - Compile`GetElement[g2, j]) g2 - 
      Compile`GetElement[g2, i]
     ]
    ],
   CompilationTarget -> "C",
   RuntimeOptions -> "Speed"
   ];

Test:

r1 = sumfunc[1, 1, 1, 1]; // RepeatedTiming // First
r2 = sumcompile[msg, q, g1, g2, 1, 1, 1., 1.]; // RepeatedTiming // First
Abs[r1 - r2]/r1

0.000414

3.17*10^-6

5.30516*10^-9

Hm. The precision is good but not very good. This might be caused by the fact that you use constants that range over many orders of magnitude.

Similar idea apply to your function bxycompile.