stage_builder.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
// Copyright 2024 XCENA Inc.
 
#pragma once
 
#include <atomic>
#include <cstdint>
#include <future>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <vector>
#include <xpti/xpti.hpp>
 
#include "pxl/direction.hpp"
#include "pxl/launch_result.hpp"
#include "pxl/map.hpp"
#include "pxl/module.hpp"
#include "pxl/ndarray.hpp"
#include "pxl/stream.hpp"
#include "pxl/task_count.hpp"
#include "pxl/type.hpp"
 
namespace pxl
{
 
namespace impl
{
class DeviceHandle;
 
struct KnobConfig
{
    std::optional<uint32_t> batchSize;
    std::optional<uint32_t> clusterBitmap;
    std::optional<pxl::LocalityMode> locality;
 
    std::optional<pxl::Map::CompletionCallback> onComplete;
    void* onCompleteArg = nullptr;
    std::optional<pxl::Map::MessageCallback> onMessage;
    void* onMessageArg = nullptr;
    std::optional<pxl::Map::ErrorCallback> onError;
    void* onErrorArg = nullptr;
 
    // stream knob — 3-state semantics:
    //   unset (nullopt):   leave Map's default stream untouched
    //   set, non-null:     route through user-supplied Stream
    //   set, null:         explicitly reset Map back to its default stream
    // (Without this distinction, stream(nullptr) was a silent no-op and there
    //  was no way to opt back out of a previously set user stream.)
    std::optional<pxl::Stream*> stream;  // nullopt = leave default; non-null = user stream; null = reset to default
    std::optional<xpti::ProfileConfig> profileConfig;
};
 
struct StageEntry
{
    std::string kernelName;
    uint32_t taskCount = 0;
    std::vector<ArgInfo_t> argList;
    std::vector<std::shared_ptr<void>> constStorage;
    std::vector<std::pair<void*, size_t>> syncTargets;
    KnobConfig knobs;
};
 
struct StageDirectionMetadata
{
    ArgDirections directions;
 
    // The shared_ptr control block owns the directions; there is no pointee.
    void operator()(void*) const noexcept
    {
    }
};
 
inline void attachStageDirections(StageEntry& entry, ArgDirections directions)
{
    if (directions.args.empty())
    {
        return;
    }
    entry.constStorage.emplace_back(nullptr, StageDirectionMetadata{std::move(directions)});
}
 
inline const ArgDirections& stageDirections(const StageEntry& entry)
{
    static const ArgDirections LegacyDirections;
    for (auto it = entry.constStorage.rbegin(); it != entry.constStorage.rend(); ++it)
    {
        if (const auto* metadata = std::get_deleter<StageDirectionMetadata>(*it))
        {
            return metadata->directions;
        }
    }
    return LegacyDirections;
}
 
struct StageGroup
{
    std::vector<StageEntry> entries;
};
 
}  // namespace impl
 
class StageGroupBuilder;
 
// [[nodiscard]] guards against building a stage chain and forgetting `.run()`.
class [[nodiscard]] StageBuilder
{
public:
    StageBuilder(std::shared_ptr<impl::DeviceHandle> device);
    ~StageBuilder();
 
    StageBuilder(StageBuilder&&) noexcept;
    StageBuilder& operator=(StageBuilder&&) noexcept;
    StageBuilder(const StageBuilder&) = delete;
    StageBuilder& operator=(const StageBuilder&) = delete;
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    StageBuilder& stage(uint32_t taskCount, Args&&... args)
    {
        const char* kernelName = nullptr;
#if defined(__PXCC_ANALYSIS__)
        kernelName = "__pxcc_placeholder__";
#else
        static_assert(sizeof(pxl::KernelTraits<Func>) > 0,
                      "Kernel function not found in KernelTraits — was it compiled with pxcc?");
        kernelName = pxl::KernelTraits<Func>::name;
#endif
 
        impl::StageEntry entry;
        entry.kernelName = kernelName;
        entry.taskCount = taskCount;
        entry.argList.reserve(sizeof...(Args));
        if constexpr (sizeof...(Args) > 0)
        {
            buildArgList(entry.argList, entry.constStorage, entry.syncTargets, std::forward<Args>(args)...);
        }
        impl::attachStageDirections(entry, detail::MakeKernelDirections<Func>());
 
        impl::StageGroup group;
        group.entries.push_back(std::move(entry));
        stages_.push_back(std::move(group));
        return *this;
    }
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<StageBuilder> stage(Args&&... args) &&;
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<StageBuilder> stage(Args&&... args) & = delete;
 
    StageGroupBuilder stage();
    StageBuilder& reserve(uint32_t n);
 
    StageBuilder& tasks(uint32_t taskCount);
 
    StageBuilder& batchSize(uint32_t n);
 
    StageBuilder& clusterBitmap(uint32_t bitmap);
 
    StageBuilder& locality(LocalityMode mode);
 
    StageBuilder& onComplete(Map::CompletionCallback cb, void* arg = nullptr);
 
    StageBuilder& onMessage(Map::MessageCallback cb, void* arg = nullptr);
 
    StageBuilder& onError(Map::ErrorCallback cb, void* arg = nullptr);
 
    StageBuilder& stream(Stream* s);
 
    StageBuilder& enableProfile(const xpti::ProfileConfig& config);
 
    LaunchResult run();
    std::future<LaunchResult> runAsync();
 
private:
    template <typename T, typename... Rest>
    void buildArgList(std::vector<ArgInfo_t>& argList, std::vector<std::shared_ptr<void>>& storage,
                      std::vector<std::pair<void*, size_t>>& syncTargets, T&& arg, Rest&&... rest)
    {
        addArg(argList, storage, syncTargets, std::forward<T>(arg));
        if constexpr (sizeof...(rest) > 0)
        {
            buildArgList(argList, storage, syncTargets, std::forward<Rest>(rest)...);
        }
    }
 
    template <typename T>
    std::enable_if_t<std::is_pointer_v<std::decay_t<T>>>
    addArg(std::vector<ArgInfo_t>& argList, std::vector<std::shared_ptr<void>>&,
           std::vector<std::pair<void*, size_t>>&, T&& ptr)
    {
        argList.emplace_back(ArgInfo_t{ArgType::DeviceMemory, static_cast<const void*>(ptr), 0});
    }
 
    template <typename T>
    std::enable_if_t<std::is_arithmetic_v<std::decay_t<T>> && !std::is_pointer_v<std::decay_t<T>> &&
                     !std::is_same_v<std::decay_t<T>, bool>>
    addArg(std::vector<ArgInfo_t>& argList, std::vector<std::shared_ptr<void>>& storage,
           std::vector<std::pair<void*, size_t>>&, T&& value)
    {
        auto stored = std::make_shared<std::decay_t<T>>(std::forward<T>(value));
        argList.emplace_back(ArgInfo_t{ArgType::Constant, stored.get(), sizeof(std::decay_t<T>)});
        storage.push_back(stored);
    }
 
    template <typename T>
    void addArg(std::vector<ArgInfo_t>& argList, std::vector<std::shared_ptr<void>>& storage,
                std::vector<std::pair<void*, size_t>>& syncTargets,
                const pxl::NDArray<T>& ndarrayArg)
    {
        // Copy the NDArray descriptor into heap storage so we don't stash the
        // caller's stack address into argList. The caller's NDArray may be a
        // function-local rvalue that dies before runAsync's worker reads it;
        // without this copy the worker would dereference a dead stack frame.
        // The underlying data buffer (the device-memory pointer behind data())
        // must remain valid until execution completes.
        auto storedNdarray = std::make_shared<pxl::NDArray<T>>(ndarrayArg);
        argList.emplace_back(ArgInfo_t{ArgType::NDArray, storedNdarray.get(), sizeof(pxl::NDArray<T>)});
        storage.push_back(storedNdarray);
        const auto byteSize = ndarrayArg.size();
        syncTargets.emplace_back(const_cast<void*>(static_cast<const void*>(ndarrayArg.data())),
                                 byteSize > 0 ? static_cast<size_t>(byteSize) : 0);
    }
 
    std::shared_ptr<impl::DeviceHandle> device_;
    std::vector<impl::StageGroup> stages_;
    uint32_t reserveCount_ = 0;
    // Atomic: same rationale as LaunchBuilder::started_ — guard against
    // accidental cross-thread reuse, return the documented "already called"
    // error rather than UB.
    std::atomic<bool> started_{false};
 
    StageGroupBuilder takeStageGroup() &&;
 
    friend class StageGroupBuilder;
    template <typename>
    friend class detail::TaskCountConfigured;
};
 
class [[nodiscard]] StageGroupBuilder
{
public:
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    StageGroupBuilder& add(uint32_t taskCount, Args&&... args)
    {
        const char* kernelName = nullptr;
#if defined(__PXCC_ANALYSIS__)
        kernelName = "__pxcc_placeholder__";
#else
        static_assert(sizeof(pxl::KernelTraits<Func>) > 0,
                      "Kernel function not found in KernelTraits — was it compiled with pxcc?");
        kernelName = pxl::KernelTraits<Func>::name;
#endif
 
        impl::StageEntry entry;
        entry.kernelName = kernelName;
        entry.taskCount = taskCount;
        entry.argList.reserve(sizeof...(Args));
        if constexpr (sizeof...(Args) > 0)
        {
            parent_.buildArgList(entry.argList, entry.constStorage, entry.syncTargets,
                                 std::forward<Args>(args)...);
        }
        impl::attachStageDirections(entry, detail::MakeKernelDirections<Func>());
        currentGroup_.entries.push_back(std::move(entry));
        return *this;
    }
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<StageGroupBuilder> add(Args&&... args) &&;
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<StageGroupBuilder> add(Args&&... args) & = delete;
 
    StageGroupBuilder& tasks(uint32_t taskCount);
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    StageBuilder& stage(uint32_t taskCount, Args&&... args)
    {
        flush();
        return parent_.stage<Func>(taskCount, std::forward<Args>(args)...);
    }
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<StageBuilder> stage(Args&&... args) &&;
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<StageBuilder> stage(Args&&... args) & = delete;
 
    StageGroupBuilder stage()
    {
        flush();
        return parent_.stage();
    }
 
    LaunchResult run()
    {
        flush();
        return parent_.run();
    }
 
    std::future<LaunchResult> runAsync()
    {
        flush();
        return parent_.runAsync();
    }
 
    StageGroupBuilder& reserve(uint32_t n)
    {
        parent_.reserve(n);
        return *this;
    }
 
    StageGroupBuilder& batchSize(uint32_t n);
 
    StageGroupBuilder& clusterBitmap(uint32_t bitmap);
 
    StageGroupBuilder& locality(LocalityMode mode);
 
    StageGroupBuilder& onComplete(Map::CompletionCallback cb, void* arg = nullptr);
 
    StageGroupBuilder& onMessage(Map::MessageCallback cb, void* arg = nullptr);
 
    StageGroupBuilder& onError(Map::ErrorCallback cb, void* arg = nullptr);
 
    StageGroupBuilder& stream(Stream* s);
 
    StageGroupBuilder& enableProfile(const xpti::ProfileConfig& config);
 
    ~StageGroupBuilder();
    StageGroupBuilder(StageGroupBuilder&&) noexcept;
 
private:
    StageGroupBuilder(StageBuilder& parent);
    // Owning constructor: takes a heap-allocated StageBuilder when the group
    // is created from a temporary parent (e.g. Launcher::stage() no-arg).
    // Without this, parent_ would dangle once the temporary expires.
    StageGroupBuilder(std::unique_ptr<StageBuilder> owned);
    void flush();
    StageGroupBuilder takeStageGroup() &&;
 
    // Heap-owned parent for the owning construction path; nullptr when the
    // group is bound to an externally-owned StageBuilder.
    std::unique_ptr<StageBuilder> owned_;
    // Always-valid reference to the active parent (owned_.get() when owning,
    // or the externally-supplied StageBuilder otherwise).
    StageBuilder& parent_;
    impl::StageGroup currentGroup_;
    bool flushed_ = false;
 
    friend class StageBuilder;
    friend class Launcher;
    template <typename>
    friend class detail::TaskCountConfigured;
};
 
inline StageGroupBuilder StageBuilder::takeStageGroup() &&
{
    return StageGroupBuilder(std::make_unique<StageBuilder>(std::move(*this)));
}
 
inline StageGroupBuilder StageGroupBuilder::takeStageGroup() &&
{
    flush();
    if (owned_)
    {
        return StageGroupBuilder(std::move(owned_));
    }
    return StageGroupBuilder(std::make_unique<StageBuilder>(std::move(parent_)));
}
 
template <auto Func, typename... Args,
          detail::EnableCompatibleKernelArguments<Func, Args...>>
detail::TaskCountRequired<StageBuilder> StageBuilder::stage(Args&&... args) &&
{
    stage<Func>(uint32_t{0}, std::forward<Args>(args)...);
    return detail::TaskCountRequired<StageBuilder>(std::move(*this));
}
 
template <auto Func, typename... Args,
          detail::EnableCompatibleKernelArguments<Func, Args...>>
detail::TaskCountRequired<StageGroupBuilder> StageGroupBuilder::add(Args&&... args) &&
{
    add<Func>(uint32_t{0}, std::forward<Args>(args)...);
    return detail::TaskCountRequired<StageGroupBuilder>(std::move(*this));
}
 
template <auto Func, typename... Args,
          detail::EnableCompatibleKernelArguments<Func, Args...>>
detail::TaskCountRequired<StageBuilder> StageGroupBuilder::stage(Args&&... args) &&
{
    auto& parent = stage<Func>(uint32_t{0}, std::forward<Args>(args)...);
    return detail::TaskCountRequired<StageBuilder>(std::move(parent));
}
 
}  // namespace pxl