kernel_call.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/optional.hpp"
#include "shambackends/DeviceBuffer.hpp"
#include <functional>
#include <optional>
 
namespace sham {
 
    namespace details {
 
        template<class T>
        const T *read_access_optional(
            shambase::opt_ref<sham::DeviceBuffer<T>> buffer, sham::EventList &depends_list) {
            if (!buffer.has_value()) {
                return nullptr;
            } else {
                return buffer.value().get().get_read_access(depends_list);
            }
        }
 
        template<class T>
        T *write_access_optional(
            shambase::opt_ref<sham::DeviceBuffer<T>> buffer, sham::EventList &depends_list) {
            if (!buffer.has_value()) {
                return nullptr;
            } else {
                return buffer.value().get().get_write_access(depends_list);
            }
        }
 
        template<class T>
        void complete_state_optional(sycl::event e, shambase::opt_ref<T> buffer) {
            if (buffer.has_value()) {
                buffer.value().get().complete_event_state(e);
            }
        }
 
        template<class Obj>
        inline auto get_read_access(Obj &o, sham::EventList &depends_list) {
            return o.get_read_access(depends_list);
        }
 
        template<class Obj>
        inline auto get_write_access(Obj &o, sham::EventList &depends_list) {
            return o.get_write_access(depends_list);
        }
        template<class Obj>
        inline auto complete_event_state(Obj &o, sycl::event e) {
            return o.complete_event_state(e);
        }
 
        template<class Obj>
        inline auto get_read_access(std::reference_wrapper<Obj> &o, sham::EventList &depends_list) {
            return o.get().get_read_access(depends_list);
        }
 
        template<class Obj>
        inline auto get_write_access(
            std::reference_wrapper<Obj> &o, sham::EventList &depends_list) {
            return o.get().get_write_access(depends_list);
        }
        template<class Obj>
        inline auto complete_event_state(std::reference_wrapper<Obj> &o, sycl::event e) {
            return o.get().complete_event_state(e);
        }
 
    } // namespace details
 
    template<class T>
    shambase::opt_ref<T> to_opt_ref(T &t) {
        return t;
    }
 
    template<class T>
    auto empty_buf_ref() {
        return shambase::opt_ref<sham::DeviceBuffer<T>>{};
    }
 
    template<class... Targ>
    struct MultiRefOpt {
        using storage_t = std::tuple<shambase::opt_ref<Targ>...>;
 
        storage_t storage;
 
        MultiRefOpt(shambase::opt_ref<Targ>... arg) : storage(arg...) {}
 
        auto get_read_access(sham::EventList &depends_list) {
            __shamrock_stack_entry();
            return std::apply(
                [&](auto &...__a) {
                    return std::tuple(details::read_access_optional(__a, depends_list)...);
                },
                storage);
        }
        auto get_write_access(sham::EventList &depends_list) {
            __shamrock_stack_entry();
            return std::apply(
                [&](auto &...__a) {
                    return std::tuple(details::write_access_optional(__a, depends_list)...);
                },
                storage);
        }
 
        void complete_event_state(sycl::event e) {
            __shamrock_stack_entry();
            std::apply(
                [&](auto &...__in) {
                    ((details::complete_state_optional(e, __in)), ...);
                },
                storage);
        }
    };
 
    namespace details {
        template<class T>
        struct mapper {
            using type = T;
        };
 
        template<class T>
        struct mapper<shambase::opt_ref<T>> {
            using type = T;
        };
    } // namespace details
 
    template<class... Targ>
    MultiRefOpt(Targ... arg) -> MultiRefOpt<typename details::mapper<Targ>::type...>;
 
    template<class... Targ>
    struct MultiRef {
        using storage_t = std::tuple<Targ &...>;
 
        storage_t storage;
 
        MultiRef(Targ &...arg) : storage(arg...) {}
 
        auto get_read_access(sham::EventList &depends_list) {
            __shamrock_stack_entry();
            return std::apply(
                [&](auto &...__a) {
                    return std::tuple(details::get_read_access(__a, depends_list)...);
                },
                storage);
        }
 
        auto get_write_access(sham::EventList &depends_list) {
            __shamrock_stack_entry();
            return std::apply(
                [&](auto &...__a) {
                    return std::tuple(details::get_write_access(__a, depends_list)...);
                },
                storage);
        }
 
        void complete_event_state(sycl::event e) {
            __shamrock_stack_entry();
            std::apply(
                [&](auto &...__in) {
                    ((details::complete_event_state(__in, e)), ...);
                },
                storage);
        }
    };
 
    namespace details {
 
        template<class index_t, class RefIn, class RefOut, class Functor>
        void typed_index_kernel_call_lambda(
            sham::DeviceQueue &q,
            RefIn in,
            RefOut in_out,
            index_t n,
            Functor &&kernel_gen,
            SourceLocation &&callsite = SourceLocation{}) {
 
            __shamrock_stack_entry_with_callsite(callsite);
 
            if (n == 0) {
                shambase::throw_with_loc<std::runtime_error>("kernel call with : n == 0");
            }
 
            sham::EventList depends_list;
 
            auto acc_in     = in.get_read_access(depends_list);
            auto acc_in_out = in_out.get_write_access(depends_list);
 
            sycl::event e;
 
            // unpack the tuples of accessors
            std::apply(
                [&](auto &...__acc_in) {
                    std::apply(
                        [&](auto &...__acc_in_out) {
                            // submit the kernel generated by the functor
                            e = q.submit(depends_list, kernel_gen(n, __acc_in..., __acc_in_out...));
                        },
                        acc_in_out);
                },
                acc_in);
 
            in.complete_event_state(e);
            in_out.complete_event_state(e);
        }
 
        template<class index_t, class RefIn, class RefOut, class Functor>
        void typed_index_kernel_call(
            sham::DeviceQueue &q,
            RefIn in,
            RefOut in_out,
            index_t n,
            Functor &&func,
            SourceLocation &&callsite = SourceLocation{}) {
 
            __shamrock_log_callsite(callsite);
 
            typed_index_kernel_call_lambda(
                q,
                in,
                in_out,
                n,
                [func
                 = std::forward<Functor>(func)](u32 n, auto... __acc_in, auto... __acc_in_out) {
                    return [=](sycl::handler &cgh) {
                        cgh.parallel_for(sycl::range<1>{n}, [=](sycl::item<1> item) {
                            shambase::check_functor_signature_deduce<void>(
                                func, index_t(item.get_linear_id()), __acc_in..., __acc_in_out...);
 
                            func(index_t(item.get_linear_id()), __acc_in..., __acc_in_out...);
                        });
                    };
                });
        }
    } // namespace details
 
    template<class RefIn, class RefOut, class Functor>
    void kernel_call(
        sham::DeviceQueue &q,
        RefIn in,
        RefOut in_out,
        u32 n,
        Functor &&func,
        SourceLocation &&callsite = SourceLocation{}) {
 
        __shamrock_log_callsite(callsite);
 
        details::typed_index_kernel_call<u32, RefIn, RefOut>(
            q, in, in_out, n, std::forward<Functor>(func));
    }
 
    template<class RefIn, class RefOut, class Functor>
    void kernel_call_u64(
        sham::DeviceQueue &q,
        RefIn in,
        RefOut in_out,
        u64 n,
        Functor &&func,
        SourceLocation &&callsite = SourceLocation{}) {
 
        __shamrock_log_callsite(callsite);
 
        details::typed_index_kernel_call<u64, RefIn, RefOut>(
            q, in, in_out, n, std::forward<Functor>(func));
    }
 
    // version where one supplies a kernel generator in the form of [&](sycl::handler &cgh) { ... }
    template<class RefIn, class RefOut, class Functor>
    void kernel_call_hndl(
        sham::DeviceQueue &q,
        RefIn in,
        RefOut in_out,
        u32 n,
        Functor &&kernel_gen,
        SourceLocation &&callsite = SourceLocation{}) {
 
        __shamrock_log_callsite(callsite);
 
        details::typed_index_kernel_call_lambda<u32, RefIn, RefOut>(
            q, in, in_out, n, std::forward<Functor>(kernel_gen));
    }
 
    template<class RefIn, class RefOut, class Functor>
    void kernel_call_hndl_u64(
        sham::DeviceQueue &q,
        RefIn in,
        RefOut in_out,
        u64 n,
        Functor &&kernel_gen,
        SourceLocation &&callsite = SourceLocation{}) {
 
        __shamrock_log_callsite(callsite);
 
        details::typed_index_kernel_call_lambda<u64, RefIn, RefOut>(
            q, in, in_out, n, std::forward<Functor>(kernel_gen));
    }
 
} // namespace sham