Add transform::VarForDynamicIndex

Use this as part of the Spirv sanitizer.
Cleans up buggy dynamic array indexing logic in the SPIR-V writer.

Fixed: tint:824
Change-Id: Ia408e49bd808bc8dbf3a1897eb47f9b33b71fdfb
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/51780
Commit-Queue: Ben Clayton <bclayton@google.com>
Commit-Queue: Antonio Maiorano <amaiorano@google.com>
Auto-Submit: Ben Clayton <bclayton@google.com>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>

src/transform/spirv.cc

@@ -28,6 +28,7 @@
 #include "src/sem/variable.h"
 #include "src/transform/external_texture_transform.h"
 #include "src/transform/manager.h"
+#include "src/transform/var_for_dynamic_index.h"
 
 TINT_INSTANTIATE_TYPEINFO(tint::transform::Spirv::Config);
 
@@ -40,6 +41,7 @@ Spirv::~Spirv() = default;
 Output Spirv::Run(const Program* in, const DataMap& data) {
   Manager manager;
   manager.Add<ExternalTextureTransform>();
+  manager.Add<VarForDynamicIndex>();
   auto transformedInput = manager.Run(in, data);
 
   auto* cfg = data.Get<Config>();

src/transform/var_for_dynamic_index.cc

@@ -0,0 +1,80 @@
+// 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"
+
+namespace tint {
+namespace transform {
+
+VarForDynamicIndex::VarForDynamicIndex() = default;
+
+VarForDynamicIndex::~VarForDynamicIndex() = default;
+
+Output VarForDynamicIndex::Run(const Program* in, const DataMap&) {
+  ProgramBuilder out;
+  CloneContext ctx(&out, in);
+
+  for (auto* node : in->ASTNodes().Objects()) {
+    if (auto* access_expr = node->As<ast::ArrayAccessorExpression>()) {
+      // Found an array accessor expression
+      auto* index_expr = access_expr->idx_expr();
+      auto* array_expr = access_expr->array();
+
+      if (index_expr->Is<ast::ScalarConstructorExpression>()) {
+        // Index expression is a literal value. As this isn't a dynamic index,
+        // we can ignore this.
+        continue;
+      }
+
+      auto* array = ctx.src->Sem().Get(array_expr);
+      if (!array->Type()->Is<sem::Array>()) {
+        // This transform currently only cares about arrays.
+        continue;
+      }
+
+      auto* stmt = array->Stmt();   // Statement that owns the expression
+      auto* block = stmt->Block();  // Block that owns the statement
+
+      // Construct a `var` declaration to hold the value in memory.
+      auto* ty = CreateASTTypeFor(&ctx, array->Type());
+      auto var_name = ctx.dst->Symbols().New("var_for_array");
+      auto* var_decl = ctx.dst->Decl(ctx.dst->Var(
+          var_name, ty, ast::StorageClass::kNone, ctx.Clone(array_expr)));
+
+      // 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(array_expr, ctx.dst->Expr(var_name));
+    }
+  }
+
+  ctx.Clone();
+
+  return Output(Program(std::move(out)));
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/var_for_dynamic_index.h

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

src/transform/var_for_dynamic_index_test.cc

@@ -0,0 +1,121 @@
+// 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/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using VarForDynamicIndexTest = TransformTest;
+
+TEST_F(VarForDynamicIndexTest, EmptyModule) {
+  auto* src = "";
+  auto* expect = "";
+
+  auto got = Run<VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexDynamic) {
+  auto* src = R"(
+fn f() {
+  var i : i32;
+  let p : array<i32, 4> = array<i32, 4>(1, 2, 3, 4);
+  let x : i32 = p[i];
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var i : i32;
+  let p : array<i32, 4> = array<i32, 4>(1, 2, 3, 4);
+  var var_for_array : array<i32, 4> = p;
+  let x : i32 = var_for_array[i];
+}
+)";
+
+  auto got = Run<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<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  let x : i32 = 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_array`), 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<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  var var_for_array : array<array<i32, 2>, 2> = p;
+  var var_for_array_1 : array<i32, 2> = var_for_array[i];
+  let x : i32 = var_for_array_1[j];
+}
+)";
+
+  auto got = Run<VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexLiteral) {
+  auto* src = R"(
+fn f() {
+  let p : array<i32, 4> = array<i32, 4>(1, 2, 3, 4);
+  let x : i32 = p[1];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(VarForDynamicIndexTest, ArrayIndexLiteralChain) {
+  auto* src = R"(
+fn f() {
+  let p : array<array<i32, 2>, 2> = array<array<i32, 2>, 2>(array<i32, 2>(1, 2), array<i32, 2>(3, 4));
+  let x : i32 = p[0][1];
+}
+)";
+
+  auto* expect = src;
+
+  auto got = Run<VarForDynamicIndex>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint