gpui_wgpu/
wgpu_atlas.rs

1use anyhow::{Context as _, Result};
2use etagere::{BucketedAtlasAllocator, size2};
3use gpui::{
4    AtlasBackend, AtlasKey, AtlasState, AtlasTextureId, AtlasTextureKind, AtlasTextureList,
5    AtlasTile, Bounds, DevicePixels, PlatformAtlas, Point, Size,
6};
7use parking_lot::Mutex;
8use std::{borrow::Cow, ops, sync::Arc};
9
10use crate::WgpuContext;
11
12fn device_size_to_etagere(size: Size<DevicePixels>) -> etagere::Size {
13    size2(size.width.0, size.height.0)
14}
15
16fn etagere_point_to_device(point: etagere::Point) -> Point<DevicePixels> {
17    Point {
18        x: DevicePixels(point.x),
19        y: DevicePixels(point.y),
20    }
21}
22
23pub struct WgpuAtlas(Mutex<AtlasState<WgpuAtlasTextures>>);
24
25struct PendingUpload {
26    id: AtlasTextureId,
27    bounds: Bounds<DevicePixels>,
28    data: Vec<u8>,
29}
30
31struct WgpuAtlasTextures {
32    device: Arc<wgpu::Device>,
33    queue: Arc<wgpu::Queue>,
34    max_texture_size: u32,
35    color_texture_format: wgpu::TextureFormat,
36    storage: WgpuAtlasStorage,
37    pending_uploads: Vec<PendingUpload>,
38}
39
40pub struct WgpuTextureInfo {
41    pub view: wgpu::TextureView,
42}
43
44impl WgpuAtlas {
45    pub fn new(
46        device: Arc<wgpu::Device>,
47        queue: Arc<wgpu::Queue>,
48        color_texture_format: wgpu::TextureFormat,
49    ) -> Self {
50        let max_texture_size = device.limits().max_texture_dimension_2d;
51        WgpuAtlas(Mutex::new(AtlasState::new(WgpuAtlasTextures {
52            device,
53            queue,
54            max_texture_size,
55            color_texture_format,
56            storage: WgpuAtlasStorage::default(),
57            pending_uploads: Vec::new(),
58        })))
59    }
60
61    pub fn from_context(context: &WgpuContext) -> Self {
62        Self::new(
63            context.device.clone(),
64            context.queue.clone(),
65            context.color_texture_format(),
66        )
67    }
68
69    pub fn before_frame(&self) {
70        let mut lock = self.0.lock();
71        lock.backend.flush_uploads();
72    }
73
74    pub fn get_texture_info(&self, id: AtlasTextureId) -> WgpuTextureInfo {
75        let lock = self.0.lock();
76        let texture = &lock.backend.storage[id];
77        WgpuTextureInfo {
78            view: texture.view.clone(),
79        }
80    }
81
82    /// Clears all cached textures and tiles, forcing them to be recreated.
83    /// Use this for incremental recovery when the device is still valid.
84    pub fn clear(&self) {
85        self.0.lock().clear(|textures| {
86            textures.storage = WgpuAtlasStorage::default();
87            textures.pending_uploads.clear();
88        });
89    }
90
91    /// Handles device lost by clearing all textures and cached tiles.
92    /// The atlas will lazily recreate textures as needed on subsequent frames.
93    pub fn handle_device_lost(&self, context: &WgpuContext) {
94        self.0.lock().clear(|textures| {
95            textures.device = context.device.clone();
96            textures.queue = context.queue.clone();
97            textures.color_texture_format = context.color_texture_format();
98            textures.storage = WgpuAtlasStorage::default();
99            textures.pending_uploads.clear();
100        });
101    }
102}
103
104impl PlatformAtlas for WgpuAtlas {
105    fn get_or_insert_with<'a>(
106        &self,
107        key: AtlasKey,
108        build: &mut dyn FnMut() -> Result<Option<(Size<DevicePixels>, Cow<'a, [u8]>)>>,
109    ) -> Result<Option<AtlasTile>> {
110        self.0.lock().get_or_insert_with(key, build)
111    }
112
113    fn remove(&self, key: &AtlasKey) {
114        self.0.lock().remove(key);
115    }
116}
117
118impl AtlasBackend for WgpuAtlasTextures {
119    fn insert(
120        &mut self,
121        kind: AtlasTextureKind,
122        size: Size<DevicePixels>,
123        bytes: &[u8],
124    ) -> Result<AtlasTile> {
125        let tile = self.allocate(size, kind).context("failed to allocate")?;
126        self.upload_texture(tile.texture_id, tile.bounds, bytes);
127        Ok(tile)
128    }
129
130    fn remove(&mut self, tile: AtlasTile) {
131        let id = tile.texture_id;
132        let Some(texture_slot) = self.storage[id.kind].textures.get_mut(id.index as usize) else {
133            return;
134        };
135
136        if let Some(mut texture) = texture_slot.take() {
137            texture.allocator.deallocate(tile.tile_id.into());
138            texture.decrement_ref_count();
139            if texture.is_unreferenced() {
140                self.pending_uploads
141                    .retain(|upload| upload.id != texture.id);
142                self.storage[id.kind]
143                    .free_list
144                    .push(texture.id.index as usize);
145            } else {
146                *texture_slot = Some(texture);
147            }
148        }
149    }
150}
151
152impl WgpuAtlasTextures {
153    fn allocate(
154        &mut self,
155        size: Size<DevicePixels>,
156        texture_kind: AtlasTextureKind,
157    ) -> Option<AtlasTile> {
158        {
159            let textures = &mut self.storage[texture_kind];
160
161            if let Some(tile) = textures
162                .iter_mut()
163                .rev()
164                .find_map(|texture| texture.allocate(size))
165            {
166                return Some(tile);
167            }
168        }
169
170        let texture = self.push_texture(size, texture_kind);
171        texture.allocate(size)
172    }
173
174    fn push_texture(
175        &mut self,
176        min_size: Size<DevicePixels>,
177        kind: AtlasTextureKind,
178    ) -> &mut WgpuAtlasTexture {
179        const DEFAULT_ATLAS_SIZE: Size<DevicePixels> = Size {
180            width: DevicePixels(1024),
181            height: DevicePixels(1024),
182        };
183        let max_texture_size = self.max_texture_size as i32;
184        let max_atlas_size = Size {
185            width: DevicePixels(max_texture_size),
186            height: DevicePixels(max_texture_size),
187        };
188
189        let size = min_size.min(&max_atlas_size).max(&DEFAULT_ATLAS_SIZE);
190        let format = match kind {
191            AtlasTextureKind::Monochrome => wgpu::TextureFormat::R8Unorm,
192            AtlasTextureKind::Subpixel | AtlasTextureKind::Polychrome => self.color_texture_format,
193        };
194
195        let texture = self.device.create_texture(&wgpu::TextureDescriptor {
196            label: Some("atlas"),
197            size: wgpu::Extent3d {
198                width: size.width.0 as u32,
199                height: size.height.0 as u32,
200                depth_or_array_layers: 1,
201            },
202            mip_level_count: 1,
203            sample_count: 1,
204            dimension: wgpu::TextureDimension::D2,
205            format,
206            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
207            view_formats: &[],
208        });
209
210        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
211
212        let texture_list = &mut self.storage[kind];
213        let index = texture_list.free_list.pop();
214
215        let atlas_texture = WgpuAtlasTexture {
216            id: AtlasTextureId {
217                index: index.unwrap_or(texture_list.textures.len()) as u32,
218                kind,
219            },
220            allocator: BucketedAtlasAllocator::new(device_size_to_etagere(size)),
221            format,
222            texture,
223            view,
224            live_atlas_keys: 0,
225        };
226
227        if let Some(ix) = index {
228            texture_list.textures[ix] = Some(atlas_texture);
229            texture_list
230                .textures
231                .get_mut(ix)
232                .and_then(|t| t.as_mut())
233                .expect("texture must exist")
234        } else {
235            texture_list.textures.push(Some(atlas_texture));
236            texture_list
237                .textures
238                .last_mut()
239                .and_then(|t| t.as_mut())
240                .expect("texture must exist")
241        }
242    }
243
244    fn upload_texture(&mut self, id: AtlasTextureId, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
245        let data = self
246            .storage
247            .get(id)
248            .map(|texture| swizzle_upload_data(bytes, texture.format))
249            .unwrap_or_else(|| bytes.to_vec());
250
251        self.pending_uploads
252            .push(PendingUpload { id, bounds, data });
253    }
254
255    fn flush_uploads(&mut self) {
256        for upload in self.pending_uploads.drain(..) {
257            let Some(texture) = self.storage.get(upload.id) else {
258                continue;
259            };
260            let bytes_per_pixel = texture.bytes_per_pixel();
261
262            self.queue.write_texture(
263                wgpu::TexelCopyTextureInfo {
264                    texture: &texture.texture,
265                    mip_level: 0,
266                    origin: wgpu::Origin3d {
267                        x: upload.bounds.origin.x.0 as u32,
268                        y: upload.bounds.origin.y.0 as u32,
269                        z: 0,
270                    },
271                    aspect: wgpu::TextureAspect::All,
272                },
273                &upload.data,
274                wgpu::TexelCopyBufferLayout {
275                    offset: 0,
276                    bytes_per_row: Some(upload.bounds.size.width.0 as u32 * bytes_per_pixel as u32),
277                    rows_per_image: None,
278                },
279                wgpu::Extent3d {
280                    width: upload.bounds.size.width.0 as u32,
281                    height: upload.bounds.size.height.0 as u32,
282                    depth_or_array_layers: 1,
283                },
284            );
285        }
286    }
287}
288
289#[derive(Default)]
290struct WgpuAtlasStorage {
291    monochrome_textures: AtlasTextureList<WgpuAtlasTexture>,
292    subpixel_textures: AtlasTextureList<WgpuAtlasTexture>,
293    polychrome_textures: AtlasTextureList<WgpuAtlasTexture>,
294}
295
296impl ops::Index<AtlasTextureKind> for WgpuAtlasStorage {
297    type Output = AtlasTextureList<WgpuAtlasTexture>;
298    fn index(&self, kind: AtlasTextureKind) -> &Self::Output {
299        match kind {
300            AtlasTextureKind::Monochrome => &self.monochrome_textures,
301            AtlasTextureKind::Subpixel => &self.subpixel_textures,
302            AtlasTextureKind::Polychrome => &self.polychrome_textures,
303        }
304    }
305}
306
307impl ops::IndexMut<AtlasTextureKind> for WgpuAtlasStorage {
308    fn index_mut(&mut self, kind: AtlasTextureKind) -> &mut Self::Output {
309        match kind {
310            AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
311            AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
312            AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
313        }
314    }
315}
316
317impl WgpuAtlasStorage {
318    fn get(&self, id: AtlasTextureId) -> Option<&WgpuAtlasTexture> {
319        self[id.kind]
320            .textures
321            .get(id.index as usize)
322            .and_then(|t| t.as_ref())
323    }
324}
325
326impl ops::Index<AtlasTextureId> for WgpuAtlasStorage {
327    type Output = WgpuAtlasTexture;
328    fn index(&self, id: AtlasTextureId) -> &Self::Output {
329        let textures = match id.kind {
330            AtlasTextureKind::Monochrome => &self.monochrome_textures,
331            AtlasTextureKind::Subpixel => &self.subpixel_textures,
332            AtlasTextureKind::Polychrome => &self.polychrome_textures,
333        };
334        textures[id.index as usize]
335            .as_ref()
336            .expect("texture must exist")
337    }
338}
339
340struct WgpuAtlasTexture {
341    id: AtlasTextureId,
342    allocator: BucketedAtlasAllocator,
343    texture: wgpu::Texture,
344    view: wgpu::TextureView,
345    format: wgpu::TextureFormat,
346    live_atlas_keys: u32,
347}
348
349impl WgpuAtlasTexture {
350    fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
351        let allocation = self.allocator.allocate(device_size_to_etagere(size))?;
352        let tile = AtlasTile {
353            texture_id: self.id,
354            tile_id: allocation.id.into(),
355            padding: 0,
356            bounds: Bounds {
357                origin: etagere_point_to_device(allocation.rectangle.min),
358                size,
359            },
360        };
361        self.live_atlas_keys += 1;
362        Some(tile)
363    }
364
365    fn bytes_per_pixel(&self) -> u8 {
366        match self.format {
367            wgpu::TextureFormat::R8Unorm => 1,
368            wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Rgba8Unorm => 4,
369            _ => 4,
370        }
371    }
372
373    fn decrement_ref_count(&mut self) {
374        self.live_atlas_keys -= 1;
375    }
376
377    fn is_unreferenced(&self) -> bool {
378        self.live_atlas_keys == 0
379    }
380}
381
382fn swizzle_upload_data(bytes: &[u8], format: wgpu::TextureFormat) -> Vec<u8> {
383    match format {
384        wgpu::TextureFormat::Rgba8Unorm => {
385            let mut data = bytes.to_vec();
386            for pixel in data.chunks_exact_mut(4) {
387                pixel.swap(0, 2);
388            }
389            data
390        }
391        _ => bytes.to_vec(),
392    }
393}
394
395#[cfg(all(test, not(target_family = "wasm")))]
396mod tests {
397    use super::*;
398    use gpui::block_on;
399    use gpui::{ImageId, RenderImageParams};
400    use std::sync::Arc;
401
402    fn test_device_and_queue() -> anyhow::Result<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> {
403        block_on(async {
404            let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
405                backends: wgpu::Backends::all(),
406                flags: wgpu::InstanceFlags::default(),
407                backend_options: wgpu::BackendOptions::default(),
408                memory_budget_thresholds: wgpu::MemoryBudgetThresholds::default(),
409                display: None,
410            });
411            let adapter = instance
412                .request_adapter(&wgpu::RequestAdapterOptions {
413                    power_preference: wgpu::PowerPreference::LowPower,
414                    compatible_surface: None,
415                    force_fallback_adapter: false,
416                })
417                .await
418                .map_err(|error| anyhow::anyhow!("failed to request adapter: {error}"))?;
419            let (device, queue) = adapter
420                .request_device(&wgpu::DeviceDescriptor {
421                    label: Some("wgpu_atlas_test_device"),
422                    required_features: wgpu::Features::empty(),
423                    required_limits: wgpu::Limits::downlevel_defaults()
424                        .using_resolution(adapter.limits())
425                        .using_alignment(adapter.limits()),
426                    memory_hints: wgpu::MemoryHints::MemoryUsage,
427                    trace: wgpu::Trace::Off,
428                    experimental_features: wgpu::ExperimentalFeatures::disabled(),
429                })
430                .await
431                .map_err(|error| anyhow::anyhow!("failed to request device: {error}"))?;
432            Ok((Arc::new(device), Arc::new(queue)))
433        })
434    }
435
436    #[test]
437    fn before_frame_skips_uploads_for_removed_texture() -> anyhow::Result<()> {
438        let (device, queue) = test_device_and_queue()?;
439
440        let atlas = WgpuAtlas::new(device, queue, wgpu::TextureFormat::Bgra8Unorm);
441        let key = AtlasKey::Image(RenderImageParams {
442            image_id: ImageId(1),
443            frame_index: 0,
444        });
445        let size = Size {
446            width: DevicePixels(1),
447            height: DevicePixels(1),
448        };
449        let mut build = || Ok(Some((size, Cow::Owned(vec![0, 0, 0, 255]))));
450
451        // Regression test: before the fix, this panicked in flush_uploads
452        atlas
453            .get_or_insert_with(key.clone(), &mut build)?
454            .expect("tile should be created");
455        atlas.remove(&key);
456        atlas.before_frame();
457        Ok(())
458    }
459
460    #[test]
461    fn remove_deallocates_tile_space_for_reuse() -> anyhow::Result<()> {
462        let (device, queue) = test_device_and_queue()?;
463        let atlas = WgpuAtlas::new(device, queue, wgpu::TextureFormat::Bgra8Unorm);
464
465        let small = Size {
466            width: DevicePixels(64),
467            height: DevicePixels(64),
468        };
469        let big = Size {
470            width: DevicePixels(700),
471            height: DevicePixels(700),
472        };
473
474        let make_key = |image_id: usize| {
475            AtlasKey::Image(RenderImageParams {
476                image_id: ImageId(image_id),
477                frame_index: 0,
478            })
479        };
480        let insert = |key: AtlasKey, size: Size<DevicePixels>| {
481            let byte_count = (size.width.0 as usize) * (size.height.0 as usize) * 4;
482            atlas
483                .get_or_insert_with(key, &mut || {
484                    Ok(Some((size, Cow::Owned(vec![0u8; byte_count]))))
485                })
486                .expect("allocation should succeed")
487                .expect("callback returns Some")
488        };
489
490        let keeper_key = make_key(1);
491        let big_key_a = make_key(2);
492        let big_key_b = make_key(3);
493
494        let keeper_tile = insert(keeper_key, small);
495        let tile_a = insert(big_key_a.clone(), big);
496        assert_eq!(keeper_tile.texture_id, tile_a.texture_id);
497
498        atlas.remove(&big_key_a);
499        let tile_b = insert(big_key_b, big);
500        assert_eq!(tile_b.texture_id, keeper_tile.texture_id);
501        Ok(())
502    }
503
504    #[test]
505    fn swizzle_upload_data_preserves_bgra_uploads() {
506        let input = vec![0x10, 0x20, 0x30, 0x40];
507        assert_eq!(
508            swizzle_upload_data(&input, wgpu::TextureFormat::Bgra8Unorm),
509            input
510        );
511    }
512
513    #[test]
514    fn swizzle_upload_data_converts_bgra_to_rgba() {
515        let input = vec![0x10, 0x20, 0x30, 0x40, 0xAA, 0xBB, 0xCC, 0xDD];
516        assert_eq!(
517            swizzle_upload_data(&input, wgpu::TextureFormat::Rgba8Unorm),
518            vec![0x30, 0x20, 0x10, 0x40, 0xCC, 0xBB, 0xAA, 0xDD]
519        );
520    }
521}
gpui_wgpu/wgpu_atlas.rs

gpui_wgpu/
wgpu_atlas.rs
