sim_lib_compute_wgpu/
segments.rs

1//! Segment planning for resident wgpu tensor buffers.
2
3/// One contiguous resident buffer segment.
4#[derive(Clone, Debug, PartialEq, Eq)]
5pub struct WgpuResidentSegment {
6    /// Segment ordinal.
7    pub index: usize,
8    /// Byte offset from the start of the logical tensor.
9    pub offset: u64,
10    /// Segment length in bytes.
11    pub bytes: u64,
12}
13
14/// Checked segment plan for a tensor payload.
15#[derive(Clone, Debug, PartialEq, Eq)]
16pub struct WgpuSegmentPlan {
17    /// Planned segments.
18    pub segments: Vec<WgpuResidentSegment>,
19}
20
21impl WgpuSegmentPlan {
22    /// Splits a payload by the stricter of the tile size and binding boundary.
23    pub fn new(total_bytes: u64, tile_bytes: u64, binding_bytes: u64) -> Self {
24        let boundary = tile_bytes.min(binding_bytes).max(1);
25        let mut segments = Vec::new();
26        let mut offset = 0;
27        while offset < total_bytes {
28            let bytes = (total_bytes - offset).min(boundary);
29            segments.push(WgpuResidentSegment {
30                index: segments.len(),
31                offset,
32                bytes,
33            });
34            offset += bytes;
35        }
36        Self { segments }
37    }
38
39    /// Returns the number of resident bytes in the plan.
40    pub fn total_bytes(&self) -> u64 {
41        self.segments.iter().map(|segment| segment.bytes).sum()
42    }
43}
sim_lib_compute_wgpu/segments.rs

sim_lib_compute_wgpu/
segments.rs
