Fix Oversized D3D11 Immediate Uniform Buffer

kMaxImmediateSizeD3D11 was defined as sizeof(RenderImmediates /
ComputeImmediates), i.e. a byte count, but every use treats it as a
number of uint32_t elements: it sizes std::array<uint32_t, N>, bounds
the element offset, and multiplies by sizeof(uint32_t) for the GPU
buffer byte size. As a result the CPU array and the GPU uniform buffer
were 4x larger than needed.

Rename the constant to kMaxImmediateSlotsD3D11 and convert the maximum
immediate struct byte size to uint32_t slots with Align, so the buffer
stays correctly sized even if a struct size is not a multiple of four.
Make IsPowerOfTwo constexpr so Align can be used in constant expressions,
and add compile-time coverage for both helpers.

Bug: 366291600
Change-Id: I7375650f0de766b154f72e8ef4f70ba266d5248a
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/329755
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Quyen Le <lehoangquyen@google.com>

src/dawn/common/Math.cpp

diff --git a/src/dawn/common/Math.cpp b/src/dawn/common/Math.cpp
index a6826b9..9579482 100644
--- a/src/dawn/common/Math.cpp
+++ b/src/dawn/common/Math.cpp
@@ -55,11 +55,6 @@
     return 1ull << (Log2(n - 1) + 1);
 }
 
-bool IsPowerOfTwo(uint64_t n) {
-    DAWN_RELEASE_ASSUME(n != 0);
-    return (n & (n - 1)) == 0;
-}
-
 bool IsPtrAligned(const void* ptr, size_t alignment) {
     DAWN_RELEASE_ASSUME(IsPowerOfTwo(alignment));
     DAWN_RELEASE_ASSUME(alignment != 0);

src/dawn/common/Math.h

diff --git a/src/dawn/common/Math.h b/src/dawn/common/Math.h
index 2bec112..7461d97 100644
--- a/src/dawn/common/Math.h
+++ b/src/dawn/common/Math.h
@@ -50,7 +50,10 @@
 // The following are not valid for 0
 uint32_t Log2(uint32_t value);
 uint32_t Log2(uint64_t value);
-bool IsPowerOfTwo(uint64_t n);
+constexpr bool IsPowerOfTwo(uint64_t n) {
+    DAWN_RELEASE_ASSUME(n != 0);
+    return (n & (n - 1)) == 0;
+}
 
 // Returns 2^exp for integrals
 template <std::integral T>

src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp

diff --git a/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp b/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp
index de84c4a..aa717b4 100644
--- a/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp
+++ b/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp
@@ -180,9 +180,9 @@
 void ScopedCommandRecordingContext::WriteUniformBufferRange(uint32_t offset,
                                                             const void* data,
                                                             size_t size) const {
-    DAWN_ASSERT(offset < CommandRecordingContext::kMaxImmediateSizeD3D11);
+    DAWN_ASSERT(offset < CommandRecordingContext::kMaxImmediateSlotsD3D11);
     DAWN_ASSERT(size <=
-                sizeof(uint32_t) * (CommandRecordingContext::kMaxImmediateSizeD3D11 - offset));
+                sizeof(uint32_t) * (CommandRecordingContext::kMaxImmediateSlotsD3D11 - offset));
     DAWN_UNSAFE_TODO(std::memcpy(&Get()->mUniformBufferData[offset], data, size));
     Get()->mUniformBufferDirty = true;
 }
@@ -383,7 +383,7 @@
     DeviceBase* device) {
     // Create a uniform buffer for user and internal Immediates.
     BufferDescriptor descriptor;
-    descriptor.size = sizeof(uint32_t) * kMaxImmediateSizeD3D11;
+    descriptor.size = sizeof(uint32_t) * kMaxImmediateSlotsD3D11;
     descriptor.usage = wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst;
     descriptor.mappedAtCreation = false;
     descriptor.label = "ImmediatesInternalBuffer";

src/dawn/native/d3d11/CommandRecordingContextD3D11.h

diff --git a/src/dawn/native/d3d11/CommandRecordingContextD3D11.h b/src/dawn/native/d3d11/CommandRecordingContextD3D11.h
index 4ee007c..a14eda6 100644
--- a/src/dawn/native/d3d11/CommandRecordingContextD3D11.h
+++ b/src/dawn/native/d3d11/CommandRecordingContextD3D11.h
@@ -35,6 +35,7 @@
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/inlined_vector.h"
 #include "src/dawn/common/Constants.h"
+#include "src/dawn/common/Math.h"
 #include "src/dawn/common/MutexProtected.h"
 #include "src/dawn/common/Ref.h"
 #include "src/dawn/common/StackAllocated.h"
@@ -92,8 +93,10 @@
     bool IsValid() const;
 
     static ResultOrError<Ref<BufferBase>> CreateInternalUniformBuffer(DeviceBase* device);
-    static constexpr uint32_t kMaxImmediateSizeD3D11 =
-        std::max(sizeof(RenderImmediates), sizeof(ComputeImmediates));
+    // The number of uint32_t elements in the immediate uniform buffer.
+    static constexpr uint32_t kMaxImmediateSlotsD3D11 =
+        Align(std::max(sizeof(RenderImmediates), sizeof(ComputeImmediates)), sizeof(uint32_t)) /
+        sizeof(uint32_t);
 
     void ReleaseKeyedMutexes();
 
@@ -116,7 +119,7 @@
 
     // The uniform buffer for built-in variables.
     Ref<GPUUsableBuffer> mUniformBuffer;
-    std::array<uint32_t, kMaxImmediateSizeD3D11> mUniformBufferData{};
+    std::array<uint32_t, kMaxImmediateSlotsD3D11> mUniformBufferData{};
     bool mUniformBufferDirty = true;
 
     absl::flat_hash_set<Ref<d3d::KeyedMutex>> mAcquiredKeyedMutexes;

src/dawn/tests/unittests/MathTests.cpp

diff --git a/src/dawn/tests/unittests/MathTests.cpp b/src/dawn/tests/unittests/MathTests.cpp
index 8447c31..feb2bd8 100644
--- a/src/dawn/tests/unittests/MathTests.cpp
+++ b/src/dawn/tests/unittests/MathTests.cpp
@@ -121,6 +121,9 @@
 
 // Tests for IsPowerOfTwo
 TEST(Math, IsPowerOfTwo) {
+    static_assert(IsPowerOfTwo(1));
+    static_assert(!IsPowerOfTwo(3));
+
     ASSERT_TRUE(IsPowerOfTwo(1));
     ASSERT_TRUE(IsPowerOfTwo(2));
     ASSERT_FALSE(IsPowerOfTwo(3));
@@ -160,6 +163,8 @@
 
 // Tests for Align
 TEST(Math, Align) {
+    static_assert(Align(77u, 4) == 80u);
+
     // 0 aligns to 0
     ASSERT_EQ(Align(0u, 4), 0u);
     ASSERT_EQ(Align(0u, 256), 0u);