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

node-graph/libraries/wgpu-executor/src/buffer.rs

@@ -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)))
+	}
+}

node-graph/libraries/wgpu-executor/src/lib.rs

@@ -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 {

node-graph/libraries/wgpu-executor/src/shader_runtime/mod.rs

@@ -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(),
-		}
-	}
-}

node-graph/libraries/wgpu-executor/src/shader_runtime/per_pixel_adjust_runtime.rs

@@ -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
 	}
 }

node-graph/libraries/wgpu-executor/src/texture_cache.rs

@@ -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
 	}
 }

node-graph/libraries/wgpu-executor/src/texture_conversion.rs

@@ -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()]);