Add transform/PadArrayElements

Replaces arrays with an explicit stride with an array to a structure holding the element padded with a `[[size]]` decoration.

Note that the HLSL writer is still not correctly emitting structure fields with a `[[size]]`, which will be fixed in a follow up change.

Bug: tint:182
Bug: tint:895
Fixed: tint:180
Fixed: tint:649
Change-Id: Ic135dfc89309ac805507e9f39392577c7f82d154
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/54582
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: David Neto <dneto@google.com>
Reviewed-by: James Price <jrprice@google.com>

src/transform/hlsl.cc

@@ -23,6 +23,7 @@
 #include "src/transform/external_texture_transform.h"
 #include "src/transform/inline_pointer_lets.h"
 #include "src/transform/manager.h"
+#include "src/transform/pad_array_elements.h"
 #include "src/transform/promote_initializers_to_const_var.h"
 #include "src/transform/simplify.h"
 #include "src/transform/wrap_arrays_in_structs.h"
@@ -52,6 +53,7 @@ Output Hlsl::Run(const Program* in, const DataMap&) {
   manager.Add<ExternalTextureTransform>();
   manager.Add<PromoteInitializersToConstVar>();
   manager.Add<WrapArraysInStructs>();
+  manager.Add<PadArrayElements>();
   data.Add<CanonicalizeEntryPointIO::Config>(
       CanonicalizeEntryPointIO::BuiltinStyle::kStructMember);
   auto out = manager.Run(in, data);

src/transform/msl.cc

@@ -27,6 +27,7 @@
 #include "src/transform/canonicalize_entry_point_io.h"
 #include "src/transform/external_texture_transform.h"
 #include "src/transform/manager.h"
+#include "src/transform/pad_array_elements.h"
 #include "src/transform/promote_initializers_to_const_var.h"
 #include "src/transform/wrap_arrays_in_structs.h"
 
@@ -43,6 +44,7 @@ Output Msl::Run(const Program* in, const DataMap&) {
   manager.Add<ExternalTextureTransform>();
   manager.Add<PromoteInitializersToConstVar>();
   manager.Add<WrapArraysInStructs>();
+  manager.Add<PadArrayElements>();
   data.Add<CanonicalizeEntryPointIO::Config>(
       CanonicalizeEntryPointIO::BuiltinStyle::kParameter);
   auto out = manager.Run(in, data);

src/transform/pad_array_elements.cc

@@ -0,0 +1,157 @@
+// 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/pad_array_elements.h"
+
+#include <unordered_map>
+#include <utility>
+
+#include "src/program_builder.h"
+#include "src/sem/array.h"
+#include "src/sem/expression.h"
+#include "src/utils/get_or_create.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using ArrayBuilder = std::function<ast::Array*()>;
+
+/// PadArray returns a function that constructs a new array in `ctx.dst` with
+/// the element type padded to account for the explicit stride. PadArray will
+/// recursively pad arrays-of-arrays. The new array element type will be added
+/// to module-scope type declarations of `ctx.dst`.
+/// @param ctx the CloneContext
+/// @param create_ast_type_for Transform::CreateASTTypeFor()
+/// @param padded_arrays a map of src array type to the new array name
+/// @param array the array type
+/// @return the new AST array
+template <typename CREATE_AST_TYPE_FOR>
+ArrayBuilder PadArray(
+    CloneContext& ctx,
+    CREATE_AST_TYPE_FOR&& create_ast_type_for,
+    std::unordered_map<const sem::Array*, ArrayBuilder>& padded_arrays,
+    const sem::Array* array) {
+  if (array->IsStrideImplicit()) {
+    // We don't want to wrap arrays that have an implicit stride
+    return nullptr;
+  }
+
+  return utils::GetOrCreate(padded_arrays, array, [&] {
+    // Generate a unique name for the array element type
+    auto name = ctx.dst->Symbols().New("tint_padded_array_element");
+
+    // Examine the element type. Is it also an array?
+    ast::Type* el_ty = nullptr;
+    if (auto* el_array = array->ElemType()->As<sem::Array>()) {
+      // Array of array - call PadArray() on the element type
+      if (auto p =
+              PadArray(ctx, create_ast_type_for, padded_arrays, el_array)) {
+        el_ty = p();
+      }
+    }
+
+    // If the element wasn't a padded array, just create the typical AST type
+    // for it
+    if (el_ty == nullptr) {
+      el_ty = create_ast_type_for(&ctx, array->ElemType());
+    }
+
+    // Structure() will create and append the ast::Struct to the
+    // global declarations of `ctx.dst`. As we haven't finished building the
+    // current module-scope statement or function, this will be placed
+    // immediately before the usage.
+    ctx.dst->Structure(
+        name,
+        {ctx.dst->Member("el", el_ty, {ctx.dst->MemberSize(array->Stride())})});
+
+    auto* dst = ctx.dst;
+    return [=] {
+      return dst->ty.array(dst->create<ast::TypeName>(name), array->Count());
+    };
+  });
+}
+
+}  // namespace
+
+PadArrayElements::PadArrayElements() = default;
+
+PadArrayElements::~PadArrayElements() = default;
+
+Output PadArrayElements::Run(const Program* in, const DataMap&) {
+  ProgramBuilder out;
+  CloneContext ctx(&out, in);
+
+  auto& sem = ctx.src->Sem();
+
+  std::unordered_map<const sem::Array*, ArrayBuilder> padded_arrays;
+  auto pad = [&](const sem::Array* array) {
+    return PadArray(ctx, CreateASTTypeFor, padded_arrays, array);
+  };
+
+  // Replace all array types with their corresponding padded array type
+  ctx.ReplaceAll([&](ast::Type* ast_type) -> ast::Type* {
+    auto* type = ctx.src->TypeOf(ast_type);
+    if (auto* array = type->UnwrapRef()->As<sem::Array>()) {
+      if (auto p = pad(array)) {
+        return p();
+      }
+    }
+    return nullptr;
+  });
+
+  // Fix up array accessors so `a[1]` becomes `a[1].el`
+  ctx.ReplaceAll(
+      [&](ast::ArrayAccessorExpression* accessor) -> ast::Expression* {
+        if (auto* array = tint::As<sem::Array>(
+                sem.Get(accessor->array())->Type()->UnwrapRef())) {
+          if (pad(array)) {
+            // Array element is wrapped in a structure. Emit a member accessor
+            // to get to the actual array element.
+            auto* idx = ctx.CloneWithoutTransform(accessor);
+            return ctx.dst->MemberAccessor(idx, "el");
+          }
+        }
+        return nullptr;
+      });
+
+  // Fix up array constructors so `A(1,2)` becomes
+  // `A(padded(1), padded(2))`
+  ctx.ReplaceAll([&](ast::TypeConstructorExpression* ctor) -> ast::Expression* {
+    if (auto* array =
+            tint::As<sem::Array>(sem.Get(ctor)->Type()->UnwrapRef())) {
+      if (auto p = pad(array)) {
+        auto* arr_ty = p();
+        auto el_typename = arr_ty->type()->As<ast::TypeName>()->name();
+
+        ast::ExpressionList args;
+        args.reserve(ctor->values().size());
+        for (auto* arg : ctor->values()) {
+          args.emplace_back(ctx.dst->Construct(
+              ctx.dst->create<ast::TypeName>(el_typename), ctx.Clone(arg)));
+        }
+
+        return ctx.dst->Construct(arr_ty, args);
+      }
+    }
+    return nullptr;
+  });
+
+  ctx.Clone();
+
+  return Output(Program(std::move(out)));
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/pad_array_elements.h

@@ -0,0 +1,51 @@
+// 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_PAD_ARRAY_ELEMENTS_H_
+#define SRC_TRANSFORM_PAD_ARRAY_ELEMENTS_H_
+
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// PadArrayElements is a transform that replaces array types with an explicit
+/// stride that is larger than the implicit stride, with an array of a new
+/// structure type. This structure holds with a single field of the element
+/// type, decorated with a [[size]] decoration to pad the structure to the
+/// required array stride. The new array types have no explicit stride,
+/// structure size is equal to the desired stride.
+/// Array index expressions and constructors are also adjusted to deal with this
+/// structure element type.
+/// This transform helps with backends that cannot directly return arrays or use
+/// them as parameters.
+class PadArrayElements : public Transform {
+ public:
+  /// Constructor
+  PadArrayElements();
+
+  /// Destructor
+  ~PadArrayElements() override;
+
+  /// Runs the transform on `program`, returning the transformation result.
+  /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
+  /// @returns the transformation result
+  Output Run(const Program* program, const DataMap& data = {}) override;
+};
+
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_PAD_ARRAY_ELEMENTS_H_

src/transform/pad_array_elements_test.cc

@@ -0,0 +1,379 @@
+// 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/pad_array_elements.h"
+
+#include <utility>
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using PadArrayElementsTest = TransformTest;
+
+TEST_F(PadArrayElementsTest, EmptyModule) {
+  auto* src = "";
+  auto* expect = "";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, ImplicitArrayStride) {
+  auto* src = R"(
+var<private> arr : array<i32, 4>;
+)";
+  auto* expect = src;
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, ArrayAsGlobal) {
+  auto* src = R"(
+var<private> arr : [[stride(8)]] array<i32, 4>;
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  [[size(8)]]
+  el : i32;
+};
+
+var<private> arr : array<tint_padded_array_element, 4>;
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, ArrayFunctionVar) {
+  auto* src = R"(
+fn f() {
+  var arr : [[stride(16)]] array<i32, 4>;
+  arr = [[stride(16)]] array<i32, 4>();
+  arr = [[stride(16)]] array<i32, 4>(1, 2, 3, 4);
+  let x = arr[3];
+}
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  [[size(16)]]
+  el : i32;
+};
+
+fn f() {
+  var arr : array<tint_padded_array_element, 4>;
+  arr = array<tint_padded_array_element, 4>();
+  arr = array<tint_padded_array_element, 4>(tint_padded_array_element(1), tint_padded_array_element(2), tint_padded_array_element(3), tint_padded_array_element(4));
+  let x = arr[3].el;
+}
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, ArrayAsParam) {
+  auto* src = R"(
+fn f(a : [[stride(12)]] array<i32, 4>) -> i32 {
+  return a[2];
+}
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  [[size(12)]]
+  el : i32;
+};
+
+fn f(a : array<tint_padded_array_element, 4>) -> i32 {
+  return a[2].el;
+}
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+// TODO(crbug.com/tint/781): Cannot parse the stride on the return array type.
+TEST_F(PadArrayElementsTest, DISABLED_ArrayAsReturn) {
+  auto* src = R"(
+fn f() -> [[stride(8)]] array<i32, 4> {
+  return array<i32, 4>(1, 2, 3, 4);
+}
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  el : i32;
+  [[size(4)]]
+  padding : u32;
+};
+
+fn f() -> array<tint_padded_array_element, 4> {
+  return array<tint_padded_array_element, 4>(tint_padded_array_element(1, 0u), tint_padded_array_element(2, 0u), tint_padded_array_element(3, 0u), tint_padded_array_element(4, 0u));
+}
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, ArrayAlias) {
+  auto* src = R"(
+type Array = [[stride(16)]] array<i32, 4>;
+
+fn f() {
+  var arr : Array;
+  arr = Array();
+  arr = Array(1, 2, 3, 4);
+  let vals : Array = Array(1, 2, 3, 4);
+  arr = vals;
+  let x = arr[3];
+}
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  [[size(16)]]
+  el : i32;
+};
+
+type Array = array<tint_padded_array_element, 4>;
+
+fn f() {
+  var arr : array<tint_padded_array_element, 4>;
+  arr = array<tint_padded_array_element, 4>();
+  arr = array<tint_padded_array_element, 4>(tint_padded_array_element(1), tint_padded_array_element(2), tint_padded_array_element(3), tint_padded_array_element(4));
+  let vals : array<tint_padded_array_element, 4> = array<tint_padded_array_element, 4>(tint_padded_array_element(1), tint_padded_array_element(2), tint_padded_array_element(3), tint_padded_array_element(4));
+  arr = vals;
+  let x = arr[3].el;
+}
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, ArraysInStruct) {
+  auto* src = R"(
+struct S {
+  a : [[stride(8)]] array<i32, 4>;
+  b : [[stride(8)]] array<i32, 8>;
+  c : [[stride(8)]] array<i32, 4>;
+  d : [[stride(12)]] array<i32, 8>;
+};
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  [[size(8)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_1 {
+  [[size(8)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_2 {
+  [[size(12)]]
+  el : i32;
+};
+
+struct S {
+  a : array<tint_padded_array_element, 4>;
+  b : array<tint_padded_array_element_1, 8>;
+  c : array<tint_padded_array_element, 4>;
+  d : array<tint_padded_array_element_2, 8>;
+};
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, ArraysOfArraysInStruct) {
+  auto* src = R"(
+struct S {
+  a : [[stride(512)]] array<i32, 4>;
+  b : [[stride(512)]] array<[[stride(32)]] array<i32, 4>, 4>;
+  c : [[stride(512)]] array<[[stride(64)]] array<[[stride(8)]] array<i32, 4>, 4>, 4>;
+};
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  [[size(512)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_2 {
+  [[size(32)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_1 {
+  [[size(512)]]
+  el : array<tint_padded_array_element_2, 4>;
+};
+
+struct tint_padded_array_element_5 {
+  [[size(8)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_4 {
+  [[size(64)]]
+  el : array<tint_padded_array_element_5, 4>;
+};
+
+struct tint_padded_array_element_3 {
+  [[size(512)]]
+  el : array<tint_padded_array_element_4, 4>;
+};
+
+struct S {
+  a : array<tint_padded_array_element, 4>;
+  b : array<tint_padded_array_element_1, 4>;
+  c : array<tint_padded_array_element_3, 4>;
+};
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, AccessArraysOfArraysInStruct) {
+  auto* src = R"(
+struct S {
+  a : [[stride(512)]] array<i32, 4>;
+  b : [[stride(512)]] array<[[stride(32)]] array<i32, 4>, 4>;
+  c : [[stride(512)]] array<[[stride(64)]] array<[[stride(8)]] array<i32, 4>, 4>, 4>;
+};
+
+fn f(s : S) -> i32 {
+  return s.a[2] + s.b[1][2] + s.c[3][1][2];
+}
+)";
+  auto* expect = R"(
+struct tint_padded_array_element {
+  [[size(512)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_2 {
+  [[size(32)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_1 {
+  [[size(512)]]
+  el : array<tint_padded_array_element_2, 4>;
+};
+
+struct tint_padded_array_element_5 {
+  [[size(8)]]
+  el : i32;
+};
+
+struct tint_padded_array_element_4 {
+  [[size(64)]]
+  el : array<tint_padded_array_element_5, 4>;
+};
+
+struct tint_padded_array_element_3 {
+  [[size(512)]]
+  el : array<tint_padded_array_element_4, 4>;
+};
+
+struct S {
+  a : array<tint_padded_array_element, 4>;
+  b : array<tint_padded_array_element_1, 4>;
+  c : array<tint_padded_array_element_3, 4>;
+};
+
+fn f(s : S) -> i32 {
+  return ((s.a[2].el + s.b[1].el[2].el) + s.c[3].el[1].el[2].el);
+}
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(PadArrayElementsTest, DeclarationOrder) {
+  auto* src = R"(
+type T0 = i32;
+
+type T1 = [[stride(8)]] array<i32, 1>;
+
+type T2 = i32;
+
+fn f1(a : [[stride(8)]] array<i32, 2>) {
+}
+
+type T3 = i32;
+
+fn f2() {
+  var v : [[stride(8)]] array<i32, 3>;
+}
+)";
+  auto* expect = R"(
+type T0 = i32;
+
+struct tint_padded_array_element {
+  [[size(8)]]
+  el : i32;
+};
+
+type T1 = array<tint_padded_array_element, 1>;
+
+type T2 = i32;
+
+struct tint_padded_array_element_1 {
+  [[size(8)]]
+  el : i32;
+};
+
+fn f1(a : array<tint_padded_array_element_1, 2>) {
+}
+
+type T3 = i32;
+
+struct tint_padded_array_element_2 {
+  [[size(8)]]
+  el : i32;
+};
+
+fn f2() {
+  var v : array<tint_padded_array_element_2, 3>;
+}
+)";
+
+  auto got = Run<PadArrayElements>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint

src/transform/transform_test.cc

@@ -78,7 +78,7 @@ TEST_F(CreateASTTypeForTest, Vector) {
 
 TEST_F(CreateASTTypeForTest, ArrayImplicitStride) {
   auto* arr = create([](ProgramBuilder& b) {
-    return b.create<sem::Array>(b.create<sem::F32>(), 2, 4, 4, 32u, true);
+    return b.create<sem::Array>(b.create<sem::F32>(), 2, 4, 4, 32u, 32u);
   });
   ASSERT_TRUE(arr->Is<ast::Array>());
   ASSERT_TRUE(arr->As<ast::Array>()->type()->Is<ast::F32>());
@@ -88,7 +88,7 @@ TEST_F(CreateASTTypeForTest, ArrayImplicitStride) {
 
 TEST_F(CreateASTTypeForTest, ArrayNonImplicitStride) {
   auto* arr = create([](ProgramBuilder& b) {
-    return b.create<sem::Array>(b.create<sem::F32>(), 2, 4, 4, 32u, false);
+    return b.create<sem::Array>(b.create<sem::F32>(), 2, 4, 4, 64u, 32u);
   });
   ASSERT_TRUE(arr->Is<ast::Array>());
   ASSERT_TRUE(arr->As<ast::Array>()->type()->Is<ast::F32>());
@@ -100,7 +100,7 @@ TEST_F(CreateASTTypeForTest, ArrayNonImplicitStride) {
                 ->decorations()[0]
                 ->As<ast::StrideDecoration>()
                 ->stride(),
-            32u);
+            64u);
 }
 
 TEST_F(CreateASTTypeForTest, Struct) {