[Perf] CUDA + AMDGPU launchers: pin stream to null at launcher scope; drop redundant pre/post-kernel stream_synchronize

Author: duburcqa

quadrants/rhi/cuda/cuda_stream_pin.h

@@ -0,0 +1,46 @@
+#pragma once
+
+// RAII guard that pins `CUDAContext::stream_` to the legacy default stream (nullptr) for the duration of the scope,
+// then restores the prior `stream_` on destruction. Pinning makes the consumer + producer streams match unconditionally
+// so the cross-stream-visibility break that CUDA's stream-ordered allocator (`cuMemAllocAsync` pool) shows on
+// pre-Ampere hardware (Turing T4 faults at `cuLaunchKernel` with `illegal-address`) is closed by construction. The
+// destructor does NOT call `stream_synchronize(nullptr)`: the helpers using this guard always issue a synchronous DtoH
+// inside the scope which already drains the null stream, and the main launcher's kernel-only fast path (no host-side
+// sync DtoH) deliberately stays asynchronous so kernel launches do not pay a forced per-launch host-GPU round-trip.
+// Restoring the prior stream while null-stream work is still in flight is safe: subsequent operations on the prior
+// stream that depend on the null-stream work will surface the standard cross-stream visibility constraints anyway, and
+// the next launch through this guard re-pins to null and serialises with prior null-stream work via legacy default
+// stream semantics. Used at the entry of every per-launch CUDA host path that issues HtoD or kernel work that needs
+// the same-stream invariant: `KernelLauncher::launch_llvm_kernel`, the adstack helpers
+// `dispatch_max_reducers_for_tasks` / `publish_adstack_metadata` / `publish_per_task_bound_count_device`. AMDGPU does
+// not need an equivalent because `AMDGPUContext::launch` passes `nullptr` for the stream parameter unconditionally;
+// HIP kernels always run on the null stream regardless of any context-level stream setting.
+
+#include "quadrants/platform/cuda/detect_cuda.h"
+
+#if defined(QD_WITH_CUDA)
+#include "quadrants/rhi/cuda/cuda_context.h"
+#include "quadrants/rhi/cuda/cuda_driver.h"
+#endif
+
+namespace quadrants::lang {
+
+#if defined(QD_WITH_CUDA)
+struct CudaDefaultStreamPinGuard {
+  bool engaged{false};
+  void *prev_stream{nullptr};
+  explicit CudaDefaultStreamPinGuard(bool engage) : engaged(engage) {
+    if (engaged) {
+      prev_stream = CUDAContext::get_instance().get_stream();
+      CUDAContext::get_instance().set_stream(nullptr);
+    }
+  }
+  ~CudaDefaultStreamPinGuard() {
+    if (engaged) {
+      CUDAContext::get_instance().set_stream(prev_stream);
+    }
+  }
+};
+#endif
+
+}  // namespace quadrants::lang

quadrants/runtime/amdgpu/kernel_launcher.cpp

@@ -258,9 +258,14 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
       }
     }
   }
-  if (transfers.size() > 0) {
-    AMDGPUDriver::get_instance().stream_synchronize(nullptr);
-  }
+  // No explicit `stream_synchronize(nullptr)` needed here. `AMDGPUContext::launch` passes `nullptr` as the stream
+  // parameter to `hipModuleLaunchKernel` unconditionally (see `rhi/amdgpu/amdgpu_context.cpp::launch`), so kernels
+  // always run on the null stream regardless of any context-level stream setting; with the HtoD copies above being
+  // sync `hipMemcpyHtoD` (host-blocking) on the same null stream, ordering is preserved without an explicit barrier
+  // even on hardware where stream-ordered allocator pools have stricter cross-stream visibility rules. Mirrors the
+  // CUDA launcher cleanup; CUDA needs a `CudaDefaultStreamPinGuard` to enforce same-stream invariant because
+  // `CUDAContext::launch` reads `CUDAContext::stream_` (mutable, may be non-null), while HIP's launch path is
+  // unconditionally null-stream.
   char *host_result_buffer = (char *)ctx.get_context().result_buffer;
   if (ctx.result_buffer_size > 0) {
     // Malloc_Async and Free_Async are available after ROCm 5.4

quadrants/runtime/cuda/kernel_launcher.cpp

@@ -1,6 +1,7 @@
 #include "quadrants/runtime/cuda/kernel_launcher.h"
 #include "quadrants/runtime/cuda/cuda_utils.h"
 #include "quadrants/rhi/cuda/cuda_context.h"
+#include "quadrants/rhi/cuda/cuda_stream_pin.h"
 #include "quadrants/runtime/llvm/llvm_runtime_executor.h"
 #include "quadrants/program/adstack_size_expr_eval.h"
 #include "quadrants/program/program.h"
@@ -208,6 +209,18 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
 
   CUDAContext::get_instance().make_current();
 
+  // Pin `CUDAContext::stream_` to the legacy default stream (nullptr) for the duration of the launch. The HtoD copies
+  // below are sync `cuMemcpyHtoD` (legacy null stream), the kernel dispatches via `cuda_module->launch` go through
+  // `CUDAContext::stream_` (now nullptr), and the post-kernel async + sync DtoH calls also land on the null stream.
+  // With every operation pinned to one stream, the cross-stream-visibility break that pre-Ampere shows for
+  // `cuMemAllocAsync`-pool buffers (Turing T4 faults at `cuLaunchKernel` with `illegal-address` when consumer +
+  // producer streams differ) is eliminated unconditionally, and the explicit `stream_synchronize(nullptr)` barriers
+  // between phases below collapse to no-ops paid for by the surrounding sync DtoH calls. Same RAII pattern as the
+  // Same RAII pattern as the `CudaDefaultStreamPinGuard` defined in `rhi/cuda/cuda_stream_pin.h`; engaging it here at
+  // launcher scope covers the HtoD/kernel/DtoH chain that the per-helper guards inside the adstack subsystem do not
+  // span.
+  CudaDefaultStreamPinGuard cuda_pin(/*engage=*/true);
+
   // |transfers| is only used for external arrays whose data is originally on
   // host. They are first transferred onto device and that device pointer is
   // stored in |device_ptrs| below. |transfers| saves its original pointer so
@@ -309,9 +322,10 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
       }
     }
   }
-  if (transfers.size() > 0) {
-    CUDADriver::get_instance().stream_synchronize(nullptr);
-  }
+  // No explicit `stream_synchronize(nullptr)` needed here. The HtoD copies above are sync `cuMemcpyHtoD` (host-blocking
+  // by definition); with the launcher's `CudaDefaultStreamPinGuard` ensuring the subsequent kernel dispatch lands on
+  // the same null stream, ordering between HtoD and kernel is preserved without an explicit barrier even on pre-Ampere
+  // pool buffers - the cross-stream visibility break only fires when producer + consumer streams differ.
   char *host_result_buffer = (char *)ctx.get_context().result_buffer;
   if (ctx.result_buffer_size > 0) {
     ctx.get_context().result_buffer = (uint64 *)device_result_buffer;
@@ -380,9 +394,11 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
     CUDADriver::get_instance().memcpy_device_to_host_async(host_result_buffer, device_result_buffer,
                                                            ctx.result_buffer_size, nullptr);
   }
-  // copy data back to host
+  // copy data back to host. With the launcher-scope `CudaDefaultStreamPinGuard` the kernel ran on the null stream, so
+  // sync `cuMemcpyDtoH` (host-blocking, on null stream) sees the kernel's writes without an explicit cross-stream
+  // barrier. The first sync DtoH below also drains the prior async `memcpy_device_to_host_async(host_result_buffer,
+  // ...)` queued on the same null stream, so an explicit `stream_synchronize(nullptr)` here would be redundant.
   if (transfers.size() > 0) {
-    CUDADriver::get_instance().stream_synchronize(nullptr);
     for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
       auto &idx = itr->first;
       CUDADriver::get_instance().memcpy_device_to_host(itr->second.first, (void *)device_ptrs[idx],

quadrants/runtime/llvm/adstack_lazy_claim/bound_eval.cpp

@@ -27,7 +27,7 @@
 #include "quadrants/rhi/amdgpu/amdgpu_context.h"
 #include "quadrants/rhi/amdgpu/amdgpu_driver.h"
 
-#include "quadrants/runtime/llvm/adstack_lazy_claim/cuda_stream_pin.h"
+#include "quadrants/rhi/cuda/cuda_stream_pin.h"
 
 namespace quadrants::lang {
 

quadrants/runtime/llvm/adstack_lazy_claim/cuda_stream_pin.h

@@ -31,7 +31,7 @@
 #include "quadrants/rhi/amdgpu/amdgpu_context.h"
 #endif
 
-#include "quadrants/runtime/llvm/adstack_lazy_claim/cuda_stream_pin.h"
+#include "quadrants/rhi/cuda/cuda_stream_pin.h"
 
 namespace quadrants::lang {