Readback: Add support for texture depth/array layers (#17479)

# Objective

Fixes #16963

## Solution

I am - no pun intended - somewhat out of my depth here but this worked
in my testing. The validation error is gone and the data read from the
GPU looks sensible. I'd greatly appreciate if somebody more familiar
with the matter could double-check this.

## References

Relevant documentation in
[WebGPU](https://gpuweb.github.io/gpuweb/#gputexelcopybufferlayout) and
[wgpu](https://github.com/gfx-rs/wgpu/blob/v23/wgpu-types/src/lib.rs#L6350).

## Testing

<details><summary>Example code for testing</summary>
<p>

```rust
use bevy::{
    image::{self as bevy_image, TextureFormatPixelInfo},
    prelude::*,
    render::{
        render_asset::RenderAssetUsages,
        render_resource::{
            Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
        },
    },
};

fn main() {
    let mut app = App::new();
    app.add_plugins(DefaultPlugins)
        .add_systems(Startup, setup)
        .add_systems(Update, readback_system);
    app.run();
}

#[derive(Resource)]
struct ImageResource(Handle<Image>);

const TEXTURE_HEIGHT: u32 = 64;
const TEXTURE_WIDTH: u32 = 32;
const TEXTURE_LAYERS: u32 = 4;
const FORMAT: TextureFormat = TextureFormat::Rgba8Uint;

fn setup(mut commands: Commands, mut images: ResMut<Assets<Image>>) {
    let layer_pixel_count = (TEXTURE_WIDTH * TEXTURE_HEIGHT) as usize;
    let layer_size = layer_pixel_count * FORMAT.pixel_size();
    let data: Vec<u8> = (0..TEXTURE_LAYERS as u8)
        .flat_map(|layer| (0..layer_size).map(move |_| layer))
        .collect();
    let image_size = data.len();
    println!("{image_size}");
    let image = Image {
        data,
        texture_descriptor: TextureDescriptor {
            label: Some("image"),
            size: Extent3d {
                width: TEXTURE_WIDTH,
                height: TEXTURE_HEIGHT,
                depth_or_array_layers: TEXTURE_LAYERS,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: TextureDimension::D2,
            format: FORMAT,
            usage: TextureUsages::COPY_DST | TextureUsages::COPY_SRC,
            view_formats: &[],
        },
        sampler: bevy_image::ImageSampler::Default,
        texture_view_descriptor: None,
        asset_usage: RenderAssetUsages::RENDER_WORLD,
    };

    commands.insert_resource(ImageResource(images.add(image)));
}

fn readback_system(
    mut commands: Commands,
    keys: Res<ButtonInput<KeyCode>>,
    image: Res<ImageResource>,
) {
    if !keys.just_pressed(KeyCode::KeyR) {
        return;
    }

    commands
        .spawn(bevy::render::gpu_readback::Readback::Texture(
            image.0.clone(),
        ))
        .observe(
            |trigger: Trigger<bevy::render::gpu_readback::ReadbackComplete>,
             mut commands: Commands| {
                info!("readback complete");

                println!("{:#?}", &trigger.0);

                commands.entity(trigger.observer()).despawn();
            },
        );
}

```

</p>
</details>

crates/bevy_render/src/gpu_readback.rs

@@ -240,16 +240,11 @@ fn prepare_buffers(
         match readback {
             Readback::Texture(image) => {
                 if let Some(gpu_image) = gpu_images.get(image) {
-                    let layout = layout_data(
-                        gpu_image.size.width,
-                        gpu_image.size.height,
-                        gpu_image.texture_format,
-                    );
+                    let layout = layout_data(gpu_image.size, gpu_image.texture_format);
                     let buffer = buffer_pool.get(
                         &render_device,
                         get_aligned_size(
-                            gpu_image.size.width,
-                            gpu_image.size.height,
+                            gpu_image.size,
                             gpu_image.texture_format.pixel_size() as u32,
                         ) as u64,
                     );
@@ -355,20 +350,32 @@ pub(crate) const fn align_byte_size(value: u32) -> u32 {
 }
 
 /// Get the size of a image when the size of each row has been rounded up to [`wgpu::COPY_BYTES_PER_ROW_ALIGNMENT`].
-pub(crate) const fn get_aligned_size(width: u32, height: u32, pixel_size: u32) -> u32 {
-    height * align_byte_size(width * pixel_size)
+pub(crate) const fn get_aligned_size(extent: Extent3d, pixel_size: u32) -> u32 {
+    extent.height * align_byte_size(extent.width * pixel_size) * extent.depth_or_array_layers
 }
 
 /// Get a [`ImageDataLayout`] aligned such that the image can be copied into a buffer.
-pub(crate) fn layout_data(width: u32, height: u32, format: TextureFormat) -> ImageDataLayout {
+pub(crate) fn layout_data(extent: Extent3d, format: TextureFormat) -> ImageDataLayout {
     ImageDataLayout {
-        bytes_per_row: if height > 1 {
+        bytes_per_row: if extent.height > 1 || extent.depth_or_array_layers > 1 {
             // 1 = 1 row
-            Some(get_aligned_size(width, 1, format.pixel_size() as u32))
+            Some(get_aligned_size(
+                Extent3d {
+                    width: extent.width,
+                    height: 1,
+                    depth_or_array_layers: 1,
+                },
+                format.pixel_size() as u32,
+            ))
         } else {
             None
         },
-        rows_per_image: None,
-        ..Default::default()
+        rows_per_image: if extent.depth_or_array_layers > 1 {
+            let (_, block_dimension_y) = format.block_dimensions();
+            Some(extent.height / block_dimension_y)
+        } else {
+            None
+        },
+        offset: 0,
     }
 }

crates/bevy_render/src/view/window/screenshot.rs

@@ -366,8 +366,7 @@ fn prepare_screenshot_state(
     let texture_view = texture.create_view(&Default::default());
     let buffer = render_device.create_buffer(&wgpu::BufferDescriptor {
         label: Some("screenshot-transfer-buffer"),
-        size: gpu_readback::get_aligned_size(size.width, size.height, format.pixel_size() as u32)
-            as u64,
+        size: gpu_readback::get_aligned_size(size, format.pixel_size() as u32) as u64,
         usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
         mapped_at_creation: false,
     });
@@ -585,17 +584,18 @@ fn render_screenshot(
     texture_view: &wgpu::TextureView,
 ) {
     if let Some(prepared_state) = &prepared.get(entity) {
+        let extent = Extent3d {
+            width,
+            height,
+            depth_or_array_layers: 1,
+        };
         encoder.copy_texture_to_buffer(
             prepared_state.texture.as_image_copy(),
             wgpu::ImageCopyBuffer {
                 buffer: &prepared_state.buffer,
-                layout: gpu_readback::layout_data(width, height, texture_format),
-            },
-            Extent3d {
-                width,
-                height,
-                ..Default::default()
+                layout: gpu_readback::layout_data(extent, texture_format),
             },
+            extent,
         );
 
         if let Some(pipeline) = pipelines.get_render_pipeline(prepared_state.pipeline_id) {