#pragma once

#include <cstdint>
#include <utility>

#include "cuda_runtime_api.h"

#include <ATen/cuda/ATenCUDAGeneral.h>

/*
* A CUDAStream interface. See CUDAStream.cpp for implementation details.
*
* Includes the CUDAStream convenience class and a pointer-based stream API.
*
* The ATen/cuda/CUDAContext interface should be preferred when working with streams.
*/

/*
* Stream pool note.
*
* A CUDAStream is an abstraction of an actual cuStream on the GPU. CUDAStreams
* are backed by cuStreams, but they use several pools to minimize the costs
* associated with creating, retaining, and destroying cuStreams.
*
* There are three pools per device, and a device's pools are lazily created.
*
* The first pool contains only the default stream. When the default stream
* is requested it's returned.
*
* The second pool is the "low priority" or "default priority" streams. In
* HIP builds there is no distinction between streams in this pool and streams
* in the third pool (below). There are 32 of these streams per device, and
* when a stream is requested one of these streams is returned round-robin.
* That is, the first stream requested is at index 0, the second at index 1...
* to index 31, then index 0 again.
*
* This means that if 33 low priority streams are requested, the first and
* last streams requested are actually the same stream (under the covers)
* and kernels enqueued on them cannot run concurrently.
*
* The third pool is the "high priority" streams. The third pool acts like
* the second pool except the streams are created with a higher priority.
*
* These pools suggest that stream users should prefer many short-lived streams,
* as the cost of acquiring and releasing streams is effectively zero. If
* many longer-lived streams are required in performance critical scenarios
* then the functionality here may need to be extended to allow, for example,
* "reserving" a subset of the pool so that other streams do not accidentally
* overlap the performance critical streams.
*/

struct CUDAStreamInternals;

namespace at {
namespace cuda {

struct CUDAEvent;

namespace detail {

// Pointer-based API (for internal use)
AT_CUDA_API CUDAStreamInternals* CUDAStream_getDefaultStream(int64_t device = -1);

AT_CUDA_API CUDAStreamInternals* CUDAStream_getStreamFromPool(
  const bool isHighPriority = false
, int64_t device = -1);

AT_CUDA_API CUDAStreamInternals* CUDAStream_getCurrentStream(int64_t device = -1);

AT_CUDA_API void CUDAStream_setStream(CUDAStreamInternals* internals);
AT_CUDA_API void CUDAStream_uncheckedSetStream(CUDAStreamInternals* internals);

AT_CUDA_API cudaStream_t CUDAStream_stream(CUDAStreamInternals*);
AT_CUDA_API int64_t CUDAStream_device(CUDAStreamInternals*);

} // namespace detail

// RAII for a CUDA stream
// Allows use as a cudaStream_t, copying, moving, and metadata access.
struct AT_CUDA_API CUDAStream {

  // Constructors
  CUDAStream() = default;
  /* implicit */ CUDAStream(CUDAStreamInternals* internals_in)
  : internals_{internals_in} { }

  // Returns true if the CUDAStream is not null.
  explicit operator bool() const noexcept { return internals_ != nullptr; }

  // Implicit conversion to cudaStream_t
  operator cudaStream_t() const { return detail::CUDAStream_stream(internals_); }

  // Less than operator (to allow use in sets)
  friend bool operator<(const CUDAStream& left, const CUDAStream& right) {
    return left.internals_ < right.internals_;
  }

  // Getters
  int64_t device() const { return detail::CUDAStream_device(internals_); }
  cudaStream_t stream() const { return detail::CUDAStream_stream(internals_); }
  CUDAStreamInternals* internals() const { return internals_; }

  void synchronize_with(const CUDAEvent& event) const;

private:
  CUDAStreamInternals* internals_ = nullptr;
};

} // namespace cuda
} // namespace at
