runtime.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/backend/cuda/device.hpp>
#include <uni20/common/trace.hpp>
 
#include <cuda_runtime_api.h>
 
#include <atomic>
#include <cstddef>
#include <memory>
#include <mutex>
#include <optional>
#include <span>
#include <type_traits>
#include <typeindex>
#include <unordered_map>
#include <utility>
#include <vector>
 
namespace uni20::cuda
{
 
class Completion;
class DeviceResources;
class StreamPool;
template <typename T>
  requires std::is_trivially_copyable_v<T>
class CudaBuffer;
 
namespace detail
{
class AllocationReclaimer;
class StreamWaiter;
template <typename T> class CudaAllocation;
} // namespace detail
 
class ScopedDevice {
  public:
    explicit ScopedDevice(int device);
    ScopedDevice(ScopedDevice const&) = delete;
    ScopedDevice& operator=(ScopedDevice const&) = delete;
    ~ScopedDevice();
 
    [[nodiscard]] int device() const noexcept { return device_; }
 
  private:
    int device_;
    int previous_device_;
    bool restore_;
};
 
class Stream {
  public:
    class State;
 
    Stream() = default;
    Stream(Stream const&) noexcept = default;
    Stream& operator=(Stream const&) noexcept = default;
    Stream(Stream&& other) noexcept = default;
    Stream& operator=(Stream&& other) noexcept = default;
    ~Stream() = default;
 
    [[nodiscard]] explicit operator bool() const noexcept { return static_cast<bool>(state_); }
 
    [[nodiscard]] int device() const noexcept;
 
    [[nodiscard]] cudaStream_t native_handle() const noexcept;
 
    void synchronize() const;
 
    [[nodiscard]] Completion record_completion() const;
 
    void wait_on(Completion const& completion) const;
 
  private:
    explicit Stream(std::shared_ptr<State> state) noexcept;
 
    std::shared_ptr<State> state_;
 
    friend class StreamPool;
};
 
class Completion {
  public:
    Completion() = default;
 
    [[nodiscard]] static Completion record(int device, cudaStream_t stream);
 
    [[nodiscard]] explicit operator bool() const noexcept { return static_cast<bool>(state_); }
 
    [[nodiscard]] int device() const noexcept;
 
    [[nodiscard]] bool ready() const;
 
    void synchronize() const;
 
  private:
    class State;
 
    explicit Completion(std::shared_ptr<State> state);
 
    std::shared_ptr<State> state_;
 
    friend class Stream;
    friend class Stream::State;
    friend class detail::AllocationReclaimer;
};
 
namespace detail
{
 
class StreamWaiter {
  public:
    using notify_type = void (*)(StreamWaiter&, Stream) noexcept;
 
    explicit StreamWaiter(notify_type notify) noexcept : notify_(notify) { CHECK(notify_ != nullptr); }
    StreamWaiter(StreamWaiter const&) = delete;
    StreamWaiter& operator=(StreamWaiter const&) = delete;
 
  private:
    void notify(Stream stream) noexcept { notify_(*this, std::move(stream)); }
 
    notify_type notify_;
    StreamWaiter* next_ = nullptr;
    bool queued_ = false;
 
    friend class ::uni20::cuda::StreamPool;
};
 
} // namespace detail
 
class StreamPool {
  public:
    struct Config
    {
        int device;
        std::size_t stream_count;
        unsigned int stream_flags = cudaStreamNonBlocking;
    };
 
    explicit StreamPool(Config config);
    StreamPool(StreamPool const&) = delete;
    StreamPool& operator=(StreamPool const&) = delete;
    ~StreamPool();
 
    [[nodiscard]] std::optional<Stream> try_acquire();
 
    [[nodiscard]] Stream acquire();
 
    [[nodiscard]] std::optional<Stream> acquire_or_enqueue(detail::StreamWaiter& waiter) noexcept;
 
    [[nodiscard]] int device() const noexcept;
 
    [[nodiscard]] std::size_t size() const noexcept;
 
    [[nodiscard]] std::size_t idle_stream_count() const noexcept;
 
    [[nodiscard]] std::size_t leased_stream_count() const noexcept;
 
    [[nodiscard]] std::size_t pending_stream_count() const noexcept;
 
    void synchronize();
 
  private:
    class Impl;
    std::unique_ptr<Impl> impl_;
 
    friend class Stream::State;
};
 
class DeviceResources {
  public:
    struct Config
    {
        Device device;
        std::size_t stream_count;
        unsigned int stream_flags = cudaStreamNonBlocking;
    };
 
    explicit DeviceResources(Config config);
    DeviceResources(DeviceResources const&) = delete;
    DeviceResources& operator=(DeviceResources const&) = delete;
    DeviceResources(DeviceResources&&) = delete;
    DeviceResources& operator=(DeviceResources&&) = delete;
    ~DeviceResources() noexcept;
 
    [[nodiscard]] Device device() const noexcept { return device_; }
 
    [[nodiscard]] StreamPool& streams() noexcept { return streams_; }
    [[nodiscard]] StreamPool const& streams() const noexcept { return streams_; }
 
    template <class Resource, class... Args> [[nodiscard]] Resource& provider_resource(Args&&... args)
    {
      std::lock_guard lock(provider_resources_mutex_);
      auto const key = std::type_index(typeid(Resource));
      auto found = provider_resources_.find(key);
      if (found != provider_resources_.end())
      {
        return static_cast<ProviderResourceModel<Resource>&>(*found->second).value;
      }
 
      auto resource = std::make_unique<ProviderResourceModel<Resource>>(std::forward<Args>(args)...);
      Resource& result = resource->value;
      provider_resources_.emplace(key, std::move(resource));
      return result;
    }
 
    [[nodiscard]] std::size_t live_allocation_count() const noexcept
    {
      return live_allocation_count_.load(std::memory_order_acquire);
    }
 
  private:
    struct ProviderResourceBase
    {
        virtual ~ProviderResourceBase() = default;
    };
 
    template <class Resource> struct ProviderResourceModel final : ProviderResourceBase
    {
        template <class... Args> explicit ProviderResourceModel(Args&&... args) : value(std::forward<Args>(args)...) {}
 
        Resource value;
    };
 
    void register_allocation() noexcept { live_allocation_count_.fetch_add(1, std::memory_order_relaxed); }
    void unregister_allocation() noexcept;
    void defer_allocation_release_after(std::span<Completion const> completions) noexcept;
    void release_allocation(void* allocation, bool blocking) noexcept;
    void destroy_provider_resources() noexcept;
 
    Device device_;
    StreamPool streams_;
    std::unique_ptr<detail::AllocationReclaimer> allocation_reclaimer_;
    std::mutex provider_resources_mutex_;
    std::unordered_map<std::type_index, std::unique_ptr<ProviderResourceBase>> provider_resources_;
    std::atomic<std::size_t> live_allocation_count_ = 0;
 
    template <typename T>
      requires std::is_trivially_copyable_v<T>
    friend class CudaBuffer;
    template <typename> friend class detail::CudaAllocation;
};
 
struct RuntimeConfig
{
    std::vector<int> device_ordinals{};
    std::optional<int> default_device{};
    std::size_t streams_per_device = 8;
    unsigned int stream_flags = cudaStreamNonBlocking;
};
 
class Runtime {
  public:
    Runtime(Runtime const&) = delete;
    Runtime& operator=(Runtime const&) = delete;
    Runtime(Runtime&&) = delete;
    Runtime& operator=(Runtime&&) = delete;
    ~Runtime() noexcept;
 
    [[nodiscard]] bool has_device(int device) const noexcept;
 
    [[nodiscard]] DeviceResources& device_resources(int device) const;
 
    [[nodiscard]] DeviceResources& default_device_resources() const;
 
    [[nodiscard]] std::optional<int> default_device() const noexcept { return default_device_; }
 
  private:
    struct construction_key
    {};
 
    explicit Runtime(RuntimeConfig config, construction_key);
 
    std::vector<std::unique_ptr<DeviceResources>> device_resources_;
    std::optional<int> default_device_;
 
    friend Runtime initialize(RuntimeConfig);
};
 
[[nodiscard]] Runtime initialize(RuntimeConfig config = {});
 
[[nodiscard]] bool is_initialized() noexcept;
 
[[nodiscard]] Runtime& runtime();
 
[[nodiscard]] DeviceResources& device_resources(int device);
 
[[nodiscard]] DeviceResources& device_resources();
 
} // namespace uni20::cuda