From 4f1bc8e2b8210d2eb596d0a38d7c58e06adc7863 Mon Sep 17 00:00:00 2001
From: Brezak <bezak.adam@proton.me>
Date: Mon, 30 Dec 2024 06:51:37 +0100
Subject: [PATCH] Don't overalign aligned values in
 `gpu_readback::align_byte_size` (#17007)

# Objective

Fix alignment calculations in our rendering code.
Fixes #16992

The `gpu_readback::align_byte_size` function incorrectly rounds aligned
values to the next alignment.
If we assume the alignment to be 256 (because that's what wgpu says it
its) the function would align 0 to 256, 256 to 512, etc...

## Solution

Forward the `gpu_readback::align_byte_size` to
`RenderDevice::align_copy_bytes_per_row` so we don't implement the same
method twice.
Simplify `RenderDevice::align_copy_bytes_per_row`.

## Testing

Ran the code provided in #16992 to see if the issue has been solved +
added a test to check if `align_copy_bytes_per_row` returns the correct
values.
---
 crates/bevy_render/src/gpu_readback.rs        | 11 +++++---
 .../bevy_render/src/renderer/render_device.rs | 28 +++++++++++++++++--
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_render/src/gpu_readback.rs b/crates/bevy_render/src/gpu_readback.rs
index bb0092edf1..0b2afaea73 100644
--- a/crates/bevy_render/src/gpu_readback.rs
+++ b/crates/bevy_render/src/gpu_readback.rs
@@ -27,7 +27,7 @@ use bevy_utils::{default, tracing::warn, HashMap};
 use encase::internal::ReadFrom;
 use encase::private::Reader;
 use encase::ShaderType;
-use wgpu::{CommandEncoder, COPY_BYTES_PER_ROW_ALIGNMENT};
+use wgpu::CommandEncoder;
 
 /// A plugin that enables reading back gpu buffers and textures to the cpu.
 pub struct GpuReadbackPlugin {
@@ -349,14 +349,17 @@ fn map_buffers(mut readbacks: ResMut<GpuReadbacks>) {
 
 // Utils
 
-pub(crate) fn align_byte_size(value: u32) -> u32 {
-    value + (COPY_BYTES_PER_ROW_ALIGNMENT - (value % COPY_BYTES_PER_ROW_ALIGNMENT))
+/// Round up a given value to be a multiple of [`wgpu::COPY_BYTES_PER_ROW_ALIGNMENT`].
+pub(crate) const fn align_byte_size(value: u32) -> u32 {
+    RenderDevice::align_copy_bytes_per_row(value as usize) as u32
 }
 
-pub(crate) fn get_aligned_size(width: u32, height: u32, pixel_size: 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)
 }
 
+/// 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 {
     ImageDataLayout {
         bytes_per_row: if height > 1 {
diff --git a/crates/bevy_render/src/renderer/render_device.rs b/crates/bevy_render/src/renderer/render_device.rs
index 407d1b361b..827d0d2ca3 100644
--- a/crates/bevy_render/src/renderer/render_device.rs
+++ b/crates/bevy_render/src/renderer/render_device.rs
@@ -231,10 +231,16 @@ impl RenderDevice {
         buffer.map_async(map_mode, callback);
     }
 
-    pub fn align_copy_bytes_per_row(row_bytes: usize) -> usize {
+    // Rounds up `row_bytes` to be a multiple of [`wgpu::COPY_BYTES_PER_ROW_ALIGNMENT`].
+    pub const fn align_copy_bytes_per_row(row_bytes: usize) -> usize {
         let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize;
-        let padded_bytes_per_row_padding = (align - row_bytes % align) % align;
-        row_bytes + padded_bytes_per_row_padding
+
+        // If row_bytes is aligned calculate a value just under the next aligned value.
+        // Otherwise calculate a value greater than the next aligned value.
+        let over_aligned = row_bytes + align - 1;
+
+        // Round the number *down* to the nearest aligned value.
+        (over_aligned / align) * align
     }
 
     pub fn get_supported_read_only_binding_type(
@@ -248,3 +254,19 @@ impl RenderDevice {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn align_copy_bytes_per_row() {
+        // Test for https://github.com/bevyengine/bevy/issues/16992
+        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize;
+
+        assert_eq!(RenderDevice::align_copy_bytes_per_row(0), 0);
+        assert_eq!(RenderDevice::align_copy_bytes_per_row(1), align);
+        assert_eq!(RenderDevice::align_copy_bytes_per_row(align + 1), align * 2);
+        assert_eq!(RenderDevice::align_copy_bytes_per_row(align), align);
+    }
+}