Use TBB to parallelise unary functions of large containers

[andrjohns]

Description

Vectorised unary functions (e.g., exp, log) currently use either looping or Eigen's vectorised functions (for arithmetic types only) to evaluate containers. We could instead use the TBB's parallel_for functionality to evaluate the container in parallel.

A minimal example of this would be applying exp to every element of a large Eigen vector:

  using Eigen::VectorXd;
  VectorXd m_d = VectorXd::Random(10000000);
  VectorXd m_res(10000000);

  tbb::task_scheduler_init init;
  tbb::parallel_for(
    tbb::blocked_range<size_t>(0, m_d.size()), 
    [&m_d,&m_res](const tbb::blocked_range<size_t>& r) {
      for (size_t i = r.begin(); i < r.end(); ++i)
        m_res[i] = std::exp(m_d[i]);
    });

Benchmarking shows a performance of ~2x that of Eigen's vectorised exp, and ~5x that of a simple loop (benchmarking code at end of post):

Run on (16 X 3700 MHz CPU s)
CPU Caches:
  L1 Data 32 KiB (x8)
  L1 Instruction 64 KiB (x8)
  L2 Unified 512 KiB (x8)
  L3 Unified 8192 KiB (x2)
Load Average: 1.53, 0.87, 0.79
-----------------------------------------------------
Benchmark           Time             CPU   Iterations
-----------------------------------------------------
ParExp       10259973 ns     10225812 ns           71
EigExp       23918425 ns     23917527 ns           28
LoopExp      53930647 ns     53929832 ns           10

This could provide large performance improvements for containers of autodiff types (where Eigen's vectorised functions can't be used), or for containers of arithmetic types where a vectorised function isn't available in Eigen.

I still need to do more testing to identify the size at which parallelisation becomes beneficial, but things look pretty promising

Current Version:

v3.2.0

Code used for benchmarking:

#include <stan/math/mix.hpp>
#include <benchmark/benchmark.h>
#include <tbb/parallel_for.h>
#include <tbb/blocked_range.h>

static void ParExp(benchmark::State& state) {
  using Eigen::VectorXd;
  VectorXd m_d = VectorXd::Random(10000000);
  VectorXd m_res(10000000);


  for (auto _ : state) {
    tbb::task_scheduler_init init;
    tbb::parallel_for(
      tbb::blocked_range<size_t>(0, m_d.size()), 
      [&m_d,&m_res](const tbb::blocked_range<size_t>& r) {
        for (size_t i = r.begin(); i < r.end(); ++i)
          m_res[i] = std::exp(m_d[i]);
      });
  }
}
BENCHMARK(ParExp);

static void EigExp(benchmark::State& state) {
  using Eigen::VectorXd;
  VectorXd m_d = VectorXd::Random(10000000);
  VectorXd m_res(10000000);

  for (auto _ : state) {
    m_res.array() = m_d.array().exp();
  }
}
BENCHMARK(EigExp);

static void LoopExp(benchmark::State& state) {
  using Eigen::VectorXd;
  VectorXd m_d = VectorXd::Random(10000000);
  VectorXd m_res(10000000);

  for (auto _ : state) {
    for(size_t i = 0; i < m_d.size(); ++i) {
      m_res[i] = std::exp(m_d[i]);
    }
  }
}
BENCHMARK(LoopExp);

BENCHMARK_MAIN();

[rok-cesnovar]

very nice! If its not too much work, want to try out lgamma also?

Also is it possible to use Eigen's vectorized form inside the parallel for on a slice of the vector?

[bob-carpenter]

How many cores did the evaluation use?

How is this going to work when other things are being sent to TBB around this, like map-rect and its descendants?

How are the autodiff types synchronized? The benchmarks looked to be only for double values.

I'd suggest putting the random generation code outside the blocks being tested so it's not included in their timing and all functions are called on the same inputs.

[wds15]

How do you suggest to integrate this into math? I mean, would you fire off the parallel version when the container size is sufficiently large... or would you provide additional function definitions which would also allow to specify a grainsize (which essentially also controls this chunking).

Other than that this plays well with the other map_rect / reduce_sum functions as we just throw tasks at the TBB and everything is designed such that we can use nested parallelism.

In all honesty, I think what would be of great interest is a variadic map (just like reduce_sum, but a 1-1 map thing). The issue here is to figure out how to concatenate outputs.

Another thing would be to provide parallel lpdfs / lpmfs which kick in for large data sets. Maybe the glm functions are actually even better suited for that given that more work is happening. Right now the glm functions are accelerated on the GPU, but multiple CPUs can help as well.

Given that you are the master, @andrjohns, of the apply* functions: What I think can be cool is to use parallel_reduce from the TBB to speed up vectorised unary reduce functions. This could be implemented on the back of reduce_sum presumably.

[t4c1]

+1 to using Eigen expressions within the parallel for instead of loop + std::exp

[andrjohns]

@rok-cesnovar

If its not too much work, want to try out lgamma also?

That saw some really nice speed-ups too:

ParLgamma_dbl    30094603 ns     30093781 ns           23
LoopLgamma_dbl  318332120 ns    318325288 ns            2

Also is it possible to use Eigen's vectorized form inside the parallel for on a slice of the vector?

I'm not sure, to be honest, since parallel_for has to be passed a loop for it to parallelise. One approach would be do something like:

m_res.segment(i, i+10) = m_d.segment(i, i+10).array().exp();

But would have to ensure that each iteration doesn't overlap, somehow. @t4c1 any ideas here?

@bob-carpenter

How many cores did the evaluation use?

This was using the default, and letting the TBB automatically choose (I have 8 physical cores available). It looks like the parallelised exp is slower than Eigen's (but faster than a loop) for 2 cores:

ParExp_dbl       27989277 ns     27987591 ns           25
EigExp_dbl       23901004 ns     23899449 ns           28
LoopExp_dbl      53516822 ns     53513638 ns           10

But faster for 4 cores:

ParExp_dbl       14673224 ns     14671461 ns           46
EigExp_dbl       23816895 ns     23814757 ns           27
LoopExp_dbl      53724174 ns     53722104 ns           10

How is this going to work when other things are being sent to TBB around this, like map-rect and its descendants?

As Sebastion mentioned, the TBB handles this kind of scheduling.

How are the autodiff types synchronized? The benchmarks looked to be only for double values.

I can't test autodiff types yet, as the perf-math library is segfaulting when assigning values to vars (i.e., just var x = 5.0; segfaults). I can post benchmarks once I've chased that down.

I'd suggest putting the random generation code outside the blocks being tested so it's not included in their timing and all functions are called on the same inputs.

The generation was already separate from the timing, but I moved to a single vector of inputs shared by all runs, and the timings were consistent.

@wds15

How do you suggest to integrate this into math?

To be honest, I hadn't thought that far ahead (got caught up in whether I could, not whether I should). My (likely simplistic) perspective was to use the parallel version depending on container size (similar to the GPU code). I feel like there's some value to 'naive' parallelism which users can just take advantage of without changing their code.

As for your other suggestions, I need to get a bit more familiar with the options in the TBB before I can talk about other applications with much confidence. I just saw the unary functions as a good first step since they can be easily represented as parallelisable loop. We can even build this into binary scalar functions, so things like lbeta(Matrix, Matrix) could be parallelised in a similar way. But I'm happy to go another direction that would generalise better.

[rok-cesnovar]

But would have to ensure that each iteration doesn't overlap

Doesn't TBB ensure that the blocked ranges don't overlap?

parallel version depending on container size (similar to the GPU code)

I would actually like to get rid of that as the values are arbitrary and would need to be modified every time some performance improvement is made on the sequential or parallel version (for the GPU Cholesky decomposition N for switching is still 1250 though recent improvements have lowered that to at least half that).

[andrjohns]

@rok-cesnovar & @t4c1 Spoke too soon, parallel_for gives you the first index and upper bound for the loop:

(size_t i = r.begin(); i < r.end(); ++i)

So I can use that to segment the Eigen vector and use Eigen's exp:

  m_res.segment(r.begin()+1, r.end()-r.begin()-1).array()
    = m_d.segment(r.begin()+1, r.end()-r.begin()-1).array().exp();

While there's no noticeable difference when using all 8-cores for the parallelism:

ParExp_dbl        9824678 ns      9816867 ns           72
ParExp_Eig_dbl    9719649 ns      9693837 ns           72
EigExp_dbl       23886641 ns     23887389 ns           27

It performs much better when dropping down to 2 cores:

ParExp_dbl       33495995 ns     33495759 ns           20
ParExp_Eig_dbl   12505482 ns     12504933 ns           52
EigExp_dbl       23954809 ns     23953778 ns           28

Which will be good for getting benefits for those with lower core-count CPUs

[t4c1]

My guess is that on 8 cores exp is so fast, it saturates memory bus.

[wds15]

Is your 8 core run improving in speedup with larger container sizes? Otherwise the grainsize could be important to tune a bit to tell the TBB what are good chunk sizes here.

My only concern with parallelising these little functions is that it may not matter after all in a big stan program. Think about Amdahl's law. You have to run a large percentage of your code in parallel to get any useful scaling. Pimping exp which is just one of many operations, is not very promising.

This is why I suggested to use the TBB for the glm functions.

I really don't want to discourage you, but maybe think about an entire Stan program where this will be useful.

[andrjohns]

I see what you mean, but I feel like this would benefit a wider range of users than the glm functions. Something like this can be built into the apply_scalar_unary and apply_scalar_binary frameworks, so that all unary and binary functions get parallelised. This would also speed up the other functions and distributions that use them, so I think this would parallelise more of a given program than the individual glms.

[wds15]

I hope you are right, but I am very doubtful about this. So maybe you can hack up a quick and dirty prototype and provide an end-to-end performance test for a Stan model?

The issue is that you end up running your program is super small slices in parallel and that is hard for me to imagine to be fast. Hence, a prototype showing the merits of this approach is really what's needed from my view to proceed.

(I really appreciate more TBB use - don't get me wrong - but we don't want to waste our resources obviously)

[rok-cesnovar]

I think it would suffice to see what effect does parallelization of unary/binary functions have on a gradient of an _lpdf function. That would probably give us enough information. Not sure which would be a good candidate.

[andrjohns]

Thanks for the comments both! Either way, this will take some benchmarking to see if it's beneficial. I'll work up a rough implementation for unary and binary functions and then we can figure out what some good models for testing would be

[wds15]

Benchmarking some lpdfs which use this is a start, but not sufficient I am afraid. While a lot of work is going on in the lpdfs, these often do not constitute most of the work as I have seen. That is I keep on encouraging people to move more stuff into reduce_sum.

This is why I would target the glm functions which should indeed represent the bulk of work in a Stan program in case you use that.

... but let's find out!

[bob-carpenter]

Other than that this plays well with the other map_rect / reduce_sum functions as we just throw tasks at the TBB and everything is designed such that we can use nested parallelism.

I'd like to see is an extension of the existing evaluation to more parallel top level jobs to see what happens when cores saturate:

• 32 / 16 / 8 / 4 / 2 / 1 parallel top-level jobs
• parallel execution of jobs vs. serial execution of the jobs (don't know if we can control lower-level parallelism granularity)

The current eval was for only 1 top-level job. What I want is a proxy for running parallel chains vs. parallelizing within a chain.