molgfx_render/scene_gpu/brick_atlas/
upload.rs

1//! Fence-safe upload, publication and eviction over fixed storage.
2
3use super::page_table::{GpuPageTable, same_generation};
4use super::types::{
5    BrickAtlasConfig, BrickAtlasError, BrickAtlasKind, BrickAtlasMetrics, BrickAtlasPoll,
6    BrickAtlasUpload,
7};
8use molgfx_core::{BrickAddress, BrickCatalog, BrickDescriptor, BrickId, BrickValueRange};
9use molgfx_gpu::{
10    BufferDesc, BufferUsage, Device, FenceValue, Queue, TextureDesc, TextureDimension,
11    TextureFormat, TextureUsage, TextureViewDesc, TextureWrite, UploadRing,
12};
13use molgfx_semantic::BrickWorkingSet;
14
15#[derive(Clone, Copy, Debug, Default)]
16enum Pending {
17    #[default]
18    Vacant,
19    Upload {
20        descriptor: BrickDescriptor,
21        fence: FenceValue,
22    },
23    Eviction {
24        brick: BrickId,
25        fence: FenceValue,
26    },
27}
28
29/// Fixed-capacity physical atlas and generational sparse page table.
30#[derive(Debug)]
31pub struct GpuBrickAtlas<D: Device> {
32    config: BrickAtlasConfig,
33    texture: D::Texture,
34    view: D::TextureView,
35    page_buffer: D::Buffer,
36    page_table: GpuPageTable,
37    working_set: BrickWorkingSet,
38    published: Vec<Option<BrickDescriptor>>,
39    pending: Vec<Pending>,
40    upload_ring: UploadRing,
41    metrics: BrickAtlasMetrics,
42}
43
44impl<D: Device> GpuBrickAtlas<D> {
45    /// Allocates the complete atlas, page table and staging storage once.
46    ///
47    /// # Errors
48    ///
49    /// Rejects zero, overflowing or device-incompatible capacities.
50    pub fn new(
51        device: &D,
52        queue: &D::Queue,
53        catalog: &BrickCatalog,
54        config: BrickAtlasConfig,
55    ) -> Result<Self, BrickAtlasError> {
56        let [width, height, brick_depth] = config.stored_shape.map(u32::from);
57        if width == 0 || height == 0 || brick_depth == 0 || config.resident_capacity == 0 {
58            return Err(BrickAtlasError::InvalidConfiguration);
59        }
60        for descriptor in catalog.descriptors() {
61            if descriptor.metadata.shape.stored() != config.stored_shape
62                || !kind_matches(config.kind, descriptor.metadata.range)
63            {
64                return Err(BrickAtlasError::InvalidConfiguration);
65            }
66        }
67        let Ok(capacity) = u32::try_from(config.resident_capacity) else {
68            return Err(BrickAtlasError::InvalidConfiguration);
69        };
70        let Some(depth) = brick_depth.checked_mul(capacity) else {
71            return Err(BrickAtlasError::InvalidConfiguration);
72        };
73        let limit = device.capabilities().max_texture_dim_3d;
74        if width > limit || height > limit || depth > limit {
75            return Err(molgfx_gpu::GpuError::LimitExceeded {
76                resource: "sparse brick atlas",
77                limit: u64::from(limit),
78            }
79            .into());
80        }
81        let page_table = GpuPageTable::new(config.resident_capacity)?;
82        if page_table.byte_len() > device.capabilities().max_storage_buffer_bytes {
83            return Err(molgfx_gpu::GpuError::LimitExceeded {
84                resource: "sparse brick page table",
85                limit: device.capabilities().max_storage_buffer_bytes,
86            }
87            .into());
88        }
89        let format = texture_format(config.kind);
90        let texture = device.create_texture(&TextureDesc {
91            label: atlas_label(config.kind),
92            width,
93            height,
94            depth,
95            dimension: TextureDimension::D3,
96            format,
97            usage: TextureUsage::TEXTURE_BINDING.union(TextureUsage::COPY_DST),
98        })?;
99        let view = device.create_texture_view(&texture, &TextureViewDesc::default());
100        let page_buffer = device.create_buffer(&BufferDesc {
101            label: "sparse brick generational page table",
102            size: page_table.byte_len(),
103            usage: BufferUsage::STORAGE.union(BufferUsage::COPY_DST),
104        })?;
105        queue.write_buffer(&page_buffer, 0, page_table.bytes());
106        let atlas_bytes = u64::from(width)
107            .checked_mul(u64::from(height))
108            .and_then(|value| value.checked_mul(u64::from(depth)))
109            .and_then(|value| value.checked_mul(4))
110            .ok_or(BrickAtlasError::InvalidConfiguration)?;
111        let page_table_bytes = page_table.byte_len();
112        Ok(Self {
113            config,
114            texture,
115            view,
116            page_buffer,
117            page_table,
118            working_set: BrickWorkingSet::new(config.resident_capacity)?,
119            published: vec![None; config.resident_capacity],
120            pending: vec![Pending::Vacant; config.resident_capacity.saturating_mul(2)],
121            upload_ring: UploadRing::new(config.uploads)?,
122            metrics: BrickAtlasMetrics {
123                atlas_bytes,
124                page_table_bytes,
125                page_table_writes: 1,
126                ..BrickAtlasMetrics::default()
127            },
128        })
129    }
130
131    /// Resets only the bounded per-frame upload allowance.
132    pub fn begin_frame(&mut self) {
133        self.upload_ring.begin_epoch();
134    }
135
136    /// Stages one brick and submits an ordered fence for its publication.
137    ///
138    /// The provider allocation is borrowed only for this call. Its bytes are
139    /// copied once into the persistent upload ring and never into scene state.
140    ///
141    /// # Errors
142    ///
143    /// Returns typed shape, semantic, lifecycle or upload backpressure.
144    pub fn stage(
145        &mut self,
146        device: &D,
147        queue: &D::Queue,
148        upload: BrickAtlasUpload<'_>,
149    ) -> Result<FenceValue, BrickAtlasError> {
150        self.validate(upload)?;
151        let id = upload.descriptor.metadata.id;
152        if self.eviction_pending(id) {
153            return Err(BrickAtlasError::EvictionPending { brick: id });
154        }
155        let Some(pending_index) = self
156            .pending
157            .iter()
158            .position(|entry| matches!(entry, Pending::Vacant))
159        else {
160            return Err(BrickAtlasError::LifecycleCapacity {
161                capacity: self.pending.len(),
162            });
163        };
164        let reservation = self.upload_ring.reserve(upload.bytes.len())?;
165        self.upload_ring
166            .bytes_mut(reservation)?
167            .copy_from_slice(upload.bytes);
168        self.upload_ring.commit(reservation.ticket())?;
169        self.upload_ring.ensure_submittable(reservation.ticket())?;
170        let slot = match self.working_set.commit(upload.descriptor) {
171            Ok(slot) => slot.get(),
172            Err(error) => {
173                let _ = self.upload_ring.cancel(reservation.ticket());
174                return Err(error.into());
175            }
176        };
177        let offset = reservation.offset();
178        let end = offset.saturating_add(reservation.len());
179        let Some(staged) = self.upload_ring.staging_bytes().get(offset..end) else {
180            let _ = self.upload_ring.cancel(reservation.ticket());
181            return Err(BrickAtlasError::InvalidConfiguration);
182        };
183        let [width, height, depth] = self.config.stored_shape.map(u32::from);
184        queue.write_texture(
185            &self.texture,
186            &TextureWrite {
187                origin: [0, 0, slot.saturating_mul(depth)],
188                size: [width, height, depth],
189                bytes_per_row: width.saturating_mul(4),
190                rows_per_image: height,
191                data: staged,
192            },
193        );
194        let fence = queue.submit_tracked(device.create_command_encoder());
195        self.upload_ring.submit(reservation.ticket(), fence)?;
196        self.pending[pending_index] = Pending::Upload {
197            descriptor: upload.descriptor,
198            fence,
199        };
200        self.refresh_metrics();
201        Ok(fence)
202    }
203
204    /// Invalidates a mapping and releases its slot only after an ordered fence.
205    ///
206    /// # Errors
207    ///
208    /// Returns a typed absent-page, duplicate-eviction or capacity error.
209    pub fn request_eviction(
210        &mut self,
211        device: &D,
212        queue: &D::Queue,
213        brick: BrickId,
214    ) -> Result<FenceValue, BrickAtlasError> {
215        if self.eviction_pending(brick) {
216            return Err(BrickAtlasError::EvictionPending { brick });
217        }
218        let Some(page) = self.working_set.get(brick).copied() else {
219            return Err(molgfx_semantic::BrickWorkingSetError::NotResident { brick }.into());
220        };
221        let Some(pending_index) = self
222            .pending
223            .iter()
224            .position(|entry| matches!(entry, Pending::Vacant))
225        else {
226            return Err(BrickAtlasError::LifecycleCapacity {
227                capacity: self.pending.len(),
228            });
229        };
230        let Ok(slot) = usize::try_from(page.slot.get()) else {
231            return Err(BrickAtlasError::InvalidConfiguration);
232        };
233        if let Some(target) = self.published.get_mut(slot) {
234            *target = None;
235        }
236        self.flush_page_table(queue)?;
237        let fence = queue.submit_tracked(device.create_command_encoder());
238        self.pending[pending_index] = Pending::Eviction { brick, fence };
239        self.refresh_metrics();
240        Ok(fence)
241    }
242
243    /// Publishes current generations and releases fence-safe evictions.
244    ///
245    /// # Errors
246    ///
247    /// Device loss or an inconsistent working-set transition is returned.
248    pub fn poll(
249        &mut self,
250        device: &D,
251        queue: &D::Queue,
252    ) -> Result<BrickAtlasPoll, BrickAtlasError> {
253        let completed = queue.completed_fence(device)?;
254        let _ = self.upload_ring.retire(completed);
255        let mut result = BrickAtlasPoll {
256            completed_fence: completed,
257            ..BrickAtlasPoll::default()
258        };
259        let mut table_changed = false;
260        for index in 0..self.pending.len() {
261            match self.pending[index] {
262                Pending::Upload { descriptor, fence } if fence <= completed => {
263                    let current = self.working_set.get(descriptor.metadata.id).copied();
264                    let publish = current.is_some_and(|page| {
265                        same_generation(
266                            descriptor,
267                            BrickDescriptor {
268                                chunk: page.chunk,
269                                metadata: page.metadata,
270                            },
271                        )
272                    }) && !self.eviction_pending(descriptor.metadata.id);
273                    if publish {
274                        if let Some(page) = current
275                            && let Ok(slot) = usize::try_from(page.slot.get())
276                            && let Some(target) = self.published.get_mut(slot)
277                        {
278                            *target = Some(descriptor);
279                            result.uploads_published += 1;
280                            table_changed = true;
281                        }
282                    } else {
283                        result.stale_completions += 1;
284                        self.metrics.stale_completions =
285                            self.metrics.stale_completions.saturating_add(1);
286                    }
287                    self.pending[index] = Pending::Vacant;
288                }
289                Pending::Eviction { brick, fence } if fence <= completed => {
290                    let page = self.working_set.evict(brick)?;
291                    let Ok(slot) = usize::try_from(page.slot.get()) else {
292                        return Err(BrickAtlasError::InvalidConfiguration);
293                    };
294                    if let Some(target) = self.published.get_mut(slot) {
295                        *target = None;
296                    }
297                    self.pending[index] = Pending::Vacant;
298                    result.evictions_completed += 1;
299                    table_changed = true;
300                }
301                Pending::Vacant | Pending::Upload { .. } | Pending::Eviction { .. } => {}
302            }
303        }
304        if table_changed {
305            self.flush_page_table(queue)?;
306        }
307        self.refresh_metrics();
308        Ok(result)
309    }
310
311    /// Resolves an exact global logical address to an atlas-local slot.
312    #[must_use]
313    pub fn resolve(&self, address: BrickAddress) -> Option<u32> {
314        self.published.iter().enumerate().find_map(|(slot, value)| {
315            value
316                .filter(|descriptor| descriptor.metadata.address == address)
317                .and_then(|_| u32::try_from(slot).ok())
318        })
319    }
320
321    /// Physical atlas view used by scalar, segmentation and field shaders.
322    #[must_use]
323    pub const fn view(&self) -> &D::TextureView {
324        &self.view
325    }
326
327    /// Persistent generational hash table consumed by sparse shader lookup.
328    #[must_use]
329    pub const fn page_buffer(&self) -> &D::Buffer {
330        &self.page_buffer
331    }
332
333    /// Hash mask supplied to sparse lookup uniforms.
334    #[must_use]
335    pub const fn page_mask(&self) -> u32 {
336        self.page_table.mask()
337    }
338
339    /// Common physical shape, including halo voxels.
340    #[must_use]
341    pub const fn stored_shape(&self) -> [u16; 3] {
342        self.config.stored_shape
343    }
344
345    /// Current bounded memory and lifecycle counters.
346    #[must_use]
347    pub const fn metrics(&self) -> BrickAtlasMetrics {
348        self.metrics
349    }
350
351    fn validate(&self, upload: BrickAtlasUpload<'_>) -> Result<(), BrickAtlasError> {
352        let metadata = upload.descriptor.metadata;
353        if metadata.shape.stored() != self.config.stored_shape {
354            return Err(BrickAtlasError::ShapeMismatch { brick: metadata.id });
355        }
356        let expected = usize::try_from(metadata.shape.voxel_count())
357            .ok()
358            .and_then(|count| count.checked_mul(4))
359            .ok_or(BrickAtlasError::InvalidConfiguration)?;
360        if upload.bytes.len() != expected {
361            return Err(BrickAtlasError::PayloadSize {
362                brick: metadata.id,
363                expected,
364                received: upload.bytes.len(),
365            });
366        }
367        if !kind_matches(self.config.kind, metadata.range) {
368            return Err(BrickAtlasError::KindMismatch { brick: metadata.id });
369        }
370        Ok(())
371    }
372
373    fn flush_page_table(&mut self, queue: &D::Queue) -> Result<(), BrickAtlasError> {
374        self.page_table.rebuild(&self.published, self.config.kind)?;
375        queue.write_buffer(&self.page_buffer, 0, self.page_table.bytes());
376        self.metrics.page_table_writes = self.metrics.page_table_writes.saturating_add(1);
377        Ok(())
378    }
379
380    fn eviction_pending(&self, brick: BrickId) -> bool {
381        self.pending.iter().any(
382            |entry| matches!(entry, Pending::Eviction { brick: pending, .. } if *pending == brick),
383        )
384    }
385
386    fn refresh_metrics(&mut self) {
387        self.metrics.resident_pages = self
388            .published
389            .iter()
390            .filter(|value| value.is_some())
391            .count();
392        self.metrics.pending_uploads = self
393            .pending
394            .iter()
395            .filter(|entry| matches!(entry, Pending::Upload { .. }))
396            .count();
397        self.metrics.pending_evictions = self
398            .pending
399            .iter()
400            .filter(|entry| matches!(entry, Pending::Eviction { .. }))
401            .count();
402    }
403}
404
405const fn kind_matches(kind: BrickAtlasKind, range: BrickValueRange) -> bool {
406    matches!(
407        (kind, range),
408        (
409            BrickAtlasKind::Scalar | BrickAtlasKind::Surface,
410            BrickValueRange::Scalar { .. }
411        ) | (
412            BrickAtlasKind::Segmentation,
413            BrickValueRange::Segmentation { .. }
414        ) | (BrickAtlasKind::Occupancy, BrickValueRange::Occupancy { .. })
415    )
416}
417
418const fn texture_format(kind: BrickAtlasKind) -> TextureFormat {
419    match kind {
420        BrickAtlasKind::Segmentation => TextureFormat::R32Uint,
421        BrickAtlasKind::Scalar | BrickAtlasKind::Occupancy | BrickAtlasKind::Surface => {
422            TextureFormat::R32Float
423        }
424    }
425}
426
427const fn atlas_label(kind: BrickAtlasKind) -> &'static str {
428    match kind {
429        BrickAtlasKind::Scalar => "sparse scalar brick atlas",
430        BrickAtlasKind::Segmentation => "sparse segmentation brick atlas",
431        BrickAtlasKind::Occupancy => "sparse occupancy brick atlas",
432        BrickAtlasKind::Surface => "sparse surface brick atlas",
433    }
434}
molgfx_render/scene_gpu/brick_atlas/upload.rs

molgfx_render/scene_gpu/brick_atlas/
upload.rs
