Fix a bug with D3D12 buffer alignment computation

We should use D3D12_CONSTANT_BUFFER_DATA_PLACEMENT_ALIGNMENT
if the usage *contains* Uniform usage, not if it is exactly
equal.

Fixed: dawn:1085
Change-Id: I540081e550c9e88efeb08169cecdc40b68d36d14
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/62242
Commit-Queue: Austin Eng <enga@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn_native/DawnNative.cpp b/src/dawn_native/DawnNative.cpp
index 15fd697..b5e1ba40 100644
--- a/src/dawn_native/DawnNative.cpp
+++ b/src/dawn_native/DawnNative.cpp
@@ -13,6 +13,8 @@
 // limitations under the License.
 
 #include "dawn_native/DawnNative.h"
+
+#include "dawn_native/Buffer.h"
 #include "dawn_native/Device.h"
 #include "dawn_native/Instance.h"
 #include "dawn_native/Texture.h"
@@ -230,4 +232,8 @@
         return object->GetLabel().c_str();
     }
 
+    uint64_t GetAllocatedSizeForTesting(WGPUBuffer buffer) {
+        return reinterpret_cast<const BufferBase*>(buffer)->GetAllocatedSize();
+    }
+
 }  // namespace dawn_native
diff --git a/src/dawn_native/d3d12/BindGroupD3D12.cpp b/src/dawn_native/d3d12/BindGroupD3D12.cpp
index 4832cdd..47c501a 100644
--- a/src/dawn_native/d3d12/BindGroupD3D12.cpp
+++ b/src/dawn_native/d3d12/BindGroupD3D12.cpp
@@ -69,7 +69,8 @@
                     switch (bindingInfo.buffer.type) {
                         case wgpu::BufferBindingType::Uniform: {
                             D3D12_CONSTANT_BUFFER_VIEW_DESC desc;
-                            desc.SizeInBytes = Align(binding.size, 256);
+                            desc.SizeInBytes =
+                                Align(binding.size, D3D12_CONSTANT_BUFFER_DATA_PLACEMENT_ALIGNMENT);
                             desc.BufferLocation =
                                 ToBackend(binding.buffer)->GetVA() + binding.offset;
 
diff --git a/src/dawn_native/d3d12/BufferD3D12.cpp b/src/dawn_native/d3d12/BufferD3D12.cpp
index 67ebc47..39fcb83 100644
--- a/src/dawn_native/d3d12/BufferD3D12.cpp
+++ b/src/dawn_native/d3d12/BufferD3D12.cpp
@@ -82,12 +82,16 @@
         }
 
         size_t D3D12BufferSizeAlignment(wgpu::BufferUsage usage) {
-            switch (usage) {
-                case wgpu::BufferUsage::Uniform:
-                    return D3D12_CONSTANT_BUFFER_DATA_PLACEMENT_ALIGNMENT;
-                default:
-                    return 1;
+            if ((usage & wgpu::BufferUsage::Uniform) != 0) {
+                // D3D buffers are always resource size aligned to 64KB. However, D3D12's validation
+                // forbids binding a CBV to an unaligned size. To prevent, one can always safely
+                // align the buffer size to the CBV data alignment as other buffer usages
+                // ignore it (no size check). The validation will still enforce bound checks with
+                // the unaligned size returned by GetSize().
+                // https://docs.microsoft.com/en-us/windows/win32/direct3d12/uploading-resources#buffer-alignment
+                return D3D12_CONSTANT_BUFFER_DATA_PLACEMENT_ALIGNMENT;
             }
+            return 1;
         }
     }  // namespace
 
@@ -103,12 +107,6 @@
     }
 
     MaybeError Buffer::Initialize(bool mappedAtCreation) {
-        // D3D buffers are always resource size aligned to 64KB. However, D3D12's validation forbids
-        // binding a CBV to an unaligned size. To prevent, one can always safely align the buffer
-        // desc size to the CBV data alignment as other buffer usages ignore it (no size check).
-        // The validation will still enforce bound checks with the unaligned size returned by
-        // GetSize().
-        // https://docs.microsoft.com/en-us/windows/win32/direct3d12/uploading-resources#buffer-alignment
         // Allocate at least 4 bytes so clamped accesses are always in bounds.
         uint64_t size = std::max(GetSize(), uint64_t(4u));
         size_t alignment = D3D12BufferSizeAlignment(GetUsage());
diff --git a/src/include/dawn_native/DawnNative.h b/src/include/dawn_native/DawnNative.h
index a0a6229..4281e05 100644
--- a/src/include/dawn_native/DawnNative.h
+++ b/src/include/dawn_native/DawnNative.h
@@ -250,6 +250,8 @@
 
     DAWN_NATIVE_EXPORT const char* GetObjectLabelForTesting(void* objectHandle);
 
+    DAWN_NATIVE_EXPORT uint64_t GetAllocatedSizeForTesting(WGPUBuffer buffer);
+
 }  // namespace dawn_native
 
 #endif  // DAWNNATIVE_DAWNNATIVE_H_
diff --git a/src/tests/BUILD.gn b/src/tests/BUILD.gn
index 8aa80f8..0c5964c 100644
--- a/src/tests/BUILD.gn
+++ b/src/tests/BUILD.gn
@@ -439,6 +439,7 @@
   }
 
   sources += [
+    "white_box/BufferAllocatedSizeTests.cpp",
     "white_box/InternalResourceUsageTests.cpp",
     "white_box/InternalStorageBufferBindingTests.cpp",
     "white_box/QueryInternalShaderTests.cpp",
diff --git a/src/tests/white_box/BufferAllocatedSizeTests.cpp b/src/tests/white_box/BufferAllocatedSizeTests.cpp
new file mode 100644
index 0000000..15c1437
--- /dev/null
+++ b/src/tests/white_box/BufferAllocatedSizeTests.cpp
@@ -0,0 +1,85 @@
+// Copyright 2021 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tests/DawnTest.h"
+
+#include "common/Math.h"
+#include "dawn_native/DawnNative.h"
+
+#include <algorithm>
+
+class BufferAllocatedSizeTests : public DawnTest {
+  protected:
+    wgpu::Buffer CreateBuffer(wgpu::BufferUsage usage, uint64_t size) {
+        wgpu::BufferDescriptor desc = {};
+        desc.usage = usage;
+        desc.size = size;
+        return device.CreateBuffer(&desc);
+    }
+
+    void SetUp() override {
+        DawnTest::SetUp();
+        DAWN_TEST_UNSUPPORTED_IF(UsesWire());
+    }
+};
+
+// Test expected allocated size for buffers with uniform usage
+TEST_P(BufferAllocatedSizeTests, UniformUsage) {
+    // Some backends have a minimum buffer size, so make sure
+    // we allocate above that.
+    constexpr uint32_t kMinBufferSize = 4u;
+
+    uint32_t requiredBufferAlignment = 1u;
+    if (IsD3D12()) {
+        requiredBufferAlignment = 256u;
+    } else if (IsMetal()) {
+        requiredBufferAlignment = 16u;
+    } else if (IsVulkan()) {
+        requiredBufferAlignment = 4u;
+    }
+
+    // Test uniform usage
+    {
+        const uint32_t bufferSize = kMinBufferSize;
+        wgpu::Buffer buffer = CreateBuffer(wgpu::BufferUsage::Uniform, bufferSize);
+        EXPECT_EQ(dawn_native::GetAllocatedSizeForTesting(buffer.Get()),
+                  Align(bufferSize, requiredBufferAlignment));
+    }
+
+    // Test uniform usage and with size just above requiredBufferAlignment allocates to the next
+    // multiple of |requiredBufferAlignment|
+    {
+        const uint32_t bufferSize = std::max(1u + requiredBufferAlignment, kMinBufferSize);
+        wgpu::Buffer buffer =
+            CreateBuffer(wgpu::BufferUsage::Uniform | wgpu::BufferUsage::Storage, bufferSize);
+        EXPECT_EQ(dawn_native::GetAllocatedSizeForTesting(buffer.Get()),
+                  Align(bufferSize, requiredBufferAlignment));
+    }
+
+    // Test uniform usage and another usage
+    {
+        const uint32_t bufferSize = kMinBufferSize;
+        wgpu::Buffer buffer =
+            CreateBuffer(wgpu::BufferUsage::Uniform | wgpu::BufferUsage::Storage, bufferSize);
+        EXPECT_EQ(dawn_native::GetAllocatedSizeForTesting(buffer.Get()),
+                  Align(bufferSize, requiredBufferAlignment));
+    }
+}
+
+DAWN_INSTANTIATE_TEST(BufferAllocatedSizeTests,
+                      D3D12Backend(),
+                      MetalBackend(),
+                      OpenGLBackend(),
+                      OpenGLESBackend(),
+                      VulkanBackend());