// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once
#include <cuda.h>
#include <cuda_runtime.h>
#include "lite/backends/cuda/cuda_utils.h"
#include "lite/backends/cuda/target_wrapper.h"

namespace paddle {
namespace lite {

// CudaStreamGuard is a encapsulation of cudaStream_t, which can accept external
// stream or internally created stream
//
//   std::unique_ptr<CudaStreamGuard> sm;
//
//   external stream: exec_stream
//     sm.reset(new CudaStreamGuard(exec_stream));
//   internal stream
//     sm.reset(new CudaStreamGuard());
//   get cudaStream_t
//     sm->stream();
class CudaStreamGuard {
 public:
  explicit CudaStreamGuard(cudaStream_t stream)
      : stream_(stream), owned_(false) {}
  CudaStreamGuard() : owned_(true) {
    lite::TargetWrapperCuda::CreateStream(&stream_);
  }
  ~CudaStreamGuard() {
    if (owned_) {
      lite::TargetWrapperCuda::DestroyStream(stream_);
    }
  }
  cudaStream_t* stream() { return &stream_; }
  bool owned() { return owned_; }

 private:
  cudaStream_t stream_;
  bool owned_{false};
};

}  // namespace lite
}  // namespace paddle
