1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
//! One backend stream, and everything that hangs off it.
//!
//! A stream owns its own memory: allocations are per stream, not per device,
//! so a buffer resolves to the stream that created it and nowhere else. What
//! it carries beside the driver handle is the state that has to move in step
//! with the work queued on it — the deferred frees, the capture window, and
//! the per-launch info buffers a capture may not allocate inside.
use cubecl_core::{
MemoryConfiguration,
ir::MemoryDeviceProperties,
server::{BufferBinding, Handle, ServerError},
};
use cubecl_server::storage::PINNED_MEMORY_ALIGNMENT;
use cubecl_server::{
logging::ServerLogger,
memory_management::{
ErrorGraph, FailureId, MemoryAllocationMode, MemoryManagement, MemoryManagementOptions,
drop_queue::{self, FlushingPolicy, PendingDropQueue},
},
metadata_cache::{MetadataCachePolicy, MetadataInfoCache},
stream::{EventStreamBackend, StreamCapture, StreamMemory},
};
use std::sync::Arc;
use cubecl_server::driver::checked;
use crate::compute::{cpu::PinnedMemoryStorage, events::Fence, gpu::GpuStorage};
#[derive(Debug)]
pub struct Stream {
pub(crate) sys: cubecl_hip_sys::hipStream_t,
pub memory_management_gpu: MemoryManagement<GpuStorage>,
pub memory_management_cpu: MemoryManagement<PinnedMemoryStorage>,
pub drop_queue: drop_queue::PendingDropQueue<Fence>,
/// This stream's graph capture (see [`StreamCapture`]): its position in
/// the lifecycle, and the memory its recorded launches were given. Enforces
/// the ordered `graph_prepare` → `begin_capture` → `end_capture`
/// transitions and gates the deferral of fenced drop-queue flushes while a
/// capture is actively recording.
pub capturing: StreamCapture,
/// Reusable per-launch info buffers (kernel shapes/strides/scalars), keyed
/// by kernel and the exact info bytes. Admission and least-recently-used
/// eviction are decided by the cache's [`MetadataCachePolicy`]; the launch
/// path sets its [`CacheMode`] from the capture lifecycle, so during graph
/// capture every buffer is cached and none is evicted mid-capture. See
/// [`StreamCapture::cache_mode`].
pub info_cache: MetadataInfoCache<Handle>,
}
impl StreamMemory for Stream {
fn failure(&self, binding: &BufferBinding) -> Option<FailureId> {
self.memory_management_gpu
.failure(&binding.memory, binding.range())
}
fn taint(&mut self, binding: &BufferBinding, failure: FailureId, failures: &mut ErrorGraph) {
self.memory_management_gpu
.taint(&binding.memory, binding.range(), failure, failures)
}
fn written(&mut self, binding: &BufferBinding, failures: &mut ErrorGraph) {
self.memory_management_gpu
.written(&binding.memory, binding.range(), failures)
}
}
#[derive(new, Debug)]
pub struct HipStreamBackend {
mem_props: MemoryDeviceProperties,
mem_config: MemoryConfiguration,
mem_alignment: usize,
is_integrated: bool,
logger: Arc<ServerLogger>,
/// Programmatic main-GPU pool layout (see
/// [`Server::install_memory_pools`](cubecl_server::server::Server::install_memory_pools)):
/// streams created after it is set build their GPU pools from it instead
/// of the runtime default. Auxiliary pools are unaffected.
#[new(default)]
gpu_pools_override: Option<MemoryConfiguration>,
}
impl HipStreamBackend {
/// The layout streams build their main-GPU pools with, and the properties
/// to resolve it against.
pub(crate) fn gpu_pools(&self) -> (MemoryConfiguration, MemoryDeviceProperties) {
let config = self
.gpu_pools_override
.clone()
.unwrap_or_else(|| self.mem_config.clone());
(config, self.mem_props.clone())
}
/// Set the main-GPU pool layout for streams created from now on.
pub(crate) fn set_gpu_pools(&mut self, config: MemoryConfiguration) {
self.gpu_pools_override = Some(config);
}
}
impl EventStreamBackend for HipStreamBackend {
type Stream = Stream;
type Event = Fence;
fn create_stream(&self) -> Self::Stream {
// SAFETY: Calling HIP FFI to create a non-blocking stream. The stream handle is
// initialized by HIP on success (asserted below) and stored for the lifetime of
// this `Stream`.
let stream = unsafe {
let mut stream: cubecl_hip_sys::hipStream_t = std::ptr::null_mut();
let stream_status = cubecl_hip_sys::hipStreamCreateWithFlags(
&mut stream,
cubecl_hip_sys::hipStreamNonBlocking,
);
// Fatal: the pool hands out streams by value and every operation
// on this backend is issued against one.
checked("hipStreamCreateWithFlags", stream_status).expect("the pool needs a stream");
stream
};
let storage = GpuStorage::new(self.mem_alignment);
// The main GPU pool honors the programmatic pool override when one was
// installed (`install_memory_pools`). The pinned pool below is left
// alone: the override targets GPU activations, and the other pools
// have deliberate configurations that must not be overridden.
let (gpu_config, gpu_props) = self.gpu_pools();
let memory_management_gpu = MemoryManagement::from_configuration(
storage,
&gpu_props,
gpu_config,
self.logger.clone(),
MemoryManagementOptions::new("Main GPU Memory"),
);
// We use the same page size and memory pools configuration for CPU pinned memory, since we
// expect the CPU to have at least the same amount of RAM as GPU memory.
// The host was never measured, so this pool states no capacity.
let memory_management_cpu = MemoryManagement::from_configuration(
PinnedMemoryStorage::new(stream),
&MemoryDeviceProperties::new(
self.mem_props.max_page_size,
PINNED_MEMORY_ALIGNMENT as u64,
),
self.mem_config.clone(),
self.logger.clone(),
MemoryManagementOptions::new("Pinned CPU Memory").mode(MemoryAllocationMode::Auto),
);
Stream {
sys: stream,
memory_management_gpu,
memory_management_cpu,
capturing: StreamCapture::default(),
info_cache: MetadataInfoCache::new(MetadataCachePolicy::default()),
drop_queue: PendingDropQueue::new(FlushingPolicy {
max_bytes_count: match self.is_integrated {
// Integrated GPUs (APUs) share memory and IOMMU with the CPU.
// Flushing more frequently prevents the GPU from reaching 100%
// utilization, which avoids transient voltage droops and IOMMU
// TLB invalidation races that cause GPU hangs on 0→100% transitions.
//
// 16 was found empirically to be a good balance between stability
// and performance, 32 still exhibited intermittent hangs.
//
// In practice the performance difference is negligible since integrated
// GPUs are typically thermally constrained anyway.
true => 16,
false => 64,
},
..Default::default()
}),
}
}
fn flush(stream: &mut Self::Stream, _failures: &mut ErrorGraph) -> Self::Event {
Fence::new(stream.sys)
}
fn wait_event(stream: &mut Self::Stream, event: Self::Event) {
event.wait_async(stream.sys);
}
fn wait_event_sync(event: Self::Event) -> Result<(), ServerError> {
event.wait_sync()
}
fn handle_cursor(stream: &Self::Stream, binding: &BufferBinding) -> u64 {
// The slice cursor the sync logic compares against the origin stream's `last_synced`
// to decide whether to wait. A freed/reallocated slice falls back to `u64::MAX`,
// which conservatively forces a wait.
stream
.memory_management_gpu
.get_cursor(binding.memory.clone())
.unwrap_or(u64::MAX)
}
}