[Perf] Per-launch host sync: drop wait_idle on SPIR-V, pin stream and drop stream_synchronize on CUDA/AMDGPU

quadrants/rhi/amdgpu/amdgpu_stream_pin.h

@@ -0,0 +1,42 @@
+#pragma once
+
+// RAII guard that pins `AMDGPUContext::stream_` to the legacy default stream (nullptr) for the duration of the scope,
+// then restores the prior `stream_` on destruction. Symmetric counterpart to `CudaDefaultStreamPinGuard` in
+// `rhi/cuda/cuda_stream_pin.h`; engaged at the same launcher-scope callsite (`KernelLauncher::launch_llvm_kernel` on
+// AMDGPU) on the default-stream fast path so the HtoD / kernel / DtoH chain lands on a single stream without an
+// explicit `stream_synchronize` between phases. AMDGPU has not been observed to hit the cross-stream pool-visibility
+// fault that motivated the CUDA guard on pre-Ampere Turing T4, but `AMDGPUContext::launch` now forwards `stream_` to
+// `hipModuleLaunchKernel` (post-streams plumbing), so the same same-stream-invariant rationale applies and keeping the
+// guard symmetric with CUDA hardens against future driver / hardware combos that tighten cross-stream pool semantics.
+// The launcher engages this guard only when entry `stream_ == nullptr` and all tasks have
+// `stream_parallel_group_id == 0`; on the explicit-stream path the streams feature requires the kernel to inherit
+// the user-supplied `stream_` and the guard stays disengaged.
+
+#include "quadrants/platform/amdgpu/detect_amdgpu.h"
+
+#if defined(QD_WITH_AMDGPU)
+#include "quadrants/rhi/amdgpu/amdgpu_context.h"
+#include "quadrants/rhi/amdgpu/amdgpu_driver.h"
+#endif
+
+namespace quadrants::lang {
+
+#if defined(QD_WITH_AMDGPU)
+struct AmdgpuDefaultStreamPinGuard {
+  bool engaged{false};
+  void *prev_stream{nullptr};
+  explicit AmdgpuDefaultStreamPinGuard(bool engage) : engaged(engage) {
+    if (engaged) {
+      prev_stream = AMDGPUContext::get_instance().get_stream();
+      AMDGPUContext::get_instance().set_stream(nullptr);
+    }
+  }
+  ~AmdgpuDefaultStreamPinGuard() {
+    if (engaged) {
+      AMDGPUContext::get_instance().set_stream(prev_stream);
+    }
+  }
+};
+#endif
+
+}  // namespace quadrants::lang

quadrants/rhi/cuda/cuda_stream_pin.h

@@ -0,0 +1,50 @@
+#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`. The launcher
+// engages the guard only on the default-stream fast path (entry `stream_ == nullptr` with all tasks on
+// `stream_parallel_group_id == 0`); on the explicit-stream path the streams feature requires the kernel to inherit
+// the user-supplied `stream_` and the guard stays disengaged. A symmetric `AmdgpuDefaultStreamPinGuard` lives in
+// `rhi/amdgpu/amdgpu_stream_pin.h` for consistency on AMDGPU - the streams plumbing now passes `stream_` through to
+// `hipModuleLaunchKernel`, so the same same-stream-invariant rationale applies, even though the cross-stream pool
+// fault that motivated the CUDA guard has not been observed on AMDGPU hardware.
+
+#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

@@ -2,6 +2,7 @@
 
 #include "quadrants/runtime/amdgpu/kernel_launcher.h"
 #include "quadrants/rhi/amdgpu/amdgpu_context.h"
+#include "quadrants/rhi/amdgpu/amdgpu_stream_pin.h"
 #include "quadrants/program/adstack_size_expr_eval.h"
 #include "quadrants/program/launch_context_builder.h"
 #include "quadrants/program/program.h"
@@ -230,6 +231,23 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
 
   auto *active_stream = AMDGPUContext::get_instance().get_stream();
 
+  // Default-stream fast path: every HtoD / DtoH / kernel-dispatch in this launcher already routes through
+  // `active_stream`, so when entry `stream_ == nullptr` AND every offloaded task launches on the same `active_stream`
+  // (i.e. `stream_parallel_group_id == 0`), the entire chain serialises on the legacy default stream. The
+  // `stream_synchronize` barriers between phases collapse to no-ops the surrounding sync DtoH (host-blocking on
+  // pageable host memory) already drains, and the `AmdgpuDefaultStreamPinGuard` re-pins `AMDGPUContext::stream_` to
+  // nullptr defensively across the launch in case an inner helper temporarily swaps it. Outside the fast path -
+  // user-supplied stream OR any task on `stream_parallel_group_id != 0` (per-group acquired streams differ from
+  // `active_stream`) - the cross-stream barriers below are load-bearing for HtoD / kernel / DtoH visibility and the
+  // pin guard would silently override the user-requested stream at the kernel-launch site, so the guard stays
+  // disengaged and main's sync semantics remain untouched. Symmetric with the CUDA launcher; the pre-Ampere pool
+  // fault that motivated the CUDA pin has not been observed on AMDGPU, but `AMDGPUContext::launch` now forwards
+  // `stream_` to `hipModuleLaunchKernel` so the same same-stream-invariant rationale applies.
+  const bool all_sgid_zero = std::all_of(offloaded_tasks.begin(), offloaded_tasks.end(),
+                                         [](const OffloadedTask &t) { return t.stream_parallel_group_id == 0; });
+  const bool default_stream_path = (active_stream == nullptr) && all_sgid_zero;
+  AmdgpuDefaultStreamPinGuard amdgpu_pin(/*engage=*/default_stream_path);
+
   char *device_result_buffer{nullptr};
   // Here we have to guarantee the result_result_buffer isn't nullptr
   // It is interesting - The code following
@@ -314,7 +332,11 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
       }
     }
   }
-  if (transfers.size() > 0) {
+  // On the default-stream fast path the post-HtoD `stream_synchronize` is redundant: HtoD goes on the null stream and
+  // the subsequent `amdgpu_module->launch` reads `AMDGPUContext::stream_` (pinned to nullptr) so the kernel dispatch
+  // serialises with the HtoD by null-stream ordering. Outside the fast path the barrier remains load-bearing because
+  // HtoD on `active_stream` is async and per-group launches read it from a different stream.
+  if (transfers.size() > 0 && !default_stream_path) {
     AMDGPUDriver::get_instance().stream_synchronize(active_stream);
   }
   char *host_result_buffer = (char *)ctx.get_context().result_buffer;
@@ -377,18 +399,32 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
     AMDGPUDriver::get_instance().memcpy_device_to_host_async(host_result_buffer, device_result_buffer,
                                                              ctx.result_buffer_size, active_stream);
   }
+  // Copy data back to host. On the default-stream fast path the kernel ran on the null stream, so a sync
+  // `hipMemcpyDtoH` (host-blocking, on null stream) sees the kernel's writes without an explicit cross-stream barrier
+  // and host-drains the prior async `memcpy_device_to_host_async(host_result_buffer, ...)` queued on the same null
+  // stream - the explicit `stream_synchronize(nullptr)` calls collapse to no-ops. Outside the fast path the barriers
+  // remain load-bearing for cross-stream visibility (per-group kernel writes vs `active_stream` DtoH).
   if (transfers.size() > 0) {
-    AMDGPUDriver::get_instance().stream_synchronize(active_stream);
+    if (!default_stream_path) {
+      AMDGPUDriver::get_instance().stream_synchronize(active_stream);
+    }
     for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
       auto &idx = itr->first;
-      AMDGPUDriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
-                                                               ctx.array_runtime_sizes[idx.arg_id], active_stream);
+      if (default_stream_path) {
+        AMDGPUDriver::get_instance().memcpy_device_to_host(itr->second.first, (void *)device_ptrs[idx],
+                                                           ctx.array_runtime_sizes[idx.arg_id]);
+      } else {
+        AMDGPUDriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
+                                                                 ctx.array_runtime_sizes[idx.arg_id], active_stream);
+      }
+    }
+    if (!default_stream_path) {
+      AMDGPUDriver::get_instance().stream_synchronize(active_stream);
     }
-    AMDGPUDriver::get_instance().stream_synchronize(active_stream);
     for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
       executor->deallocate_memory_on_device(itr->second.second);
     }
-  } else if (ctx.result_buffer_size > 0) {
+  } else if (ctx.result_buffer_size > 0 && !default_stream_path) {
     AMDGPUDriver::get_instance().stream_synchronize(active_stream);
   }
   // Persistent scratch: no per-launch free for the per-handle `arg_buffer` / `runtime_context` or the launcher-global

quadrants/runtime/cuda/kernel_launcher.cpp

@@ -4,6 +4,7 @@
 #include "quadrants/runtime/cuda/cuda_utils.h"
 #include "quadrants/rhi/cuda/cuda_context.h"
 #include "quadrants/rhi/cuda/cuda_driver.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"
@@ -285,6 +286,23 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
 
   auto *active_stream = CUDAContext::get_instance().get_stream();
 
+  // Default-stream fast path: every HtoD / DtoH / kernel-dispatch in this launcher already routes through
+  // `active_stream`, so when entry `stream_ == nullptr` AND every offloaded task launches on the same `active_stream`
+  // (i.e. `stream_parallel_group_id == 0`), the entire chain serialises on the legacy default stream. The
+  // `stream_synchronize` barriers between phases collapse to no-ops the surrounding sync DtoH (host-blocking on
+  // pageable host memory) already drains, and the `CudaDefaultStreamPinGuard` re-pins `CUDAContext::stream_` to
+  // nullptr defensively across the launch in case an inner helper temporarily swaps it. Outside the fast path -
+  // user-supplied stream OR any task on `stream_parallel_group_id != 0` (per-group acquired streams differ from
+  // `active_stream`) - the cross-stream barriers below are load-bearing for HtoD / kernel / DtoH visibility and the
+  // pin guard would silently override the user-requested stream at the kernel-launch site, so the guard stays
+  // disengaged and main's sync semantics remain untouched. Closes the cross-stream-visibility break that
+  // `cuMemAllocAsync`-pool buffers show on pre-Ampere hardware (Turing T4 faults at `cuLaunchKernel` with
+  // `illegal-address` when consumer + producer streams differ) on the default-stream path.
+  const bool all_sgid_zero = std::all_of(offloaded_tasks.begin(), offloaded_tasks.end(),
+                                         [](const OffloadedTask &t) { return t.stream_parallel_group_id == 0; });
+  const bool default_stream_path = (active_stream == nullptr) && all_sgid_zero;
+  CudaDefaultStreamPinGuard cuda_pin(/*engage=*/default_stream_path);
+
   char *device_result_buffer{nullptr};
   // Launcher-global persistent `result_buffer`. See `kernel_launcher.h` for why this one is shared across handles
   // (kernel writes + synchronous host readback before any other reader runs). `arg_buffer` and `runtime_context`
@@ -370,7 +388,11 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
       }
     }
   }
-  if (transfers.size() > 0) {
+  // On the default-stream fast path the post-HtoD `stream_synchronize` is redundant: HtoD goes on the null stream and
+  // the subsequent `cuda_module->launch` reads `CUDAContext::stream_` (pinned to nullptr) so the kernel dispatch
+  // serialises with the HtoD by null-stream ordering. Outside the fast path the barrier remains load-bearing because
+  // HtoD on `active_stream` is async and per-group launches read it from a different stream.
+  if (transfers.size() > 0 && !default_stream_path) {
     CUDADriver::get_instance().stream_synchronize(active_stream);
   }
   char *host_result_buffer = (char *)ctx.get_context().result_buffer;
@@ -459,19 +481,32 @@ 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, active_stream);
   }
-  // copy data back to host
+  // Copy data back to host. On the default-stream fast path the kernel ran on the null stream, so a sync
+  // `cuMemcpyDtoH` (host-blocking, on null stream) sees the kernel's writes without an explicit cross-stream barrier
+  // and host-drains the prior async `memcpy_device_to_host_async(host_result_buffer, ...)` queued on the same null
+  // stream - the explicit `stream_synchronize(nullptr)` calls collapse to no-ops. Outside the fast path the barriers
+  // remain load-bearing for cross-stream visibility (per-group kernel writes vs `active_stream` DtoH).
   if (transfers.size() > 0) {
-    CUDADriver::get_instance().stream_synchronize(active_stream);
+    if (!default_stream_path) {
+      CUDADriver::get_instance().stream_synchronize(active_stream);
+    }
     for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
       auto &idx = itr->first;
-      CUDADriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
-                                                             ctx.array_runtime_sizes[idx.arg_id], active_stream);
+      if (default_stream_path) {
+        CUDADriver::get_instance().memcpy_device_to_host(itr->second.first, (void *)device_ptrs[idx],
+                                                         ctx.array_runtime_sizes[idx.arg_id]);
+      } else {
+        CUDADriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
+                                                               ctx.array_runtime_sizes[idx.arg_id], active_stream);
+      }
+    }
+    if (!default_stream_path) {
+      CUDADriver::get_instance().stream_synchronize(active_stream);
     }
-    CUDADriver::get_instance().stream_synchronize(active_stream);
     for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
       executor->deallocate_memory_on_device(itr->second.second);
     }
-  } else if (ctx.result_buffer_size > 0) {
+  } else if (ctx.result_buffer_size > 0 && !default_stream_path) {
     CUDADriver::get_instance().stream_synchronize(active_stream);
   }
   // Free per-call ephemeral buffers (explicit-stream path).  The free is stream-ordered: it won't execute until all

quadrants/runtime/gfx/runtime.cpp

@@ -210,11 +210,14 @@ class HostDeviceContextBlitter {
 
     if (require_sync) {
       if (readback_sizes.size()) {
+        // `readback_data` internally allocates a staging buffer, records a `buffer_copy` cmdlist, and submits it via
+        // `submit_synced` (which on every backend drains the compute stream's pending cmdbufs). Both Metal's
+        // `command_sync` and Vulkan's `vkQueueWaitIdle` drain everything queued, so the kernel cmdlist submitted just
+        // above completes as part of the same wait the readback already pays. A separate `wait_idle()` here would
+        // double the host-GPU round-trip per readback launch with no extra correctness; the `command_complete_sema`
+        // dependency edge is what the readback's submit honors on Vulkan, and Metal ignores `wait_semaphores` but
+        // drains in submission order so the kernel still completes before the staging buffer's `buffer_copy` runs.
         StreamSemaphore command_complete_sema = device_->get_compute_stream()->submit(cmdlist);
-
-        device_->wait_idle();
-
-        // In this case `readback_data` syncs
         QD_ASSERT(device_->readback_data(readback_dev_ptrs.data(), readback_host_ptrs.data(), readback_sizes.data(),
                                          int(readback_sizes.size()), {command_complete_sema}) == RhiResult::success);
       } else {

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/metadata_publish.cpp

@@ -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 {