libcruft-util/cruft/util/job/dispatch.hpp

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * Copyright 2018 Danny Robson <danny@nerdcruft.net>
 */

#pragma once

#include "queue.hpp"

#include "../extent.hpp"
#include "../region.hpp"
#include "../concepts.hpp"

#include <vector>

namespace cruft::job {
    /// call a function across all elements of a container using the supplied
    /// job queue.
    ///
    /// threads will have work sizes dictated by a supplied extent.
    ///
    /// returns a cookie that will block at destruction until all jobs have
    /// completed. it will take ownership of an forwarding-reference function
    /// if one is supplied.
    ///
    /// TODO: extend to 1d and 3d
    ///
    /// \param area  The total size of the work unit
    /// \param chunk The size of the subdivided units to be send to threads.
    template <
        typename ContainerT,
        typename FunctionT,
        typename ...ArgsT,
        size_t DimensionV
    >
    requires requires (
        ContainerT t,
        cruft::point<DimensionV,int> p,
        FunctionT f,
        ArgsT ...args
    ) {
        { t[p] = std::invoke (f, p, args...) };
    }
    auto
    dispatch (
        cruft::job::queue &q,
        ContainerT &data,
        cruft::region<DimensionV,int> total_area,
        cruft::extent<DimensionV,int> chunk,
        FunctionT &&func,
        ArgsT ...args
    ) {
        auto chunked_func = [&func] (
            ContainerT &inner_data,
            auto param,
            ArgsT... inner_args
        ) {
            for (cruft::point<DimensionV,int> p: param)
                inner_data[p] = std::invoke (func, p, inner_args...);
        };

        std::vector<queue::cookie> cookies;

        extent2i const count = (total_area.e + chunk - 1) / chunk;

        for (auto index: count.step ()) {
            cruft::point<DimensionV,int> base = index * chunk + total_area.p.template as<cruft::vector> ();
            region<DimensionV,int> const oversized_work { base, chunk };
            auto const constrained_work = intersection (oversized_work, total_area);
            cookies.push_back (
                q.submit (
                    chunked_func,
                    std::ref (data),
                    constrained_work.step (),
                    args...
                )
            );
        }

        return std::tuple (std::forward<FunctionT> (func), std::move (cookies));
    }
}