Improve texture caching by allowing weak refs and more formats (#4447)

* Destroy GPU textures when the last reference drops

* Request pooled textures synchronously

* Pool cached textures by format

* Add the owned GPU buffer type

* Route the remaining GPU allocations through the executor

* Add weak texture parking to the texture cache

* Raise the texture cache budget to 1GB for native and 512MB for wasm
This commit is contained in:
Timon
2026-08-27 15:19:00 +00:00
committed by GitHub
parent 8c48d5acce
commit 96cc520c4b
13 changed files with 213 additions and 169 deletions

View File

@@ -140,7 +140,7 @@ mod cpu {
pub use gpu::GPU;
#[cfg(feature = "wgpu")]
pub use gpu::Texture;
pub use gpu::{Texture, TextureWeakRef};
#[cfg(feature = "wgpu")]
mod gpu {
@@ -149,37 +149,57 @@ mod gpu {
use std::sync::Arc;
#[derive(Clone, Debug, PartialEq, Eq, Hash, DynAny)]
pub struct Texture(Arc<wgpu::Texture>);
pub struct Texture(Arc<TextureInner>);
#[derive(Debug, PartialEq, Eq, Hash)]
struct TextureInner(wgpu::Texture);
impl Drop for TextureInner {
fn drop(&mut self) {
self.0.destroy();
}
}
impl Texture {
pub fn is_shared(&self) -> bool {
Arc::strong_count(&self.0) > 1
}
pub fn is_weakly_shared(&self) -> bool {
Arc::weak_count(&self.0) > 0
}
pub fn downgrade(&self) -> TextureWeakRef {
TextureWeakRef(Arc::downgrade(&self.0))
}
}
#[derive(Clone, Debug)]
pub struct TextureWeakRef(std::sync::Weak<TextureInner>);
impl TextureWeakRef {
pub fn upgrade(&self) -> Option<Texture> {
self.0.upgrade().map(Texture)
}
}
impl Deref for Texture {
type Target = wgpu::Texture;
fn deref(&self) -> &Self::Target {
&self.0
&self.0.0
}
}
impl AsRef<wgpu::Texture> for Texture {
fn as_ref(&self) -> &wgpu::Texture {
&self.0
}
}
impl From<Arc<wgpu::Texture>> for Texture {
fn from(texture: Arc<wgpu::Texture>) -> Self {
Self(texture)
&self.0.0
}
}
impl From<wgpu::Texture> for Texture {
fn from(texture: wgpu::Texture) -> Self {
Self(Arc::new(texture))
}
}
impl From<Texture> for Arc<wgpu::Texture> {
fn from(texture: Texture) -> Self {
texture.0
Self(Arc::new(TextureInner(texture)))
}
}

View File

@@ -0,0 +1,34 @@
use std::ops::Deref;
use std::sync::Arc;
#[derive(Clone, Debug)]
pub struct Buffer(Arc<BufferInner>);
#[derive(Debug)]
struct BufferInner(wgpu::Buffer);
impl Drop for BufferInner {
fn drop(&mut self) {
self.0.destroy();
}
}
impl Deref for Buffer {
type Target = wgpu::Buffer;
fn deref(&self) -> &Self::Target {
&self.0.0
}
}
impl AsRef<wgpu::Buffer> for Buffer {
fn as_ref(&self) -> &wgpu::Buffer {
&self.0.0
}
}
impl From<wgpu::Buffer> for Buffer {
fn from(buffer: wgpu::Buffer) -> Self {
Self(Arc::new(BufferInner(buffer)))
}
}

View File

@@ -1,3 +1,4 @@
mod buffer;
mod context;
mod pipeline;
pub mod shader_runtime;
@@ -12,16 +13,18 @@ use core_types::color::SRGBA8;
use futures::lock::Mutex;
use glam::UVec2;
use graphene_application_io::{ApplicationIo, EditorApi};
use raster_types::Texture;
use std::sync::Arc;
use vello::{AaConfig, AaSupport, RenderParams, Renderer, RendererOptions, Scene};
use wgpu::util::DeviceExt;
use wgpu::{Origin3d, TextureAspect};
pub use buffer::Buffer;
pub use context::Context as WgpuContext;
pub use context::ContextBuilder as WgpuContextBuilder;
pub use pipeline::AsyncPipeline as AsyncWgpuPipeline;
pub use pipeline::Pipeline as WgpuPipeline;
pub use pipeline::PipelineCache as WgpuPipelineCache;
pub use raster_types::Texture;
pub use rendering::RenderContext;
pub use wgpu::Backends as WgpuBackends;
pub use wgpu::Features as WgpuFeatures;
@@ -30,7 +33,10 @@ pub use wgpu_sync::Instance as WgpuInstance;
pub use wgpu_sync::Queue as WgpuQueue;
pub use wgpu_sync::Surface as WgpuSurface;
const TEXTURE_CACHE_SIZE: u64 = 256 * 1024 * 1024; // 256 MiB
#[cfg(not(target_family = "wasm"))]
const TEXTURE_CACHE_SIZE: u64 = 1024 * 1024 * 1024; // 1GB
#[cfg(target_family = "wasm")]
const TEXTURE_CACHE_SIZE: u64 = 512 * 1024 * 1024; // 512MB
#[derive(dyn_any::DynAny, Clone)]
pub struct WgpuExecutor {
@@ -41,16 +47,12 @@ impl WgpuExecutor {
pub fn context(&self) -> &WgpuContext {
&self.inner.context
}
pub fn shader_runtime(&self) -> &ShaderRuntime {
&self.inner.shader_runtime
}
}
#[derive(dyn_any::DynAny)]
pub struct WgpuExecutorInner {
context: WgpuContext,
texture_cache: Mutex<TextureCache>,
texture_cache: std::sync::Mutex<TextureCache>,
vello_renderer: Mutex<Renderer>,
shader_runtime: ShaderRuntime,
}
@@ -69,7 +71,7 @@ impl<'a, T: ApplicationIo<Executor = WgpuExecutor>> From<&'a EditorApi<T>> for &
impl WgpuExecutor {
pub async fn render_vello_scene(&self, scene: &Scene, size: UVec2, context: &RenderContext, background: Option<Color>) -> Result<Texture> {
let texture = self.request_texture(size).await;
let texture = self.request_texture(size);
let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());
@@ -109,8 +111,20 @@ impl WgpuExecutor {
pipeline.init::<P>(self);
}
pub async fn request_texture(&self, size: UVec2) -> Texture {
self.inner.texture_cache.lock().await.request_texture(&self.context().device, size)
pub fn request_texture(&self, size: UVec2) -> Texture {
self.request_texture_with_format(size, wgpu::TextureFormat::Rgba8Unorm)
}
pub fn request_texture_with_format(&self, size: UVec2, format: wgpu::TextureFormat) -> Texture {
self.inner.texture_cache.lock().unwrap().request_texture(&self.context().device, size, format)
}
pub fn create_buffer(&self, desc: &wgpu::BufferDescriptor) -> Buffer {
self.context().device.create_buffer(desc).into()
}
pub fn create_buffer_init(&self, desc: &wgpu::util::BufferInitDescriptor) -> Buffer {
self.context().device.create_buffer_init(desc).into()
}
}
@@ -134,7 +148,7 @@ impl WgpuExecutor {
let texture_cache = TextureCache::new(TEXTURE_CACHE_SIZE);
let shader_runtime = ShaderRuntime::new(&context);
let shader_runtime = ShaderRuntime::default();
Some(Self {
inner: Arc::new(WgpuExecutorInner {

View File

@@ -1,20 +1,10 @@
use crate::WgpuContext;
use crate::shader_runtime::per_pixel_adjust_runtime::PerPixelAdjustShaderRuntime;
pub mod per_pixel_adjust_runtime;
pub const FULLSCREEN_VERTEX_SHADER_NAME: &str = "fullscreen_vertex_fullscreen_vertex";
#[derive(Default)]
pub struct ShaderRuntime {
context: WgpuContext,
per_pixel_adjust: PerPixelAdjustShaderRuntime,
}
impl ShaderRuntime {
pub fn new(context: &WgpuContext) -> Self {
Self {
context: context.clone(),
per_pixel_adjust: PerPixelAdjustShaderRuntime::new(),
}
}
}

View File

@@ -1,16 +1,17 @@
use crate::WgpuContext;
use crate::shader_runtime::{FULLSCREEN_VERTEX_SHADER_NAME, ShaderRuntime};
use crate::shader_runtime::FULLSCREEN_VERTEX_SHADER_NAME;
use crate::{Buffer, WgpuContext, WgpuExecutor};
use core_types::list::{Item, List};
use core_types::shaders::buffer_struct::BufferStruct;
use futures::lock::Mutex;
use glam::UVec2;
use raster_types::{GPU, Raster};
use std::borrow::Cow;
use std::collections::HashMap;
use wgpu::util::{BufferInitDescriptor, DeviceExt};
use wgpu::util::BufferInitDescriptor;
use wgpu::{
BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingResource, BindingType, Buffer, BufferBinding, BufferBindingType, BufferUsages, ColorTargetState, Face,
BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingResource, BindingType, BufferBinding, BufferBindingType, BufferUsages, ColorTargetState, Face,
FragmentState, FrontFace, LoadOp, Operations, PipelineLayoutDescriptor, PolygonMode, PrimitiveState, PrimitiveTopology, RenderPassColorAttachment, RenderPassDescriptor, RenderPipelineDescriptor,
ShaderModuleDescriptor, ShaderSource, ShaderStages, StoreOp, TextureDescriptor, TextureDimension, TextureFormat, TextureSampleType, TextureViewDescriptor, TextureViewDimension, VertexState,
ShaderModuleDescriptor, ShaderSource, ShaderStages, StoreOp, TextureFormat, TextureSampleType, TextureViewDescriptor, TextureViewDimension, VertexState,
};
pub struct PerPixelAdjustShaderRuntime {
@@ -32,22 +33,21 @@ impl PerPixelAdjustShaderRuntime {
}
}
impl ShaderRuntime {
impl WgpuExecutor {
pub async fn run_per_pixel_adjust<T: BufferStruct>(&self, shaders: &Shaders<'_>, textures: List<Raster<GPU>>, args: Option<&T>) -> List<Raster<GPU>> {
let mut cache = self.per_pixel_adjust.pipeline_cache.lock().await;
let mut cache = self.inner.shader_runtime.per_pixel_adjust.pipeline_cache.lock().await;
let pipeline = cache
.entry(shaders.fragment_shader_name.to_owned())
.or_insert_with(|| PerPixelAdjustGraphicsPipeline::new(&self.context, shaders));
.or_insert_with(|| PerPixelAdjustGraphicsPipeline::new(self.context(), shaders));
let arg_buffer = args.map(|args| {
let device = &self.context.device;
device.create_buffer_init(&BufferInitDescriptor {
self.create_buffer_init(&BufferInitDescriptor {
label: Some(&format!("{} arg buffer", pipeline.name.as_str())),
usage: BufferUsages::STORAGE,
contents: bytemuck::bytes_of(&T::write(*args)),
})
});
pipeline.dispatch(&self.context, textures, arg_buffer)
pipeline.dispatch(self, textures, arg_buffer)
}
}
@@ -160,9 +160,9 @@ impl PerPixelAdjustGraphicsPipeline {
}
}
pub fn dispatch(&self, context: &WgpuContext, textures: List<Raster<GPU>>, arg_buffer: Option<Buffer>) -> List<Raster<GPU>> {
pub fn dispatch(&self, executor: &WgpuExecutor, textures: List<Raster<GPU>>, arg_buffer: Option<Buffer>) -> List<Raster<GPU>> {
assert_eq!(self.has_uniform, arg_buffer.is_some());
let device = &context.device;
let device = &executor.context().device;
let name = self.name.as_str();
let mut cmd = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
@@ -203,16 +203,7 @@ impl PerPixelAdjustGraphicsPipeline {
entries,
});
let tex_out = device.create_texture(&TextureDescriptor {
label: Some(&format!("{name} texture out")),
size: tex_in.size(),
mip_level_count: 1,
sample_count: 1,
dimension: TextureDimension::D2,
format,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::RENDER_ATTACHMENT,
view_formats: &[format],
});
let tex_out = executor.request_texture_with_format(UVec2::new(tex_in.width(), tex_in.height()), format);
let view_out = tex_out.create_view(&TextureViewDescriptor::default());
let mut rp = cmd.begin_render_pass(&RenderPassDescriptor {
@@ -237,7 +228,7 @@ impl PerPixelAdjustGraphicsPipeline {
Item::from_parts(Raster::new_gpu(tex_out), attributes)
})
.collect::<List<_>>();
context.queue.submit([cmd.finish()]);
executor.context().queue.submit([cmd.finish()]);
out
}
}

View File

@@ -1,11 +1,10 @@
use glam::UVec2;
use raster_types::Texture;
use std::collections::VecDeque;
use std::sync::Arc;
pub(crate) struct TextureCache {
/// Always sorted oldest-first by insertion/last-use order.
textures: VecDeque<Arc<wgpu::Texture>>,
textures: VecDeque<Texture>,
max_free_bytes: u64,
}
@@ -17,49 +16,53 @@ impl TextureCache {
}
}
pub fn request_texture(&mut self, device: &wgpu::Device, size: UVec2) -> Texture {
pub fn request_texture(&mut self, device: &wgpu::Device, size: UVec2, format: wgpu::TextureFormat) -> Texture {
let size = size.max(UVec2::ONE);
if let Some(pos) = self
.textures
.iter()
.position(|texture| UVec2::new(texture.width(), texture.height()) == size && Arc::strong_count(texture) == 1)
.position(|texture| UVec2::new(texture.width(), texture.height()) == size && texture.format() == format && !texture.is_shared() && !texture.is_weakly_shared())
{
let entry = self.textures.remove(pos).unwrap();
let texture = entry.clone();
self.textures.push_back(entry);
return texture.into();
return texture;
}
let incoming_bytes = size.x as u64 * size.y as u64 * 4;
let incoming_bytes = size.x as u64 * size.y as u64 * format.block_copy_size(None).unwrap_or(4) as u64;
self.evict_until_fits(incoming_bytes);
let texture = Arc::new(device.create_texture(&wgpu::TextureDescriptor {
label: Some(&format!("cached_texture_{}x{}", size.x, size.y)),
size: wgpu::Extent3d {
width: size.x,
height: size.y,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8Unorm,
usage: wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::RENDER_ATTACHMENT,
view_formats: &[],
}));
let texture: Texture = device
.create_texture(&wgpu::TextureDescriptor {
label: Some(&format!("cached_{}x{}", size.x, size.y)),
size: wgpu::Extent3d {
width: size.x,
height: size.y,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format,
usage: {
let common = wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::RENDER_ATTACHMENT;
match format {
wgpu::TextureFormat::Rgba8Unorm => common | wgpu::TextureUsages::STORAGE_BINDING,
_ => common,
}
},
view_formats: &[],
})
.into();
self.textures.push_back(texture.clone());
texture.into()
texture
}
fn total_free_bytes(&self) -> u64 {
self.textures
.iter()
.filter(|texture| Arc::strong_count(texture) == 1)
.map(|texture| texture.memory_size_estimate())
.sum()
self.textures.iter().filter(|texture| !texture.is_shared()).map(|texture| texture.memory_size_estimate()).sum()
}
fn evict_until_fits(&mut self, incoming_bytes: u64) {
@@ -70,18 +73,19 @@ impl TextureCache {
return;
}
self.textures.retain(|texture| {
if free_bytes + incoming_bytes <= max_free_bytes {
return true;
}
if Arc::strong_count(texture) == 1 {
free_bytes -= texture.memory_size_estimate();
texture.destroy();
false
} else {
true
}
});
for parked in [false, true] {
self.textures.retain(|texture| {
if free_bytes + incoming_bytes <= max_free_bytes {
return true;
}
if !texture.is_shared() && texture.is_weakly_shared() == parked {
free_bytes -= texture.memory_size_estimate();
false
} else {
true
}
});
}
}
}
@@ -91,6 +95,6 @@ trait TextureMemoryCostEstimateExt {
impl TextureMemoryCostEstimateExt for wgpu::Texture {
fn memory_size_estimate(&self) -> u64 {
self.width() as u64 * self.height() as u64 * 4
self.width() as u64 * self.height() as u64 * self.format().block_copy_size(None).unwrap_or(4) as u64
}
}

View File

@@ -1,40 +1,33 @@
use crate::WgpuExecutor;
use crate::{Buffer, WgpuExecutor};
use core_types::Color;
use core_types::color::SRGBA8;
use core_types::list::{Item, List};
use core_types::ops::Convert;
use core_types::transform::Footprint;
use raster_types::Image;
use raster_types::{CPU, GPU, Raster};
use wgpu::util::{DeviceExt, TextureDataOrder};
use wgpu::{Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages};
use raster_types::{CPU, GPU, Raster, Texture};
use wgpu::{Extent3d, TextureFormat};
/// Uploads CPU image data to a GPU texture
///
/// Creates a new WGPU texture with RGBA8UnormSrgb format and uploads the provided
/// image data. The texture is configured for binding, copying, and source operations.
fn upload_to_texture(device: &wgpu::Device, queue: &wgpu::Queue, image: &Raster<CPU>) -> wgpu::Texture {
fn upload_to_texture(executor: &WgpuExecutor, queue: &wgpu::Queue, image: &Raster<CPU>) -> Texture {
let rgba8_data: Vec<SRGBA8> = image.data.iter().map(|x| (*x).into()).collect();
device.create_texture_with_data(
queue,
&TextureDescriptor {
label: Some("upload_to_texture staging texture"),
size: Extent3d {
width: image.width,
height: image.height,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: TextureDimension::D2,
format: TextureFormat::Rgba8UnormSrgb,
usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST | TextureUsages::COPY_SRC,
view_formats: &[],
},
TextureDataOrder::LayerMajor,
let texture = executor.request_texture_with_format(glam::UVec2::new(image.width, image.height), TextureFormat::Rgba8UnormSrgb);
queue.write_texture(
texture.as_image_copy(),
bytemuck::cast_slice(rgba8_data.as_slice()),
)
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(4 * image.width),
rows_per_image: Some(image.height),
},
Extent3d {
width: image.width,
height: image.height,
depth_or_array_layers: 1,
},
);
texture
}
/// Converts a Raster<GPU> texture to Raster<CPU> by downloading the underlying texture data.
@@ -44,7 +37,7 @@ fn upload_to_texture(device: &wgpu::Device, queue: &wgpu::Queue, image: &Raster<
/// - 4 bytes-per-pixel RGBA8
/// - Texture has COPY_SRC usage
struct RasterGpuToRasterCpuConverter {
buffer: wgpu::Buffer,
buffer: Buffer,
width: u32,
height: u32,
unpadded_bytes_per_row: u32,
@@ -52,7 +45,7 @@ struct RasterGpuToRasterCpuConverter {
_source: raster_types::Texture,
}
impl RasterGpuToRasterCpuConverter {
fn new(device: &wgpu::Device, encoder: &mut wgpu::CommandEncoder, data_gpu: Raster<GPU>) -> Self {
fn new(executor: &WgpuExecutor, encoder: &mut wgpu::CommandEncoder, data_gpu: Raster<GPU>) -> Self {
let texture = data_gpu.data();
let width = texture.width();
let height = texture.height();
@@ -62,7 +55,7 @@ impl RasterGpuToRasterCpuConverter {
let padded_bytes_per_row = unpadded_bytes_per_row.div_ceil(align) * align;
let buffer_size = padded_bytes_per_row as u64 * height as u64;
let buffer = device.create_buffer(&wgpu::BufferDescriptor {
let buffer = executor.create_buffer(&wgpu::BufferDescriptor {
label: Some("texture_download_buffer"),
size: buffer_size,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
@@ -151,13 +144,12 @@ impl<'i> Convert<List<Raster<GPU>>, &'i WgpuExecutor> for List<Raster<GPU>> {
/// Converts a `List<Raster<CPU>>` to `List<Raster<GPU>>` by uploading each image to a texture
impl<'i> Convert<List<Raster<GPU>>, &'i WgpuExecutor> for List<Raster<CPU>> {
async fn convert(self, _: Footprint, executor: &'i WgpuExecutor) -> List<Raster<GPU>> {
let device = &executor.context().device;
let queue = executor.context().queue.lock();
let list = self
.into_iter()
.map(|row| {
let (image, attributes) = row.into_parts();
let texture = upload_to_texture(device, &queue, &image);
let texture = upload_to_texture(executor, &queue, &image);
Item::from_parts(Raster::new_gpu(texture), attributes)
})
@@ -171,9 +163,8 @@ impl<'i> Convert<List<Raster<GPU>>, &'i WgpuExecutor> for List<Raster<CPU>> {
/// Converts single CPU raster to GPU by uploading to texture
impl<'i> Convert<Raster<GPU>, &'i WgpuExecutor> for Raster<CPU> {
async fn convert(self, _: Footprint, executor: &'i WgpuExecutor) -> Raster<GPU> {
let device = &executor.context().device;
let queue = executor.context().queue.lock();
let texture = upload_to_texture(device, &queue, &self);
let texture = upload_to_texture(executor, &queue, &self);
queue.submit([]);
Raster::new_gpu(texture)
@@ -202,7 +193,7 @@ impl<'i> Convert<List<Raster<CPU>>, &'i WgpuExecutor> for List<Raster<GPU>> {
for row in self {
let (element, attributes) = row.into_parts();
converters.push(RasterGpuToRasterCpuConverter::new(device, &mut encoder, element));
converters.push(RasterGpuToRasterCpuConverter::new(executor, &mut encoder, element));
rows_meta.push(Item::from_parts((), attributes));
}
@@ -239,7 +230,7 @@ impl<'i> Convert<Raster<CPU>, &'i WgpuExecutor> for Raster<GPU> {
label: Some("single_texture_download_encoder"),
});
let converter = RasterGpuToRasterCpuConverter::new(device, &mut encoder, self);
let converter = RasterGpuToRasterCpuConverter::new(executor, &mut encoder, self);
queue.submit([encoder.finish()]);