// Copyright (c) 2026 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.

#include <c10/cuda/CUDAStream.h>

#include <atomic>
#include <memory>
#include <mutex>
#include <vector>

#include "paddle/phi/api/include/context_pool.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
#include "paddle/phi/backends/gpu/gpu_info.h"
#endif

namespace c10::cuda {

namespace {

#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)

constexpr int kStreamsPerPool = 32;

std::once_flag g_init_once;
c10::DeviceIndex g_num_gpus = -1;

struct DevicePools {
#ifdef PADDLE_WITH_HIP
  std::vector<hipStream_t> low_priority;
  std::vector<hipStream_t> high_priority;
#else
  std::vector<cudaStream_t> low_priority;
  std::vector<cudaStream_t> high_priority;
#endif
  std::atomic<uint32_t> lp_counter{0};
  std::atomic<uint32_t> hp_counter{0};
  std::once_flag init_flag;
};

std::vector<std::unique_ptr<DevicePools>> g_pools;

void initGlobalState() {
  std::call_once(g_init_once, []() {
    g_num_gpus =
        static_cast<c10::DeviceIndex>(phi::backends::gpu::GetGPUDeviceCount());
    g_pools.resize(g_num_gpus);
    for (auto& ptr : g_pools) {
      ptr = std::make_unique<DevicePools>();
    }
  });
}

void initDevicePools(c10::DeviceIndex device_index) {
  phi::backends::gpu::GPUDeviceGuard guard(device_index);
  int lo_pri = 0, hi_pri = 0;
#ifdef PADDLE_WITH_HIP
  C10_CUDA_CHECK(hipDeviceGetStreamPriorityRange(&lo_pri, &hi_pri));
#else
  C10_CUDA_CHECK(cudaDeviceGetStreamPriorityRange(&lo_pri, &hi_pri));
#endif

  auto& pool = *g_pools[device_index];
  pool.low_priority.resize(kStreamsPerPool);
  pool.high_priority.resize(kStreamsPerPool);

  for (int i = 0; i < kStreamsPerPool; ++i) {
#ifdef PADDLE_WITH_HIP
    C10_CUDA_CHECK(hipStreamCreateWithPriority(
        &pool.low_priority[i], hipStreamNonBlocking, lo_pri));
    C10_CUDA_CHECK(hipStreamCreateWithPriority(
        &pool.high_priority[i], hipStreamNonBlocking, hi_pri));
#else
    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
        &pool.low_priority[i], cudaStreamNonBlocking, lo_pri));
    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
        &pool.high_priority[i], cudaStreamNonBlocking, hi_pri));
#endif
  }
}

inline void check_gpu(c10::DeviceIndex device_index) {
  TORCH_CHECK(device_index >= 0 && device_index < g_num_gpus,
              "Device index value ",
              static_cast<int>(device_index),
              " is out of index range [0, ",
              static_cast<int>(g_num_gpus),
              ")");
}

inline phi::GPUContext* getMutableGPUContext(c10::DeviceIndex device_index) {
  return static_cast<phi::GPUContext*>(
      paddle::experimental::DeviceContextPool::Instance().GetMutable(
          phi::GPUPlace(device_index)));
}

#ifdef PADDLE_WITH_HIP
inline hipStream_t getPaddleCurrentStream(c10::DeviceIndex device_index) {
  auto* current_stream =
      paddle::GetCurrentCUDAStream(phi::GPUPlace(device_index));
  return current_stream == nullptr ? nullptr : current_stream->raw_stream();
}
#else
inline cudaStream_t getPaddleCurrentStream(c10::DeviceIndex device_index) {
  auto* current_stream =
      paddle::GetCurrentCUDAStream(phi::GPUPlace(device_index));
  return current_stream == nullptr ? nullptr : current_stream->raw_stream();
}
#endif

#endif  // defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)

}  // namespace

#ifdef PADDLE_WITH_HIP
inline CUDAStream make_cuda_stream(hipStream_t raw,
                                   c10::DeviceIndex device_index) {
#else
inline CUDAStream make_cuda_stream(cudaStream_t raw,
                                   c10::DeviceIndex device_index) {
#endif
  c10::StreamId sid =
      static_cast<c10::StreamId>(reinterpret_cast<intptr_t>(raw));
  return CUDAStream(
      c10::Stream(c10::Stream::UNSAFE,
                  c10::Device(c10::DeviceType::CUDA, device_index),
                  sid));
}

CUDAStream getStreamFromPool(const int priority,
                             c10::DeviceIndex device_index) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  initGlobalState();
  if (device_index == -1) {
    device_index =
        static_cast<c10::DeviceIndex>(phi::backends::gpu::GetCurrentDeviceId());
  }
  check_gpu(device_index);

  std::call_once(
      g_pools[device_index]->init_flag, initDevicePools, device_index);

  const uint32_t idx = (priority < 0 ? g_pools[device_index]->hp_counter++
                                     : g_pools[device_index]->lp_counter++) %
                       kStreamsPerPool;
#ifdef PADDLE_WITH_HIP
  hipStream_t raw = (priority < 0 ? g_pools[device_index]->high_priority[idx]
                                  : g_pools[device_index]->low_priority[idx]);
#else
  cudaStream_t raw = (priority < 0 ? g_pools[device_index]->high_priority[idx]
                                   : g_pools[device_index]->low_priority[idx]);
#endif

  return make_cuda_stream(raw, device_index);
#else
  TORCH_CHECK(false, "getStreamFromPool is not supported without CUDA/HIP");
  return getDefaultCUDAStream(device_index);
#endif
}

CUDAStream getStreamFromPool(const bool isHighPriority,
                             c10::DeviceIndex device_index) {
  return getStreamFromPool(isHighPriority ? -1 : 0, device_index);
}

#ifdef PADDLE_WITH_HIP
CUDAStream getStreamFromExternal(hipStream_t ext_stream,
                                 c10::DeviceIndex device_index) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  initGlobalState();
  check_gpu(device_index);
#endif
  return make_cuda_stream(ext_stream, device_index);
}
#else
CUDAStream getStreamFromExternal(cudaStream_t ext_stream,
                                 c10::DeviceIndex device_index) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  initGlobalState();
  check_gpu(device_index);
#endif
  return make_cuda_stream(ext_stream, device_index);
}
#endif

CUDAStream getDefaultCUDAStream(c10::DeviceIndex device_index) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  initGlobalState();
  if (device_index == -1) {
    device_index =
        static_cast<c10::DeviceIndex>(phi::backends::gpu::GetCurrentDeviceId());
  }
  check_gpu(device_index);
#endif
  return CUDAStream(c10::Stream(
      c10::Stream::DEFAULT, c10::Device(c10::DeviceType::CUDA, device_index)));
}

CUDAStream getCurrentCUDAStream(c10::DeviceIndex device_index) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  initGlobalState();
  if (device_index == -1) {
    device_index =
        static_cast<c10::DeviceIndex>(phi::backends::gpu::GetCurrentDeviceId());
  }
  check_gpu(device_index);
  auto raw = getPaddleCurrentStream(device_index);
  if (raw == nullptr) {
    return getDefaultCUDAStream(device_index);
  }
  return make_cuda_stream(raw, device_index);
#else
  return getDefaultCUDAStream(device_index);
#endif
}

void setCurrentCUDAStream(CUDAStream stream) {
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  initGlobalState();
  c10::DeviceIndex idx = stream.unwrap().device_index();
  check_gpu(idx);
  getMutableGPUContext(idx)->SetStream(stream.stream());
#else
  (void)stream;
#endif
}

}  // namespace c10::cuda