Add DuplicateStorageStruct transform

This transform is currently required by Dawn for it's GLSL backend so
that SPIRV-Cross does not end up renaming the structures internally,
which it does to fix aliasing problems. We can remove this transform
if/once Dawn using binding numbers rather than names for uniform/storage
buffer data.

Bug: tint:386
Bug: tint:808

CloneContext: allow insert before/after of a node marked for removal
Change-Id: I3ff3b37bca62db07d5c759250dd4777e279b7a4f
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/51403
Commit-Queue: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/BUILD.gn

@@ -509,6 +509,8 @@ libtint_source_set("libtint_core_all_src") {
     "transform/canonicalize_entry_point_io.h",
     "transform/decompose_storage_access.cc",
     "transform/decompose_storage_access.h",
+    "transform/duplicate_storage_structs.cc",
+    "transform/duplicate_storage_structs.h",
     "transform/external_texture_transform.cc",
     "transform/external_texture_transform.h",
     "transform/first_index_offset.cc",

src/CMakeLists.txt

@@ -273,6 +273,8 @@ set(TINT_LIB_SRCS
   transform/canonicalize_entry_point_io.h
   transform/decompose_storage_access.cc
   transform/decompose_storage_access.h
+  transform/duplicate_storage_structs.cc
+  transform/duplicate_storage_structs.h
   transform/external_texture_transform.cc
   transform/external_texture_transform.h
   transform/first_index_offset.cc
@@ -800,6 +802,7 @@ if(${TINT_BUILD_TESTS})
       transform/calculate_array_length_test.cc
       transform/canonicalize_entry_point_io_test.cc
       transform/decompose_storage_access_test.cc
+      transform/duplicate_storage_structs_test.cc
       transform/external_texture_transform_test.cc
       transform/first_index_offset_test.cc
       transform/renamer_test.cc

src/transform/duplicate_storage_structs.cc

@@ -0,0 +1,126 @@
+// 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/duplicate_storage_structs.h"
+
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "src/program_builder.h"
+#include "src/sem/variable.h"
+#include "src/utils/get_or_create.h"
+#include "src/utils/hash.h"
+
+namespace {
+struct StructAccessPair {
+  tint::sem::Struct const* str;
+  tint::ast::AccessControl::Access access;
+
+  bool operator==(const StructAccessPair& rhs) const {
+    return str == rhs.str && access == rhs.access;
+  }
+};
+}  // namespace
+
+namespace std {
+/// Custom std::hash specialization for StructAccessPair so
+/// StructAccessPairs can be used as keys for std::unordered_map and
+/// std::unordered_set.
+template <>
+class hash<StructAccessPair> {
+ public:
+  /// @param struct_access_pair the StructAccessPair to create a hash for
+  /// @return the hash value
+  inline std::size_t operator()(
+      const StructAccessPair& struct_access_pair) const {
+    return tint::utils::Hash(struct_access_pair.str, struct_access_pair.access);
+  }
+};
+}  // namespace std
+
+namespace tint {
+namespace transform {
+DuplicateStorageStructs::DuplicateStorageStructs() = default;
+DuplicateStorageStructs::~DuplicateStorageStructs() = default;
+
+Output DuplicateStorageStructs::Run(const Program* in, const DataMap&) {
+  ProgramBuilder out;
+  CloneContext ctx(&out, in);
+
+  // Create struct duplicates of storage buffers per access type.
+  // We do this by finding all storage variables, and creating a copy of the
+  // Struct they point to per struct/access pair.
+
+  std::unordered_map<StructAccessPair, ast::Struct*> sap_to_new_struct;
+
+  for (auto* node : in->ASTNodes().Objects()) {
+    if (auto* ast_var = node->As<ast::Variable>()) {
+      auto* var = in->Sem().Get(ast_var);
+      if (auto* str = var->Type()->UnwrapRef()->As<sem::Struct>()) {
+        // Skip non-storage structs
+        if (!str->UsedAs(ast::StorageClass::kStorage)) {
+          continue;
+        }
+
+        auto sap = StructAccessPair{str, var->AccessControl()};
+
+        auto* new_str = utils::GetOrCreate(sap_to_new_struct, sap, [&] {
+          // For this ast::Struct/Access pair, create a clone of the existing
+          // ast::Struct with a new name.
+          auto* old_str = str->Declaration();
+          std::string old_name = ctx.src->Symbols().NameFor(old_str->name());
+
+          auto* new_struct = [&] {
+            Symbol new_name = ctx.dst->Symbols().New(old_name);
+            auto new_members = ctx.Clone(old_str->members());
+            auto new_decorations = ctx.Clone(old_str->decorations());
+            return ctx.dst->create<ast::Struct>(new_name, new_members,
+                                                new_decorations);
+          }();
+
+          // Insert it after the original struct
+          ctx.InsertAfter(ctx.src->AST().GlobalDeclarations(),
+                          str->Declaration(), new_struct);
+
+          // Remove the original struct
+          ctx.Remove(ctx.src->AST().GlobalDeclarations(), old_str);
+
+          return new_struct;
+        });
+
+        // Replace the existing variable with a new one who's type is an
+        // AccessControl<TypeName<"New Struct">>.
+        auto* new_var = [&] {
+          auto new_name = ctx.Clone(ast_var->symbol());
+          auto* new_ty = ctx.dst->ty.access(
+              var->AccessControl(), ctx.dst->ty.type_name(new_str->name()));
+          auto* new_constructor = ctx.Clone(ast_var->constructor());
+          auto new_decorations = ctx.Clone(ast_var->decorations());
+          return ctx.dst->create<ast::Variable>(
+              new_name, ast_var->declared_storage_class(), new_ty,
+              ast_var->is_const(), new_constructor, new_decorations);
+        }();
+        ctx.Replace(ast_var, new_var);
+      }
+    }
+  }
+
+  ctx.Clone();
+
+  return Output(Program(std::move(out)));
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/duplicate_storage_structs.h

@@ -0,0 +1,54 @@
+// 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_DUPLICATE_STORAGE_STRUCTS_H_
+#define SRC_TRANSFORM_DUPLICATE_STORAGE_STRUCTS_H_
+
+#include <string>
+#include <unordered_map>
+
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// DuplicateStorageStructs is a Transform that creates duplicates of storage
+/// structures, one per access control type.
+///
+/// This transform is currently required by Dawn for it's GLSL backend so that
+/// SPIRV-Cross does not end up renaming the structures internally, which it
+/// does to fix aliasing problems (see
+/// https://github.com/KhronosGroup/SPIRV-Cross/pull/799). We can remove this
+/// transform if/once Dawn using binding numbers rather than names for
+/// uniform/storage buffer data. See https://crbug.com/tint/386.
+///
+class DuplicateStorageStructs : public Transform {
+ public:
+  /// Constructor
+  DuplicateStorageStructs();
+
+  /// Destructor
+  ~DuplicateStorageStructs() override;
+
+  /// Runs the transform on `program`, returning the transformation result.
+  /// @param program the source program to transform
+  /// @param data optional extra transform-specific data
+  /// @returns the transformation result
+  Output Run(const Program* program, const DataMap& data = {}) override;
+};
+
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_DUPLICATE_STORAGE_STRUCTS_H_

src/transform/duplicate_storage_structs_test.cc

@@ -0,0 +1,275 @@
+// 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/duplicate_storage_structs.h"
+
+#include "src/transform/test_helper.h"
+namespace tint {
+namespace transform {
+namespace {
+
+using DuplicateStorageStructsTransformTest = TransformTest;
+
+TEST_F(DuplicateStorageStructsTransformTest, Simple) {
+  auto* src = R"(
+[[block]]
+  struct S {
+  a : f32;
+};
+
+[[group(0), binding(0)]] var<storage> r : [[access(read)]] S;
+[[group(0), binding(1)]] var<storage> w : [[access(write)]] S;
+
+[[stage(compute)]]
+  fn main() {
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct S_1 {
+  a : f32;
+};
+
+[[block]]
+struct S_2 {
+  a : f32;
+};
+
+[[group(0), binding(0)]] var<storage> r : [[access(read)]] S_1;
+
+[[group(0), binding(1)]] var<storage> w : [[access(write)]] S_2;
+
+[[stage(compute)]]
+fn main() {
+}
+)";
+
+  auto got = Run<DuplicateStorageStructs>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DuplicateStorageStructsTransformTest, MultipleStorageStructs) {
+  auto* src = R"(
+[[block]]
+  struct S1 {
+  a : f32;
+};
+
+[[block]]
+  struct S2 {
+  a : i32;
+};
+
+[[group(0), binding(0)]] var<storage> r1 : [[access(read)]] S1;
+[[group(0), binding(1)]] var<storage> w1 : [[access(write)]] S1;
+
+[[group(0), binding(2)]] var<storage> r2 : [[access(read)]] S2;
+[[group(0), binding(3)]] var<storage> w2 : [[access(write)]] S2;
+
+[[stage(compute)]]
+  fn main() {
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct S1_1 {
+  a : f32;
+};
+
+[[block]]
+struct S1_2 {
+  a : f32;
+};
+
+[[block]]
+struct S2_1 {
+  a : i32;
+};
+
+[[block]]
+struct S2_2 {
+  a : i32;
+};
+
+[[group(0), binding(0)]] var<storage> r1 : [[access(read)]] S1_1;
+
+[[group(0), binding(1)]] var<storage> w1 : [[access(write)]] S1_2;
+
+[[group(0), binding(2)]] var<storage> r2 : [[access(read)]] S2_1;
+
+[[group(0), binding(3)]] var<storage> w2 : [[access(write)]] S2_2;
+
+[[stage(compute)]]
+fn main() {
+}
+)";
+
+  auto got = Run<DuplicateStorageStructs>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DuplicateStorageStructsTransformTest, MultipleStorageVariables) {
+  auto* src = R"(
+[[block]]
+  struct S {
+  a : f32;
+};
+
+[[group(0), binding(0)]] var<storage> r1 : [[access(read)]] S;
+[[group(0), binding(1)]] var<storage> r2 : [[access(read)]] S;
+[[group(0), binding(2)]] var<storage> r3 : [[access(read)]] S;
+[[group(0), binding(3)]] var<storage> w1 : [[access(write)]] S;
+[[group(0), binding(4)]] var<storage> w2 : [[access(write)]] S;
+[[group(0), binding(5)]] var<storage> w3 : [[access(write)]] S;
+[[group(0), binding(3)]] var<storage> rw1 : [[access(read_write)]] S;
+[[group(0), binding(4)]] var<storage> rw2 : [[access(read_write)]] S;
+[[group(0), binding(5)]] var<storage> rw3 : [[access(read_write)]] S;
+
+[[stage(compute)]]
+  fn main() {
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct S_1 {
+  a : f32;
+};
+
+[[block]]
+struct S_2 {
+  a : f32;
+};
+
+[[block]]
+struct S_3 {
+  a : f32;
+};
+
+[[group(0), binding(0)]] var<storage> r1 : [[access(read)]] S_1;
+
+[[group(0), binding(1)]] var<storage> r2 : [[access(read)]] S_1;
+
+[[group(0), binding(2)]] var<storage> r3 : [[access(read)]] S_1;
+
+[[group(0), binding(3)]] var<storage> w1 : [[access(write)]] S_2;
+
+[[group(0), binding(4)]] var<storage> w2 : [[access(write)]] S_2;
+
+[[group(0), binding(5)]] var<storage> w3 : [[access(write)]] S_2;
+
+[[group(0), binding(3)]] var<storage> rw1 : [[access(read_write)]] S_3;
+
+[[group(0), binding(4)]] var<storage> rw2 : [[access(read_write)]] S_3;
+
+[[group(0), binding(5)]] var<storage> rw3 : [[access(read_write)]] S_3;
+
+[[stage(compute)]]
+fn main() {
+}
+)";
+
+  auto got = Run<DuplicateStorageStructs>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(DuplicateStorageStructsTransformTest, MixedStructTypes) {
+  auto* src = R"(
+
+[[block]]
+struct S_1 {
+  a : f32;
+};
+
+[[block]]
+struct S_2 {
+  a : f32;
+};
+
+[[block]]
+struct S_3 {
+  a : f32;
+};
+
+[[group(0), binding(0)]] var<storage> r1 : [[access(read)]] S_1;
+
+[[group(0), binding(1)]] var<storage> r2 : [[access(read)]] S_1;
+
+[[group(0), binding(2)]] var<storage> r3 : [[access(read)]] S_1;
+
+[[group(0), binding(3)]] var<storage> w1 : [[access(write)]] S_2;
+
+[[group(0), binding(4)]] var<storage> w2 : [[access(write)]] S_2;
+
+[[group(0), binding(5)]] var<storage> w3 : [[access(write)]] S_2;
+
+[[group(0), binding(3)]] var<storage> rw1 : [[access(read_write)]] S_3;
+
+[[group(0), binding(4)]] var<storage> rw2 : [[access(read_write)]] S_3;
+
+[[group(0), binding(5)]] var<storage> rw3 : [[access(read_write)]] S_3;
+
+[[stage(compute)]]
+fn main() {
+}
+)";
+
+  auto* expect = R"(
+[[block]]
+struct S_1_1 {
+  a : f32;
+};
+
+[[block]]
+struct S_2_1 {
+  a : f32;
+};
+
+[[block]]
+struct S_3_1 {
+  a : f32;
+};
+
+[[group(0), binding(0)]] var<storage> r1 : [[access(read)]] S_1_1;
+
+[[group(0), binding(1)]] var<storage> r2 : [[access(read)]] S_1_1;
+
+[[group(0), binding(2)]] var<storage> r3 : [[access(read)]] S_1_1;
+
+[[group(0), binding(3)]] var<storage> w1 : [[access(write)]] S_2_1;
+
+[[group(0), binding(4)]] var<storage> w2 : [[access(write)]] S_2_1;
+
+[[group(0), binding(5)]] var<storage> w3 : [[access(write)]] S_2_1;
+
+[[group(0), binding(3)]] var<storage> rw1 : [[access(read_write)]] S_3_1;
+
+[[group(0), binding(4)]] var<storage> rw2 : [[access(read_write)]] S_3_1;
+
+[[group(0), binding(5)]] var<storage> rw3 : [[access(read_write)]] S_3_1;
+
+[[stage(compute)]]
+fn main() {
+}
+)";
+
+  auto got = Run<DuplicateStorageStructs>(src);
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint

src/transform/spirv.cc

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

test/BUILD.gn

@@ -257,9 +257,9 @@ tint_unittests_source_set("tint_unittests_core_src") {
     "../src/resolver/intrinsic_test.cc",
     "../src/resolver/is_host_shareable_test.cc",
     "../src/resolver/is_storeable_test.cc",
+    "../src/resolver/pipeline_overridable_constant_test.cc",
     "../src/resolver/ptr_ref_test.cc",
     "../src/resolver/ptr_ref_validation_test.cc",
-    "../src/resolver/pipeline_overridable_constant_test.cc",
     "../src/resolver/resolver_test.cc",
     "../src/resolver/resolver_test_helper.cc",
     "../src/resolver/resolver_test_helper.h",
@@ -299,6 +299,7 @@ tint_unittests_source_set("tint_unittests_core_src") {
     "../src/transform/calculate_array_length_test.cc",
     "../src/transform/canonicalize_entry_point_io_test.cc",
     "../src/transform/decompose_storage_access_test.cc",
+    "../src/transform/duplicate_storage_structs_test.cc",
     "../src/transform/external_texture_transform_test.cc",
     "../src/transform/first_index_offset_test.cc",
     "../src/transform/renamer_test.cc",