Fix GPU out-of-memory crash by reusing overlay textures (#3614)

* Refactor TargetTexture into proper abstraction with ensure_size() method

* Remove redundant overlays_texture field, use view() directly

* Use if-let syntax in render_vello_scene_to_target_texture to avoid explicit unwrap

* Implement TargetTexture::new() constructor to avoid dummy textures

* fix compile error

* cleanup

* Avoid cloning texture view

---------

Co-authored-by: Timon <me@timon.zip>
Co-authored-by: Dennis Kobert <dennis@kobert.dev>

desktop/src/render/state.rs

@@ -1,6 +1,7 @@
-use crate::window::Window;
+use std::borrow::Cow;
 
-use crate::wrapper::{Color, WgpuContext, WgpuExecutor};
+use crate::window::Window;
+use crate::wrapper::{Color, TargetTexture, WgpuContext, WgpuExecutor};
 
 #[derive(derivative::Derivative)]
 #[derivative(Debug)]
@@ -17,7 +18,7 @@ pub(crate) struct RenderState {
 	viewport_scale: [f32; 2],
 	viewport_offset: [f32; 2],
 	viewport_texture: Option<wgpu::Texture>,
-	overlays_texture: Option<wgpu::Texture>,
+	overlays_texture: Option<TargetTexture>,
 	ui_texture: Option<wgpu::Texture>,
 	bind_group: Option<wgpu::BindGroup>,
 	#[derivative(Debug = "ignore")]
@@ -208,11 +209,6 @@ impl RenderState {
 		self.update_bindgroup();
 	}
 
-	pub(crate) fn bind_overlays_texture(&mut self, overlays_texture: wgpu::Texture) {
-		self.overlays_texture = Some(overlays_texture);
-		self.update_bindgroup();
-	}
-
 	pub(crate) fn bind_ui_texture(&mut self, bind_ui_texture: wgpu::Texture) {
 		self.ui_texture = Some(bind_ui_texture);
 		self.update_bindgroup();
@@ -236,12 +232,15 @@ impl RenderState {
 			return;
 		};
 		let size = glam::UVec2::new(viewport_texture.width(), viewport_texture.height());
-		let texture = futures::executor::block_on(self.executor.render_vello_scene_to_texture(&scene, size, &Default::default(), Color::TRANSPARENT));
-		let Ok(texture) = texture else {
-			tracing::error!("Error rendering overlays");
+		let result = futures::executor::block_on(
+			self.executor
+				.render_vello_scene_to_target_texture(&scene, size, &Default::default(), Color::TRANSPARENT, &mut self.overlays_texture),
+		);
+		if let Err(e) = result {
+			tracing::error!("Error rendering overlays: {:?}", e);
 			return;
-		};
-		self.bind_overlays_texture(texture);
+		}
+		self.update_bindgroup();
 	}
 
 	pub(crate) fn render(&mut self, window: &Window) -> Result<(), RenderError> {
@@ -312,7 +311,11 @@ impl RenderState {
 
 	fn update_bindgroup(&mut self) {
 		let viewport_texture_view = self.viewport_texture.as_ref().unwrap_or(&self.transparent_texture).create_view(&wgpu::TextureViewDescriptor::default());
-		let overlays_texture_view = self.overlays_texture.as_ref().unwrap_or(&self.transparent_texture).create_view(&wgpu::TextureViewDescriptor::default());
+		let overlays_texture_view = self
+			.overlays_texture
+			.as_ref()
+			.map(|target| Cow::Borrowed(target.view()))
+			.unwrap_or_else(|| Cow::Owned(self.transparent_texture.create_view(&wgpu::TextureViewDescriptor::default())));
 		let ui_texture_view = self.ui_texture.as_ref().unwrap_or(&self.transparent_texture).create_view(&wgpu::TextureViewDescriptor::default());
 
 		let bind_group = self.context.device.create_bind_group(&wgpu::BindGroupDescriptor {
@@ -324,7 +327,7 @@ impl RenderState {
 				},
 				wgpu::BindGroupEntry {
 					binding: 1,
-					resource: wgpu::BindingResource::TextureView(&overlays_texture_view),
+					resource: wgpu::BindingResource::TextureView(&overlays_texture_view.as_ref()),
 				},
 				wgpu::BindGroupEntry {
 					binding: 2,

desktop/wrapper/src/lib.rs

@@ -6,6 +6,7 @@ pub use graphite_editor::consts::FILE_EXTENSION;
 // TODO: Remove usage of this reexport in desktop create and remove this line
 pub use graphene_std::Color;
 
+pub use wgpu_executor::TargetTexture;
 pub use wgpu_executor::WgpuContext;
 pub use wgpu_executor::WgpuContextBuilder;
 pub use wgpu_executor::WgpuExecutor;

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

@@ -48,12 +48,58 @@ pub struct Surface {
 	pub blitter: TextureBlitter,
 }
 
+#[derive(Clone, Debug)]
 pub struct TargetTexture {
 	texture: wgpu::Texture,
 	view: wgpu::TextureView,
 	size: UVec2,
 }
 
+impl TargetTexture {
+	/// Creates a new TargetTexture with the specified size.
+	pub fn new(device: &wgpu::Device, size: UVec2) -> Self {
+		let size = size.max(UVec2::ONE);
+		let texture = device.create_texture(&wgpu::TextureDescriptor {
+			label: None,
+			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,
+			usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_SRC,
+			format: VELLO_SURFACE_FORMAT,
+			view_formats: &[],
+		});
+		let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+		Self { texture, view, size }
+	}
+
+	/// Ensures the texture has the specified size, creating a new one if needed.
+	/// This allows reusing the same texture across frames when the size hasn't changed.
+	pub fn ensure_size(&mut self, device: &wgpu::Device, size: UVec2) {
+		let size = size.max(UVec2::ONE);
+		if self.size == size {
+			return;
+		}
+
+		*self = Self::new(device, size);
+	}
+
+	/// Returns a reference to the texture view for rendering.
+	pub fn view(&self) -> &wgpu::TextureView {
+		&self.view
+	}
+
+	/// Returns a reference to the underlying texture.
+	pub fn texture(&self) -> &wgpu::Texture {
+		&self.texture
+	}
+}
+
 #[cfg(target_family = "wasm")]
 pub type Window = web_sys::HtmlCanvasElement;
 #[cfg(not(target_family = "wasm"))]
@@ -71,32 +117,14 @@ impl WgpuExecutor {
 		self.render_vello_scene_to_target_texture(scene, size, context, background, &mut output).await?;
 		Ok(output.unwrap().texture)
 	}
+	pub async fn render_vello_scene_to_target_texture(&self, scene: &Scene, size: UVec2, context: &RenderContext, background: Color, output: &mut Option<TargetTexture>) -> Result<()> {
+		// Initialize (lazily) if this is the first call
+		if output.is_none() {
+			*output = Some(TargetTexture::new(&self.context.device, size));
+		}
 
-	async fn render_vello_scene_to_target_texture(&self, scene: &Scene, size: UVec2, context: &RenderContext, background: Color, output: &mut Option<TargetTexture>) -> Result<()> {
-		let size = size.max(UVec2::ONE);
-		let target_texture = if let Some(target_texture) = output
-			&& target_texture.size == size
-		{
-			target_texture
-		} else {
-			let texture = self.context.device.create_texture(&wgpu::TextureDescriptor {
-				label: None,
-				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,
-				usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_SRC,
-				format: VELLO_SURFACE_FORMAT,
-				view_formats: &[],
-			});
-			let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
-			*output = Some(TargetTexture { texture, view, size });
-			output.as_mut().unwrap()
-		};
+		if let Some(target_texture) = output.as_mut() {
+			target_texture.ensure_size(&self.context.device, size);
 
 			let [r, g, b, a] = background.to_rgba8_srgb();
 			let render_params = RenderParams {
@@ -117,11 +145,12 @@ impl WgpuExecutor {
 					};
 					renderer.override_image(&image_brush.image, Some(texture_view));
 				}
-			renderer.render_to_texture(&self.context.device, &self.context.queue, scene, &target_texture.view, &render_params)?;
+				renderer.render_to_texture(&self.context.device, &self.context.queue, scene, target_texture.view(), &render_params)?;
 				for (image_brush, _) in context.resource_overrides.iter() {
 					renderer.override_image(&image_brush.image, None);
 				}
 			}
+		}
 		Ok(())
 	}