7
$\begingroup$

In the course of writing a custom solver for a system of coupled PDEs that NDSolve doesn't handle particularly well, I need to do many matrix multiplications of small vectors (3 dimensions, in this particular case) by small matrices ($ 3 \times 3 $). The vectors and matrices are both laid out in lists (so can be thought of as rank 2 and rank 3 tensors, respectively), meaning the operation can be expressed like so:

ClearAll[manyDots];
manyDots[m_,x_] := MapThread[Dot,{m,x}];

Here's the symbolic output of a simple case:

Clear[a,b,c,u,v,w];
With[{
  m = Array[#, {2, 2}] & /@ {a, b, c},
  x = Array[#, {2}] & /@ {u, v, w}
  },
 manyDots[m, x]]

(* {{a[1, 1] u[1] + a[1, 2] u[2], 
     a[2, 1] u[1] + a[2, 2] u[2]}, 
    {b[1, 1] v[1] + b[1, 2] v[2], 
     b[2, 1] v[1] + b[2, 2] v[2]}, 
    {c[1, 1] w[1] + c[1, 2] w[2], 
     c[2, 1] w[1] + c[2, 2] w[2]}} *)

Now for some test data:

ClearAll[mTest, xTest];
mTest = RandomReal[{}, {1000, 3, 3}];
xTest = RandomReal[{}, {1000, 3}];

Happily, these come out of RandomReal nicely packed:

Developer`PackedArrayQ /@ {mTest, xTest}

Still, my naïve implementation is not terribly fast:

slow = First@RepeatedTiming[manyDots[mTest, xTest];, 1.0]
(* 0.00066 *)

We're doing 1000 Dot operations, and each one involves about a dozen floating point adds or multiplies:

ops = Count[Array[a, {3, 3}].Array[v, {3}], _Plus | _Times, Infinity]
(* 12 *)

ops*1000/slow
(* 1.8*10^7 *)

Surely we can do better than 18 million FLOPS.

The next step is to try compiling the function:

ClearAll[compiledManyDots];
compiledManyDots = Compile[{{m, _Real, 3}, {x, _Real, 2}},
   MapThread[Dot, {m, x}],
   CompilationTarget -> "C"];

This is about an order of magnitude faster:

compiled = First@RepeatedTiming[compiledManyDots[mTest, xTest];, 1.0]
(* 0.000083 *)

slow/compiled
(* 8. *)

ops*1000/compiled
(* 1.4*10^8 *)

Still, I'd think it should be possible to squeeze another small integer factor out of that by being clever. So far I haven't tried much except a tedious exercise of reimplementing Dot inline in the compiled function body using Do, which didn't help.

$\endgroup$

2 Answers 2

6
$\begingroup$

There's NDSolve`FEM`MapThreadDot[mTest, xTest], which seems fairly fast on Henrik's example:

mTest = RandomReal[{}, {1000000, 3, 3}];
xTest = RandomReal[{}, {1000000, 3}];
compiled1 = First@RepeatedTiming[a1 = compiledManyDots[mTest, xTest];]
compiled2 = First@RepeatedTiming[a2 = compiledManyDots2[mTest, xTest];]
femMTD =    First@RepeatedTiming[a3 = NDSolve`FEM`MapThreadDot[mTest, xTest];]
Max@Abs@Differences@{a1, a2, a3}
(*
  0.13   <-- compiledManyDots
  0.054  <-- compiledManyDots2
  0.031  <-- NDSolve`FEM`MapThreadDot

  0.
*)
$\endgroup$
4
  • $\begingroup$ Wow Michael, that's amazing! $\endgroup$ Feb 25, 2018 at 14:50
  • 1
    $\begingroup$ @HenrikSchumacher Thanks. I imagine they've minimized the overhead in shipping the dot products off to BLAS. Also, BLAS may take advantage of the FMA instruction and your compiled version may not -- that alone could explain the difference. I can't imagine a more efficient approach than yours (using 'Compile' in Mathematica, that is). $\endgroup$
    – Michael E2
    Feb 25, 2018 at 16:08
  • 1
    $\begingroup$ Indeed, FMA is a very good point... $\endgroup$ Feb 25, 2018 at 22:03
  • $\begingroup$ This is awesome. Sped up my solver by a factor of two just from dropping it in! $\endgroup$
    – Pillsy
    Feb 25, 2018 at 22:40
4
$\begingroup$

The best I can achieve is the following which is on my machine (quad core) aboute twice as fast:

ClearAll[compiledManyDots2];
Quiet@Block[{m, mm, x, xx},
  mm = Table[Compile`GetElement[m, i, j], {i, 1, 3}, {j, 1, 3}];
  xx = Table[Compile`GetElement[x, i], {i, 1, 3}];
  product = mm.xx;
  With[{code = mm.xx},
   compiledManyDots2 = Compile[{{m, _Real, 2}, {x, _Real, 1}},
      code,
      CompilationTarget -> "C",
      RuntimeAttributes -> {Listable},
      Parallelization -> True,
      RuntimeOptions -> "Speed"
      ];
   ]
  ]

mTest = RandomReal[{}, {1000000, 3, 3}];
xTest = RandomReal[{}, {1000000, 3}];
compiled1 = First@RepeatedTiming[a1 = compiledManyDots[mTest, xTest];]
compiled2 = First@RepeatedTiming[a2 = compiledManyDots2[mTest, xTest];]
Max[Abs[a1 - a2]]

0.095

0.051

0.

The C code of the compuational kernel obtained from

s = ExportString[compiledManyDots2, "C"];
StringTake[s, StringPosition[s, "DLLEXPORT"][[-1, 1]] ;;]

contains no bound checks (which is good) but some redundant local variable declarations and indexing. I am not sure if this gets optimized away by the compiler. The actual threading of this kernel over the jobs is done by the Listable atrtibute and thus out of my reach.

$\endgroup$
3
  • $\begingroup$ I like the way you generate the code symbolically; it seems the speedup is considerably bigger with larger arrays, suggesting there's a fairly big cost around the call to the function. $\endgroup$
    – Pillsy
    Feb 25, 2018 at 13:58
  • $\begingroup$ There is certainly some managing overhead for distributing the jobs over the cores of the CPU. Towards symbolical calculations: In particular for nonlinear operations, e.g. on local dofs of a FE mesh, creating symbolical code and let Compile and the actual C compiler do all the optimization is often much easier and more efficient(!) than hand coded approaches. The drawback is that you might need different code for similar problems but with different dimensions... $\endgroup$ Feb 25, 2018 at 14:07
  • $\begingroup$ Moreover, time measurements, even with RepeatedTiming, in the range of percentages of milliseconds are not very trustedworthy... $\endgroup$ Feb 25, 2018 at 14:13

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.