spirv: Add block decorations with a transform

Any struct which is used as the store type of a buffer variable needs
to have a block decoration. If that struct is nested inside an array
or another struct, we wrap it inside another struct first.

This removes the SPIR-V backend's reliance on the [[block]] attribute,
which will soon be deprecated and removed.

Bug: tint:1324
Change-Id: Ib6ad54f24a3e4a090da9faeed699f266abcb66ba
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/72082
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/transform/add_spirv_block_decoration.cc

@@ -0,0 +1,130 @@
+// 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/add_spirv_block_decoration.h"
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
+#include "src/program_builder.h"
+#include "src/sem/variable.h"
+#include "src/utils/map.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::transform::AddSpirvBlockDecoration);
+TINT_INSTANTIATE_TYPEINFO(
+    tint::transform::AddSpirvBlockDecoration::SpirvBlockDecoration);
+
+namespace tint {
+namespace transform {
+
+AddSpirvBlockDecoration::AddSpirvBlockDecoration() = default;
+
+AddSpirvBlockDecoration::~AddSpirvBlockDecoration() = default;
+
+void AddSpirvBlockDecoration::Run(CloneContext& ctx, const DataMap&, DataMap&) {
+  auto& sem = ctx.src->Sem();
+
+  // Collect the set of structs that are nested in other types.
+  std::unordered_set<const sem::Struct*> nested_structs;
+  for (auto* node : ctx.src->ASTNodes().Objects()) {
+    if (auto* arr = sem.Get<sem::Array>(node->As<ast::Array>())) {
+      if (auto* nested_str = arr->ElemType()->As<sem::Struct>()) {
+        nested_structs.insert(nested_str);
+      }
+    } else if (auto* str = sem.Get<sem::Struct>(node->As<ast::Struct>())) {
+      for (auto* member : str->Members()) {
+        if (auto* nested_str = member->Type()->As<sem::Struct>()) {
+          nested_structs.insert(nested_str);
+        }
+      }
+    }
+  }
+
+  // A map from a struct in the source program to a block-decorated wrapper that
+  // contains it in the destination program.
+  std::unordered_map<const sem::Struct*, const ast::Struct*> wrapper_structs;
+
+  // Process global variables that are buffers.
+  for (auto* var : ctx.src->AST().GlobalVariables()) {
+    auto* sem_var = sem.Get<sem::GlobalVariable>(var);
+    if (var->declared_storage_class != ast::StorageClass::kStorage &&
+        var->declared_storage_class != ast::StorageClass::kUniform) {
+      continue;
+    }
+
+    auto* str = sem.Get<sem::Struct>(var->type);
+    if (!str) {
+      // TODO(jrprice): We'll need to wrap these too, when WGSL supports this.
+      TINT_ICE(Transform, ctx.dst->Diagnostics())
+          << "non-struct buffer types are not yet supported";
+      continue;
+    }
+
+    if (nested_structs.count(str)) {
+      const char* kInnerStructMemberName = "inner";
+
+      // This struct is nested somewhere else, so we need to wrap it first.
+      auto* wrapper = utils::GetOrCreate(wrapper_structs, str, [&]() {
+        auto* block =
+            ctx.dst->ASTNodes().Create<SpirvBlockDecoration>(ctx.dst->ID());
+        auto wrapper_name =
+            ctx.src->Symbols().NameFor(str->Declaration()->name) + "_block";
+        auto* ret = ctx.dst->create<ast::Struct>(
+            ctx.dst->Symbols().New(wrapper_name),
+            ast::StructMemberList{ctx.dst->Member(kInnerStructMemberName,
+                                                  CreateASTTypeFor(ctx, str))},
+            ast::DecorationList{block});
+        ctx.InsertAfter(ctx.src->AST().GlobalDeclarations(), str->Declaration(),
+                        ret);
+        return ret;
+      });
+      ctx.Replace(var->type, ctx.dst->ty.Of(wrapper));
+
+      // Insert a member accessor to get the original struct from the wrapper at
+      // any usage of the original variable.
+      for (auto* user : sem_var->Users()) {
+        ctx.Replace(user->Declaration(),
+                    ctx.dst->MemberAccessor(ctx.Clone(var->symbol),
+                                            kInnerStructMemberName));
+      }
+    } else {
+      // Add a block decoration to this struct directly.
+      auto* block =
+          ctx.dst->ASTNodes().Create<SpirvBlockDecoration>(ctx.dst->ID());
+      ctx.InsertFront(str->Declaration()->decorations, block);
+    }
+  }
+
+  ctx.Clone();
+}
+
+AddSpirvBlockDecoration::SpirvBlockDecoration::SpirvBlockDecoration(
+    ProgramID pid)
+    : Base(pid) {}
+AddSpirvBlockDecoration::SpirvBlockDecoration::~SpirvBlockDecoration() =
+    default;
+std::string AddSpirvBlockDecoration::SpirvBlockDecoration::InternalName()
+    const {
+  return "spirv_block";
+}
+
+const AddSpirvBlockDecoration::SpirvBlockDecoration*
+AddSpirvBlockDecoration::SpirvBlockDecoration::Clone(CloneContext* ctx) const {
+  return ctx->dst->ASTNodes()
+      .Create<AddSpirvBlockDecoration::SpirvBlockDecoration>(ctx->dst->ID());
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/add_spirv_block_decoration.h

@@ -0,0 +1,74 @@
+// 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_ADD_SPIRV_BLOCK_DECORATION_H_
+#define SRC_TRANSFORM_ADD_SPIRV_BLOCK_DECORATION_H_
+
+#include <string>
+
+#include "src/ast/internal_decoration.h"
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// AddSpirvBlockDecoration is a transform that adds an
+/// [[internal(spirv_block)]] attribute to any structure that is used as the
+/// store type of a buffer. If that structure is nested inside another structure
+/// or an array, then it is wrapped inside another structure which gets the
+/// [[internal(spirv_block)]] attribute instead.
+class AddSpirvBlockDecoration
+    : public Castable<AddSpirvBlockDecoration, Transform> {
+ public:
+  /// SpirvBlockDecoration is an InternalDecoration that is used to decorate a
+  // structure that needs a SPIR-V block decoration.
+  class SpirvBlockDecoration
+      : public Castable<SpirvBlockDecoration, ast::InternalDecoration> {
+   public:
+    /// Constructor
+    /// @param program_id the identifier of the program that owns this node
+    explicit SpirvBlockDecoration(ProgramID program_id);
+    /// Destructor
+    ~SpirvBlockDecoration() override;
+
+    /// @return a short description of the internal decoration which will be
+    /// displayed as `[[internal(<name>)]]`
+    std::string InternalName() const override;
+
+    /// Performs a deep clone of this object using the CloneContext `ctx`.
+    /// @param ctx the clone context
+    /// @return the newly cloned object
+    const SpirvBlockDecoration* Clone(CloneContext* ctx) const override;
+  };
+
+  /// Constructor
+  AddSpirvBlockDecoration();
+
+  /// Destructor
+  ~AddSpirvBlockDecoration() 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_ADD_SPIRV_BLOCK_DECORATION_H_

src/transform/add_spirv_block_decoration_test.cc

@@ -0,0 +1,462 @@
+// 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/add_spirv_block_decoration.h"
+
+#include <memory>
+#include <utility>
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using AddSpirvBlockDecorationTest = TransformTest;
+
+TEST_F(AddSpirvBlockDecorationTest, EmptyModule) {
+  auto* src = "";
+  auto* expect = "";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Noop_UsedForPrivateVar) {
+  auto* src = R"(
+struct S {
+  f : f32;
+};
+
+var<private> p : S;
+
+[[stage(fragment)]]
+fn main() {
+  p.f = 1.0;
+}
+)";
+  auto* expect = src;
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Noop_UsedForShaderIO) {
+  auto* src = R"(
+struct S {
+  [[location(0)]]
+  f : f32;
+};
+
+[[stage(fragment)]]
+fn main() -> S {
+  return S();
+}
+)";
+  auto* expect = src;
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Basic) {
+  auto* src = R"(
+struct S {
+  f : f32;
+};
+
+[[group(0), binding(0)]]
+var<uniform> u : S;
+
+[[stage(fragment)]]
+fn main() {
+  let f = u.f;
+}
+)";
+  auto* expect = R"(
+[[internal(spirv_block)]]
+struct S {
+  f : f32;
+};
+
+[[group(0), binding(0)]] var<uniform> u : S;
+
+[[stage(fragment)]]
+fn main() {
+  let f = u.f;
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Nested_OuterBuffer_InnerNotBuffer) {
+  auto* src = R"(
+struct Inner {
+  f : f32;
+};
+
+struct Outer {
+  i : Inner;
+};
+
+[[group(0), binding(0)]]
+var<uniform> u : Outer;
+
+[[stage(fragment)]]
+fn main() {
+  let f = u.i.f;
+}
+)";
+  auto* expect = R"(
+struct Inner {
+  f : f32;
+};
+
+[[internal(spirv_block)]]
+struct Outer {
+  i : Inner;
+};
+
+[[group(0), binding(0)]] var<uniform> u : Outer;
+
+[[stage(fragment)]]
+fn main() {
+  let f = u.i.f;
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Nested_OuterBuffer_InnerBuffer) {
+  auto* src = R"(
+struct Inner {
+  f : f32;
+};
+
+struct Outer {
+  i : Inner;
+};
+
+[[group(0), binding(0)]]
+var<uniform> u0 : Outer;
+
+[[group(0), binding(1)]]
+var<uniform> u1 : Inner;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.i.f;
+  let f1 = u1.f;
+}
+)";
+  auto* expect = R"(
+struct Inner {
+  f : f32;
+};
+
+[[internal(spirv_block)]]
+struct Inner_block {
+  inner : Inner;
+};
+
+[[internal(spirv_block)]]
+struct Outer {
+  i : Inner;
+};
+
+[[group(0), binding(0)]] var<uniform> u0 : Outer;
+
+[[group(0), binding(1)]] var<uniform> u1 : Inner_block;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.i.f;
+  let f1 = u1.inner.f;
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Nested_OuterNotBuffer_InnerBuffer) {
+  auto* src = R"(
+struct Inner {
+  f : f32;
+};
+
+struct Outer {
+  i : Inner;
+};
+
+var<private> p : Outer;
+
+[[group(0), binding(1)]]
+var<uniform> u : Inner;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = p.i.f;
+  let f1 = u.f;
+}
+)";
+  auto* expect = R"(
+struct Inner {
+  f : f32;
+};
+
+[[internal(spirv_block)]]
+struct Inner_block {
+  inner : Inner;
+};
+
+struct Outer {
+  i : Inner;
+};
+
+var<private> p : Outer;
+
+[[group(0), binding(1)]] var<uniform> u : Inner_block;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = p.i.f;
+  let f1 = u.inner.f;
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Nested_InnerUsedForMultipleBuffers) {
+  auto* src = R"(
+struct Inner {
+  f : f32;
+};
+
+struct S {
+  i : Inner;
+};
+
+[[group(0), binding(0)]]
+var<uniform> u0 : S;
+
+[[group(0), binding(1)]]
+var<uniform> u1 : Inner;
+
+[[group(0), binding(2)]]
+var<uniform> u2 : Inner;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.i.f;
+  let f1 = u1.f;
+  let f2 = u2.f;
+}
+)";
+  auto* expect = R"(
+struct Inner {
+  f : f32;
+};
+
+[[internal(spirv_block)]]
+struct Inner_block {
+  inner : Inner;
+};
+
+[[internal(spirv_block)]]
+struct S {
+  i : Inner;
+};
+
+[[group(0), binding(0)]] var<uniform> u0 : S;
+
+[[group(0), binding(1)]] var<uniform> u1 : Inner_block;
+
+[[group(0), binding(2)]] var<uniform> u2 : Inner_block;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.i.f;
+  let f1 = u1.inner.f;
+  let f2 = u2.inner.f;
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, StructInArray) {
+  auto* src = R"(
+struct S {
+  f : f32;
+};
+
+[[group(0), binding(0)]]
+var<uniform> u : S;
+
+[[stage(fragment)]]
+fn main() {
+  let f = u.f;
+  let a = array<S, 4>();
+}
+)";
+  auto* expect = R"(
+struct S {
+  f : f32;
+};
+
+[[internal(spirv_block)]]
+struct S_block {
+  inner : S;
+};
+
+[[group(0), binding(0)]] var<uniform> u : S_block;
+
+[[stage(fragment)]]
+fn main() {
+  let f = u.inner.f;
+  let a = array<S, 4>();
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, StructInArray_MultipleBuffers) {
+  auto* src = R"(
+struct S {
+  f : f32;
+};
+
+[[group(0), binding(0)]]
+var<uniform> u0 : S;
+
+[[group(0), binding(1)]]
+var<uniform> u1 : S;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.f;
+  let f1 = u1.f;
+  let a = array<S, 4>();
+}
+)";
+  auto* expect = R"(
+struct S {
+  f : f32;
+};
+
+[[internal(spirv_block)]]
+struct S_block {
+  inner : S;
+};
+
+[[group(0), binding(0)]] var<uniform> u0 : S_block;
+
+[[group(0), binding(1)]] var<uniform> u1 : S_block;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.inner.f;
+  let f1 = u1.inner.f;
+  let a = array<S, 4>();
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(AddSpirvBlockDecorationTest, Aliases_Nested_OuterBuffer_InnerBuffer) {
+  auto* src = R"(
+struct Inner {
+  f : f32;
+};
+
+type MyInner = Inner;
+
+struct Outer {
+  i : MyInner;
+};
+
+type MyOuter = Outer;
+
+[[group(0), binding(0)]]
+var<uniform> u0 : MyOuter;
+
+[[group(0), binding(1)]]
+var<uniform> u1 : MyInner;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.i.f;
+  let f1 = u1.f;
+}
+)";
+  auto* expect = R"(
+struct Inner {
+  f : f32;
+};
+
+[[internal(spirv_block)]]
+struct Inner_block {
+  inner : Inner;
+};
+
+type MyInner = Inner;
+
+[[internal(spirv_block)]]
+struct Outer {
+  i : MyInner;
+};
+
+type MyOuter = Outer;
+
+[[group(0), binding(0)]] var<uniform> u0 : MyOuter;
+
+[[group(0), binding(1)]] var<uniform> u1 : Inner_block;
+
+[[stage(fragment)]]
+fn main() {
+  let f0 = u0.i.f;
+  let f1 = u1.inner.f;
+}
+)";
+
+  auto got = Run<AddSpirvBlockDecoration>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint