cubecl_server/memory_management/
mod.rs

1//! Memory pools and the bookkeeping behind a [`Handle`](crate::server::Handle).
2//!
3//! The value types a client sees — reports, configuration, the handle itself —
4//! are `cubecl-runtime`'s and re-exported here.
5
6pub use cubecl_runtime::memory_management::*;
7
8pub(crate) mod memory_pool;
9
10mod error_graph;
11mod taint;
12
13/// Export utilities to keep track of CPU buffers when performing async data copies.
14#[cfg(multi_threading)]
15pub mod drop_queue;
16
17pub use error_graph::*;
18pub use taint::*;
19
20mod dynamic;
21pub(crate) use dynamic::*;
22
23/// Dynamic memory management strategy.
24mod memory_manage;
25pub use memory_manage::*;
26
27/// Moving live allocations off outdated pages.
28pub mod relocation;
29
30/// Release memory after a failed wait only when the device is poisoned and can no longer use it.
31/// Otherwise completion is unknown, so keep the allocations reserved by leaking their owners.
32pub(crate) fn release_or_leak<T>(memory: T, err: &crate::server::ServerError) {
33    if err.is_device_poisoned() {
34        core::mem::drop(memory);
35    } else {
36        core::mem::forget(memory);
37        log::error!("leaking memory the device may still be using, after a failed wait: {err}");
38    }
39}
cubecl_server/memory_management/mod.rs

cubecl_server/memory_management/
mod.rs
