Account for vec3 padding in Metal (flutter#181563)

Re-landing of flutter#181340
incorporating flutter#181550.

Vec3 uniforms on the metal backend are vec4-aligned. This PR accounts
for that padding.

---------

Co-authored-by: gaaclarke <30870216+gaaclarke@users.noreply.github.com>

Author: walley892

engine/src/flutter/impeller/compiler/compiler_test.cc

@@ -106,6 +106,7 @@ bool CompilerTestBase::CanCompileAndReflect(
       entry_point_name);
 
   Reflector::Options reflector_options;
+  reflector_options.target_platform = GetParam();
   reflector_options.header_file_name = ReflectionHeaderName(fixture_name);
   reflector_options.shader_name = "shader_name";
 

engine/src/flutter/impeller/compiler/reflector.cc

@@ -26,6 +26,7 @@
 #include "impeller/geometry/matrix.h"
 #include "impeller/geometry/scalar.h"
 #include "impeller/runtime_stage/runtime_stage.h"
+#include "runtime_stage_types_flatbuffers.h"
 #include "spirv_common.hpp"
 
 namespace impeller {
@@ -372,6 +373,27 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
     uniform_description.columns = spir_type.columns;
     uniform_description.bit_width = spir_type.width;
     uniform_description.array_elements = GetArrayElements(spir_type);
+
+    if (TargetPlatformIsMetal(options_.target_platform) &&
+        uniform_description.type == spirv_cross::SPIRType::BaseType::Float) {
+      // Metal aligns float3 to 16 bytes.
+      // Metal aligns float3x3 COLUMNS to 16 bytes.
+      // For float3: Size 12. Padding 4. Stride 16.
+      // For float3x3: Size 36. Padding 12 (4 per col). Stride 48.
+
+      if (spir_type.vecsize == 3 &&
+          (spir_type.columns == 1 || spir_type.columns == 3)) {
+        for (size_t c = 0; c < spir_type.columns; c++) {
+          for (size_t v = 0; v < 3; v++) {
+            uniform_description.padding_layout.push_back(
+                fb::PaddingType::kFloat);
+          }
+          uniform_description.padding_layout.push_back(
+              fb::PaddingType::kPadding);
+        }
+      }
+    }
+
     FML_CHECK(data->backend != RuntimeStageBackend::kVulkan ||
               spir_type.basetype ==
                   spirv_cross::SPIRType::BaseType::SampledImage)
@@ -396,7 +418,7 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
     size_t binding =
         compiler_->get_decoration(ubo.id, spv::Decoration::DecorationBinding);
     auto members = ReadStructMembers(ubo.type_id);
-    std::vector<uint8_t> struct_layout;
+    std::vector<fb::PaddingType> padding_layout;
     size_t float_count = 0;
 
     for (size_t i = 0; i < members.size(); i += 1) {
@@ -407,7 +429,7 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
           size_t padding_count =
               (member.size + sizeof(float) - 1) / sizeof(float);
           while (padding_count > 0) {
-            struct_layout.push_back(0);
+            padding_layout.push_back(fb::PaddingType::kPadding);
             padding_count--;
           }
           break;
@@ -418,18 +440,18 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
             // and 0 layout property per byte of padding
             for (auto i = 0; i < member.array_elements; i++) {
               for (auto j = 0u; j < member.size / sizeof(float); j++) {
-                struct_layout.push_back(1);
+                padding_layout.push_back(fb::PaddingType::kFloat);
               }
               for (auto j = 0u; j < member.element_padding / sizeof(float);
                    j++) {
-                struct_layout.push_back(0);
+                padding_layout.push_back(fb::PaddingType::kPadding);
               }
             }
           } else {
             size_t member_float_count = member.byte_length / sizeof(float);
             float_count += member_float_count;
             while (member_float_count > 0) {
-              struct_layout.push_back(1);
+              padding_layout.push_back(fb::PaddingType::kFloat);
               member_float_count--;
             }
           }
@@ -446,7 +468,7 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
         .location = binding,
         .binding = binding,
         .type = spirv_cross::SPIRType::Struct,
-        .struct_layout = std::move(struct_layout),
+        .padding_layout = std::move(padding_layout),
         .struct_float_count = float_count,
     });
   }

engine/src/flutter/impeller/compiler/runtime_stage_data.cc

@@ -332,8 +332,8 @@ std::unique_ptr<fb::RuntimeStageT> RuntimeStageData::CreateStageFlatbuffer(
       desc->array_elements = uniform.array_elements.value();
     }
 
-    for (const auto& byte_type : uniform.struct_layout) {
-      desc->struct_layout.push_back(static_cast<fb::StructByteType>(byte_type));
+    for (const auto& byte_type : uniform.padding_layout) {
+      desc->padding_layout.push_back(static_cast<fb::PaddingType>(byte_type));
     }
     desc->struct_float_count = uniform.struct_float_count;
 

engine/src/flutter/impeller/compiler/types.h

@@ -12,6 +12,7 @@
 #include <optional>
 #include <string>
 
+#include "runtime_stage_types_flatbuffers.h"
 #include "shaderc/shaderc.hpp"
 #include "spirv_cross.hpp"
 #include "spirv_msl.hpp"
@@ -55,7 +56,11 @@ struct UniformDescription {
   size_t columns = 0u;
   size_t bit_width = 0u;
   std::optional<size_t> array_elements = std::nullopt;
-  std::vector<uint8_t> struct_layout = {};
+  /// The layout of padding bytes in the uniform buffer.
+  /// The format matches the values in the flatbuffer
+  /// UniformDescription::padding_layout.
+  /// \see RuntimeEffectContents::EmplaceUniform
+  std::vector<fb::PaddingType> padding_layout = {};
   size_t struct_float_count = 0u;
 };
 

engine/src/flutter/impeller/core/runtime_types.cc

@@ -6,14 +6,37 @@
 
 namespace impeller {
 
-size_t RuntimeUniformDescription::GetSize() const {
-  size_t size = dimensions.rows * dimensions.cols * bit_width / 8u;
+size_t RuntimeUniformDescription::GetDartSize() const {
+  size_t size = 0;
+  if (!padding_layout.empty()) {
+    for (impeller::RuntimePaddingType byte_type : padding_layout) {
+      if (byte_type == RuntimePaddingType::kFloat) {
+        size += sizeof(float);
+      }
+    }
+  } else {
+    size = dimensions.rows * dimensions.cols * bit_width / 8u;
+  }
+  if (array_elements.value_or(0) > 0) {
+    // Covered by check on the line above.
+    // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
+    size *= array_elements.value();
+  }
+  return size;
+}
+
+size_t RuntimeUniformDescription::GetGPUSize() const {
+  size_t size = 0;
+  if (padding_layout.empty()) {
+    size = dimensions.rows * dimensions.cols * bit_width / 8u;
+  } else {
+    size = sizeof(float) * padding_layout.size();
+  }
   if (array_elements.value_or(0) > 0) {
     // Covered by check on the line above.
     // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
     size *= array_elements.value();
   }
-  size += sizeof(float) * struct_layout.size();
   return size;
 }
 

engine/src/flutter/impeller/core/runtime_types.h

@@ -38,6 +38,11 @@ struct RuntimeUniformDimensions {
   size_t cols = 0;
 };
 
+enum class RuntimePaddingType : uint8_t {
+  kPadding = 0,
+  kFloat = 1,
+};
+
 struct RuntimeUniformDescription {
   std::string name;
   size_t location = 0u;
@@ -47,11 +52,16 @@ struct RuntimeUniformDescription {
   RuntimeUniformDimensions dimensions = {};
   size_t bit_width = 0u;
   std::optional<size_t> array_elements;
-  std::vector<uint8_t> struct_layout = {};
+  std::vector<RuntimePaddingType> padding_layout = {};
   size_t struct_float_count = 0u;
 
-  /// @brief  Computes the total number of bytes that this uniform requires.
-  size_t GetSize() const;
+  /// @brief  Computes the total number of bytes that this uniform requires for
+  /// representation in the Dart float buffer.
+  size_t GetDartSize() const;
+
+  /// @brief  Computes the total number of bytes that this uniform requires for
+  /// representation in the GPU.
+  size_t GetGPUSize() const;
 };
 
 }  // namespace impeller

engine/src/flutter/impeller/display_list/dl_runtime_effect_impeller.cc

@@ -38,7 +38,7 @@ size_t DlRuntimeEffectImpeller::uniform_size() const {
 
   size_t total = 0;
   for (const auto& uniform : runtime_stage_->GetUniforms()) {
-    total += uniform.GetSize();
+    total += uniform.GetGPUSize();
   }
   return total;
 }

engine/src/flutter/impeller/entity/contents/runtime_effect_contents.cc

@@ -24,35 +24,48 @@
 
 namespace impeller {
 
-namespace {
-constexpr char kPaddingType = 0;
-constexpr char kFloatType = 1;
-}  // namespace
-
 // static
-BufferView RuntimeEffectContents::EmplaceVulkanUniform(
-    const std::shared_ptr<const std::vector<uint8_t>>& input_data,
+BufferView RuntimeEffectContents::EmplaceUniform(
+    const uint8_t* source_data,
     HostBuffer& data_host_buffer,
-    const RuntimeUniformDescription& uniform,
-    size_t minimum_uniform_alignment) {
-  // TODO(jonahwilliams): rewrite this to emplace directly into
-  // HostBuffer.
-  std::vector<float> uniform_buffer;
-  uniform_buffer.reserve(uniform.struct_layout.size());
-  size_t uniform_byte_index = 0u;
-  for (char byte_type : uniform.struct_layout) {
-    if (byte_type == kPaddingType) {
-      uniform_buffer.push_back(0.f);
-    } else {
-      FML_DCHECK(byte_type == kFloatType);
-      uniform_buffer.push_back(reinterpret_cast<const float*>(
-          input_data->data())[uniform_byte_index++]);
-    }
+    const RuntimeUniformDescription& uniform) {
+  size_t minimum_uniform_alignment =
+      data_host_buffer.GetMinimumUniformAlignment();
+  size_t alignment = std::max(uniform.bit_width / 8, minimum_uniform_alignment);
+
+  if (uniform.padding_layout.empty()) {
+    return data_host_buffer.Emplace(source_data, uniform.GetGPUSize(),
+                                    alignment);
   }
 
+  // If the uniform has a padding layout, we need to repack the data.
+  // We can do this by using the EmplaceProc to write directly to the
+  // HostBuffer.
   return data_host_buffer.Emplace(
-      reinterpret_cast<const void*>(uniform_buffer.data()),
-      sizeof(float) * uniform_buffer.size(), minimum_uniform_alignment);
+      uniform.GetGPUSize(), alignment,
+      [&uniform, source_data](uint8_t* destination) {
+        size_t count = uniform.array_elements.value_or(1);
+        if (count == 0) {
+          // Make sure to run at least once.
+          count = 1;
+        }
+        size_t uniform_byte_index = 0u;
+        size_t struct_float_index = 0u;
+        auto* float_destination = reinterpret_cast<float*>(destination);
+        auto* float_source = reinterpret_cast<const float*>(source_data);
+
+        for (size_t i = 0; i < count; i++) {
+          for (RuntimePaddingType byte_type : uniform.padding_layout) {
+            if (byte_type == RuntimePaddingType::kPadding) {
+              float_destination[struct_float_index++] = 0.f;
+            } else {
+              FML_DCHECK(byte_type == RuntimePaddingType::kFloat);
+              float_destination[struct_float_index++] =
+                  float_source[uniform_byte_index++];
+            }
+          }
+        }
+      });
 }
 
 void RuntimeEffectContents::SetRuntimeStage(
@@ -284,12 +297,8 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
               << "Uniform " << uniform.name
               << " had unexpected type kFloat for Vulkan backend.";
 
-          size_t alignment =
-              std::max(uniform.bit_width / 8,
-                       data_host_buffer.GetMinimumUniformAlignment());
-          BufferView buffer_view =
-              data_host_buffer.Emplace(uniform_data_->data() + buffer_offset,
-                                       uniform.GetSize(), alignment);
+          BufferView buffer_view = EmplaceUniform(
+              uniform_data_->data() + buffer_offset, data_host_buffer, uniform);
 
           ShaderUniformSlot uniform_slot;
           uniform_slot.name = uniform.name.c_str();
@@ -298,7 +307,7 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
                                    DescriptorType::kUniformBuffer, uniform_slot,
                                    std::move(metadata), std::move(buffer_view));
           buffer_index++;
-          buffer_offset += uniform.GetSize();
+          buffer_offset += uniform.GetDartSize();
           buffer_location++;
           break;
         }
@@ -309,12 +318,10 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
           uniform_slot.binding = uniform.location;
           uniform_slot.name = uniform.name.c_str();
 
-          pass.BindResource(ShaderStage::kFragment,
-                            DescriptorType::kUniformBuffer, uniform_slot,
-                            nullptr,
-                            EmplaceVulkanUniform(
-                                uniform_data_, data_host_buffer, uniform,
-                                data_host_buffer.GetMinimumUniformAlignment()));
+          pass.BindResource(
+              ShaderStage::kFragment, DescriptorType::kUniformBuffer,
+              uniform_slot, nullptr,
+              EmplaceUniform(uniform_data_->data(), data_host_buffer, uniform));
         }
       }
     }

engine/src/flutter/impeller/entity/contents/runtime_effect_contents.h

@@ -37,11 +37,17 @@ class RuntimeEffectContents final : public ColorSourceContents {
   bool BootstrapShader(const ContentContext& renderer) const;
 
   // Visible for testing
-  static BufferView EmplaceVulkanUniform(
-      const std::shared_ptr<const std::vector<uint8_t>>& input_data,
-      HostBuffer& host_buffer,
-      const RuntimeUniformDescription& uniform,
-      size_t minimum_uniform_alignment);
+  /// Copies the uniform data into the host buffer.
+  ///
+  /// If the `uniform` has a `padding_layout`, it is used to repack the data.
+  ///
+  /// @param source_data The pointer to the start of the uniform data in the
+  ///        source.
+  /// @param host_buffer The host buffer to emplace the uniform data into.
+  /// @param uniform The description of the uniform being emplaced.
+  static BufferView EmplaceUniform(const uint8_t* source_data,
+                                   HostBuffer& host_buffer,
+                                   const RuntimeUniformDescription& uniform);
 
  private:
   bool RegisterShader(const ContentContext& renderer) const;

engine/src/flutter/impeller/entity/entity_unittests.cc

@@ -1922,12 +1922,9 @@ TEST_P(EntityTest, RuntimeEffectSetsRightSizeWhenUniformIsStruct) {
   uniform_data->resize(sizeof(FragUniforms));
   memcpy(uniform_data->data(), &frag_uniforms, sizeof(FragUniforms));
 
-  auto buffer_view = RuntimeEffectContents::EmplaceVulkanUniform(
-      uniform_data, GetContentContext()->GetTransientsDataBuffer(),
-      runtime_stage->GetUniforms()[0],
-      GetContentContext()
-          ->GetTransientsDataBuffer()
-          .GetMinimumUniformAlignment());
+  auto buffer_view = RuntimeEffectContents::EmplaceUniform(
+      uniform_data->data(), GetContentContext()->GetTransientsDataBuffer(),
+      runtime_stage->GetUniforms()[0]);
 
   // 16 bytes:
   //   8 bytes for iResolution