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 "shambase/type_traits.hpp"
#include "shambackends/DeviceBuffer.hpp"
#include "shambackends/kernel_call/MultiRef.hpp"
#include "shambackends/kernel_call/MultiRefOpt.hpp"
#include "shambackends/kernel_call/buffer_access_utils.hpp"
#include <functional>
#include <optional>
 
namespace sham {
 
    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);
        }

        // Helper types for kernel signature checking (and clean error messages)

        template<class index_t, class RefIn, class RefOut>
        struct kernel_gen_args {
 
            using args_types = decltype(std::tuple_cat(
                std::tuple<index_t>{},
 
                std::declval<decltype(std::declval<RefIn>().get_read_access(
                    std::declval<sham::EventList &>()))>(),
 
                std::declval<decltype(std::declval<RefOut>().get_write_access(
                    std::declval<sham::EventList &>()))>()));
        };

        template<typename Tuple>
        struct tuple_to_signature;
 
        template<typename... Ts>
        struct tuple_to_signature<std::tuple<Ts...>> {
            using type = void(Ts...);
        };

        template<typename Signature>
        struct expected_kernel_signature;
 
        template<typename Ret, typename... Ts>
        struct expected_kernel_signature<Ret(Ts...)> {};

        template<class index_t, class RefIn, class RefOut>
        using kernel_expected_signature = expected_kernel_signature<typename tuple_to_signature<
            typename kernel_gen_args<index_t, RefIn, RefOut>::args_types>::type>;

        template<typename F, typename Signature>
        struct is_kernel_invocable;
 
        template<typename F, typename Ret, typename... Ts>
        struct is_kernel_invocable<F, expected_kernel_signature<Ret(Ts...)>>
            : std::bool_constant<std::invocable<F, Ts...>> {};

        template<typename F, typename Signature>
        concept kernel_invocable = is_kernel_invocable<F, Signature>::value;

        template<typename Signature>
        struct expected_kernel_fn_type;
 
        template<typename Ret, typename... Ts>
        struct expected_kernel_fn_type<expected_kernel_signature<Ret(Ts...)>> {
            using type = Ret(Ts...);
        };

        template<class ExpectedFnSignature>
        inline constexpr bool matches_expected_kernel_signature = false;

        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{}) {
 
            using expected_sig = kernel_expected_signature<index_t, RefIn, RefOut>;
 
            if constexpr (is_kernel_invocable<Functor, expected_sig>::value) {
                __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) {
                                func(index_t(item.get_linear_id()), __acc_in..., __acc_in_out...);
                            });
                        };
                    });
            } else {
                static_assert(
                    matches_expected_kernel_signature<
                        typename expected_kernel_fn_type<expected_sig>::type>,
                    "Kernel functor is not invocable with the expected signature; see "
                    "matches_expected_kernel_signature<void(...)> for the required argument list.");
            }
        }
    } // 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