[spirv-writer] Infer and emit matrix layout

WGSL matrices are always column major, with tightly packed columns.
We need to infer that layout and decorate the containing structure
members, if those members require a layout.

Bug: tint:200
Change-Id: Ieceb460ed6c7eeb1244beecbf4e0eee0b8b8e373
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/26180
Reviewed-by: Sarah Mashayekhi <sarahmashay@google.com>

src/writer/spirv/builder.cc

@@ -128,6 +128,17 @@ uint32_t IndexFromName(char name) {
   return std::numeric_limits<uint32_t>::max();
 }
 
+/// Returns the matrix type that is |type| or that is wrapped by
+/// one or more levels of an arrays inside of |type|.
+/// @param type the given type, which must not be null
+/// @returns the nested matrix type, or nullptr if none
+ast::type::MatrixType* GetNestedMatrixType(ast::type::Type* type) {
+  while (type->IsArray()) {
+    type = type->AsArray()->type();
+  }
+  return type->IsMatrix() ? type->AsMatrix() : nullptr;
+}
+
 }  // namespace
 
 Builder::AccessorInfo::AccessorInfo() : source_id(0), source_type(nullptr) {}
@@ -2019,18 +2030,40 @@ uint32_t Builder::GenerateStructMember(uint32_t struct_id,
   push_debug(spv::Op::OpMemberName, {Operand::Int(struct_id), Operand::Int(idx),
                                      Operand::String(member->name())});
 
+  bool has_layout = false;
   for (const auto& deco : member->decorations()) {
     if (deco->IsOffset()) {
       push_annot(spv::Op::OpMemberDecorate,
                  {Operand::Int(struct_id), Operand::Int(idx),
                   Operand::Int(SpvDecorationOffset),
                   Operand::Int(deco->AsOffset()->offset())});
+      has_layout = true;
     } else {
       error_ = "unknown struct member decoration";
       return 0;
     }
   }
 
+  if (has_layout) {
+    // Infer and emit matrix layout.
+    auto* matrix_type = GetNestedMatrixType(member->type());
+    if (matrix_type) {
+      push_annot(spv::Op::OpMemberDecorate,
+                 {Operand::Int(struct_id), Operand::Int(idx),
+                  Operand::Int(SpvDecorationColMajor)});
+      if (!matrix_type->type()->IsF32()) {
+        error_ = "matrix scalar element type must be f32";
+        return 0;
+      }
+      const auto scalar_elem_size = 4;
+      const auto effective_row_count = (matrix_type->rows() == 2) ? 2 : 4;
+      push_annot(spv::Op::OpMemberDecorate,
+                 {Operand::Int(struct_id), Operand::Int(idx),
+                  Operand::Int(SpvDecorationMatrixStride),
+                  Operand::Int(effective_row_count * scalar_elem_size)});
+    }
+  }
+
   return GenerateTypeIfNeeded(member->type());
 }
 

src/writer/spirv/builder_type_test.cc

@@ -403,6 +403,173 @@ OpMemberDecorate %1 1 Offset 8
 )");
 }
 
+TEST_F(BuilderTest_Type, GenerateStruct_NonLayout_Matrix) {
+  // Don't infer layout for matrix when there is no offset.
+  ast::type::F32Type f32;
+  ast::type::MatrixType glsl_mat2x2(&f32, 2, 2);
+  ast::type::MatrixType glsl_mat2x3(&f32, 3, 2);  // 2 columns, 3 rows
+  ast::type::MatrixType glsl_mat4x4(&f32, 4, 4);
+
+  ast::StructMemberDecorationList empty_a;
+  ast::StructMemberDecorationList empty_b;
+  ast::StructMemberDecorationList empty_c;
+  ast::StructMemberList members;
+  members.push_back(std::make_unique<ast::StructMember>("a", &glsl_mat2x2,
+                                                        std::move(empty_a)));
+  members.push_back(std::make_unique<ast::StructMember>("b", &glsl_mat2x3,
+                                                        std::move(empty_b)));
+  members.push_back(std::make_unique<ast::StructMember>("c", &glsl_mat4x4,
+                                                        std::move(empty_c)));
+
+  auto s = std::make_unique<ast::Struct>(ast::StructDecoration::kNone,
+                                         std::move(members));
+  ast::type::StructType s_type(std::move(s));
+
+  ast::Module mod;
+  Builder b(&mod);
+  auto id = b.GenerateTypeIfNeeded(&s_type);
+  ASSERT_FALSE(b.has_error()) << b.error();
+  EXPECT_EQ(id, 1u);
+
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%4 = OpTypeFloat 32
+%3 = OpTypeVector %4 2
+%2 = OpTypeMatrix %3 2
+%6 = OpTypeVector %4 3
+%5 = OpTypeMatrix %6 2
+%8 = OpTypeVector %4 4
+%7 = OpTypeMatrix %8 4
+%1 = OpTypeStruct %2 %5 %7
+)");
+  EXPECT_EQ(DumpInstructions(b.debug()), R"(OpMemberName %1 0 "a"
+OpMemberName %1 1 "b"
+OpMemberName %1 2 "c"
+)");
+  EXPECT_EQ(DumpInstructions(b.annots()), "");
+}
+
+TEST_F(BuilderTest_Type, GenerateStruct_DecoratedMembers_LayoutMatrix) {
+  // We have to infer layout for matrix when it also has an offset.
+  ast::type::F32Type f32;
+  ast::type::MatrixType glsl_mat2x2(&f32, 2, 2);
+  ast::type::MatrixType glsl_mat2x3(&f32, 3, 2);  // 2 columns, 3 rows
+  ast::type::MatrixType glsl_mat4x4(&f32, 4, 4);
+
+  ast::StructMemberDecorationList a_decos;
+  a_decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(0));
+  ast::StructMemberDecorationList b_decos;
+  b_decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(16));
+  ast::StructMemberDecorationList c_decos;
+  c_decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(48));
+
+  ast::StructMemberList members;
+  members.push_back(std::make_unique<ast::StructMember>("a", &glsl_mat2x2,
+                                                        std::move(a_decos)));
+  members.push_back(std::make_unique<ast::StructMember>("b", &glsl_mat2x3,
+                                                        std::move(b_decos)));
+  members.push_back(std::make_unique<ast::StructMember>("c", &glsl_mat4x4,
+                                                        std::move(c_decos)));
+
+  auto s = std::make_unique<ast::Struct>(ast::StructDecoration::kNone,
+                                         std::move(members));
+  ast::type::StructType s_type(std::move(s));
+
+  ast::Module mod;
+  Builder b(&mod);
+  auto id = b.GenerateTypeIfNeeded(&s_type);
+  ASSERT_FALSE(b.has_error()) << b.error();
+  EXPECT_EQ(id, 1u);
+
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%4 = OpTypeFloat 32
+%3 = OpTypeVector %4 2
+%2 = OpTypeMatrix %3 2
+%6 = OpTypeVector %4 3
+%5 = OpTypeMatrix %6 2
+%8 = OpTypeVector %4 4
+%7 = OpTypeMatrix %8 4
+%1 = OpTypeStruct %2 %5 %7
+)");
+  EXPECT_EQ(DumpInstructions(b.debug()), R"(OpMemberName %1 0 "a"
+OpMemberName %1 1 "b"
+OpMemberName %1 2 "c"
+)");
+  EXPECT_EQ(DumpInstructions(b.annots()), R"(OpMemberDecorate %1 0 Offset 0
+OpMemberDecorate %1 0 ColMajor
+OpMemberDecorate %1 0 MatrixStride 8
+OpMemberDecorate %1 1 Offset 16
+OpMemberDecorate %1 1 ColMajor
+OpMemberDecorate %1 1 MatrixStride 16
+OpMemberDecorate %1 2 Offset 48
+OpMemberDecorate %1 2 ColMajor
+OpMemberDecorate %1 2 MatrixStride 16
+)");
+}
+
+TEST_F(BuilderTest_Type, GenerateStruct_DecoratedMembers_LayoutArraysOfMatrix) {
+  // We have to infer layout for matrix when it also has an offset.
+  // The decoration goes on the struct member, even if the matrix is buried
+  // in levels of arrays.
+  ast::type::F32Type f32;
+
+  ast::type::MatrixType glsl_mat2x2(&f32, 2, 2);
+  ast::type::ArrayType arr_mat2x2(&glsl_mat2x2, 1);  // Singly nested array
+
+  ast::type::MatrixType glsl_mat2x3(&f32, 3, 2);  // 2 columns, 3 rows
+  ast::type::ArrayType arr_mat2x3(&glsl_mat2x3, 1);
+  ast::type::ArrayType arr_arr_mat2x2(&arr_mat2x3, 1);  // Doubly nested array
+
+  ast::type::MatrixType glsl_mat4x4(&f32, 4, 4);
+  ast::type::ArrayType rtarr_mat4x4(&glsl_mat4x4);  // Runtime array
+
+  ast::StructMemberDecorationList a_decos;
+  a_decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(0));
+  ast::StructMemberDecorationList b_decos;
+  b_decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(16));
+  ast::StructMemberDecorationList c_decos;
+  c_decos.push_back(std::make_unique<ast::StructMemberOffsetDecoration>(48));
+
+  ast::StructMemberList members;
+  members.push_back(std::make_unique<ast::StructMember>("a", &glsl_mat2x2,
+                                                        std::move(a_decos)));
+  members.push_back(std::make_unique<ast::StructMember>("b", &glsl_mat2x3,
+                                                        std::move(b_decos)));
+  members.push_back(std::make_unique<ast::StructMember>("c", &glsl_mat4x4,
+                                                        std::move(c_decos)));
+
+  auto s = std::make_unique<ast::Struct>(ast::StructDecoration::kNone,
+                                         std::move(members));
+  ast::type::StructType s_type(std::move(s));
+
+  ast::Module mod;
+  Builder b(&mod);
+  auto id = b.GenerateTypeIfNeeded(&s_type);
+  ASSERT_FALSE(b.has_error()) << b.error();
+  EXPECT_EQ(id, 1u);
+
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%4 = OpTypeFloat 32
+%3 = OpTypeVector %4 2
+%2 = OpTypeMatrix %3 2
+%6 = OpTypeVector %4 3
+%5 = OpTypeMatrix %6 2
+%8 = OpTypeVector %4 4
+%7 = OpTypeMatrix %8 4
+%1 = OpTypeStruct %2 %5 %7
+)");
+  EXPECT_EQ(DumpInstructions(b.debug()), R"(OpMemberName %1 0 "a"
+OpMemberName %1 1 "b"
+OpMemberName %1 2 "c"
+)");
+  EXPECT_EQ(DumpInstructions(b.annots()), R"(OpMemberDecorate %1 0 Offset 0
+OpMemberDecorate %1 0 ColMajor
+OpMemberDecorate %1 0 MatrixStride 8
+OpMemberDecorate %1 1 Offset 16
+OpMemberDecorate %1 1 ColMajor
+OpMemberDecorate %1 1 MatrixStride 16
+OpMemberDecorate %1 2 Offset 48
+OpMemberDecorate %1 2 ColMajor
+OpMemberDecorate %1 2 MatrixStride 16
+)");
+}
+
 TEST_F(BuilderTest_Type, GenerateU32) {
   ast::type::U32Type u32;