[msl-writer] Add struct offset support.

This CL adds support for injecting padding into structs in the MSL
backend.

Bug: tint:8
Change-Id: I83631a71ce4a2f00b61974ee2c0c7ca1b97f3028
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/24601
Reviewed-by: David Neto <dneto@google.com>

CMakeLists.txt

@@ -92,7 +92,7 @@ if(${TINT_BUILD_DOCS})
   else()
     message("Doxygen not found. Skipping documentation")
   endif(DOXYGEN_FOUND)
-endif(TINT_BUILD_DOCS)
+endif()
 
 if(MSVC)
   # We don't want to have to copy the C Runtime DLL everywhere the executable

src/writer/msl/generator_impl.cc

@@ -36,6 +36,7 @@
 #include "src/ast/member_accessor_expression.h"
 #include "src/ast/return_statement.h"
 #include "src/ast/sint_literal.h"
+#include "src/ast/struct_member_offset_decoration.h"
 #include "src/ast/switch_statement.h"
 #include "src/ast/type/alias_type.h"
 #include "src/ast/type/array_type.h"
@@ -70,6 +71,14 @@ bool last_is_break_or_fallthrough(const ast::StatementList& stmts) {
   return stmts.back()->IsBreak() || stmts.back()->IsFallthrough();
 }
 
+uint32_t adjust_for_alignment(uint32_t count, uint32_t alignment) {
+  const auto spill = count % alignment;
+  if (spill == 0) {
+    return count;
+  }
+  return count + alignment - spill;
+}
+
 }  // namespace
 
 GeneratorImpl::GeneratorImpl() = default;
@@ -132,6 +141,94 @@ bool GeneratorImpl::Generate(const ast::Module& module) {
   return true;
 }
 
+uint32_t GeneratorImpl::calculate_largest_alignment(
+    ast::type::StructType* type) {
+  auto* stct = type->AsStruct()->impl();
+  uint32_t largest_alignment = 0;
+  for (const auto& mem : stct->members()) {
+    auto align = calculate_alignment_size(mem->type());
+    if (align == 0) {
+      return 0;
+    }
+    if (!mem->type()->IsStruct()) {
+      largest_alignment = std::max(largest_alignment, align);
+    } else {
+      largest_alignment =
+          std::max(largest_alignment,
+                   calculate_largest_alignment(mem->type()->AsStruct()));
+    }
+  }
+  return largest_alignment;
+}
+
+uint32_t GeneratorImpl::calculate_alignment_size(ast::type::Type* type) {
+  if (type->IsAlias()) {
+    return calculate_alignment_size(type->AsAlias()->type());
+  }
+  if (type->IsArray()) {
+    auto* ary = type->AsArray();
+    // TODO(dsinclair): Handle array stride and adjust for alignment.
+    uint32_t type_size = calculate_alignment_size(ary->type());
+    return ary->size() * type_size;
+  }
+  if (type->IsBool()) {
+    return 1;
+  }
+  if (type->IsPointer()) {
+    return 0;
+  }
+  if (type->IsF32() || type->IsI32() || type->IsU32()) {
+    return 4;
+  }
+  if (type->IsMatrix()) {
+    auto* mat = type->AsMatrix();
+    // TODO(dsinclair): Handle MatrixStride
+    // https://github.com/gpuweb/gpuweb/issues/773
+    uint32_t type_size = calculate_alignment_size(mat->type());
+    return mat->rows() * mat->columns() * type_size;
+  }
+  if (type->IsStruct()) {
+    auto* stct = type->AsStruct()->impl();
+    uint32_t count = 0;
+    uint32_t largest_alignment = 0;
+    // Offset decorations in WGSL must be in increasing order.
+    for (const auto& mem : stct->members()) {
+      for (const auto& deco : mem->decorations()) {
+        if (deco->IsOffset()) {
+          count = deco->AsOffset()->offset();
+        }
+      }
+      auto align = calculate_alignment_size(mem->type());
+      if (align == 0) {
+        return 0;
+      }
+      if (!mem->type()->IsStruct()) {
+        largest_alignment = std::max(largest_alignment, align);
+      } else {
+        largest_alignment =
+            std::max(largest_alignment,
+                     calculate_largest_alignment(mem->type()->AsStruct()));
+      }
+
+      // Round up to the alignment size
+      count = adjust_for_alignment(count, align);
+      count += align;
+    }
+    // Round struct up to largest align size
+    count = adjust_for_alignment(count, largest_alignment);
+    return count;
+  }
+  if (type->IsVector()) {
+    auto* vec = type->AsVector();
+    uint32_t type_size = calculate_alignment_size(vec->type());
+    if (vec->size() == 2) {
+      return 2 * type_size;
+    }
+    return 4 * type_size;
+  }
+  return 0;
+}
+
 bool GeneratorImpl::EmitAliasType(const ast::type::AliasType* alias) {
   make_indent();
   out_ << "typedef ";
@@ -1254,15 +1351,37 @@ bool GeneratorImpl::EmitType(ast::type::Type* type, const std::string& name) {
     out_ << "struct {" << std::endl;
 
     increment_indent();
+    uint32_t current_offset = 0;
+    uint32_t pad_count = 0;
     for (const auto& mem : str->members()) {
       make_indent();
-      // TODO(dsinclair): Member decorations?
+      for (const auto& deco : mem->decorations()) {
-      // if (!mem->decorations().empty()) {
+        if (deco->IsOffset()) {
-      // }
+          uint32_t offset = deco->AsOffset()->offset();
+          if (offset != current_offset) {
+            out_ << "int8_t pad_" << pad_count << "["
+                 << (offset - current_offset) << "];" << std::endl;
+            pad_count++;
+            make_indent();
+          }
+          current_offset = offset;
+        } else {
+          error_ = "unsupported member decoration: " + deco->to_str();
+          return false;
+        }
+      }
 
       if (!EmitType(mem->type(), mem->name())) {
         return false;
       }
+      auto size = calculate_alignment_size(mem->type());
+      if (size == 0) {
+        error_ =
+            "unable to calculate byte size for: " + mem->type()->type_name();
+        return false;
+      }
+      current_offset += size;
+
       // Array member name will be output with the type
       if (!mem->type()->IsArray()) {
         out_ << " " << namer_.NameFor(mem->name());

src/writer/msl/generator_impl.h

@@ -22,6 +22,7 @@
 #include "src/ast/literal.h"
 #include "src/ast/module.h"
 #include "src/ast/scalar_constructor_expression.h"
+#include "src/ast/type/struct_type.h"
 #include "src/ast/type_constructor_expression.h"
 #include "src/scope_stack.h"
 #include "src/writer/msl/namer.h"
@@ -43,6 +44,16 @@ class GeneratorImpl : public TextGenerator {
   /// @returns true on successful generation; false otherwise
   bool Generate(const ast::Module& module);
 
+  /// Calculates the alignment size of the given |type|. This returns 0
+  /// for pointers as the size is unknown.
+  /// @param type the type to calculate the alignment size for
+  /// @returns the number of bytes used to align |type| or 0 on error
+  uint32_t calculate_alignment_size(ast::type::Type* type);
+  /// Calculates the largest alignment seen within a struct
+  /// @param type the struct to calculate
+  /// @returns the largest alignment value
+  uint32_t calculate_largest_alignment(ast::type::StructType* type);
+
   /// Handles generating an alias
   /// @param alias the alias to generate
   /// @returns true if the alias was emitted

src/writer/msl/generator_impl_test.cc

@@ -22,6 +22,19 @@
 #include "src/ast/identifier_expression.h"
 #include "src/ast/module.h"
 #include "src/ast/pipeline_stage.h"
+#include "src/ast/struct.h"
+#include "src/ast/struct_member.h"
+#include "src/ast/struct_member_offset_decoration.h"
+#include "src/ast/type/alias_type.h"
+#include "src/ast/type/array_type.h"
+#include "src/ast/type/bool_type.h"
+#include "src/ast/type/f32_type.h"
+#include "src/ast/type/i32_type.h"
+#include "src/ast/type/matrix_type.h"
+#include "src/ast/type/pointer_type.h"
+#include "src/ast/type/struct_type.h"
+#include "src/ast/type/u32_type.h"
+#include "src/ast/type/vector_type.h"
 #include "src/ast/type/void_type.h"
 #include "src/writer/msl/namer.h"
 
@@ -108,6 +121,200 @@ INSTANTIATE_TEST_SUITE_P(
                     MslBuiltinData{ast::Builtin::kGlobalInvocationId,
                                    "thread_position_in_grid"}));
 
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_alias) {
+  ast::type::F32Type f32;
+  ast::type::AliasType alias("a", &f32);
+  GeneratorImpl g;
+  EXPECT_EQ(4u, g.calculate_alignment_size(&alias));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_array) {
+  ast::type::F32Type f32;
+  ast::type::ArrayType ary(&f32, 4);
+  GeneratorImpl g;
+  EXPECT_EQ(4u * 4u, g.calculate_alignment_size(&ary));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_bool) {
+  ast::type::BoolType bool_type;
+  GeneratorImpl g;
+  EXPECT_EQ(1u, g.calculate_alignment_size(&bool_type));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_f32) {
+  ast::type::F32Type f32;
+  GeneratorImpl g;
+  EXPECT_EQ(4u, g.calculate_alignment_size(&f32));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_i32) {
+  ast::type::I32Type i32;
+  GeneratorImpl g;
+  EXPECT_EQ(4u, g.calculate_alignment_size(&i32));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_matrix) {
+  ast::type::F32Type f32;
+  ast::type::MatrixType mat(&f32, 3, 2);
+  GeneratorImpl g;
+  EXPECT_EQ(4u * 3u * 2u, g.calculate_alignment_size(&mat));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_pointer) {
+  ast::type::BoolType bool_type;
+  ast::type::PointerType ptr(&bool_type, ast::StorageClass::kPrivate);
+  GeneratorImpl g;
+  EXPECT_EQ(0u, g.calculate_alignment_size(&ptr));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_struct) {
+  ast::type::I32Type i32;
+  ast::type::F32Type f32;
+
+  ast::StructMemberDecorationList decos;
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(4));
+
+  ast::StructMemberList members;
+  members.push_back(
+      std::make_unique<ast::StructMember>("a", &i32, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(32));
+  members.push_back(
+      std::make_unique<ast::StructMember>("b", &f32, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(128));
+  members.push_back(
+      std::make_unique<ast::StructMember>("c", &f32, std::move(decos)));
+
+  auto str = std::make_unique<ast::Struct>();
+  str->set_members(std::move(members));
+
+  ast::type::StructType s(std::move(str));
+
+  GeneratorImpl g;
+  EXPECT_EQ(132u, g.calculate_alignment_size(&s));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_struct_of_struct) {
+  ast::type::I32Type i32;
+  ast::type::F32Type f32;
+  ast::type::VectorType fvec(&f32, 3);
+
+  ast::StructMemberDecorationList decos;
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(0));
+
+  ast::StructMemberList members;
+  members.push_back(
+      std::make_unique<ast::StructMember>("a", &i32, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(16));
+  members.push_back(
+      std::make_unique<ast::StructMember>("b", &fvec, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(32));
+  members.push_back(
+      std::make_unique<ast::StructMember>("c", &f32, std::move(decos)));
+
+  auto inner_str = std::make_unique<ast::Struct>();
+  inner_str->set_members(std::move(members));
+
+  ast::type::StructType inner_s(std::move(inner_str));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(0));
+  members.push_back(
+      std::make_unique<ast::StructMember>("d", &f32, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(32));
+  members.push_back(
+      std::make_unique<ast::StructMember>("e", &inner_s, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(64));
+  members.push_back(
+      std::make_unique<ast::StructMember>("f", &f32, std::move(decos)));
+
+  auto outer_str = std::make_unique<ast::Struct>();
+  outer_str->set_members(std::move(members));
+
+  ast::type::StructType outer_s(std::move(outer_str));
+
+  GeneratorImpl g;
+  EXPECT_EQ(80u, g.calculate_alignment_size(&outer_s));
+}
+
+TEST_F(MslGeneratorImplTest, calculate_alignment_size_u32) {
+  ast::type::U32Type u32;
+  GeneratorImpl g;
+  EXPECT_EQ(4u, g.calculate_alignment_size(&u32));
+}
+
+struct MslVectorSizeData {
+  uint32_t elements;
+  uint32_t byte_size;
+};
+inline std::ostream& operator<<(std::ostream& out, MslVectorSizeData data) {
+  out << data.elements;
+  return out;
+}
+using MslVectorSizeBoolTest = testing::TestWithParam<MslVectorSizeData>;
+TEST_P(MslVectorSizeBoolTest, calculate) {
+  auto param = GetParam();
+
+  ast::type::BoolType bool_type;
+  ast::type::VectorType vec(&bool_type, param.elements);
+  GeneratorImpl g;
+  EXPECT_EQ(param.byte_size, g.calculate_alignment_size(&vec));
+}
+INSTANTIATE_TEST_SUITE_P(MslGeneratorImplTest,
+                         MslVectorSizeBoolTest,
+                         testing::Values(MslVectorSizeData{2u, 2u},
+                                         MslVectorSizeData{3u, 4u},
+                                         MslVectorSizeData{4u, 4u}));
+
+using MslVectorSizeI32Test = testing::TestWithParam<MslVectorSizeData>;
+TEST_P(MslVectorSizeI32Test, calculate) {
+  auto param = GetParam();
+
+  ast::type::I32Type i32;
+  ast::type::VectorType vec(&i32, param.elements);
+  GeneratorImpl g;
+  EXPECT_EQ(param.byte_size, g.calculate_alignment_size(&vec));
+}
+INSTANTIATE_TEST_SUITE_P(MslGeneratorImplTest,
+                         MslVectorSizeI32Test,
+                         testing::Values(MslVectorSizeData{2u, 8u},
+                                         MslVectorSizeData{3u, 16u},
+                                         MslVectorSizeData{4u, 16u}));
+
+using MslVectorSizeU32Test = testing::TestWithParam<MslVectorSizeData>;
+TEST_P(MslVectorSizeU32Test, calculate) {
+  auto param = GetParam();
+
+  ast::type::U32Type u32;
+  ast::type::VectorType vec(&u32, param.elements);
+  GeneratorImpl g;
+  EXPECT_EQ(param.byte_size, g.calculate_alignment_size(&vec));
+}
+INSTANTIATE_TEST_SUITE_P(MslGeneratorImplTest,
+                         MslVectorSizeU32Test,
+                         testing::Values(MslVectorSizeData{2u, 8u},
+                                         MslVectorSizeData{3u, 16u},
+                                         MslVectorSizeData{4u, 16u}));
+
+using MslVectorSizeF32Test = testing::TestWithParam<MslVectorSizeData>;
+TEST_P(MslVectorSizeF32Test, calculate) {
+  auto param = GetParam();
+
+  ast::type::F32Type f32;
+  ast::type::VectorType vec(&f32, param.elements);
+  GeneratorImpl g;
+  EXPECT_EQ(param.byte_size, g.calculate_alignment_size(&vec));
+}
+INSTANTIATE_TEST_SUITE_P(MslGeneratorImplTest,
+                         MslVectorSizeF32Test,
+                         testing::Values(MslVectorSizeData{2u, 8u},
+                                         MslVectorSizeData{3u, 16u},
+                                         MslVectorSizeData{4u, 16u}));
+
 }  // namespace
 }  // namespace msl
 }  // namespace writer

src/writer/msl/generator_impl_type_test.cc

@@ -176,8 +176,7 @@ TEST_F(MslGeneratorImplTest, DISABLED_EmitType_Pointer) {
   EXPECT_EQ(g.result(), "float*");
 }
 
-// TODO(dsinclair): How to translate offsets?
+TEST_F(MslGeneratorImplTest, EmitType_Struct) {
-TEST_F(MslGeneratorImplTest, DISABLED_EmitType_Struct) {
   ast::type::I32Type i32;
   ast::type::F32Type f32;
 
@@ -203,6 +202,42 @@ TEST_F(MslGeneratorImplTest, DISABLED_EmitType_Struct) {
 })");
 }
 
+TEST_F(MslGeneratorImplTest, EmitType_Struct_InjectPadding) {
+  ast::type::I32Type i32;
+  ast::type::F32Type f32;
+
+  ast::StructMemberDecorationList decos;
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(4));
+
+  ast::StructMemberList members;
+  members.push_back(
+      std::make_unique<ast::StructMember>("a", &i32, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(32));
+  members.push_back(
+      std::make_unique<ast::StructMember>("b", &f32, std::move(decos)));
+
+  decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(128));
+  members.push_back(
+      std::make_unique<ast::StructMember>("c", &f32, std::move(decos)));
+
+  auto str = std::make_unique<ast::Struct>();
+  str->set_members(std::move(members));
+
+  ast::type::StructType s(std::move(str));
+
+  GeneratorImpl g;
+  ASSERT_TRUE(g.EmitType(&s, "")) << g.error();
+  EXPECT_EQ(g.result(), R"(struct {
+  int8_t pad_0[4];
+  int a;
+  int8_t pad_1[24];
+  float b;
+  int8_t pad_2[92];
+  float c;
+})");
+}
+
 TEST_F(MslGeneratorImplTest, EmitType_Struct_NameCollision) {
   ast::type::I32Type i32;
   ast::type::F32Type f32;