hlsl: Pointer support

Add `transform::InlinePointerLets` - a Transform that moves all usage of function-scope  `let` statements of a pointer type into their places of usage.

Make the HLSL writer transform pointer parameters to `inout`.

Fixed: tint:183
Change-Id: I0a7552fa6cd31c7b7691e64feae3170a81cc6c49
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/51281
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: David Neto <dneto@google.com>

src/transform/hlsl.cc

@@ -27,6 +27,7 @@
 #include "src/transform/canonicalize_entry_point_io.h"
 #include "src/transform/decompose_storage_access.h"
 #include "src/transform/external_texture_transform.h"
+#include "src/transform/inline_pointer_lets.h"
 #include "src/transform/manager.h"
 
 namespace tint {
@@ -41,6 +42,7 @@ Output Hlsl::Run(const Program* in, const DataMap& data) {
   manager.Add<DecomposeStorageAccess>();
   manager.Add<CalculateArrayLength>();
   manager.Add<ExternalTextureTransform>();
+  manager.Add<InlinePointerLets>();
   auto out = manager.Run(in, data);
   if (!out.program.IsValid()) {
     return out;

src/transform/inline_pointer_lets.cc

@@ -0,0 +1,191 @@
+// 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/inline_pointer_lets.h"
+
+#include <memory>
+#include <unordered_map>
+#include <utility>
+
+#include "src/program_builder.h"
+#include "src/sem/block_statement.h"
+#include "src/sem/function.h"
+#include "src/sem/statement.h"
+#include "src/sem/variable.h"
+#include "src/utils/scoped_assignment.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+/// Traverses the expression `expr` looking for non-literal array indexing
+/// expressions that would affect the computed address of a pointer expression.
+/// The function-like argument `cb` is called for each found.
+/// @param program the program that owns all the expression nodes
+/// @param expr the expression to traverse
+/// @param cb a function-like object with the signature
+/// `void(const ast::Expression*)`, which is called for each array index
+/// expression
+template <typename F>
+void CollectSavedArrayIndices(const Program* program,
+                              ast::Expression* expr,
+                              F&& cb) {
+  if (auto* a = expr->As<ast::ArrayAccessorExpression>()) {
+    CollectSavedArrayIndices(program, a->array(), cb);
+
+    if (!a->idx_expr()->Is<ast::ScalarConstructorExpression>()) {
+      cb(a->idx_expr());
+    }
+    return;
+  }
+
+  if (auto* m = expr->As<ast::MemberAccessorExpression>()) {
+    CollectSavedArrayIndices(program, m->structure(), cb);
+    return;
+  }
+
+  if (auto* u = expr->As<ast::UnaryOpExpression>()) {
+    CollectSavedArrayIndices(program, u->expr(), cb);
+    return;
+  }
+
+  // Note: Other ast::Expression types can be safely ignored as they cannot be
+  // used to generate a reference or pointer.
+  // See https://gpuweb.github.io/gpuweb/wgsl/#forming-references-and-pointers
+}
+
+// PtrLet represents a `let` declaration of a pointer type.
+struct PtrLet {
+  // A map of ptr-let initializer sub-expression to the name of generated
+  // variable that holds the saved value of this sub-expression, when resolved
+  // at the point of the ptr-let declaration.
+  std::unordered_map<const ast::Expression*, Symbol> saved_vars;
+};
+
+}  // namespace
+
+InlinePointerLets::InlinePointerLets() = default;
+
+InlinePointerLets::~InlinePointerLets() = default;
+
+Output InlinePointerLets::Run(const Program* in, const DataMap&) {
+  ProgramBuilder out;
+  CloneContext ctx(&out, in);
+
+  // If not null, current_ptr_let is the current PtrLet being operated on.
+  PtrLet* current_ptr_let = nullptr;
+  // A map of the AST `let` variable to the PtrLet
+  std::unordered_map<const ast::Variable*, std::unique_ptr<PtrLet>> ptr_lets;
+
+  // Register the ast::Expression transform handler.
+  // This performs two different transformations:
+  // * Identifiers that resolve to the pointer-typed `let` declarations are
+  // replaced with the inlined (and recursively transformed) initializer
+  // expression for the `let` declaration.
+  // * Sub-expressions inside the pointer-typed `let` initializer expression
+  // that have been hoisted to a saved variable are replaced with the saved
+  // variable identifier.
+  ctx.ReplaceAll([&](ast::Expression* expr) -> ast::Expression* {
+    if (current_ptr_let) {
+      // We're currently processing the initializer expression of a
+      // pointer-typed `let` declaration. Look to see if we need to swap this
+      // Expression with a saved variable.
+      auto it = current_ptr_let->saved_vars.find(expr);
+      if (it != current_ptr_let->saved_vars.end()) {
+        return ctx.dst->Expr(it->second);
+      }
+    }
+    if (auto* ident = expr->As<ast::IdentifierExpression>()) {
+      if (auto* vu = in->Sem().Get<sem::VariableUser>(ident)) {
+        auto* var = vu->Variable()->Declaration();
+        auto it = ptr_lets.find(var);
+        if (it != ptr_lets.end()) {
+          // We've found an identifier that resolves to a `let` declaration.
+          // We need to replace this identifier with the initializer expression
+          // of the `let` declaration. Clone the initializer expression to make
+          // a copy. Note that this will call back into this ReplaceAll()
+          // handler for sub-expressions of the initializer.
+          auto* ptr_let = it->second.get();
+          // TINT_SCOPED_ASSIGNMENT provides a stack of PtrLet*, this is
+          // required to handle the 'chaining' of inlined `let`s.
+          TINT_SCOPED_ASSIGNMENT(current_ptr_let, ptr_let);
+          return ctx.Clone(var->constructor());
+        }
+      }
+    }
+    return nullptr;
+  });
+
+  // Find all the pointer-typed `let` declarations.
+  // Note that these must be function-scoped, as module-scoped `let`s are not
+  // permitted.
+  for (auto* node : in->ASTNodes().Objects()) {
+    if (auto* let = node->As<ast::VariableDeclStatement>()) {
+      if (!let->variable()->is_const()) {
+        continue;  // Not a `let` declaration. Ignore.
+      }
+
+      auto* var = in->Sem().Get(let->variable());
+      if (!var->Type()->Is<sem::Pointer>()) {
+        continue;  // Not a pointer type. Ignore.
+      }
+
+      // We're dealing with a pointer-typed `let` declaration.
+      auto ptr_let = std::make_unique<PtrLet>();
+      TINT_SCOPED_ASSIGNMENT(current_ptr_let, ptr_let.get());
+
+      auto* block = ctx.src->Sem().Get(let)->Block()->Declaration();
+
+      // Scan the initializer expression for array index expressions that need
+      // to be hoist to temporary "saved" variables.
+      CollectSavedArrayIndices(
+          ctx.src, var->Declaration()->constructor(),
+          [&](ast::Expression* idx_expr) {
+            // We have a sub-expression that needs to be saved.
+            // Create a new variable
+            auto saved_name = ctx.dst->Symbols().New(
+                ctx.src->Symbols().NameFor(var->Declaration()->symbol()) +
+                "_save");
+            auto* saved = ctx.dst->Decl(
+                ctx.dst->Const(saved_name, nullptr, ctx.Clone(idx_expr)));
+            // Place this variable after the pointer typed let. Order here is
+            // important as order-of-operations needs to be preserved.
+            // CollectSavedArrayIndices() visits the LHS of an array accessor
+            // before the index expression.
+            // Note that repeated calls to InsertAfter() with the same `after`
+            // argument will result in nodes to inserted in the order the calls
+            // are made (last call is inserted last).
+            ctx.InsertAfter(block->statements(), let, saved);
+            // Record the substitution of `idx_expr` to the saved variable with
+            // the symbol `saved_name`. This will be used by the ReplaceAll()
+            // handler above.
+            ptr_let->saved_vars.emplace(idx_expr, saved_name);
+          });
+
+      // Record the pointer-typed `let` declaration.
+      // This will be used by the ReplaceAll() handler above.
+      ptr_lets.emplace(let->variable(), std::move(ptr_let));
+      // As the original `let` declaration will be fully inlined, there's no
+      // need for the original declaration to exist. Remove it.
+      ctx.Remove(block->statements(), let);
+    }
+  }
+
+  ctx.Clone();
+
+  return Output(Program(std::move(out)));
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/inline_pointer_lets.h

@@ -0,0 +1,52 @@
+// 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_INLINE_POINTER_LETS_H_
+#define SRC_TRANSFORM_INLINE_POINTER_LETS_H_
+
+#include <string>
+#include <unordered_map>
+
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// InlinePointerLets is a Transform that moves all usage of function-scope
+/// `let` statements of a pointer type into their places of usage.
+///
+/// Parameters of a pointer type are not adjusted.
+///
+/// Note: InlinePointerLets does not operate on module-scope `let`s, as these
+/// cannot be pointers: https://gpuweb.github.io/gpuweb/wgsl/#module-constants
+/// `A module-scope let-declared constant must be of atomic-free plain type.`
+class InlinePointerLets : public Transform {
+ public:
+  /// Constructor
+  InlinePointerLets();
+
+  /// Destructor
+  ~InlinePointerLets() 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_INLINE_POINTER_LETS_H_

src/transform/inline_pointer_lets_test.cc

@@ -0,0 +1,323 @@
+// 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/inline_pointer_lets.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using InlinePointerLetsTest = TransformTest;
+
+TEST_F(InlinePointerLetsTest, EmptyModule) {
+  auto* src = "";
+  auto* expect = "";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(InlinePointerLetsTest, Basic) {
+  auto* src = R"(
+fn f() {
+  var v : i32;
+  let p : ptr<function, i32> = &v;
+  let x : i32 = *p;
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var v : i32;
+  let x : i32 = *(&(v));
+}
+)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(InlinePointerLetsTest, ComplexChain) {
+  auto* src = R"(
+fn f() {
+  var m : mat4x4<f32>;
+  let mp : ptr<function, mat4x4<f32>> = &m;
+  let vp : ptr<function, vec4<f32>> = &(*mp)[2];
+  let fp : ptr<function, f32> = &(*vp)[1];
+  let f : f32 = *fp;
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var m : mat4x4<f32>;
+  let f : f32 = *(&(*(&(*(&(m))[2]))[1]));
+}
+)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(InlinePointerLetsTest, Param) {
+  auto* src = R"(
+fn x(p : ptr<function, i32>) -> i32 {
+  return *p;
+}
+
+fn f() {
+  var v : i32;
+  let p : ptr<function, i32> = &v;
+  var r : i32 = x(p);
+}
+)";
+
+  auto* expect = R"(
+fn x(p : ptr<function, i32>) -> i32 {
+  return *(p);
+}
+
+fn f() {
+  var v : i32;
+  var r : i32 = x(&(v));
+}
+)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(InlinePointerLetsTest, SavedVars) {
+  auto* src = R"(
+struct S {
+  i : i32;
+};
+
+fn arr() {
+  var a : array<S, 2>;
+  var i : i32 = 0;
+  var j : i32 = 0;
+  let p : ptr<function, i32> = &a[i + j].i;
+  i = 2;
+  *p = 4;
+}
+
+fn vec() {
+  var v : vec3<f32>;
+  var i : i32 = 0;
+  var j : i32 = 0;
+  let p : ptr<function, f32> = &v[i + j];
+  i = 2;
+  *p = 4.0;
+}
+
+fn mat() {
+  var m : mat3x3<f32>;
+  var i : i32 = 0;
+  var j : i32 = 0;
+  let p : ptr<function, vec3<f32>> = &m[i + j];
+  i = 2;
+  *p = vec3<f32>(4.0, 5.0, 6.0);
+}
+)";
+
+  auto* expect = R"(
+struct S {
+  i : i32;
+};
+
+fn arr() {
+  var a : array<S, 2>;
+  var i : i32 = 0;
+  var j : i32 = 0;
+  let p_save = (i + j);
+  i = 2;
+  *(&(a[p_save].i)) = 4;
+}
+
+fn vec() {
+  var v : vec3<f32>;
+  var i : i32 = 0;
+  var j : i32 = 0;
+  let p_save_1 = (i + j);
+  i = 2;
+  *(&(v[p_save_1])) = 4.0;
+}
+
+fn mat() {
+  var m : mat3x3<f32>;
+  var i : i32 = 0;
+  var j : i32 = 0;
+  let p_save_2 = (i + j);
+  i = 2;
+  *(&(m[p_save_2])) = vec3<f32>(4.0, 5.0, 6.0);
+}
+)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(InlinePointerLetsTest, DontSaveLiterals) {
+  auto* src = R"(
+fn f() {
+  var arr : array<i32, 2>;
+  let p1 : ptr<function, i32> = &arr[1];
+  *p1 = 4;
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var arr : array<i32, 2>;
+  *(&(arr[1])) = 4;
+}
+)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(InlinePointerLetsTest, SavedVarsChain) {
+  auto* src = R"(
+fn f() {
+  var arr : array<array<i32, 2>, 2>;
+  let i : i32 = 0;
+  let j : i32 = 1;
+  let p : ptr<function, array<i32, 2>> = &arr[i];
+  let q : ptr<function, i32> = &(*p)[j];
+  *q = 12;
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var arr : array<array<i32, 2>, 2>;
+  let i : i32 = 0;
+  let j : i32 = 1;
+  let p_save = i;
+  let q_save = j;
+  *(&(*(&(arr[p_save]))[q_save])) = 12;
+}
+)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(InlinePointerLetsTest, MultiSavedVarsInSinglePtrLetExpr) {
+  auto* src = R"(
+fn x() -> i32 {
+  return 1;
+}
+
+fn y() -> i32 {
+  return 1;
+}
+
+fn z() -> i32 {
+  return 1;
+}
+
+struct Inner {
+  a : array<i32, 2>;
+};
+
+struct Outer {
+  a : array<Inner, 2>;
+};
+
+fn f() {
+  var arr : array<Outer, 2>;
+  let p : ptr<function, i32> = &arr[x()].a[y()].a[z()];
+  *p = 1;
+  *p = 2;
+}
+)";
+
+  auto* expect = R"(
+fn x() -> i32 {
+  return 1;
+}
+
+fn y() -> i32 {
+  return 1;
+}
+
+fn z() -> i32 {
+  return 1;
+}
+
+struct Inner {
+  a : array<i32, 2>;
+};
+
+struct Outer {
+  a : array<Inner, 2>;
+};
+
+fn f() {
+  var arr : array<Outer, 2>;
+  let p_save = x();
+  let p_save_1 = y();
+  let p_save_2 = z();
+  *(&(arr[p_save].a[p_save_1].a[p_save_2])) = 1;
+  *(&(arr[p_save].a[p_save_1].a[p_save_2])) = 2;
+}
+)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+// TODO(crbug.com/tint/819): Enable when we support inter-scope shadowing.
+TEST_F(InlinePointerLetsTest, DISABLED_ModificationAfterInline) {
+  auto* src = R"(
+fn x(p : ptr<function, i32>) -> i32 {
+  return *p;
+}
+
+fn f() {
+  var i : i32 = 1;
+  let p : ptr<function, i32> = &i;
+  if (true) {
+    var i : i32 = 2;
+    x(p);
+  }
+}
+)";
+
+  auto* expect = R"(<TODO>)";
+
+  auto got = Run<InlinePointerLets>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+}  // namespace
+}  // namespace transform
+}  // namespace tint