10
$\begingroup$

Consider the following matrices:

phRotMatrix[
   ph_] = {{Cos[ph], Sin[ph], 0}, {-Sin[ph], Cos[ph], 0}, {0, 0, 1}};
thRotMatrix[
   th_] = {{1, 0, 0}, {0, Cos[th], Sin[th]}, {0, -Sin[th], Cos[th]}};

and some arbitrary vector vec. I would like to compile the product phRotMatrix[ph].thRotMatrix[th].vec. This is how I do:

CompileProd = 
 Compile[{{rotphMatr, _Real, 2}, {rotthMatr, _Real, 2}, {vec, _Real, 
    1}}, rotphMatr.rotthMatr.vec, CompilationTarget -> "C", 
  RuntimeOptions -> "Speed"]

However, it does not provide any speedup in comparison to the ordinary evaluation:

thrand = RandomReal[{0, Pi}];
phrand = RandomReal[{-Pi, Pi}];
(CompileProd[phRotMatrix[phrand], thRotMatrix[thrand], {1, 2, 3}]; // 
  AbsoluteTiming)
(phRotMatrix[phrand].thRotMatrix[thrand].{1, 2, 3}; // AbsoluteTiming)

{0.0000304,Null}

{0.0000161,Null}

Could you please tell me what I do wrong?

$\endgroup$
3
  • 2
    $\begingroup$ If you are using numerical evaluations, I would assume that Dot and other matrix operations already rely on heavily optimized routines from standard libraries that use fast compiled code. I would therefore be surprised (and disappointed!) if a user-compiled routine were any faster. If you are doing symbolic evaluation, I am not sure that this can be achieved within Compile; Have you checked if you get kicked back out to the main evaluator instead? $\endgroup$
    – MarcoB
    Dec 8, 2021 at 22:42
  • 1
    $\begingroup$ You probably cannot improve timing on a single dot product like you've shown. Are you using your function in a loop? It's generally better to map / thread a function over a large number of items in a list (cache locality + simd reasons). And generate batches of random values using the second argument to RandomReal, not single random values at a time in a loop. $\endgroup$
    – flinty
    Dec 8, 2021 at 23:16
  • $\begingroup$ @flinty : yes, I will use it in something like a loop, see my question mathematica.stackexchange.com/questions/259431/… I will call BlockThreeBodyPhaseSpaceNew2[E1, E3, m, m1, m2, m3] many times for different values of E1, E3. $\endgroup$ Dec 8, 2021 at 23:19

2 Answers 2

14
$\begingroup$

Compiling such linear algebra operatations does make some sense because we can give the compiler valuable information, e.g., that we use only $3 \times 3$ matrices. This allows the compiler to optimize the code by, e.g., loop unrolling and enabling or disabling vectorization. But you have to give the compiler some symbolic code to optimize. Otherwise it will resolve calls to Dot by calling some BLAS routines (that are typically quite efficient for general and large matrices, but not necessarily for matrices of fixed small size).

This is how I would do it from within Mathematica. (I don't want to write C/C++ code directly here.)

A = Table[Compile`GetElement[a, i, j], {i, 1, 3}, {j, 1, 3}];
B = Table[Compile`GetElement[b, i, j], {i, 1, 3}, {j, 1, 3}];
V = Table[Compile`GetElement[v, i], {i, 1, 3}];
dotThread = With[{code = A.B.V},
   
   Compile[{{a, _Real, 2}, {b, _Real, 2}, {v, _Real, 1}},
    code,
    CompilationTarget -> "C",
    RuntimeAttributes -> Listable,
    Parallelization -> True,
    RuntimeOptions -> "Speed"
    ]
   ];

Now compare:

phRotMatrix[ph_] = {{Cos[ph], Sin[ph], 0}, {-Sin[ph], Cos[ph], 0}, {0, 0, 1}};
thRotMatrix[th_] = {{1, 0, 0}, {0, Cos[th], Sin[th]}, {0, -Sin[th], Cos[th]}};

CompileProd = 
  Compile[{{rotphMatr, _Real, 2}, {rotthMatr, _Real, 2}, {vec, _Real, 
     1}}, rotphMatr.rotthMatr.vec, CompilationTarget -> "C", 
   RuntimeOptions -> "Speed"];

n = 1000000;
phlist = RandomReal[{-Pi, Pi}, n];
thlist = RandomReal[{-Pi, Pi}, n];
veclist = RandomReal[{-1, 1}, {n, 3}];
Alist = phRotMatrix /@ phlist; // AbsoluteTiming // First
Blist = thRotMatrix /@ thlist; // AbsoluteTiming // First
result = MapThread[CompileProd, {Alist, Blist, veclist}]; // AbsoluteTiming // First
result2 = dotThread[Alist, Blist, veclist]; // AbsoluteTiming // First
Max[Abs[result - result2]]

2.7701

2.73302

1.99988

0.811031

0.

So this made the dot product only twice as fast. Hm. Not so convincing. The reason is probably that the parallelization incurs some overhead.

Here is a way to do all the work in a single thread:

A = Table[Compile`GetElement[a, k, i, j], {i, 1, 3}, {j, 1, 3}];
B = Table[Compile`GetElement[b, k, i, j], {i, 1, 3}, {j, 1, 3}];
V = Table[Compile`GetElement[v, k, i], {i, 1, 3}];
dotMany = With[{code = A.B.V},
   Compile[{{a, _Real, 3}, {b, _Real, 3}, {v, _Real, 2}, {first, _Integer}, {last, _Integer}},
    Table[code, {k, first, last}],
    CompilationTarget -> "C",
    RuntimeAttributes -> Listable,
    Parallelization -> True,
    RuntimeOptions -> "Speed"
    ]
   ];

result3 = dotMany[Alist, Blist, veclist, 1, Length[veclist]]; // 
  AbsoluteTiming // First
Max[Abs[result - result3]]

0.476497

0.

This is already 4 times as fast as the original CompileProd.

We could try to put this function on 4 different threads as follows:

threadcount = 4;
{first, last} = Transpose[ Partition[Round[Subdivide[1, Length[veclist] + 1, threadcount]], 2, 1]] - {0, 1};
result4 = dotMany[Alist, Blist, veclist, first, last]; // 
  AbsoluteTiming // First
Max[Abs[Flatten[result] - Flatten[result4]]]

0.901562

0.

Hm. How unfortunate. It became slower for some reason...

Even more can be done.

Dense matrix-vector multiplications are heavily memory bound. As your matrices are rotation matrices, it should be cheaper to send only the angles. So instead of $3 \times 3 + 3 \times 3 + 3 = 21$, you send only $1 + 1 + 3 = 5$ doubles per matrix-matrix-vector tuple. Moreover, your rotation matrices have many zeros. This means we can save quite many multiplications and additions. So the following should be more efficient:

V = Table[Compile`GetElement[v, k, i], {i, 1, 3}];
rotateMany = 
  With[{code = 
     phRotMatrix[Compile`GetElement[ph, k]].thRotMatrix[Compile`GetElement[th, k]].V},
   Compile[{{ph, _Real, 1}, {th, _Real, 1}, {v, _Real,2}, {first, _Integer}, {last, _Integer}},
    Table[code, {k, first, last}],
    CompilationTarget -> "C",
    RuntimeAttributes -> Listable,
    Parallelization -> True,
    RuntimeOptions -> "Speed"
    ]
   ];

Running this on a single thread:

result5 = rotateMany[phlist, thlist, veclist, 1, Length[veclist]]; // 
  AbsoluteTiming // First
Max[Abs[Flatten[result] - Flatten[result5]]]

0.073883

4.44089*10^-16

Wow, that got much better! And on 4 thread?

result6 = rotateMany[phlist, thlist, veclist, first, last]; //   AbsoluteTiming // First
Max[Abs[Flatten[result] - Flatten[result6]]]

0.024217

4.44089*10^-16

This is almost 100 times faster than CompileProd. And it even generates the rotation matrices on the fly, so it even spares us from generating them in the first place. This also shows that the bottleneck with my first multi-thread attempts was caused by memory bandwidth problems.

$\endgroup$
9
$\begingroup$

This is too long for a comment, but also not a complete answer. I'm tossing out an idea in case this was not already considered.

One needs to recall two things:

  1. the product of two rotation matrices is also a rotation matrix
  2. any $3\times3$ rotation matrix can be expressed in terms of an axis and angle

In your specific case, the product you are interested in corresponds to the angle

tt = ArcCos[Cos[th] Cos[ph/2]^2 - Sin[ph/2]^2]

and the (unnormalized) axis

vv = {-2 Cos[ph/2]^2 Sin[th], Sin[ph] Sin[th], -2 Cos[th/2]^2 Sin[ph]}

from which you can apply the Rodrigues rotation formula

vn = vv/Sqrt[vv . vv];
vec = vec Cos[tt] + Cross[vn, vec] Sin[tt] + vn (vn . vec) (1 - Cos[tt])

where you should retain common subexpressions, and simplify Cos[tt] and Sin[tt] so that you wouldn't need to evaluate ArcCos[]. I suspect compiling a suitably optimized version of this would give good performance.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.