compute.h

Go to the documentation of this file.

 
#ifndef NNGN_COMPUTE_COMPUTE_H
#define NNGN_COMPUTE_COMPUTE_H
 
#include <array>
#include <cassert>
#include <cstring>
#include <memory>
#include <numeric>
#include <ranges>
#include <span>
#include <string>
#include <string_view>
#include <vector>
 
#include "utils/concepts.h"
#include "utils/def.h"
#include "utils/ranges.h"
#include "utils/utils.h"
 
namespace nngn {
 
struct Compute {
    enum class Backend : u8 {
        PSEUDOCOMP,
        OPENCL_BACKEND,
    };
    enum class DeviceType : u8 {
        CPU = 1u << 0, GPU = 1u << 1
    };
    struct Version { u32 major = {}, minor = {}; };
    struct OpenCLParameters {
        Version version = {};
        bool debug = {};
        DeviceType preferred_device = DeviceType::CPU;
    };
    enum class Type : u8 {
        NONE,
        LOCAL,
        SCALAR_BEGIN,
        BYTE = SCALAR_BEGIN, INT, UINT, FLOAT,
        VECTOR_BEGIN,
        BYTEV = VECTOR_BEGIN, INTV, UINTV, FLOATV,
        VECTOR_END = FLOATV,
        DATA,
        HANDLE_BEGIN,
        BUFFER = HANDLE_BEGIN, IMAGE, SAMPLER,
        HANDLE_END = SAMPLER,
        N,
    };
    enum MemFlag : u8 {
        READ_WRITE = 1u << 0, WRITE_ONLY = 1u << 1, READ_ONLY = 1u << 2,
    };
    enum ExecFlag : u8 {
        BLOCKING = 1u << 0,
    };
    enum Limit : u8 {
        COMPUTE_UNITS, WORK_GROUP_SIZE, LOCAL_MEMORY, N,
    };
    enum ProfInfo : u8 {
        QUEUED = 1u << 0, SUBMIT = 1u << 1, START = 1u << 2, END = 1u << 3,
        PROF_INFO_MAX = END, PROF_INFO_ALL = QUEUED | SUBMIT | START | END,
    };
    struct Handle {
        u32 id = {};
        constexpr explicit operator bool() const { return this->id; }
    };
    struct Buffer : Handle { static constexpr auto type = Type::BUFFER; };
    struct Image : Handle { static constexpr auto type = Type::IMAGE; };
    struct Sampler : Handle { static constexpr auto type = Type::SAMPLER; };
    struct Program : Handle {};
    struct Kernel : Handle {};
    struct Event;
    struct Events {
        std::size_t n_wait = {};
        const Event *const *wait_list = {};
        Event *const *events = {};
    };
    struct DataArg {
        std::size_t s = {};
        const std::byte *p = {};
        template<typename T>
        explicit constexpr DataArg(const T *t) : s(sizeof(T)), p(as_bytes(t)) {}
        constexpr auto begin() const { return this->p; }
        constexpr auto end() const { return this->p + this->s; }
    };
    template<typename T> static constexpr Type arg_type = Type::NONE;
    static constexpr bool is_vector_type(Type t);
    static constexpr bool is_handle_type(Type t);
    template<Type t> static constexpr Type to_vector_type();
    static std::unique_ptr<Compute> create(
        Backend b, const void *params = nullptr);
    NNGN_VIRTUAL(Compute)
    virtual bool init() = 0;
    virtual size_t n_platforms() const = 0;
    virtual size_t n_devices() const = 0;
    virtual void get_limits(u64 *p) const = 0;
    std::vector<u64> get_limits() const;
    virtual std::string platform_name() const = 0;
    virtual std::string device_name() const = 0;
    virtual Buffer create_buffer(
        MemFlag flags, std::size_t n, const std::byte *p) = 0;
    virtual bool read_buffer(
        Buffer b, std::size_t off, std::size_t n, std::byte *p,
        Events events) const = 0;
    virtual bool fill_buffer(
        Buffer b, std::size_t off, std::size_t n, std::byte v,
        Events events) const = 0;
    virtual bool fill_buffer(
        Buffer b, std::size_t off, std::size_t n,
        std::size_t pattern_size, const std::byte *p,
        Events events) const = 0;
    virtual bool write_buffer(
        Buffer b, std::size_t off, std::size_t n, const std::byte *p,
        Events events) const = 0;
    virtual bool write_buffer_rect(
        Buffer b,
        std::array<std::size_t, 3> buffer_origin,
        std::array<std::size_t, 3> host_origin,
        std::array<std::size_t, 3> region,
        std::size_t buffer_row_pitch, std::size_t buffer_slice_pitch,
        std::size_t host_row_pitch, std::size_t host_slice_pitch,
        const std::byte *p, Events events) const = 0;
    virtual void *map_buffer(
        Buffer b, MemFlag flags, std::size_t off, std::size_t n,
        Events events) const = 0;
    virtual bool unmap_buffer(Buffer b, void *p, Events events) const = 0;
    template<typename ...Ts>
    bool write_struct(Buffer b, Events events, Ts &&...ts) const;
    virtual bool release_buffer(Buffer b) = 0;
    virtual Image create_image(
        Type type, std::size_t w, std::size_t h, MemFlag flags,
        const std::byte *p) = 0;
    virtual bool read_image(
        Image i, std::size_t w, std::size_t h, std::byte *p, Events events)
        const = 0;
    virtual bool fill_image(
        Image i, std::size_t w, std::size_t h, const void *v,
        Events events) const = 0;
    virtual bool release_image(Image i) = 0;
    virtual Sampler create_sampler() = 0;
    virtual bool release_sampler(Sampler s) = 0;
    virtual Program create_program(std::string_view src, const char *opts) = 0;
    virtual bool release_program(Program p) = 0;
    virtual Kernel create_kernel(
        Program program, const char *func,
        std::size_t len, const Type *types,
        const std::size_t *sizes, const std::byte *const *data,
        Events events) = 0;
    virtual bool release_kernel(Kernel k) = 0;
    virtual std::size_t n_events(std::size_t n, const Type *types) const = 0;
    virtual bool prof_info(
        ProfInfo info, std::size_t n, const Event *const *events, u64 *out)
        const = 0;
    virtual bool wait(std::size_t n, const Event *const *v) const = 0;
    virtual bool release_events(std::size_t n, const Event *const *v) const = 0;
    virtual bool execute(
        Kernel kernel, ExecFlag flags,
        u32 n_dim, const std::size_t *global_size,
        const std::size_t *local_size, Events events) const = 0;
    virtual bool execute(
        Program program, const std::string &func, ExecFlag flags,
        u32 n_dim, const std::size_t *global_size,
        const std::size_t *local_size, std::size_t len, const Type *types,
        const std::size_t *sizes, const std::byte *const *data,
        Events events) const = 0;
    template<typename ...Ts>
    Kernel create_kernel(
        Program program, const char *func, Events events, Ts &&...ts);
    template<typename ...Ts>
    bool execute(
        Program program, const std::string &func, ExecFlag flags,
        u32 n_dim, const std::size_t *global_size,
        const std::size_t *local_size, Events events, Ts &&...ts);
};
 
template<Compute::Backend>
std::unique_ptr<Compute> compute_create_backend(const void *params);
 
template<typename T>
static constexpr Compute::Type arg_type = Compute::Type::NONE;
template<>
constexpr auto Compute::arg_type<std::byte> = Compute::Type::BYTE;
template<>
constexpr auto Compute::arg_type<i32> = Compute::Type::INT;
template<>
constexpr auto Compute::arg_type<u32> = Compute::Type::UINT;
template<>
constexpr auto Compute::arg_type<float> = Compute::Type::FLOAT;
 
template<std::derived_from<Compute::Handle> T>
constexpr auto Compute::arg_type<T> = T::type;
 
template<std::ranges::range T>
constexpr auto Compute::arg_type<T> =
    Compute::to_vector_type<Compute::arg_type<std::ranges::range_value_t<T>>>();
 
inline constexpr bool Compute::is_vector_type(Type t)
    { return Type::VECTOR_BEGIN <= t && t <= Type::VECTOR_END; }
inline constexpr bool Compute::is_handle_type(Type t)
    { return Type::HANDLE_BEGIN <= t && t <= Type::HANDLE_END; }
 
namespace detail {
 
inline auto arg_size(const std::byte&) { return sizeof(std::byte); }
inline auto arg_ptr(const std::byte &b) { return as_bytes(&b); }
 
inline auto arg_size(const Compute::Handle &t) { return sizeof(t.id); }
inline auto arg_ptr(const Compute::Handle &t) { return as_bytes(&t.id); }
 
template<arithmetic T> auto arg_size(const T&) { return sizeof(T); }
auto arg_ptr(const arithmetic auto &t) { return as_bytes(&t); }
 
auto arg_size(const std::ranges::sized_range auto &r) {
    // TODO https://bugs.llvm.org/show_bug.cgi?id=39663
    // return std::span{r}.size_bytes();
    auto s = std::span{r};
    return s.size_bytes();
}
auto arg_ptr(const std::ranges::range auto &r)
    { return std::as_bytes(std::span{r}).data(); }
 
}
 
template<Compute::Type t>
inline constexpr auto Compute::to_vector_type() -> Type {
    using T = std::underlying_type_t<Compute::Type>;
    if constexpr(Type::SCALAR_BEGIN <= t && t < Type::VECTOR_BEGIN)
        return static_cast<Compute::Type>(
            static_cast<T>(Compute::Type::VECTOR_BEGIN)
            + static_cast<T>(t)
            - static_cast<T>(Compute::Type::BYTE));
    else
        static_assert(nngn::always_false<std::integral_constant<Type, t>>{});
}
 
inline std::vector<u64> Compute::get_limits() const {
    std::vector<u64> ret(Limit::N);
    this->get_limits(ret.data());
    return ret;
}
 
template<typename ...Ts>
bool Compute::write_struct(Buffer b, Events events, Ts &&...ts) const {
    const std::array sizes = {detail::arg_size(ts)...};
    std::vector<std::byte> data(std::reduce(cbegin(sizes), cend(sizes)));
    assert(!data.empty());
    auto copy = [p = data.data(), s = sizes.data()](const auto &x) mutable
        { std::memcpy(p, detail::arg_ptr(x), *s); p += *s++; };
    (..., copy(ts));
    return this->write_buffer(b, 0, data.size(), data.data(), events);
}
 
template<typename ...Ts>
auto Compute::create_kernel(
    Program program, const char *func, Events events, Ts &&...ts
) -> Kernel {
    constexpr auto n = sizeof...(Ts);
    std::array<Type, n> types = {};
    std::array<std::size_t, n> sizes = {};
    std::array<const std::byte*, n> data = {};
    auto f =
        [tp = types.data(), sp = sizes.data(), dp = data.data()]
        <typename T>(const T &t) mutable
    {
        *tp++ = Compute::arg_type<T>;
        *sp++ = detail::arg_size(t);
        *dp++ = as_bytes(detail::arg_ptr(t));
    };
    (..., f(ts));
    return this->create_kernel(
        program, func, n, types.data(), sizes.data(), data.data(), events);
}
 
template<typename ...Ts>
bool Compute::execute(
        Program program, const std::string &func, ExecFlag flags,
        u32 n_dim, const std::size_t *global_size,
        const std::size_t *local_size, Events events, Ts &&...ts) {
    constexpr auto n = sizeof...(Ts);
    std::array<Type, n> types = {};
    std::array<std::size_t, n> sizes = {};
    std::array<const std::byte*, n> data = {};
    auto f =
        [tp = types.data(), sp = sizes.data(), dp = data.data()]
        <typename T>(const T &t) mutable
    {
        *tp++ = Compute::arg_type<T>;
        *sp++ = detail::arg_size(t);
        *dp++ = as_bytes(detail::arg_ptr(t));
    };
    (..., f(ts));
    return this->execute(
        program, func, flags, n_dim, global_size, local_size,
        n, types.data(), sizes.data(), data.data(), events);
}
 
}
 
#endif