direction.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 XCENA Inc.
 
#pragma once
 
#include <cstddef>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace pxl
{
 
template <typename T>
class NDArray;
 
enum class ArgDir : unsigned char
{
    InOut,
    Input,
    Output,
};
 
struct ArgDirections
{
    // Empty means the legacy InOut policy for every argument.
    std::vector<ArgDir> args;
};
 
namespace detail
{
 
template <typename T>
struct KernelFunctionTraits;
 
template <typename R, typename... Args>
struct KernelFunctionTraits<R (*)(Args...)>
{
    static constexpr std::size_t arity = sizeof...(Args);
    using parameters = std::tuple<Args...>;
};
 
template <typename R, typename... Args>
struct KernelFunctionTraits<R (*)(Args...) noexcept> : KernelFunctionTraits<R (*)(Args...)>
{
};
 
template <auto Func>
inline constexpr std::size_t KernelArityV = KernelFunctionTraits<decltype(Func)>::arity;
 
template <auto Func, std::size_t Index>
using KernelParamT =
    std::tuple_element_t<Index, typename KernelFunctionTraits<decltype(Func)>::parameters>;
 
template <typename T, typename = void>
struct KernelParamDirection
{
    static constexpr ArgDir value = ArgDir::InOut;
};
 
template <typename T>
struct KernelParamDirection<
    T, std::void_t<typename T::element_type, typename T::kernel_buffer_direction>>
{
private:
    using direction_tag = typename T::kernel_buffer_direction;
    static constexpr bool reads = direction_tag::reads_existing_data;
    static constexpr bool writes = direction_tag::writes_result;
 
public:
    static_assert(reads || writes, "A directional kernel buffer needs read or write access");
    static constexpr ArgDir value = reads ? (writes ? ArgDir::InOut : ArgDir::Input)
                                          : ArgDir::Output;
};
 
template <typename T>
inline constexpr ArgDir KernelParamDirectionV =
    KernelParamDirection<std::remove_cv_t<std::remove_reference_t<T>>>::value;
 
template <typename T, typename = void>
struct IsDirectionalKernelParam : std::false_type
{
};
 
template <typename T>
struct IsDirectionalKernelParam<
    T, std::void_t<typename T::element_type, typename T::kernel_buffer_direction>>
    : std::true_type
{
};
 
template <typename T>
inline constexpr bool IsDirectionalKernelParamV =
    IsDirectionalKernelParam<std::remove_cv_t<std::remove_reference_t<T>>>::value;
 
template <typename T>
struct ActualKernelBuffer
{
    static constexpr bool valid = false;
    using element_type = void;
    using pointer = void*;
};
 
template <typename T>
struct ActualKernelBuffer<T*>
{
    static constexpr bool valid = true;
    using element_type = T;
    using pointer = T*;
};
 
template <typename T>
struct ActualKernelBuffer<pxl::NDArray<T>>
{
    static constexpr bool valid = true;
    using element_type = T;
    using pointer = T*;
};
 
template <typename Param, typename Actual, bool = IsDirectionalKernelParamV<Param>>
struct KernelArgumentCompatible : std::true_type
{
};
 
template <typename Param, typename Actual>
struct KernelArgumentCompatible<Param, Actual, true>
{
private:
    using parameter_type = std::remove_cv_t<std::remove_reference_t<Param>>;
    using parameter_element = typename parameter_type::element_type;
    using direction_tag = typename parameter_type::kernel_buffer_direction;
    using actual_buffer = ActualKernelBuffer<std::decay_t<Actual>>;
    using actual_element = typename actual_buffer::element_type;
    static constexpr bool writes = direction_tag::writes_result;
    using wire_element = std::conditional_t<writes, parameter_element,
                                            std::add_const_t<parameter_element>>;
 
public:
    static constexpr bool value =
        !std::is_reference_v<Param> &&
        actual_buffer::valid &&
        !std::is_void_v<std::remove_cv_t<parameter_element>> &&
        !std::is_function_v<parameter_element> &&
        (!writes || !std::is_const_v<parameter_element>) &&
        std::is_same_v<std::remove_cv_t<actual_element>,
                       std::remove_cv_t<parameter_element>> &&
        std::is_convertible_v<typename actual_buffer::pointer, wire_element*>;
};
 
template <typename Param, typename Actual>
inline constexpr bool KernelArgumentCompatibleV =
    KernelArgumentCompatible<Param, Actual>::value;
 
template <auto Func, typename ActualTuple, std::size_t... Indices>
constexpr bool KernelArgumentsCompatible(std::index_sequence<Indices...>)
{
    return (KernelArgumentCompatibleV<
                KernelParamT<Func, Indices>,
                std::tuple_element_t<Indices, ActualTuple>> &&
            ...);
}
 
template <auto Func, typename... Actual>
constexpr bool KernelArgumentsCompatible()
{
    if constexpr (sizeof...(Actual) != KernelArityV<Func>)
    {
        return false;
    }
    else
    {
        return KernelArgumentsCompatible<Func, std::tuple<Actual...>>(
            std::make_index_sequence<KernelArityV<Func>>{});
    }
}
 
template <auto Func, typename... Actual>
inline constexpr bool KernelArgumentsCompatibleV =
    KernelArgumentsCompatible<Func, Actual...>();
 
template <auto Func, typename... Actual>
using EnableCompatibleKernelArguments =
    std::enable_if_t<KernelArgumentsCompatibleV<Func, Actual...>, int>;
 
template <auto Func, std::size_t... Indices>
ArgDirections MakeKernelDirections(std::index_sequence<Indices...>)
{
    ArgDirections directions;
    if constexpr (((KernelParamDirectionV<KernelParamT<Func, Indices>> == ArgDir::InOut) && ...))
    {
        return directions;
    }
    directions.args = {KernelParamDirectionV<KernelParamT<Func, Indices>>...};
    return directions;
}
 
template <auto Func>
ArgDirections MakeKernelDirections()
{
    return MakeKernelDirections<Func>(std::make_index_sequence<KernelArityV<Func>>{});
}
 
}  // namespace detail
 
namespace impl
{
 
inline ArgDir directionAt(const std::vector<ArgDir>& directions, std::size_t index) noexcept
{
    return index < directions.size() ? directions[index] : ArgDir::InOut;
}
 
}  // namespace impl
 
}  // namespace pxl