Re-allow dynamic indexing of 'let' arrays and matrices

Spec change: https://github.com/gpuweb/gpuweb/pull/2427
Reverses: tint:867

This reverts and fixes commits:
 b6fdcc54df
 10442eff7d

Added a bunch of end-to-end tests.

Fixed: tint:1352
Change-Id: I34968243bbec1cab838c8ba50a6f027146bbfd06
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/75401
Reviewed-by: David Neto <dneto@google.com>
Reviewed-by: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>

src/transform/var_for_dynamic_index.cc

@@ -0,0 +1,91 @@
+// 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/var_for_dynamic_index.h"
+
+#include <utility>
+
+#include "src/program_builder.h"
+#include "src/sem/array.h"
+#include "src/sem/block_statement.h"
+#include "src/sem/expression.h"
+#include "src/sem/statement.h"
+#include "src/transform/for_loop_to_loop.h"
+
+namespace tint {
+namespace transform {
+
+VarForDynamicIndex::VarForDynamicIndex() = default;
+
+VarForDynamicIndex::~VarForDynamicIndex() = default;
+
+void VarForDynamicIndex::Run(CloneContext& ctx, const DataMap&, DataMap&) {
+  ProgramBuilder out;
+  if (!Requires<ForLoopToLoop>(ctx)) {
+    return;
+  }
+
+  auto& sem = ctx.src->Sem();
+
+  for (auto* node : ctx.src->ASTNodes().Objects()) {
+    if (auto* access_expr = node->As<ast::IndexAccessorExpression>()) {
+      // Found an array accessor expression
+      auto* index_expr = access_expr->index;
+      auto* object_expr = access_expr->object;
+
+      if (sem.Get(index_expr)->ConstantValue()) {
+        // Index expression resolves to a compile time value.
+        // As this isn't a dynamic index, we can ignore this.
+        continue;
+      }
+
+      auto* indexed = sem.Get(object_expr);
+      if (!indexed->Type()->IsAnyOf<sem::Array, sem::Matrix>()) {
+        // This transform currently only cares about array and matrices.
+        continue;
+      }
+
+      // Construct a `var` declaration to hold the value in memory.
+      // TODO(bclayton): group multiple accesses in the same object.
+      // e.g. arr[i] + arr[i+1] // Don't create two vars for this
+      auto var_name = ctx.dst->Symbols().New("var_for_index");
+      auto* var_decl = ctx.dst->Decl(
+          ctx.dst->Var(var_name, nullptr, ctx.Clone(object_expr)));
+
+      // Statement that owns the expression
+      auto* stmt = indexed->Stmt();
+      // Block that owns the statement
+      auto* block = stmt->Parent()->As<sem::BlockStatement>();
+      if (!block) {
+        TINT_ICE(Transform, ctx.dst->Diagnostics())
+            << "statement parent was not a block";
+        continue;
+      }
+
+      // Insert the `var` declaration before the statement that performs the
+      // indexing. Note that for indexing chains, AST node ordering guarantees
+      // that the inner-most index variable will be placed first in the block.
+      ctx.InsertBefore(block->Declaration()->statements, stmt->Declaration(),
+                       var_decl);
+
+      // Replace the original index expression with the new `var`.
+      ctx.Replace(object_expr, ctx.dst->Expr(var_name));
+    }
+  }
+
+  ctx.Clone();
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/var_for_dynamic_index.h

@@ -0,0 +1,49 @@
+// 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_VAR_FOR_DYNAMIC_INDEX_H_
+#define SRC_TRANSFORM_VAR_FOR_DYNAMIC_INDEX_H_
+
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// A transform that extracts array and matrix values that are dynamically
+/// indexed to a temporary `var` local before performing the index. This
+/// transform is used by the SPIR-V writer as there is no SPIR-V instruction
+/// that can dynamically index a non-pointer composite.
+/// Requires the ForLoopToLoop transform to be run first.
+class VarForDynamicIndex : public Transform {
+ public:
+  /// Constructor
+  VarForDynamicIndex();
+
+  /// Destructor
+  ~VarForDynamicIndex() 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_VAR_FOR_DYNAMIC_INDEX_H_

src/transform/var_for_dynamic_index_test.cc

@@ -0,0 +1,327 @@
+// 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/var_for_dynamic_index.h"
+#include "src/transform/for_loop_to_loop.h"
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using VarForDynamicIndexTest = TransformTest;
+
+TEST_F(VarForDynamicIndexTest, EmptyModule) {
+  auto* src = "";
+  auto* expect = "";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexDynamic) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  let p = array<i32, 4>(1, 2, 3, 4);
+  let x = p[i];
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  let p = array<i32, 4>(1, 2, 3, 4);
+  var var_for_index = p;
+  let x = var_for_index[i];
+}
+)";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, MatrixIndexDynamic) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  let x = p[i];
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  var var_for_index = p;
+  let x = var_for_index[i];
+}
+)";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexDynamicChain) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  var j : i32;
+  let p = array<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  let x = p[i][j];
+}
+)";
+
+  // TODO(bclayton): Optimize this case:
+  // This output is not as efficient as it could be.
+  // We only actually need to hoist the inner-most array to a `var`
+  // (`var_for_index`), as later indexing operations will be working with
+  // references, not values.
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  var j : i32;
+  let p = array<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  var var_for_index = p;
+  var var_for_index_1 = var_for_index[i];
+  let x = var_for_index_1[j];
+}
+)";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexInForLoopInit) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  let p = array<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  for(let x = p[i]; ; ) {
+    break;
+  }
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  let p = array<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  {
+    var var_for_index = p;
+    let x = var_for_index[i];
+    loop {
+      break;
+    }
+  }
+}
+)";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, MatrixIndexInForLoopInit) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  for(let x = p[i]; ; ) {
+    break;
+  }
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  {
+    var var_for_index = p;
+    let x = var_for_index[i];
+    loop {
+      break;
+    }
+  }
+}
+)";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexInForLoopCond) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  let p = array<i32, 2>(1, 2);
+  for(; p[i] < 3; ) {
+    break;
+  }
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  let p = array<i32, 2>(1, 2);
+  loop {
+    var var_for_index = p;
+    if (!((var_for_index[i] < 3))) {
+      break;
+    }
+    break;
+  }
+}
+)";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, MatrixIndexInForLoopCond) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  for(; p[i].x < 3.0; ) {
+    break;
+  }
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  loop {
+    var var_for_index = p;
+    if (!((var_for_index[i].x < 3.0))) {
+      break;
+    }
+    break;
+  }
+}
+)";
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexLiteral) {
+  auto* src = R"(
+fn f() {
+  let p = array<i32, 4>(1, 2, 3, 4);
+  let x = p[1];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, MatrixIndexLiteral) {
+  auto* src = R"(
+fn f() {
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  let x = p[1];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexConstantLet) {
+  auto* src = R"(
+fn f() {
+  let p = array<i32, 4>(1, 2, 3, 4);
+  let c = 1;
+  let x = p[c];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, MatrixIndexConstantLet) {
+  auto* src = R"(
+fn f() {
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  let c = 1;
+  let x = p[c];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexLiteralChain) {
+  auto* src = R"(
+fn f() {
+  let p = array<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  let x = p[0][1];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, MatrixIndexLiteralChain) {
+  auto* src = R"(
+fn f() {
+  let p = mat2x2(1.0, 2.0, 3.0, 4.0);
+  let x = p[0][1];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<ForLoopToLoop, VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint