launcher.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 <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/stage_builder.hpp"
#include "pxl/stream.hpp"
#include "pxl/task_count.hpp"
#include "pxl/type.hpp"
 
namespace pxl
{
 
namespace impl
{
class LaunchBuilderImpl;
class LaunchBuilderTestAccess;
class DeviceHandle;
}  // namespace impl
 
// [[nodiscard]] catches the most common new-user mistake — building the chain
// and forgetting `.run()`. Without it, the temporary is silently destroyed,
// the kernel never executes, and there is no runtime signal.
class [[nodiscard]] LaunchBuilder
{
public:
    LaunchBuilder& reserve(uint32_t n);
 
    LaunchBuilder& tasks(uint32_t taskCount);
 
    LaunchBuilder& batchSize(uint32_t n);
 
    LaunchBuilder& clusterBitmap(uint32_t bitmap);
 
    LaunchBuilder& locality(LocalityMode mode);
 
    LaunchBuilder& onComplete(Map::CompletionCallback cb, void* arg = nullptr);
 
    LaunchBuilder& onMessage(Map::MessageCallback cb, void* arg = nullptr);
 
    LaunchBuilder& onError(Map::ErrorCallback cb, void* arg = nullptr);
 
    LaunchBuilder& stream(Stream* s);
 
    LaunchBuilder& enableProfile(const xpti::ProfileConfig& config);
 
    LaunchResult run();
 
    std::future<LaunchResult> runAsync();
 
    ~LaunchBuilder();
    LaunchBuilder(LaunchBuilder&& other) noexcept;
    LaunchBuilder& operator=(LaunchBuilder&& other) noexcept;
 
    LaunchBuilder(const LaunchBuilder&) = delete;
    LaunchBuilder& operator=(const LaunchBuilder&) = delete;
 
private:
    LaunchBuilder(std::unique_ptr<impl::LaunchBuilderImpl> impl);
 
    std::unique_ptr<impl::LaunchBuilderImpl> impl_;
    // Atomic so concurrent run()/runAsync() races on the same builder are
    // diagnosed (returned as the "already called" failure) rather than UB.
    // The builder is not designed to be shared across threads, but the guard
    // is cheap and protects users from accidental sharing patterns.
    std::atomic<bool> started_{false};
 
    friend class Launcher;
    friend class impl::LaunchBuilderTestAccess;
};
 
class Launcher
{
public:
    Launcher();
 
    explicit Launcher(uint32_t deviceId);
 
    ~Launcher();
 
    Launcher(const Launcher&) = delete;
    Launcher& operator=(const Launcher&) = delete;
    Launcher(Launcher&&) noexcept;
    Launcher& operator=(Launcher&&) noexcept;
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    LaunchBuilder execute(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
 
        std::vector<ArgInfo_t> argList;
        std::vector<std::shared_ptr<void>> storage;
        std::vector<std::pair<void*, size_t>> syncTargets;
        auto directions = detail::MakeKernelDirections<Func>();
        argList.reserve(sizeof...(Args));
        if constexpr (sizeof...(Args) > 0)
        {
            buildArgList(argList, storage, syncTargets, std::forward<Args>(args)...);
        }
 
        if (impl_ == nullptr)
        {
            initFromArgs(argList, syncTargets);
        }
        if (!isValid())
        {
            return makeEmptyBuilder();
        }
 
        return execute(kernelName, taskCount,
                       std::move(argList), std::move(storage), std::move(syncTargets),
                       std::move(directions));
    }
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<LaunchBuilder> execute(Args&&... args)
    {
        return detail::TaskCountRequired<LaunchBuilder>(
            execute<Func>(uint32_t{0}, std::forward<Args>(args)...));
    }
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    StageBuilder stage(uint32_t taskCount, Args&&... args)
    {
        if (impl_ == nullptr)
        {
            std::vector<ArgInfo_t> tempArgList;
            std::vector<std::shared_ptr<void>> tempStorage;
            std::vector<std::pair<void*, size_t>> tempSyncTargets;
            tempArgList.reserve(sizeof...(Args));
            if constexpr (sizeof...(Args) > 0)
            {
                buildArgList(tempArgList, tempStorage, tempSyncTargets, std::forward<Args>(args)...);
            }
            initFromArgs(tempArgList, tempSyncTargets);
        }
        StageBuilder builder(impl_);
        // Call returns StageBuilder& to the local; discard ref and return the
        // local by value to enable NRVO (StageBuilder is move-only).
        builder.stage<Func>(taskCount, std::forward<Args>(args)...);
        return builder;
    }
 
    template <auto Func, typename... Args,
              detail::EnableCompatibleKernelArguments<Func, Args...> = 0>
    detail::TaskCountRequired<StageBuilder> stage(Args&&... args)
    {
        return detail::TaskCountRequired<StageBuilder>(
            stage<Func>(uint32_t{0}, std::forward<Args>(args)...));
    }
 
    StageGroupBuilder stage()
    {
        // Allocate the parent StageBuilder on the heap and transfer ownership
        // to the returned StageGroupBuilder. Without owning the parent here,
        // the StageGroupBuilder would hold a reference into this stack frame
        // and dangle as soon as the function returns.
        return StageGroupBuilder(std::make_unique<StageBuilder>(impl_));
    }
 
    uint32_t deviceId() const;
 
    bool isValid() const;
 
private:
    LaunchBuilder execute(const char* kernelName, uint32_t taskCount,
                          std::vector<ArgInfo_t> argList,
                          std::vector<std::shared_ptr<void>> storage,
                          std::vector<std::pair<void*, size_t>> syncTargets);
 
    LaunchBuilder execute(const char* kernelName, uint32_t taskCount,
                          std::vector<ArgInfo_t> argList,
                          std::vector<std::shared_ptr<void>> storage,
                          std::vector<std::pair<void*, size_t>> syncTargets,
                          ArgDirections directions);
 
    void initFromArgs(const std::vector<ArgInfo_t>& argList,
                      const std::vector<std::pair<void*, size_t>>& syncTargets);
    LaunchBuilder makeEmptyBuilder();
 
    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 runAsync's worker
        // thread never reads a dead caller stack frame. The data buffer behind
        // ndarrayArg.data() is owned by pxl::allocateMemory and outlives the
        // builder; only the descriptor needs heap-stable storage.
        //
        // On a device without hardware Back-Invalidate the launch synchronizes
        // this view's host cache lines, and that works at cache-line
        // granularity: a view sharing a line with host-written data outside it
        // can have those host bytes written back over the kernel's output.
        // Views that start at the allocation base, or are cache-line aligned
        // and sized, are unaffected. The same holds for a pointer argument
        // that points into the middle of an allocation.
        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);
        // NDArray::size() returns int64_t; guard against negative values before casting to size_t.
        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> impl_;
};
 
}  // namespace pxl