4
$\begingroup$

I have a question / confusion over improving the speed of MapThread[Dot,...] for lists of tensors. My problem involves taking two lists of tensors and then computing the list of dot products. I have found that I can speed this up by a factor of 2 or 3 by compiling Dot and using Listable to run over the list. However, it seems that once I reach rank 3 tensors, the compiled version is actually slower than the naive MapThread[Dot,...] implementation.

The background is that I have a list of functions together with their derivatives

$$s_\alpha,\quad \partial_i s_\alpha,\quad \partial_i \partial_j s_\alpha,$$

where $\alpha=1,\ldots,N_f$ and $i,j=1,\ldots,5$. These are lists of rank 1, 2 and 3 tensors. I then evaluate these for $N_p$ points, giving me a list of complex, machine precision tensors in packed arrays. I then need to compute various dot products, for example

$$ \partial_i \partial_j s_\alpha\, \bar{s}_\alpha\quad\text{or}\quad \partial_i \partial_j s_\alpha\, \partial_k \partial_l \bar{s}_\alpha,$$

giving a list of length $N_p$ with rank 2 (5 x 5) or rank 4 (5 x 5 x 5 x 5) tensors.

Here is a working example where I take the dot product for a list of rank 3 and rank 1 tensors (where the compiled version is indeed quicker) and for two lists of rank 3 tensors (where the compiled version is slower). I'm running this on a 16-core machine with 256 GB of RAM, so I don't think I'm running into a memory issue. The only thing I can think of is that there is some copying involved for the compiled versions and the copying becomes the limiting factor for the higher rank tensors?

Thanks in advance for any suggestions!

(*number of points and functions*)
Np = 200000;
Nf = 100;

(*value of functions and second derivatives at each point*)
s = RandomComplex[{-0.1, 0.1}, {Np, Nf}];
s$deriv2 = RandomComplex[{-0.1, 0.1}, {Np, 5, 5, Nf}];

(*define compiled function for rank 3 x rank 1 dot*)
MapThreadDot31 = 
  Compile[{{m, _Complex, 3}, {x, _Complex, 1}}, Dot[m, x], 
   RuntimeOptions -> "Speed", Parallelization -> True, 
   RuntimeAttributes -> {Listable}, CompilationTarget -> "C"];

(*define compiled function for rank 3 x rank 3 dot*)
MapThreadDot33 = 
  Compile[{{m, _Complex, 3}, {x, _Complex, 3}}, Dot[m, x], 
   RuntimeOptions -> "Speed", Parallelization -> True, 
   RuntimeAttributes -> {Listable}, CompilationTarget -> "C"];

(*compute dot product for rank 3 x rank 1 tensors*)
(*compiled version is more than 10x quicker*)
MapThreadDot31[s$deriv2, Conjugate[s]]; // AbsoluteTiming
MapThread[Dot, {s$deriv2, Conjugate[s]}]; // AbsoluteTiming

(*0.717989 s*)
(*10.6436 s*)

(*compute dot product for rank 3 x rank 3 tensors*)
(*compiled version is 3x slower*)
MapThreadDot33[s$deriv2, 
   ConjugateTranspose[s$deriv2, {1, 3, 4, 2}]]; // AbsoluteTiming
MapThread[
   Dot, {s$deriv2, 
    ConjugateTranspose[s$deriv2, {1, 3, 4, 2}]}]; // AbsoluteTiming

(*135.947 s*)
(*48.2371 s*)

Update

As a simpler example, we can compute the dot products of lists of large vectors, where I run into the same problem. Here is some sample code that computes the dot product for vectors using a number of different ways - MapThread[Dot,..], mixtures of Listable, Table (from Shadowray's comment),Parallelization, compiling to C or WVM. Some observations:

  • Listable + compiled to C + parallelize is quicker than all of the methods, until we seem to hit some size limit for the vectors, after which it becomes the slowest.
  • Listable + parallelize suffers from this problem greatly. Table + parallelize does not seem to.
  • Table does not seem to gain any speed up from parallelize, even for smaller vectors.
  • Listable without parallelize performs the same as Table with or without parallelize.

Any ideas for recovering the speed up obtained from parallelizing the Listable method? I use vectors of greatly varying length, so it would be nice if a single method worked for all of them. Thanks for the help so far!

(* Listable + C + parallel *)
MapThreadDot11$listable$C = 
  Compile[{{m, _Complex, 1}, {x, _Complex, 1}}, Dot[m, x], 
   RuntimeOptions -> "Speed", Parallelization -> True, 
   RuntimeAttributes -> {Listable}, CompilationTarget -> "C"];

(* Listable + WVM + parallel *)
MapThreadDot11$listable = 
  Compile[{{m, _Complex, 1}, {x, _Complex, 1}}, Dot[m, x], 
   RuntimeOptions -> "Speed", Parallelization -> True, 
   RuntimeAttributes -> {Listable}];

(* Listable + C *)
MapThreadDot11$listable$C$1 = 
  Compile[{{m, _Complex, 1}, {x, _Complex, 1}}, Dot[m, x], 
   RuntimeOptions -> "Speed", RuntimeAttributes -> {Listable}, 
   CompilationTarget -> "C"];

(* Listable + WVM *)
MapThreadDot11$listable$1 = 
  Compile[{{m, _Complex, 1}, {x, _Complex, 1}}, Dot[m, x], 
   RuntimeOptions -> "Speed", RuntimeAttributes -> {Listable}];

(* Table + C + parallel *)
MapThreadDot11$table$C = 
  Compile[{{m, _Complex, 2}, {x, _Complex, 2}}, 
   Table[Dot[m[[i]], x[[i]]], {i, 1, Length[m]}], 
   RuntimeOptions -> "Speed", CompilationTarget -> "C", 
   Parallelization -> True];

(* Table + WVM + parallel *)
MapThreadDot11$table = 
  Compile[{{m, _Complex, 2}, {x, _Complex, 2}}, 
   Table[Dot[m[[i]], x[[i]]], {i, 1, Length[m]}], 
   RuntimeOptions -> "Speed", CompilationTarget -> "C", 
   Parallelization -> True];

(* Table + C *)
MapThreadDot11$table$C$1 = 
  Compile[{{m, _Complex, 2}, {x, _Complex, 2}}, 
   Table[Dot[m[[i]], x[[i]]], {i, 1, Length[m]}], 
   RuntimeOptions -> "Speed", CompilationTarget -> "C"];

(* Table + WVM *)
MapThreadDot11$table$1 = 
  Compile[{{m, _Complex, 2}, {x, _Complex, 2}}, 
   Table[Dot[m[[i]], x[[i]]], {i, 1, Length[m]}], 
   RuntimeOptions -> "Speed"];

(* Define number of points and functions *)
Np = 200000;
Nf = 500;
s = RandomComplex[{-0.1, 0.1}, {Np, Nf}];

(*Timings when we let Nf vary from 500 to 4000*)
MapThread[Dot, {s, Conjugate[s]}]; // AbsoluteTiming
(*500: 2.6s, 1000: 5.5s, 2000:  11s, 3000: 16s, 4000: 17s*)
MapThreadDot11$listable$C[s, Conjugate[s]]; // AbsoluteTiming
(*500: 0.3s, 1000: 0.7s, 2000: 1.7s, 3000: 71s, 4000: 99s*)
MapThreadDot11$listable[s, Conjugate[s]]; // AbsoluteTiming
(*500: 0.3s, 1000: 0.7s, 2000: 1.7s, 3000: 88s, 4000: 67s*)
MapThreadDot11$listable$C$1[s, Conjugate[s]]; // AbsoluteTiming
(*500: 1.0s, 1000: 1.7s, 2000: 3.5s, 3000:  9s, 4000:  8s*)
MapThreadDot11$listable$1[s, Conjugate[s]]; // AbsoluteTiming
(*500: 0.9s, 1000: 1.7s, 2000: 3.6s, 3000:  7s, 4000: 11s*)
MapThreadDot11$table$C[s, Conjugate[s]]; // AbsoluteTiming
(*500: 0.9s, 1000: 2.7s, 2000: 3.6s, 3000:  6s, 4000: 11s*)
MapThreadDot11$table[s, Conjugate[s]]; // AbsoluteTiming
(*500: 0.9s, 1000: 1.7s, 2000: 3.6s, 3000:  6s, 4000:  9s*)
MapThreadDot11$table$C$1[s, Conjugate[s]]; // AbsoluteTiming
(*500: 0.7s, 1000: 1.7s, 2000: 3.6s, 3000:  8s, 4000:  9s*)
MapThreadDot11$table$1[s, Conjugate[s]]; // AbsoluteTiming
(*500: 0.9s, 1000: 2.9s, 2000: 4.2s, 3000: 10s, 4000: 10s*)

Update 2

There seems to be some sudden change in behaviour as Nf reaches 2048. This causes a sudden change in the timings, as set out below. For comparison, I've used the MapThreadDot from FEM - this is on par performance wise with the other methods suggested. While MapThreadDot11$listable$C is faster for Nf<2048, it suddenly becomes slower than the other functions for Nf>=2048. Note that this seems to be independent of Np, the number of points - changing Np to a much smaller or larger number does not change this conclusion (I have set this to 100,000 in the below example so it should run with much less RAM). Could it be that it suddenly uses a different BLAS function once the size of the vectors that we are taking the dot product of reaches 2048? It is definitely linked to parallelizing the calculation over many cores. I see this problem on the 16 core machine I've been using, but not on my 2 core laptop...

(* Listable + C + parallel *)
MapThreadDot11$listable$C = 
  Compile[{{m, _Complex, 1}, {x, _Complex, 1}}, Dot[m, x], 
   RuntimeOptions -> "Speed", Parallelization -> True, 
   RuntimeAttributes -> {Listable}, CompilationTarget -> "C"];

(* Define number of points *)
Np = 100000;

(* Timings for different values of Nf *)

Nf = 2046;
s = RandomComplex[{-0.1, 0.1}, {Np, Nf}];

NDSolve`FEM`MapThreadDot[s,Conjugate[s]]; //AbsoluteTiming
MapThreadDot11$listable$C[s, Conjugate[s]]; //AbsoluteTiming
(* Nf: 2046; FEM: 1.3s; Listable: 0.5s *)
(* Nf: 2047; FEM: 1.3s; Listable: 0.5s *)
(* Nf: 2048; FEM: 1.3s; Listable: 1.0s *)
(* Nf: 2049; FEM: 1.3s; Listable: 1.3s *)
(* Nf: 2050; FEM: 1.3s; Listable: 1.9s *)
Improve this question
$\endgroup$
12
  • 1
    $\begingroup$ If I had to guess, I'd assume this issue is the memory cost of passing those huge tensors to C $\endgroup$
    – b3m2a1
    Commented Jul 3, 2019 at 21:46
  • $\begingroup$ On my low-RAM machine (16 GB), MapThreadDot33 is still faster; however, a lot of swapping is going, so that does not mean much... $\endgroup$
    – Henrik Schumacher
    Commented Jul 3, 2019 at 23:15
  • $\begingroup$ Try this: MapThreadDot33 = Compile[{{m, _Complex, 4}, {x, _Complex, 4}}, Table[Dot[m[[i]], x[[i]]], {i, 1, Length[m]}]]; $\endgroup$
    – Ray Shadow
    Commented Jul 4, 2019 at 0:48
  • $\begingroup$ So it seems that removing the Parallelization -> True removes all difference. Compiling to WVM instead of C seems to make no difference, so it seems the performance difference is due to copying the large tensors to the many parallel kernels (16 in this case, and the problem seems to get even worse on the 36 core machines). In addition, on my laptop there is little to no difference (I would guess as it only has two cores). For the rank 1 x rank 3 tensor calculation, the parallelized compiled function is still much faster. Any ideas for recovering some of this speed? $\endgroup$
    – ala10
    Commented Jul 4, 2019 at 16:48
  • $\begingroup$ How does the MapThreadDot33 function from my previous comment perform on your system? Note that it is not listable and rely on the Table function. $\endgroup$
    – Ray Shadow
    Commented Jul 4, 2019 at 21:35

0

Reset to default

Your Answer

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