Add support for matrix construction from scalars

Use a transform to convert these to the vector form for the MSL and
SPIR-V backends.

MSL only has the scalar form from version 2.0 onwards.

Fixed: tint:1123
Change-Id: I384abd9872d9eae52a10a37cbd6aa96004692e9c
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/67360
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: James Price <jrprice@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/transform/vectorize_scalar_matrix_constructors.cc

@@ -0,0 +1,73 @@
+// Copyright 2021 The Tint 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 "src/transform/vectorize_scalar_matrix_constructors.h"
+
+#include <utility>
+
+#include "src/program_builder.h"
+#include "src/sem/expression.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::transform::VectorizeScalarMatrixConstructors);
+
+namespace tint {
+namespace transform {
+
+VectorizeScalarMatrixConstructors::VectorizeScalarMatrixConstructors() =
+    default;
+
+VectorizeScalarMatrixConstructors::~VectorizeScalarMatrixConstructors() =
+    default;
+
+void VectorizeScalarMatrixConstructors::Run(CloneContext& ctx,
+                                            const DataMap&,
+                                            DataMap&) {
+  ctx.ReplaceAll([&](const ast::TypeConstructorExpression* constructor)
+                     -> const ast::TypeConstructorExpression* {
+    // Check if this is a matrix constructor with scalar arguments.
+    auto* mat_type = ctx.src->Sem().Get(constructor->type)->As<sem::Matrix>();
+    if (!mat_type) {
+      return nullptr;
+    }
+    if (constructor->values.size() == 0) {
+      return nullptr;
+    }
+    if (!ctx.src->Sem().Get(constructor->values[0])->Type()->is_scalar()) {
+      return nullptr;
+    }
+
+    // Build a list of vector expressions for each column.
+    ast::ExpressionList columns;
+    for (uint32_t c = 0; c < mat_type->columns(); c++) {
+      // Build a list of scalar expressions for each value in the column.
+      ast::ExpressionList row_values;
+      for (uint32_t r = 0; r < mat_type->rows(); r++) {
+        row_values.push_back(
+            ctx.Clone(constructor->values[c * mat_type->rows() + r]));
+      }
+
+      // Construct the column vector.
+      auto* col = ctx.dst->vec(CreateASTTypeFor(ctx, mat_type->type()),
+                               mat_type->rows(), row_values);
+      columns.push_back(col);
+    }
+
+    return ctx.dst->Construct(CreateASTTypeFor(ctx, mat_type), columns);
+  });
+
+  ctx.Clone();
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/vectorize_scalar_matrix_constructors.h

@@ -0,0 +1,46 @@
+// Copyright 2021 The Tint 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.
+
+#ifndef SRC_TRANSFORM_VECTORIZE_SCALAR_MATRIX_CONSTRUCTORS_H_
+#define SRC_TRANSFORM_VECTORIZE_SCALAR_MATRIX_CONSTRUCTORS_H_
+
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// A transform that converts scalar matrix constructors to the vector form.
+class VectorizeScalarMatrixConstructors
+    : public Castable<VectorizeScalarMatrixConstructors, Transform> {
+ public:
+  /// Constructor
+  VectorizeScalarMatrixConstructors();
+
+  /// Destructor
+  ~VectorizeScalarMatrixConstructors() override;
+
+ protected:
+  /// Runs the transform using the CloneContext built for transforming a
+  /// program. Run() is responsible for calling Clone() on the CloneContext.
+  /// @param ctx the CloneContext primed with the input program and
+  /// ProgramBuilder
+  /// @param inputs optional extra transform-specific input data
+  /// @param outputs optional extra transform-specific output data
+  void Run(CloneContext& ctx, const DataMap& inputs, DataMap& outputs) override;
+};
+
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_VECTORIZE_SCALAR_MATRIX_CONSTRUCTORS_H_

src/transform/vectorize_scalar_matrix_constructors_test.cc

@@ -0,0 +1,114 @@
+// Copyright 2021 The Tint 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 "src/transform/vectorize_scalar_matrix_constructors.h"
+
+#include <string>
+#include <utility>
+
+#include "src/transform/test_helper.h"
+#include "src/utils/string.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using VectorizeScalarMatrixConstructorsTest =
+    TransformTestWithParam<std::pair<uint32_t, uint32_t>>;
+
+TEST_P(VectorizeScalarMatrixConstructorsTest, Basic) {
+  uint32_t cols = GetParam().first;
+  uint32_t rows = GetParam().second;
+  std::string mat_type =
+      "mat" + std::to_string(cols) + "x" + std::to_string(rows) + "<f32>";
+  std::string vec_type = "vec" + std::to_string(rows) + "<f32>";
+  std::string scalar_values;
+  std::string vector_values;
+  for (uint32_t c = 0; c < cols; c++) {
+    if (c > 0) {
+      vector_values += ", ";
+      scalar_values += ", ";
+    }
+    vector_values += vec_type + "(";
+    for (uint32_t r = 0; r < rows; r++) {
+      if (r > 0) {
+        scalar_values += ", ";
+        vector_values += ", ";
+      }
+      auto value = std::to_string(c * rows + r) + ".0";
+      scalar_values += value;
+      vector_values += value;
+    }
+    vector_values += ")";
+  }
+
+  std::string tmpl = R"(
+[[stage(fragment)]]
+fn main() {
+  let m = ${matrix}(${values});
+}
+)";
+  tmpl = utils::ReplaceAll(tmpl, "${matrix}", mat_type);
+  auto src = utils::ReplaceAll(tmpl, "${values}", scalar_values);
+  auto expect = utils::ReplaceAll(tmpl, "${values}", vector_values);
+
+  auto got = Run<VectorizeScalarMatrixConstructors>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_P(VectorizeScalarMatrixConstructorsTest, NonScalarConstructors) {
+  uint32_t cols = GetParam().first;
+  uint32_t rows = GetParam().second;
+  std::string mat_type =
+      "mat" + std::to_string(cols) + "x" + std::to_string(rows) + "<f32>";
+  std::string vec_type = "vec" + std::to_string(rows) + "<f32>";
+  std::string columns;
+  for (uint32_t c = 0; c < cols; c++) {
+    if (c > 0) {
+      columns += ", ";
+    }
+    columns += vec_type + "()";
+  }
+
+  std::string tmpl = R"(
+[[stage(fragment)]]
+fn main() {
+  let m = ${matrix}(${columns});
+}
+)";
+  tmpl = utils::ReplaceAll(tmpl, "${matrix}", mat_type);
+  auto src = utils::ReplaceAll(tmpl, "${columns}", columns);
+  auto expect = src;
+
+  auto got = Run<VectorizeScalarMatrixConstructors>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+INSTANTIATE_TEST_SUITE_P(VectorizeScalarMatrixConstructorsTest,
+                         VectorizeScalarMatrixConstructorsTest,
+                         testing::Values(std::make_pair(2, 2),
+                                         std::make_pair(2, 3),
+                                         std::make_pair(2, 4),
+                                         std::make_pair(3, 2),
+                                         std::make_pair(3, 3),
+                                         std::make_pair(3, 4),
+                                         std::make_pair(4, 2),
+                                         std::make_pair(4, 3),
+                                         std::make_pair(4, 4)));
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint