kernel_call_distrib.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once
 
#include "shambase/DistributedData.hpp"
#include "shambackends/kernel_call.hpp"
 
namespace sham {
 
    template<class... Targ>
    struct DDMultiRef {
        using storage_t = std::tuple<Targ &...>;
 
        storage_t storage;
 
        DDMultiRef(Targ &...arg) : storage(arg...) {}
 
        auto get(u64 id) {
            shamlog_debug_ln(
                "kern call",
                "called DDMultiRef.get, id =",
                id,
                SourceLocation{}.format_one_line_func());
            return std::apply(
                [id](auto &...args) {
                    return sham::MultiRef{args.get(id)...};
                },
                storage);
        }
    };
 
    template<class index_t, class RefIn, class RefOut, class Functor>
    inline void distributed_data_kernel_call(
        sham::DeviceScheduler_ptr dev_sched,
        RefIn in,
        RefOut in_out,
        const shambase::DistributedData<index_t> &thread_counts,
        Functor &&func) {
 
        auto mrefs_in
            = thread_counts.template map<decltype(in.get(0))>([&](u64 id, const index_t &n) {
                  shamlog_debug_ln("kern call", "build multi ref in for patch", id);
                  return in.get(id);
              });
 
        auto mrefs_in_out
            = thread_counts.template map<decltype(in_out.get(0))>([&](u64 id, const index_t &n) {
                  shamlog_debug_ln("kern call", "build multi ref in_out for patch", id);
                  return in_out.get(id);
              });
 
        thread_counts.for_each([&](u64 id, const index_t &n) {
            shamlog_debug_ln(
                "kern call", "calling sham::kernel_call on patch", id, " thread count", n);
            sham::kernel_call(
                dev_sched->get_queue(),
                mrefs_in.get(id),
                mrefs_in_out.get(id),
                n,
                std::forward<Functor>(func));
        });
    }
 
    // version where one supplies a kernel generator in the form of [&](sycl::handler &cgh) { ... }
    template<class index_t, class RefIn, class RefOut, class Functor>
    inline void distributed_data_kernel_call_hndl(
        sham::DeviceScheduler_ptr dev_sched,
        RefIn in,
        RefOut in_out,
        const shambase::DistributedData<index_t> &thread_counts,
        Functor &&kernel_gen) {
 
        auto mrefs_in
            = thread_counts.template map<decltype(in.get(0))>([&](u64 id, const index_t &n) {
                  shamlog_debug_ln("kern call", "build multi ref in for patch", id);
                  return in.get(id);
              });
 
        auto mrefs_in_out
            = thread_counts.template map<decltype(in_out.get(0))>([&](u64 id, const index_t &n) {
                  shamlog_debug_ln("kern call", "build multi ref in_out for patch", id);
                  return in_out.get(id);
              });
 
        thread_counts.for_each([&](u64 id, const index_t &n) {
            shamlog_debug_ln(
                "kern call", "calling sham::kernel_call_hndl on patch", id, " thread count", n);
 
            sham::kernel_call_hndl(
                dev_sched->get_queue(),
                mrefs_in.get(id),
                mrefs_in_out.get(id),
                n,
                std::forward<Functor>(kernel_gen));
        });
    }
 
} // namespace sham