cubecl_server/command/
staging.rs

1//! What a host-to-device copy should do about its source buffer.
2//!
3//! The decision is plain data with a constructor, which is the point: it is
4//! the only part of a copy that can be checked without a device, and inline in
5//! the copy it never was.
6
7use cubecl_common::bytes::AllocationProperty;
8
9/// A megabyte, for the thresholds below.
10const MB: usize = 1024 * 1024;
11
12/// Transfers up to this size go through a pinned staging buffer, which the
13/// driver can DMA from without a bounce. Above it the copy is long enough that
14/// the bounce costs less than pinning would.
15const STAGE_MAX: usize = 100 * MB;
16
17/// Above this size the drop queue is flushed after the copy, so the source is
18/// released promptly rather than waiting for the next batch to fill.
19const FLUSH_MIN: usize = 10 * MB;
20
21/// What a host-to-device copy does about its source.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct Staging {
24    /// Copy the source into pinned host memory before handing it to the
25    /// driver.
26    pub through_pinned: bool,
27    /// Flush the drop queue once the copy is enqueued, rather than leaving the
28    /// source held until the next batch fills.
29    pub flush_after: bool,
30}
31
32impl Staging {
33    /// What a copy of `size` bytes of memory allocated as `property` should do.
34    pub fn of(size: usize, property: AllocationProperty) -> Self {
35        let file_backed = matches!(property, AllocationProperty::File);
36        Self {
37            // File-backed data is staged whatever its size: the driver reads
38            // the source asynchronously, and it has to be real memory by then.
39            // Otherwise stage only what is small enough to be worth pinning,
40            // and never what is pinned already — that would be a redundant
41            // pinned-to-pinned copy.
42            through_pinned: file_backed
43                || (size < STAGE_MAX && !matches!(property, AllocationProperty::Pinned)),
44            // A large source, or one mapped from a file, is worth releasing
45            // now rather than holding until the batch fills.
46            flush_after: file_backed || size > FLUSH_MIN,
47        }
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    /// File-backed data is staged whatever its size.
56    ///
57    /// The driver reads the source after the copy is enqueued, so a mapping
58    /// that may still fault has to become real memory first. Size is why
59    /// staging is *worth* it for everything else; for a file it is why the
60    /// copy is correct.
61    #[test]
62    fn file_backed_data_is_always_staged() {
63        for size in [1, STAGE_MAX, STAGE_MAX * 4] {
64            assert!(Staging::of(size, AllocationProperty::File).through_pinned);
65        }
66    }
67
68    /// Already-pinned memory is handed to the driver as it is.
69    ///
70    /// Staging it would copy pinned memory into pinned memory, paying for a
71    /// bounce that exists to avoid one.
72    #[test]
73    fn pinned_data_is_never_restaged() {
74        for size in [1, STAGE_MAX / 2, STAGE_MAX * 4] {
75            assert!(!Staging::of(size, AllocationProperty::Pinned).through_pinned);
76        }
77    }
78
79    /// Ordinary host memory is staged only while it is small enough that
80    /// pinning costs less than the bounce it saves.
81    #[test]
82    fn plain_data_is_staged_up_to_the_threshold() {
83        assert!(Staging::of(STAGE_MAX - 1, AllocationProperty::Native).through_pinned);
84        assert!(!Staging::of(STAGE_MAX, AllocationProperty::Native).through_pinned);
85    }
86
87    /// A copy big enough to be worth releasing promptly flushes the queue,
88    /// whether or not it was staged.
89    #[test]
90    fn a_large_source_is_released_without_waiting_for_the_batch() {
91        assert!(!Staging::of(FLUSH_MIN, AllocationProperty::Native).flush_after);
92        assert!(Staging::of(FLUSH_MIN + 1, AllocationProperty::Native).flush_after);
93        assert!(Staging::of(1, AllocationProperty::File).flush_after);
94    }
95}
cubecl_server/command/staging.rs

cubecl_server/command/
staging.rs
