align getStreamFromPool with Pytorch

Author: youge325

paddle/phi/api/include/compat/c10/cuda/CUDAStream.h

@@ -21,6 +21,7 @@
 #include <algorithm>
 #include <array>
 #include <atomic>
+#include <cstdint>
 #include <functional>
 #include <mutex>
 #include <ostream>
@@ -39,8 +40,36 @@ static constexpr int max_compile_time_stream_priorities = 4;
 
 namespace detail {
 
-constexpr int kStreamsPerPool = 32;
+constexpr int kStreamsPerPoolBits = 5;
+constexpr int kStreamsPerPool = 1 << kStreamsPerPoolBits;  // 32
 constexpr int kMaxDevices = 64;
+constexpr unsigned int kDefaultFlags = cudaStreamNonBlocking;
+
+// Global stream state - inline for header-only implementation
+inline std::array<std::once_flag, kMaxDevices>& device_flags() {
+  static std::array<std::once_flag, kMaxDevices> flags;
+  return flags;
+}
+
+inline std::array<std::array<std::atomic<uint32_t>, kMaxDevices>,
+                  max_compile_time_stream_priorities>&
+priority_counters() {
+  static std::array<std::array<std::atomic<uint32_t>, kMaxDevices>,
+                    max_compile_time_stream_priorities>
+      counters;
+  return counters;
+}
+
+inline std::array<
+    std::array<std::array<cudaStream_t, kStreamsPerPool>, kMaxDevices>,
+    max_compile_time_stream_priorities>&
+streams() {
+  static std::array<
+      std::array<std::array<cudaStream_t, kStreamsPerPool>, kMaxDevices>,
+      max_compile_time_stream_priorities>
+      stream_arrays;
+  return stream_arrays;
+}
 
 inline int gpu_device_count() {
   static const int count = phi::backends::gpu::GetGPUDeviceCount();
@@ -59,32 +88,31 @@ inline void check_device_index(int device_index) {
               ")");
 }
 
-struct StreamPoolState {
-  cudaStream_t low_priority[kStreamsPerPool]{};
-  cudaStream_t high_priority[kStreamsPerPool]{};
-  std::atomic<uint32_t> lp_counter{0};
-  std::atomic<uint32_t> hp_counter{0};
-  std::once_flag init_flag;
-};
-
-inline StreamPoolState& get_pool(int device_index) {
-  check_device_index(device_index);
-  static StreamPoolState states[kMaxDevices];
-  return states[device_index];
+// Init a single CUDA stream with given priority
+inline void initSingleStream(int priority_idx,
+                             int device_index,
+                             int stream_idx) {
+  phi::backends::gpu::GPUDeviceGuard guard(device_index);
+  auto& stream = streams()[priority_idx][device_index][stream_idx];
+  // priority_idx 0 = highest priority (most negative value)
+  int pri = -(priority_idx);
+  C10_CUDA_CHECK(cudaStreamCreateWithPriority(&stream, kDefaultFlags, pri));
 }
 
-inline void init_pool(int device_index, StreamPoolState* state) {
-  phi::backends::gpu::GPUDeviceGuard guard(device_index);
-  int lo_pri = 0, hi_pri = 0;
-  C10_CUDA_CHECK(cudaDeviceGetStreamPriorityRange(&lo_pri, &hi_pri));
+// Init stream pools for a device (called once per device)
+inline void initDeviceStreamState(int device_index) {
   for (int i = 0; i < kStreamsPerPool; ++i) {
-    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
-        &state->low_priority[i], cudaStreamNonBlocking, lo_pri));
-    C10_CUDA_CHECK(cudaStreamCreateWithPriority(
-        &state->high_priority[i], cudaStreamNonBlocking, hi_pri));
+    for (int p = 0; p < max_compile_time_stream_priorities; ++p) {
+      initSingleStream(p, device_index, i);
+    }
   }
 }
 
+// Helper to get round-robin index
+inline uint32_t get_idx(std::atomic<uint32_t>* counter) {
+  return counter->fetch_add(1) % kStreamsPerPool;
+}
+
 struct TLSStreamState {
   cudaStream_t streams[kMaxDevices]{};
   bool has_stream[kMaxDevices]{};
@@ -95,6 +123,25 @@ inline TLSStreamState& get_tls() {
   return s;
 }
 
+// Global initialization flag and max priorities
+inline void initGlobalStreamState() {
+  // This is called once to initialize global state
+  int leastPriority = 0, greatestPriority = 0;
+  C10_CUDA_CHECK(
+      cudaDeviceGetStreamPriorityRange(&leastPriority, &greatestPriority));
+  // max_stream_priorities is computed at runtime but bounded by compile-time
+  // constant For simplicity, we use the compile-time max
+}
+
+inline void initCUDAStreamsOnce() {
+  static std::once_flag init_flag;
+  std::call_once(init_flag, initGlobalStreamState);
+
+  auto& tls = get_tls();
+  // Initialize TLS current streams to default (null)
+  // This is lazy - we don't need to pre-initialize all entries
+}
+
 }  // namespace detail
 
 class CUDAStream {
@@ -191,6 +238,7 @@ inline CUDAStream make_cuda_stream(cudaStream_t raw,
 }
 
 inline CUDAStream getCurrentCUDAStream(c10::DeviceIndex device_index = -1) {
+  detail::initCUDAStreamsOnce();
   if (device_index == -1) {
     device_index = phi::backends::gpu::GetCurrentDeviceId();
   }
@@ -210,37 +258,41 @@ inline CUDAStream getCurrentCUDAStream(c10::DeviceIndex device_index = -1) {
 
 inline CUDAStream getStreamFromPool(const int priority,
                                     c10::DeviceIndex device_index = -1) {
+  detail::initCUDAStreamsOnce();
   if (device_index == -1) {
     device_index = phi::backends::gpu::GetCurrentDeviceId();
   }
-  auto& state = detail::get_pool(device_index);
-  std::call_once(state.init_flag, [device_index, &state]() {
-    detail::init_pool(device_index, &state);
-  });
+  detail::check_device_index(device_index);
 
-  cudaStream_t raw;
+  // Initialize stream pools for this device (once)
+  std::call_once(detail::device_flags()[device_index],
+                 detail::initDeviceStreamState,
+                 device_index);
 
-  // Keep parity with PyTorch API shape: negative priority selects the
-  // high-priority pool, non-negative selects the low-priority pool.
-  if (priority < 0) {
-    raw = state.high_priority[state.hp_counter.fetch_add(1) %
-                              detail::kStreamsPerPool];
-  } else {
-    raw = state.low_priority[state.lp_counter.fetch_add(1) %
-                             detail::kStreamsPerPool];
-  }
+  // Map priority to priority index: higher priority = lower (more negative)
+  // value PyTorch: pri_idx = clamp(-priority, 0, max_stream_priorities - 1)
+  int pri_idx =
+      std::clamp(-priority, 0, max_compile_time_stream_priorities - 1);
+  const auto idx =
+      detail::get_idx(&detail::priority_counters()[pri_idx][device_index]);
+
+  cudaStream_t raw = detail::streams()[pri_idx][device_index][idx];
   return make_cuda_stream(raw, device_index);
 }
 
 /**
  * Get a new stream from the CUDA stream pool.
  *
- * This overload matches PyTorch's bool-based entry point and preserves the
- * single-argument form `getStreamFromPool(true)` for high-priority requests.
+ * This overload matches PyTorch's bool-based entry point.
  */
 inline CUDAStream getStreamFromPool(const bool isHighPriority = false,
                                     c10::DeviceIndex device_index = -1) {
-  return getStreamFromPool(isHighPriority ? -1 : 0, device_index);
+  // High priority: -1 (highest priority)
+  // Low priority: 0 (default priority)
+  // Using -1 to match typical CUDA priority range and ensure
+  // getStreamFromPool(true) and getStreamFromPool(-1) behave consistently
+  int priority = isHighPriority ? -1 : 0;
+  return getStreamFromPool(priority, device_index);
 }
 
 inline CUDAStream getStreamFromExternal(cudaStream_t ext_stream,
@@ -252,20 +304,18 @@ inline CUDAStream getStreamFromExternal(cudaStream_t ext_stream,
 /**
  * Set the current CUDA stream for the device of the given stream in the
  * calling thread.
- *
- * Implements per-thread, per-device current stream semantics: the change is
- * local to the calling OS thread and does not affect any shared state such as
- * Paddle's GPUContext.  Other threads continue to see their own current stream.
  */
 inline void setCurrentCUDAStream(CUDAStream stream) {
   c10::DeviceIndex idx = stream.unwrap().device_index();
   detail::check_device_index(idx);
+  detail::initCUDAStreamsOnce();
   auto& tls = detail::get_tls();
   tls.streams[idx] = stream.stream();
   tls.has_stream[idx] = true;
 }
 
 inline CUDAStream getDefaultCUDAStream(c10::DeviceIndex device_index = -1) {
+  detail::initCUDAStreamsOnce();
   if (device_index == -1) {
     device_index = phi::backends::gpu::GetCurrentDeviceId();
   }