flushing thread local buffer at end of parallel loop with TBB - c++

Question

I want to parallelize a loop (using tbb) which contains some expensive but vectorizable iterations (randomly spread). My idea was to buffer those and flush the buffer whenever it reaches the vector size. Such a buffer must be thread-local. For example,
// dummy for testing
void do_vectorized_work(size_t k, size_t*indices)
{}
// dummy for testing
bool requires_expensive_work(size_t k)
{ return (k&7)==0; }
struct buffer
{
size_t K=0, B[vector_size];
void load(size_t i)
{
B[K++]=i;
if(K==vector_size)
flush();
}
void flush()
{
do_vectorized_work(K,B);
K=0;
}
};
void do_work_in_parallel(size_t N)
{
tbb::enumerable_thread_specific<buffer> tl_buffer;
tbb::parallel_for(size_t(0),N,[&](size_t i)
{
if(requires_expensive_work(i))
tl_buffer.local().load(i);
});
}
However, this leaves the buffers non-empty, so I still have to flush each of them a final time
for(auto&b:tl_buffer)
b.flush();
but this is serial! Of course, I can also try to do this in parallel
using tl_range = typename tbb::enumerable_thread_specific<buffer>::range_type;
tbb::parallel_for(tl_buffer.range(),[](tl_range const&range)
{
for(auto r:range)
r->flush();
});
But I'm not sure this is efficient (since there are only as many buffers as there are threads). I was wondering whether it is possible to avoid this final flush after the event. I.e. is it possible to use tbb::tasks (replacing tbb::parallel_for) in such a way that each thread's final task is to flush its buffer?

Answer 1

No, a worker thread does not have complete information about whether this particular task is the last task of the given work or not (this is how work-stealing works). Thus, it is not possible to implement such a function on the level of parallel_for or the scheduler itself. Thus, I'd recommend you to go with these two approaches you describe.
There are two other things you can do about this though.
make it asynchronous. I.e. enqueue a task which will get everything flushed. It will help to remove this code from the hot path on the main thread. Just be careful if there are any dependencies which need to be set on completion of this task.
use tbb::task_scheduler_observer in order to initialize thread-specific data and release it lazily when threads get shut down or when there is no work remains for some time. The latter requires using local observer feature which is not yet officially supported but remains stable for few years already.
Example:
#define TBB_PREVIEW_LOCAL_OBSERVER 1
#include <tbb/tbb.h>
#include <assert.h>
typedef void * buffer_t;
const static int bufsz = 1024;
class thread_buffer_allocator: public tbb::task_scheduler_observer {
tbb::enumerable_thread_specific<buffer_t> _buf;
public:
thread_buffer_allocator( )
: tbb::task_scheduler_observer( /*local=*/ true ) {
observe(true); // activate the observer
}
~thread_buffer_allocator( ) {
observe(false); // deactivate the observer
for(auto &b : _buf) {
printf("destructor: cleared: %p\n", b);
free(b);
}
}
/*override*/ void on_scheduler_entry( bool worker ) {
assert(_buf.local() == nullptr);
_buf.local() = malloc(bufsz);
printf("on entry: %p\n", _buf.local());
}
/*override*/ void on_scheduler_exit( bool worker ) {
printf("on exit\n");
if(_buf.local()) {
printf("on exit: cleared %p\n", _buf.local());
free(_buf.local());
_buf.local() = nullptr;
}
}
};
int main() {
thread_buffer_allocator buffers_scope;
tbb::parallel_for(0, 1024*1024*1024, [&](auto i){
usleep(i%3);
});
return 0;
}

Answer 2

It occurred to me that this can be solved by reduction.
struct buffer
{
std::size_t K=0, B[vector_size];
void load(std::size_t i)
{
B[K++]=i;
if(K==vector_size) flush();
}
void flush()
{
do_vectorized_work(K,B);
K=0;
}
buffer(buffer const&, tbb::split)
{}
void operator()(tbb::block_range<std::size_t> const&range)
{ for(i:range) load(i); }
bool empty()
{ return K==0; }
std::size_t pop()
{ return K? B[--K] : 0; }
void join(buffer&rhs)
{ while(!rhs.empty()) load(rhs.pop()); }
};
void do_work_in_parallel(std::size_t N)
{
buffer buff;
tbb::parallel_reduce(tbb::block_range<std::size_t>(0,N,vector_size),buff);
if(!buff.empty())
buff.flush();
}