Add the async convert trait and restore the GPU texture download

node-graph/libraries/core-types/src/ops.rs

@@ -47,6 +47,14 @@ pub trait Convert<T, C>: Sized {
 	fn convert(self, footprint: Footprint, converter: C) -> T;
 }
 
+/// The genuinely asynchronous counterpart of [`Convert`], for conversions whose work completes outside
+/// the evaluation, such as the GPU-to-CPU texture readback. Consumed by the `convert_async` kernel on
+/// the async source tier; a conversion pair implements exactly one of the two traits.
+pub trait ConvertAsync<T, C>: Sized {
+	#[must_use]
+	fn convert(self, footprint: Footprint, converter: C) -> crate::runtime::SourceFuture<T>;
+}
+
 impl<T: ToString + Send> Convert<String, ()> for T {
 	/// Converts this type into a `String` using its `ToString` implementation.
 	#[inline]

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

@@ -3,7 +3,8 @@ use core_types::Color;
 use core_types::Ctx;
 use core_types::color::SRGBA8;
 use core_types::list::{Item, List};
-use core_types::ops::Convert;
+use core_types::ops::{Convert, ConvertAsync};
+use core_types::runtime::SourceFuture;
 use core_types::transform::Footprint;
 use raster_types::Image;
 use raster_types::{CPU, GPU, Raster};
@@ -85,6 +86,172 @@ impl<'i> Convert<List<Raster<CPU>>, &'i WgpuExecutor> for List<Raster<CPU>> {
 	}
 }
 
+/// Converts a Raster<GPU> texture to Raster<CPU> by downloading the underlying texture data.
+///
+/// Assumptions:
+/// - 2D texture, mip level 0
+/// - 4 bytes-per-pixel RGBA8
+/// - Texture has COPY_SRC usage
+struct RasterGpuToRasterCpuConverter {
+	buffer: wgpu::Buffer,
+	width: u32,
+	height: u32,
+	unpadded_bytes_per_row: u32,
+	padded_bytes_per_row: u32,
+	_source: raster_types::Texture,
+}
+impl RasterGpuToRasterCpuConverter {
+	fn new(device: &wgpu::Device, encoder: &mut wgpu::CommandEncoder, data_gpu: Raster<GPU>) -> Self {
+		let texture = data_gpu.data();
+		let width = texture.width();
+		let height = texture.height();
+		let bytes_per_pixel = 4; // RGBA8
+		let unpadded_bytes_per_row = width * bytes_per_pixel;
+		let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
+		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 {
+			label: Some("texture_download_buffer"),
+			size: buffer_size,
+			usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+			mapped_at_creation: false,
+		});
+
+		encoder.copy_texture_to_buffer(
+			wgpu::TexelCopyTextureInfo {
+				texture,
+				mip_level: 0,
+				origin: wgpu::Origin3d::ZERO,
+				aspect: wgpu::TextureAspect::All,
+			},
+			wgpu::TexelCopyBufferInfo {
+				buffer: &buffer,
+				layout: wgpu::TexelCopyBufferLayout {
+					offset: 0,
+					bytes_per_row: Some(padded_bytes_per_row),
+					rows_per_image: Some(height),
+				},
+			},
+			Extent3d {
+				width,
+				height,
+				depth_or_array_layers: 1,
+			},
+		);
+
+		Self {
+			buffer,
+			width,
+			height,
+			unpadded_bytes_per_row,
+			padded_bytes_per_row,
+			// Keep source texture alive
+			_source: data_gpu.texture.clone(),
+		}
+	}
+
+	async fn convert(self, device: &wgpu::Device) -> Result<Raster<CPU>, wgpu::BufferAsyncError> {
+		let buffer_slice = self.buffer.slice(..);
+		let (sender, receiver) = futures::channel::oneshot::channel();
+		buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
+			let _ = sender.send(result);
+		});
+
+		let _ = device.poll(wgpu::wgt::PollType::wait_indefinitely());
+
+		receiver.await.expect("Failed to receive map result")?;
+
+		let view = buffer_slice.get_mapped_range();
+
+		let row_stride = self.padded_bytes_per_row as usize;
+		let row_bytes = self.unpadded_bytes_per_row as usize;
+		let mut cpu_data: Vec<Color> = Vec::with_capacity((self.width * self.height) as usize);
+		for row in 0..self.height as usize {
+			let start = row * row_stride;
+			let row_slice = &view[start..start + row_bytes];
+			for px in row_slice.chunks_exact(4) {
+				// `Image<Color>` pixels are stored linear-light with associated (premultiplied) alpha
+				let srgba = SRGBA8::new(px[0], px[1], px[2], px[3]);
+				cpu_data.push(Color::from(srgba).apply_opacity(px[3] as f32 / 255.));
+			}
+		}
+
+		drop(view);
+		self.buffer.unmap();
+		let cpu_image = Image {
+			data: cpu_data,
+			width: self.width,
+			height: self.height,
+			base64_string: None,
+		};
+
+		Ok(Raster::new_cpu(cpu_image))
+	}
+}
+
+/// Converts a `List<Raster<GPU>>` to `List<Raster<CPU>>` by downloading texture data in one go then asynchronously maps all buffers and processes the results.
+impl<'i> ConvertAsync<List<Raster<CPU>>, &'i WgpuExecutor> for List<Raster<GPU>> {
+	fn convert(self, _: Footprint, executor: &'i WgpuExecutor) -> SourceFuture<List<Raster<CPU>>> {
+		let device = executor.context().device.clone();
+		let queue = executor.context().queue.lock();
+
+		let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
+			label: Some("batch_texture_download_encoder"),
+		});
+
+		let mut converters = Vec::new();
+		let mut rows_meta = Vec::new();
+
+		for row in self {
+			let (element, attributes) = row.into_parts();
+			converters.push(RasterGpuToRasterCpuConverter::new(&device, &mut encoder, element));
+			rows_meta.push(Item::from_parts((), attributes));
+		}
+
+		queue.submit([encoder.finish()]);
+
+		Box::pin(async move {
+			let mut map_futures = Vec::new();
+			for converter in converters {
+				map_futures.push(converter.convert(&device));
+			}
+
+			let map_results = futures::future::try_join_all(map_futures)
+				.await
+				.map_err(|_| "Failed to receive map result")
+				.expect("Buffer mapping communication failed");
+
+			map_results
+				.into_iter()
+				.zip(rows_meta)
+				.map(|(element, row)| {
+					let (_, attributes) = row.into_parts();
+					Item::from_parts(element, attributes)
+				})
+				.collect()
+		})
+	}
+}
+
+/// Converts single GPU raster to CPU by downloading texture data
+impl<'i> ConvertAsync<Raster<CPU>, &'i WgpuExecutor> for Raster<GPU> {
+	fn convert(self, _: Footprint, executor: &'i WgpuExecutor) -> SourceFuture<Raster<CPU>> {
+		let device = executor.context().device.clone();
+		let queue = executor.context().queue.lock();
+
+		let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
+			label: Some("single_texture_download_encoder"),
+		});
+
+		let converter = RasterGpuToRasterCpuConverter::new(&device, &mut encoder, self);
+
+		queue.submit([encoder.finish()]);
+
+		Box::pin(async move { converter.convert(&device).await.expect("Failed to download texture data") })
+	}
+}
+
 /// Uploads an raster texture from the CPU to the GPU. This is now deprecated and the Convert node should be used in the future.
 ///
 /// Accepts either individual raster data or a `List` of raster elements and converts it to the GPU format using the WgpuExecutor's device and queue.