From 58016bfb31a05c94d14c4924a4132a30bc43ec68 Mon Sep 17 00:00:00 2001
From: Antonio Maiorano <amaiorano@google.com>
Date: Wed, 14 Jul 2021 15:33:31 +0000
Subject: [PATCH] Make storage buffer length UBO size a multiple of 16 bytes

Conditional on whether TINT_EXPECTS_UBOS_TO_BE_MULTIPLE_OF_16 is
defined, which it will be in a future Tint CL.

Bug: tint:984
Change-Id: I191cd59fe6947c9113a8fad7dbfbf79903969b27
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/57960
Commit-Queue: Antonio Maiorano <amaiorano@google.com>
Reviewed-by: Austin Eng <enga@chromium.org>
---
 src/dawn_native/metal/CommandBufferMTL.mm | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/dawn_native/metal/CommandBufferMTL.mm b/src/dawn_native/metal/CommandBufferMTL.mm
index 4b389bf210..176186df3e 100644
--- a/src/dawn_native/metal/CommandBufferMTL.mm
+++ b/src/dawn_native/metal/CommandBufferMTL.mm
@@ -30,6 +30,8 @@
 #include "dawn_native/metal/TextureMTL.h"
 #include "dawn_native/metal/UtilsMetal.h"
 
+#include <tint/tint.h>
+
 namespace dawn_native { namespace metal {
 
     namespace {
@@ -272,7 +274,7 @@ namespace dawn_native { namespace metal {
 
             // The lengths of buffers are stored as 32bit integers because that is the width the
             // MSL code generated by SPIRV-Cross expects.
-            PerStage<std::array<uint32_t, kGenericMetalBufferSlots>> data;
+            PerStage<std::array<uint32_t, ((kGenericMetalBufferSlots / 4) + 1) * 4>> data;
 
             void Apply(id<MTLRenderCommandEncoder> render,
                        RenderPipeline* pipeline,
@@ -292,6 +294,10 @@ namespace dawn_native { namespace metal {
                         bufferCount += pipeline->GetVertexBufferCount();
                     }
 
+#ifdef TINT_EXPECTS_UBOS_TO_BE_MULTIPLE_OF_16
+                    bufferCount = Align(bufferCount, 4);
+                    ASSERT(bufferCount <= data[SingleShaderStage::Vertex].size());
+#endif
                     [render setVertexBytes:data[SingleShaderStage::Vertex].data()
                                     length:sizeof(uint32_t) * bufferCount
                                    atIndex:kBufferLengthBufferSlot];
@@ -300,6 +306,10 @@ namespace dawn_native { namespace metal {
                 if (stagesToApply & wgpu::ShaderStage::Fragment) {
                     uint32_t bufferCount = ToBackend(pipeline->GetLayout())
                                                ->GetBufferBindingCount(SingleShaderStage::Fragment);
+#ifdef TINT_EXPECTS_UBOS_TO_BE_MULTIPLE_OF_16
+                    bufferCount = Align(bufferCount, 4);
+                    ASSERT(bufferCount <= data[SingleShaderStage::Fragment].size());
+#endif
                     [render setFragmentBytes:data[SingleShaderStage::Fragment].data()
                                       length:sizeof(uint32_t) * bufferCount
                                      atIndex:kBufferLengthBufferSlot];
@@ -320,6 +330,10 @@ namespace dawn_native { namespace metal {
 
                 uint32_t bufferCount = ToBackend(pipeline->GetLayout())
                                            ->GetBufferBindingCount(SingleShaderStage::Compute);
+#ifdef TINT_EXPECTS_UBOS_TO_BE_MULTIPLE_OF_16
+                bufferCount = Align(bufferCount, 4);
+                ASSERT(bufferCount <= data[SingleShaderStage::Compute].size());
+#endif
                 [compute setBytes:data[SingleShaderStage::Compute].data()
                            length:sizeof(uint32_t) * bufferCount
                           atIndex:kBufferLengthBufferSlot];