API to avoid heap-allocating tasks in Scope::spawn

[davidlattimore]

For some kinds of problems / algorithms, spawning tasks into a scope is useful. It's unfortunate that each task needs a separate heap allocation. I've tried a few different approaches which I've documented here, but have returned to using Scope::spawn due to downsides of those approaches.

It seems like there could be a few options for changes / additions to the Scope::spawn API that might allow the heap allocation to be avoided.

One way would be to make the Scope generic over some type then allow pushing new work items of that type. Perhaps Scope wouldn't be the right name for that. It's probably more of a dynamically growable work queue where work items can add more work items and runs until all work is complete with no additional work items added.

Another approach might be to configure each scope with a work-item size. Spawned tasks would then need size and alignment <= that size, but could then be stored adjacent to one another in some internal storage within the scope.

Join can be used for some of these kinds of problems, but can end up pretty awkward and can easily result in stack overflows if you're not careful.

[cuviper]

Can you sketch out some concrete APIs for your ideas? Don't worry about naming, but I think this will make it easier for us to talk through what is desired and how feasible that is.

[davidlattimore]

Sure! Here's the sort of thing I had imagined:

/// Execute `init` to spawn initial work, then execute `op` for each work item.
/// `op` can spawn additional work.
pub fn typed_scope<'scope, T, INIT, OP>(init: INIT, op: OP)
    where
        INIT: FnOnce(&TypedScope<'scope, T>),
        OP: Fn(T, &TypedScope<'scope, T>);

Which might be used something like this:

    rayon::typed_scope(
        |scope| {
            for work_item in initial_work {
                scope.add_work_item(work_item);
            }
        },
        |work_item, scope| {
            process_work_item(work_item, scope);
        },
    );

Alternatively, we could assume that the user already has some kind of collection of initial work:

pub fn typed_scope<'scope, T, INIT, OP>(init: INIT, op: OP)
    where
        INIT: IntoIterator<Item = T>,
        OP: Fn(T, &TypedScope<'scope, T>);

Which could be used as:

    rayon::typed_scope(initial_work, |work_item, scope| {
        process_work_item(work_item, scope);
    });

or more succinctly:

rayon::typed_scope(initial_work, process_work_item);

In either case, the scope type would be generic over the lifetime and the type of work items that it contains.

pub struct Scope<'scope, T>...

Possibly with some bounds on T. Possibly T: Send + Sync + 'static might be fine, but I suspect T: Send + Sync + 'scope would be more flexible.

I've previously attempted to get an API somewhat like this by means of a crossbeam-channel fed into par_bridge (see the blog post linked in the top comment for details of how this was done). There are a few issues with this. The first is that any use of par_bridge where the input iterator can block, doesn't play nicely with other work occurring in parallel. The par_bridge generally ends up occupying all threads until the channel closes even if there are periods of time when there's no work available. The threads that don't have work end up blocking waiting for the Mutex in par_bridge::IterParallelProducer except for one which has the lock and is blocking waiting for data from the channel (via a call to Iterator::next).

Another issue with the approach is that the mutex is somewhat heavyweight and unnecessary since the underlying channel is safe to read concurrently from multiple threads.

When imagining the typed_scope API above, I had imagined that the scope would hold a lock-free, unordered bag of work items that multiple threads could pull from. This might be more efficient than a crossbeam-channel that tries to maintain FIFO ordering. In fact the bag of work starts to sound a bit like crossbeam_deque. I guess the tricky thing is whether it would be possible for a scope to have its own crossbeam_deque in addition to the main one that rayon uses. If it did, how could work be pulled from either depending on what was available. e.g. the scope might at several points in time only have a single work item being processed, but then that work item results in heaps more being spawned. How I'd imagine this could work is that jobs are added to the main deque to schedule processing of work from the scope. Those jobs run until they fail to find more work, then exit, allowing other work from the main deque to be processed again. The scope would probably need to track how many threads were working on the scope and if additional work is added, push additional jobs to the main deque again. I think it should be possible for the jobs pushed to the main queue to not require heap allocations, since they should only need a pointer to the scope and a pointer to the function vtable. Although, I guess care would be needed to allow the scope to safely terminate if there's still jobs in the main queue that are stuck behind other work being processed from the main queue.

I guess if it is possible to have a scope with its own crossbeam_deque, then that suggests that it might be possible to have an API more like the current API, but where the closures are stored in the scope's deque. The caveat would be that the scope would need to be configured to hold some maximum closure size. Possibly this would mean that the scope would be generic over the size of the work items, or it could be a runtime check.