#pragma once

#include <c10/util/ArrayRef.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <ATen/cuda/CUDAContext.h>

#include <vector>

namespace at { namespace cuda {

// TODO: Implement this generically in c10.  You'll need some way to get
// the number of GPUs from the GuardImpl, in that case.
class CUDAMultiStreamGuard final {
public:
  /// Calls `set_stream` on each of the streams in the list.
  /// This may be useful if you need to set different streams
  /// for different devices.
  explicit CUDAMultiStreamGuard(ArrayRef<CUDAStream> streams) : CUDAMultiStreamGuard() {
    for (const auto& s : streams) {
      setCurrentCUDAStream(s);
    }
  }

  CUDAMultiStreamGuard() {
    const size_t device_count = getNumGPUs();
    original_streams_.reserve(device_count);
    for (size_t device = 0; device < device_count; ++device) {
      original_streams_.push_back(getCurrentCUDAStream(device));
    }
  }

  CUDAMultiStreamGuard(const CUDAGuard&) = delete;
  CUDAMultiStreamGuard& operator=(const CUDAGuard&) = delete;

  // See Note [Move construction for RAII guards is tricky]
  CUDAMultiStreamGuard(CUDAGuard&& other) = delete;

  // See Note [Move assignment for RAII guards is tricky]
  CUDAMultiStreamGuard& operator=(CUDAGuard&& other) = delete;

  ArrayRef<CUDAStream> original_streams() const {
    return original_streams_;
  }

  /// Resets the CUDA stream on each device to the one that was active upon
  /// construction.
  ~CUDAMultiStreamGuard() {
    for (const auto& s : original_streams_) {
      setCurrentCUDAStream(s);
    }
  }

private:
  /// The original streams that were active on all devices.
  std::vector<CUDAStream> original_streams_;
};

}} // namespace at::cuda