cublas/execution.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/backend/cuda/resource_pool.hpp>
#include <uni20/backend/cuda/runtime.hpp>
 
#include <cublas_v2.h>
 
#include <cstddef>
#include <optional>
 
namespace uni20::cuda
{
class DeviceResources;
}
 
namespace uni20::cublas
{
 
class ExecutionAcquireAwaiter;
 
class HandleSlot {
  public:
    explicit HandleSlot(int device);
    HandleSlot(HandleSlot const&) = delete;
    HandleSlot& operator=(HandleSlot const&) = delete;
    HandleSlot(HandleSlot&& other) noexcept;
    HandleSlot& operator=(HandleSlot&& other) noexcept;
    ~HandleSlot();
 
    [[nodiscard]] int device() const noexcept { return device_; }
 
    [[nodiscard]] cublasHandle_t native_handle() const noexcept { return handle_; }
 
    void bind(cuda::Stream const& stream);
 
  private:
    void reset() noexcept;
 
    int device_ = -1;
    cublasHandle_t handle_ = nullptr;
};
 
class ExecutionLease {
  public:
    using handle_lease_type = cuda::ResourceLease<HandleSlot>;
 
    ExecutionLease() noexcept = default;
    ExecutionLease(ExecutionLease const&) = delete;
    ExecutionLease& operator=(ExecutionLease const&) = delete;
    ExecutionLease(ExecutionLease&& other) noexcept;
    ExecutionLease& operator=(ExecutionLease&& other) noexcept;
    ~ExecutionLease();
 
    [[nodiscard]] explicit operator bool() const noexcept
    {
      return static_cast<bool>(handle_) && static_cast<bool>(stream_);
    }
 
    [[nodiscard]] HandleSlot& handle() const { return handle_.get(); }
 
    [[nodiscard]] cuda::Stream const& stream() const noexcept { return stream_; }
 
    void release() noexcept;
 
  private:
    ExecutionLease(handle_lease_type handle, cuda::Stream stream);
 
    handle_lease_type handle_;
    cuda::Stream stream_;
 
    friend class ExecutionPool;
    friend class ExecutionAcquireAwaiter;
};
 
class ExecutionPool {
  public:
    ExecutionPool(cuda::StreamPool& streams, std::size_t handle_count);
    ExecutionPool(ExecutionPool const&) = delete;
    ExecutionPool& operator=(ExecutionPool const&) = delete;
    ExecutionPool(ExecutionPool&&) = delete;
    ExecutionPool& operator=(ExecutionPool&&) = delete;
 
    [[nodiscard]] std::optional<ExecutionLease> try_acquire();
 
    [[nodiscard]] ExecutionLease acquire();
 
    [[nodiscard]] int device() const noexcept { return streams_->device(); }
 
    [[nodiscard]] std::size_t handle_count() const noexcept { return handles_.size(); }
 
    [[nodiscard]] std::size_t idle_handle_count() const noexcept { return handles_.idle_count(); }
 
  private:
    [[nodiscard]] ExecutionLease make_lease(cuda::ResourceLease<HandleSlot> handle, cuda::Stream stream);
 
    cuda::StreamPool* streams_;
    cuda::ResourcePool<HandleSlot> handles_;
 
    friend class ExecutionAcquireAwaiter;
};
 
[[nodiscard]] ExecutionPool& execution_pool(cuda::DeviceResources& resources);
 
[[nodiscard]] ExecutionPool& execution_pool(cuda::DeviceResources& resources, std::size_t handle_count);
 
} // namespace uni20::cublas