Skip to content

Commit cbb9a52

Browse files
duburcqaduburcqa
committed
[Perf] CUDA + AMDGPU launchers: pin stream to null at launcher scope; drop redundant pre/post-kernel stream_synchronize
1 parent 5211549 commit cbb9a52

7 files changed

Lines changed: 178 additions & 58 deletions

File tree

‎quadrants/rhi/amdgpu/amdgpu_stream_pin.h‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#pragma once
2+
3+
// RAII guard that pins `AMDGPUContext::stream_` to the legacy default stream (nullptr) for the duration of the scope,
4+
// then restores the prior `stream_` on destruction. Symmetric counterpart to `CudaDefaultStreamPinGuard` in
5+
// `rhi/cuda/cuda_stream_pin.h`; engaged at the same launcher-scope callsite (`KernelLauncher::launch_llvm_kernel` on
6+
// AMDGPU) on the default-stream fast path so the HtoD / kernel / DtoH chain lands on a single stream without an
7+
// explicit `stream_synchronize` between phases. AMDGPU has not been observed to hit the cross-stream pool-visibility
8+
// fault that motivated the CUDA guard on pre-Ampere Turing T4, but `AMDGPUContext::launch` now forwards `stream_` to
9+
// `hipModuleLaunchKernel` (post-streams plumbing), so the same same-stream-invariant rationale applies and keeping the
10+
// guard symmetric with CUDA hardens against future driver / hardware combos that tighten cross-stream pool semantics.
11+
// The launcher engages this guard only when entry `stream_ == nullptr` and all tasks have
12+
// `stream_parallel_group_id == 0`; on the explicit-stream path the streams feature requires the kernel to inherit
13+
// the user-supplied `stream_` and the guard stays disengaged.
14+
15+
#include "quadrants/platform/amdgpu/detect_amdgpu.h"
16+
17+
#if defined(QD_WITH_AMDGPU)
18+
#include "quadrants/rhi/amdgpu/amdgpu_context.h"
19+
#include "quadrants/rhi/amdgpu/amdgpu_driver.h"
20+
#endif
21+
22+
namespace quadrants::lang {
23+
24+
#if defined(QD_WITH_AMDGPU)
25+
struct AmdgpuDefaultStreamPinGuard {
26+
bool engaged{false};
27+
void *prev_stream{nullptr};
28+
explicit AmdgpuDefaultStreamPinGuard(bool engage) : engaged(engage) {
29+
if (engaged) {
30+
prev_stream = AMDGPUContext::get_instance().get_stream();
31+
AMDGPUContext::get_instance().set_stream(nullptr);
32+
}
33+
}
34+
~AmdgpuDefaultStreamPinGuard() {
35+
if (engaged) {
36+
AMDGPUContext::get_instance().set_stream(prev_stream);
37+
}
38+
}
39+
};
40+
#endif
41+
42+
} // namespace quadrants::lang

‎quadrants/rhi/cuda/cuda_stream_pin.h‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
#pragma once
2+
3+
// RAII guard that pins `CUDAContext::stream_` to the legacy default stream (nullptr) for the duration of the scope,
4+
// then restores the prior `stream_` on destruction. Pinning makes the consumer + producer streams match unconditionally
5+
// so the cross-stream-visibility break that CUDA's stream-ordered allocator (`cuMemAllocAsync` pool) shows on
6+
// pre-Ampere hardware (Turing T4 faults at `cuLaunchKernel` with `illegal-address`) is closed by construction. The
7+
// destructor does NOT call `stream_synchronize(nullptr)`: the helpers using this guard always issue a synchronous DtoH
8+
// inside the scope which already drains the null stream, and the main launcher's kernel-only fast path (no host-side
9+
// sync DtoH) deliberately stays asynchronous so kernel launches do not pay a forced per-launch host-GPU round-trip.
10+
// Restoring the prior stream while null-stream work is still in flight is safe: subsequent operations on the prior
11+
// stream that depend on the null-stream work will surface the standard cross-stream visibility constraints anyway, and
12+
// the next launch through this guard re-pins to null and serialises with prior null-stream work via legacy default
13+
// stream semantics. Used at the entry of every per-launch CUDA host path that issues HtoD or kernel work that needs
14+
// the same-stream invariant: `KernelLauncher::launch_llvm_kernel`, the adstack helpers
15+
// `dispatch_max_reducers_for_tasks` / `publish_adstack_metadata` / `publish_per_task_bound_count_device`. The launcher
16+
// engages the guard only on the default-stream fast path (entry `stream_ == nullptr` with all tasks on
17+
// `stream_parallel_group_id == 0`); on the explicit-stream path the streams feature requires the kernel to inherit
18+
// the user-supplied `stream_` and the guard stays disengaged. A symmetric `AmdgpuDefaultStreamPinGuard` lives in
19+
// `rhi/amdgpu/amdgpu_stream_pin.h` for consistency on AMDGPU - the streams plumbing now passes `stream_` through to
20+
// `hipModuleLaunchKernel`, so the same same-stream-invariant rationale applies, even though the cross-stream pool
21+
// fault that motivated the CUDA guard has not been observed on AMDGPU hardware.
22+
23+
#include "quadrants/platform/cuda/detect_cuda.h"
24+
25+
#if defined(QD_WITH_CUDA)
26+
#include "quadrants/rhi/cuda/cuda_context.h"
27+
#include "quadrants/rhi/cuda/cuda_driver.h"
28+
#endif
29+
30+
namespace quadrants::lang {
31+
32+
#if defined(QD_WITH_CUDA)
33+
struct CudaDefaultStreamPinGuard {
34+
bool engaged{false};
35+
void *prev_stream{nullptr};
36+
explicit CudaDefaultStreamPinGuard(bool engage) : engaged(engage) {
37+
if (engaged) {
38+
prev_stream = CUDAContext::get_instance().get_stream();
39+
CUDAContext::get_instance().set_stream(nullptr);
40+
}
41+
}
42+
~CudaDefaultStreamPinGuard() {
43+
if (engaged) {
44+
CUDAContext::get_instance().set_stream(prev_stream);
45+
}
46+
}
47+
};
48+
#endif
49+
50+
} // namespace quadrants::lang

‎quadrants/runtime/amdgpu/kernel_launcher.cpp‎

Lines changed: 42 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#include "quadrants/runtime/amdgpu/kernel_launcher.h"
44
#include "quadrants/rhi/amdgpu/amdgpu_context.h"
5+
#include "quadrants/rhi/amdgpu/amdgpu_stream_pin.h"
56
#include "quadrants/program/adstack_size_expr_eval.h"
67
#include "quadrants/program/launch_context_builder.h"
78
#include "quadrants/program/program.h"
@@ -230,6 +231,23 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
230231

231232
auto *active_stream = AMDGPUContext::get_instance().get_stream();
232233

234+
// Default-stream fast path: every HtoD / DtoH / kernel-dispatch in this launcher already routes through
235+
// `active_stream`, so when entry `stream_ == nullptr` AND every offloaded task launches on the same `active_stream`
236+
// (i.e. `stream_parallel_group_id == 0`), the entire chain serialises on the legacy default stream. The
237+
// `stream_synchronize` barriers between phases collapse to no-ops the surrounding sync DtoH (host-blocking on
238+
// pageable host memory) already drains, and the `AmdgpuDefaultStreamPinGuard` re-pins `AMDGPUContext::stream_` to
239+
// nullptr defensively across the launch in case an inner helper temporarily swaps it. Outside the fast path -
240+
// user-supplied stream OR any task on `stream_parallel_group_id != 0` (per-group acquired streams differ from
241+
// `active_stream`) - the cross-stream barriers below are load-bearing for HtoD / kernel / DtoH visibility and the
242+
// pin guard would silently override the user-requested stream at the kernel-launch site, so the guard stays
243+
// disengaged and main's sync semantics remain untouched. Symmetric with the CUDA launcher; the pre-Ampere pool
244+
// fault that motivated the CUDA pin has not been observed on AMDGPU, but `AMDGPUContext::launch` now forwards
245+
// `stream_` to `hipModuleLaunchKernel` so the same same-stream-invariant rationale applies.
246+
const bool all_sgid_zero = std::all_of(offloaded_tasks.begin(), offloaded_tasks.end(),
247+
[](const OffloadedTask &t) { return t.stream_parallel_group_id == 0; });
248+
const bool default_stream_path = (active_stream == nullptr) && all_sgid_zero;
249+
AmdgpuDefaultStreamPinGuard amdgpu_pin(/*engage=*/default_stream_path);
250+
233251
char *device_result_buffer{nullptr};
234252
// Here we have to guarantee the result_result_buffer isn't nullptr
235253
// It is interesting - The code following
@@ -314,7 +332,11 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
314332
}
315333
}
316334
}
317-
if (transfers.size() > 0) {
335+
// On the default-stream fast path the post-HtoD `stream_synchronize` is redundant: HtoD goes on the null stream and
336+
// the subsequent `amdgpu_module->launch` reads `AMDGPUContext::stream_` (pinned to nullptr) so the kernel dispatch
337+
// serialises with the HtoD by null-stream ordering. Outside the fast path the barrier remains load-bearing because
338+
// HtoD on `active_stream` is async and per-group launches read it from a different stream.
339+
if (transfers.size() > 0 && !default_stream_path) {
318340
AMDGPUDriver::get_instance().stream_synchronize(active_stream);
319341
}
320342
char *host_result_buffer = (char *)ctx.get_context().result_buffer;
@@ -377,18 +399,32 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
377399
AMDGPUDriver::get_instance().memcpy_device_to_host_async(host_result_buffer, device_result_buffer,
378400
ctx.result_buffer_size, active_stream);
379401
}
402+
// Copy data back to host. On the default-stream fast path the kernel ran on the null stream, so a sync
403+
// `hipMemcpyDtoH` (host-blocking, on null stream) sees the kernel's writes without an explicit cross-stream barrier
404+
// and host-drains the prior async `memcpy_device_to_host_async(host_result_buffer, ...)` queued on the same null
405+
// stream - the explicit `stream_synchronize(nullptr)` calls collapse to no-ops. Outside the fast path the barriers
406+
// remain load-bearing for cross-stream visibility (per-group kernel writes vs `active_stream` DtoH).
380407
if (transfers.size() > 0) {
381-
AMDGPUDriver::get_instance().stream_synchronize(active_stream);
408+
if (!default_stream_path) {
409+
AMDGPUDriver::get_instance().stream_synchronize(active_stream);
410+
}
382411
for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
383412
auto &idx = itr->first;
384-
AMDGPUDriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
385-
ctx.array_runtime_sizes[idx.arg_id], active_stream);
413+
if (default_stream_path) {
414+
AMDGPUDriver::get_instance().memcpy_device_to_host(itr->second.first, (void *)device_ptrs[idx],
415+
ctx.array_runtime_sizes[idx.arg_id]);
416+
} else {
417+
AMDGPUDriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
418+
ctx.array_runtime_sizes[idx.arg_id], active_stream);
419+
}
420+
}
421+
if (!default_stream_path) {
422+
AMDGPUDriver::get_instance().stream_synchronize(active_stream);
386423
}
387-
AMDGPUDriver::get_instance().stream_synchronize(active_stream);
388424
for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
389425
executor->deallocate_memory_on_device(itr->second.second);
390426
}
391-
} else if (ctx.result_buffer_size > 0) {
427+
} else if (ctx.result_buffer_size > 0 && !default_stream_path) {
392428
AMDGPUDriver::get_instance().stream_synchronize(active_stream);
393429
}
394430
// Persistent scratch: no per-launch free for the per-handle `arg_buffer` / `runtime_context` or the launcher-global

‎quadrants/runtime/cuda/kernel_launcher.cpp‎

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "quadrants/runtime/cuda/cuda_utils.h"
55
#include "quadrants/rhi/cuda/cuda_context.h"
66
#include "quadrants/rhi/cuda/cuda_driver.h"
7+
#include "quadrants/rhi/cuda/cuda_stream_pin.h"
78
#include "quadrants/runtime/llvm/llvm_runtime_executor.h"
89
#include "quadrants/program/adstack_size_expr_eval.h"
910
#include "quadrants/program/program.h"
@@ -285,6 +286,23 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
285286

286287
auto *active_stream = CUDAContext::get_instance().get_stream();
287288

289+
// Default-stream fast path: every HtoD / DtoH / kernel-dispatch in this launcher already routes through
290+
// `active_stream`, so when entry `stream_ == nullptr` AND every offloaded task launches on the same `active_stream`
291+
// (i.e. `stream_parallel_group_id == 0`), the entire chain serialises on the legacy default stream. The
292+
// `stream_synchronize` barriers between phases collapse to no-ops the surrounding sync DtoH (host-blocking on
293+
// pageable host memory) already drains, and the `CudaDefaultStreamPinGuard` re-pins `CUDAContext::stream_` to
294+
// nullptr defensively across the launch in case an inner helper temporarily swaps it. Outside the fast path -
295+
// user-supplied stream OR any task on `stream_parallel_group_id != 0` (per-group acquired streams differ from
296+
// `active_stream`) - the cross-stream barriers below are load-bearing for HtoD / kernel / DtoH visibility and the
297+
// pin guard would silently override the user-requested stream at the kernel-launch site, so the guard stays
298+
// disengaged and main's sync semantics remain untouched. Closes the cross-stream-visibility break that
299+
// `cuMemAllocAsync`-pool buffers show on pre-Ampere hardware (Turing T4 faults at `cuLaunchKernel` with
300+
// `illegal-address` when consumer + producer streams differ) on the default-stream path.
301+
const bool all_sgid_zero = std::all_of(offloaded_tasks.begin(), offloaded_tasks.end(),
302+
[](const OffloadedTask &t) { return t.stream_parallel_group_id == 0; });
303+
const bool default_stream_path = (active_stream == nullptr) && all_sgid_zero;
304+
CudaDefaultStreamPinGuard cuda_pin(/*engage=*/default_stream_path);
305+
288306
char *device_result_buffer{nullptr};
289307
// Launcher-global persistent `result_buffer`. See `kernel_launcher.h` for why this one is shared across handles
290308
// (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
370388
}
371389
}
372390
}
373-
if (transfers.size() > 0) {
391+
// On the default-stream fast path the post-HtoD `stream_synchronize` is redundant: HtoD goes on the null stream and
392+
// the subsequent `cuda_module->launch` reads `CUDAContext::stream_` (pinned to nullptr) so the kernel dispatch
393+
// serialises with the HtoD by null-stream ordering. Outside the fast path the barrier remains load-bearing because
394+
// HtoD on `active_stream` is async and per-group launches read it from a different stream.
395+
if (transfers.size() > 0 && !default_stream_path) {
374396
CUDADriver::get_instance().stream_synchronize(active_stream);
375397
}
376398
char *host_result_buffer = (char *)ctx.get_context().result_buffer;
@@ -459,19 +481,32 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx
459481
CUDADriver::get_instance().memcpy_device_to_host_async(host_result_buffer, device_result_buffer,
460482
ctx.result_buffer_size, active_stream);
461483
}
462-
// copy data back to host
484+
// Copy data back to host. On the default-stream fast path the kernel ran on the null stream, so a sync
485+
// `cuMemcpyDtoH` (host-blocking, on null stream) sees the kernel's writes without an explicit cross-stream barrier
486+
// and host-drains the prior async `memcpy_device_to_host_async(host_result_buffer, ...)` queued on the same null
487+
// stream - the explicit `stream_synchronize(nullptr)` calls collapse to no-ops. Outside the fast path the barriers
488+
// remain load-bearing for cross-stream visibility (per-group kernel writes vs `active_stream` DtoH).
463489
if (transfers.size() > 0) {
464-
CUDADriver::get_instance().stream_synchronize(active_stream);
490+
if (!default_stream_path) {
491+
CUDADriver::get_instance().stream_synchronize(active_stream);
492+
}
465493
for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
466494
auto &idx = itr->first;
467-
CUDADriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
468-
ctx.array_runtime_sizes[idx.arg_id], active_stream);
495+
if (default_stream_path) {
496+
CUDADriver::get_instance().memcpy_device_to_host(itr->second.first, (void *)device_ptrs[idx],
497+
ctx.array_runtime_sizes[idx.arg_id]);
498+
} else {
499+
CUDADriver::get_instance().memcpy_device_to_host_async(itr->second.first, (void *)device_ptrs[idx],
500+
ctx.array_runtime_sizes[idx.arg_id], active_stream);
501+
}
502+
}
503+
if (!default_stream_path) {
504+
CUDADriver::get_instance().stream_synchronize(active_stream);
469505
}
470-
CUDADriver::get_instance().stream_synchronize(active_stream);
471506
for (auto itr = transfers.begin(); itr != transfers.end(); itr++) {
472507
executor->deallocate_memory_on_device(itr->second.second);
473508
}
474-
} else if (ctx.result_buffer_size > 0) {
509+
} else if (ctx.result_buffer_size > 0 && !default_stream_path) {
475510
CUDADriver::get_instance().stream_synchronize(active_stream);
476511
}
477512
// Free per-call ephemeral buffers (explicit-stream path). The free is stream-ordered: it won't execute until all

‎quadrants/runtime/llvm/adstack_lazy_claim/bound_eval.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
#include "quadrants/rhi/amdgpu/amdgpu_context.h"
2828
#include "quadrants/rhi/amdgpu/amdgpu_driver.h"
2929

30-
#include "quadrants/runtime/llvm/adstack_lazy_claim/cuda_stream_pin.h"
30+
#include "quadrants/rhi/cuda/cuda_stream_pin.h"
3131

3232
namespace quadrants::lang {
3333

‎quadrants/runtime/llvm/adstack_lazy_claim/cuda_stream_pin.h‎

Lines changed: 0 additions & 43 deletions
This file was deleted.

‎quadrants/runtime/llvm/adstack_lazy_claim/metadata_publish.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
#include "quadrants/rhi/amdgpu/amdgpu_context.h"
3232
#endif
3333

34-
#include "quadrants/runtime/llvm/adstack_lazy_claim/cuda_stream_pin.h"
34+
#include "quadrants/rhi/cuda/cuda_stream_pin.h"
3535

3636
namespace quadrants::lang {
3737

0 commit comments

Comments
 (0)