Add texture pool to render cache node (#3804)

* Add texture pool to render cache node

* Use direct texture copy instead of bitter and fix graphene_cli

* Remove warnings

* Fix wgpu import path

* Code review fixes

desktop/src/render/state.rs

@@ -11,13 +11,13 @@ pub(crate) struct RenderState {
 	executor: WgpuExecutor,
 	config: wgpu::SurfaceConfiguration,
 	render_pipeline: wgpu::RenderPipeline,
-	transparent_texture: wgpu::Texture,
+	transparent_texture: std::sync::Arc<wgpu::Texture>,
 	sampler: wgpu::Sampler,
 	desired_width: u32,
 	desired_height: u32,
 	viewport_scale: [f32; 2],
 	viewport_offset: [f32; 2],
-	viewport_texture: Option<wgpu::Texture>,
+	viewport_texture: Option<std::sync::Arc<wgpu::Texture>>,
 	overlays_texture: Option<TargetTexture>,
 	ui_texture: Option<wgpu::Texture>,
 	bind_group: Option<wgpu::BindGroup>,
@@ -50,7 +50,7 @@ impl RenderState {
 
 		surface.configure(&context.device, &config);
 
-		let transparent_texture = context.device.create_texture(&wgpu::TextureDescriptor {
+		let transparent_texture = std::sync::Arc::new(context.device.create_texture(&wgpu::TextureDescriptor {
 			label: Some("Transparent Texture"),
 			size: wgpu::Extent3d {
 				width: 1,
@@ -63,7 +63,7 @@ impl RenderState {
 			format: wgpu::TextureFormat::Bgra8UnormSrgb,
 			usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
 			view_formats: &[],
-		});
+		}));
 
 		// Create shader module
 		let shader = context.device.create_shader_module(wgpu::include_wgsl!("composite_shader.wgsl"));
@@ -207,7 +207,7 @@ impl RenderState {
 		}
 	}
 
-	pub(crate) fn bind_viewport_texture(&mut self, viewport_texture: wgpu::Texture) {
+	pub(crate) fn bind_viewport_texture(&mut self, viewport_texture: std::sync::Arc<wgpu::Texture>) {
 		self.viewport_texture = Some(viewport_texture);
 		self.update_bindgroup();
 	}

desktop/wrapper/src/lib.rs

@@ -58,7 +58,7 @@ impl DesktopWrapper {
 }
 
 pub enum NodeGraphExecutionResult {
-	HasRun(Option<wgpu::Texture>),
+	HasRun(Option<std::sync::Arc<wgpu::Texture>>),
 	NotRun,
 }
 

editor/src/node_graph_executor/runtime.rs

@@ -59,6 +59,9 @@ pub struct NodeRuntime {
 	/// Cached surface for Wasm viewport rendering (reused across frames)
 	#[cfg(all(target_family = "wasm", feature = "gpu"))]
 	wasm_viewport_surface: Option<wgpu_executor::WgpuSurface>,
+	/// Currently displayed texture, the runtime keeps a reference to it to avoid the texture getting destroyed while it is still in use.
+	#[cfg(all(target_family = "wasm", feature = "gpu"))]
+	current_viewport_texture: Option<ImageTexture>,
 }
 
 /// Messages passed from the editor thread to the node runtime thread.
@@ -144,6 +147,8 @@ impl NodeRuntime {
 			inspect_state: None,
 			#[cfg(all(target_family = "wasm", feature = "gpu"))]
 			wasm_viewport_surface: None,
+			#[cfg(all(target_family = "wasm", feature = "gpu"))]
+			current_viewport_texture: None,
 		}
 	}
 
@@ -275,7 +280,7 @@ impl NodeRuntime {
 								.gpu_executor()
 								.expect("GPU executor should be available when we receive a texture");
 
-							let raster_cpu = Raster::new_gpu(image_texture.texture).convert(Footprint::BOUNDLESS, executor).await;
+							let raster_cpu = Raster::new_gpu(image_texture.texture.as_ref().clone()).convert(Footprint::BOUNDLESS, executor).await;
 
 							let (data, width, height) = raster_cpu.to_flat_u8();
 
@@ -299,7 +304,7 @@ impl NodeRuntime {
 								.gpu_executor()
 								.expect("GPU executor should be available when we receive a texture");
 
-							let raster_cpu = Raster::new_gpu(image_texture.texture).convert(Footprint::BOUNDLESS, executor).await;
+							let raster_cpu = Raster::new_gpu(image_texture.texture.as_ref().clone()).convert(Footprint::BOUNDLESS, executor).await;
 
 							self.sender.send_eyedropper_preview(raster_cpu);
 							continue;
@@ -354,13 +359,22 @@ impl NodeRuntime {
 							);
 
 							let surface_texture = surface_inner.get_current_texture().expect("Failed to get surface texture");
+							self.current_viewport_texture = Some(image_texture.clone());
 
-							// Blit the rendered texture to the surface
-							surface.surface.blitter.copy(
-								&executor.context.device,
-								&mut encoder,
-								&image_texture.texture.create_view(&vello::wgpu::TextureViewDescriptor::default()),
-								&surface_texture.texture.create_view(&vello::wgpu::TextureViewDescriptor::default()),
+							encoder.copy_texture_to_texture(
+								vello::wgpu::TexelCopyTextureInfoBase {
+									texture: image_texture.texture.as_ref(),
+									mip_level: 0,
+									origin: Default::default(),
+									aspect: Default::default(),
+								},
+								vello::wgpu::TexelCopyTextureInfoBase {
+									texture: &surface_texture.texture,
+									mip_level: 0,
+									origin: Default::default(),
+									aspect: Default::default(),
+								},
+								image_texture.texture.size(),
 							);
 
 							executor.context.queue.submit([encoder.finish()]);

node-graph/graphene-cli/src/export.rs

@@ -69,7 +69,7 @@ pub async fn export_document(
 			}
 			RenderOutputType::Texture(image_texture) => {
 				// Convert GPU texture to CPU buffer
-				let gpu_raster = Raster::<GPU>::new_gpu(image_texture.texture);
+				let gpu_raster = Raster::<GPU>::new_gpu(image_texture.texture.as_ref().clone());
 				let cpu_raster: Raster<CPU> = gpu_raster.convert(Footprint::BOUNDLESS, wgpu_executor).await;
 				let (data, width, height) = cpu_raster.to_flat_u8();
 

node-graph/libraries/application-io/src/lib.rs

@@ -50,10 +50,10 @@ impl Size for web_sys::HtmlCanvasElement {
 	}
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, DynAny)]
 pub struct ImageTexture {
 	#[cfg(feature = "wgpu")]
-	pub texture: wgpu::Texture,
+	pub texture: Arc<wgpu::Texture>,
 	#[cfg(not(feature = "wgpu"))]
 	pub texture: (),
 }
@@ -89,10 +89,6 @@ impl PartialEq for ImageTexture {
 	}
 }
 
-unsafe impl StaticType for ImageTexture {
-	type Static = ImageTexture;
-}
-
 #[cfg(feature = "wgpu")]
 impl Size for ImageTexture {
 	fn size(&self) -> UVec2 {

node-graph/nodes/gstd/src/pixel_preview.rs

@@ -61,7 +61,7 @@ pub async fn pixel_preview<'a: 'n>(
 	let exec = editor_api.application_io.as_ref().unwrap().gpu_executor().unwrap();
 	let resampled = exec.resample_texture(&source_texture.texture, physical_resolution, &transform);
 
-	result.data = RenderOutputType::Texture(ImageTexture { texture: resampled });
+	result.data = RenderOutputType::Texture(ImageTexture { texture: resampled.into() });
 
 	result
 		.metadata

node-graph/nodes/gstd/src/render_cache.rs

@@ -106,6 +106,12 @@ struct TileCacheImpl {
 	timestamp: u64,
 	total_memory: usize,
 	cache_key: CacheKey,
+	texture_cache_resolution: UVec2,
+	/// Pool of textures of the same size: `texture_cache_resolution`.
+	/// Reusing textures reduces the wgpu allocation pressure,
+	/// which is a problem on web since we have to wait for
+	/// the browser to garbage collect unused textures, eating up memory.
+	texture_cache: Vec<Arc<wgpu::Texture>>,
 }
 
 #[derive(Clone, Default, dyn_any::DynAny, Debug)]
@@ -214,6 +220,36 @@ impl TileCacheImpl {
 		self.regions.clear();
 		self.total_memory = 0;
 	}
+
+	pub fn request_texture(&mut self, size: UVec2, device: &wgpu::Device) -> Arc<wgpu::Texture> {
+		if self.texture_cache_resolution != size {
+			self.texture_cache_resolution = size;
+			self.texture_cache.clear();
+		}
+		self.texture_cache.truncate(5);
+		for texture in &self.texture_cache {
+			if Arc::strong_count(texture) == 1 {
+				return Arc::clone(texture);
+			}
+		}
+		let texture = Arc::new(device.create_texture(&wgpu::TextureDescriptor {
+			label: Some("viewport_output"),
+			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::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::TEXTURE_BINDING,
+			view_formats: &[],
+		}));
+		self.texture_cache.push(texture.clone());
+
+		texture
+	}
 }
 
 impl TileCache {
@@ -224,6 +260,10 @@ impl TileCache {
 	pub fn store_regions(&self, regions: Vec<CachedRegion>) {
 		self.0.lock().unwrap().store_regions(regions);
 	}
+
+	pub fn request_texture(&self, size: UVec2, device: &wgpu::Device) -> Arc<wgpu::Texture> {
+		self.0.lock().unwrap().request_texture(size, device)
+	}
 }
 
 fn group_into_regions(tiles: &[TileCoord], max_region_area: u32) -> Vec<RenderRegion> {
@@ -413,7 +453,11 @@ pub async fn render_output_cache<'a: 'n>(
 	}
 
 	let exec = editor_api.application_io.as_ref().unwrap().gpu_executor().unwrap();
-	let (output_texture, combined_metadata) = composite_cached_regions(&all_regions, physical_resolution, &device_origin_offset, &footprint.transform, exec);
+
+	let device = &exec.context.device;
+	let output_texture = tile_cache.request_texture(physical_resolution, device);
+
+	let combined_metadata = composite_cached_regions(&all_regions, output_texture.as_ref(), &device_origin_offset, &footprint.transform, exec);
 
 	RenderOutput {
 		data: RenderOutputType::Texture(ImageTexture { texture: output_texture }),
@@ -462,7 +506,7 @@ where
 	let memory_size = (region_pixel_size.x * region_pixel_size.y) as usize * BYTES_PER_PIXEL;
 
 	CachedRegion {
-		texture: rendered_texture.texture,
+		texture: rendered_texture.texture.as_ref().clone(),
 		texture_size: region_pixel_size,
 		tiles: region.tiles.clone(),
 		metadata: result.metadata,
@@ -473,29 +517,14 @@ where
 
 fn composite_cached_regions(
 	regions: &[CachedRegion],
-	output_resolution: UVec2,
+	output_texture: &wgpu::Texture,
 	device_origin_offset: &DVec2,
 	viewport_transform: &DAffine2,
 	exec: &wgpu_executor::WgpuExecutor,
-) -> (wgpu::Texture, rendering::RenderMetadata) {
+) -> rendering::RenderMetadata {
 	let device = &exec.context.device;
 	let queue = &exec.context.queue;
-
-	// TODO: Use texture pool to reuse existing unused textures instead of allocating fresh ones every time
-	let output_texture = device.create_texture(&wgpu::TextureDescriptor {
-		label: Some("viewport_output"),
-		size: wgpu::Extent3d {
-			width: output_resolution.x,
-			height: output_resolution.y,
-			depth_or_array_layers: 1,
-		},
-		mip_level_count: 1,
-		sample_count: 1,
-		dimension: wgpu::TextureDimension::D2,
-		format: wgpu::TextureFormat::Rgba8Unorm,
-		usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC | wgpu::TextureUsages::TEXTURE_BINDING,
-		view_formats: &[],
-	});
+	let output_resolution = UVec2::new(output_texture.width(), output_texture.height());
 
 	let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("composite") });
 	let mut combined_metadata = rendering::RenderMetadata::default();
@@ -548,5 +577,5 @@ fn composite_cached_regions(
 	}
 
 	queue.submit([encoder.finish()]);
-	(output_texture, combined_metadata)
+	combined_metadata
 }

node-graph/nodes/gstd/src/render_node.rs

@@ -196,10 +196,11 @@ async fn render<'a: 'n>(ctx: impl Ctx + ExtractFootprint + ExtractVarArgs, edito
 				None
 			};
 
-			let texture = exec
-				.render_vello_scene_to_texture(&scene, physical_resolution, context, background)
-				.await
-				.expect("Failed to render Vello scene");
+			let texture = Arc::new(
+				exec.render_vello_scene_to_texture(&scene, physical_resolution, context, background)
+					.await
+					.expect("Failed to render Vello scene"),
+			);
 
 			RenderOutputType::Texture(ImageTexture { texture })
 		}