[resolver] Track pipeline stage usages for structs

This will be used by the generators to determine how to handle
location decorations.

Change-Id: Ib0e0ce852a5da3819781b402c5625a440c4c9544
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/46400
Auto-Submit: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>

src/CMakeLists.txt

@@ -483,6 +483,7 @@ if(${TINT_BUILD_TESTS})
     resolver/resolver_test_helper.h
     resolver/resolver_test.cc
     resolver/struct_layout_test.cc
+    resolver/struct_pipeline_stage_use_test.cc
     resolver/struct_storage_class_use_test.cc
     resolver/type_constructor_validation_test.cc
     resolver/type_validation_test.cc

src/resolver/resolver.cc

@@ -293,6 +293,54 @@ bool Resolver::Function(ast::Function* func) {
   variable_stack_.push_scope();
   for (auto* param : func->params()) {
     variable_stack_.set(param->symbol(), CreateVariableInfo(param));
+
+    if (auto* str =
+            param->declared_type()->UnwrapAliasIfNeeded()->As<type::Struct>()) {
+      auto* info = Structure(str);
+      if (!info) {
+        return false;
+      }
+      switch (func->pipeline_stage()) {
+        case ast::PipelineStage::kVertex:
+          info->pipeline_stage_uses.emplace(
+              semantic::PipelineStageUsage::kVertexInput);
+          break;
+        case ast::PipelineStage::kFragment:
+          info->pipeline_stage_uses.emplace(
+              semantic::PipelineStageUsage::kFragmentInput);
+          break;
+        case ast::PipelineStage::kCompute:
+          info->pipeline_stage_uses.emplace(
+              semantic::PipelineStageUsage::kComputeInput);
+          break;
+        case ast::PipelineStage::kNone:
+          break;
+      }
+    }
+  }
+
+  if (auto* str =
+          func->return_type()->UnwrapAliasIfNeeded()->As<type::Struct>()) {
+    auto* info = Structure(str);
+    if (!info) {
+      return false;
+    }
+    switch (func->pipeline_stage()) {
+      case ast::PipelineStage::kVertex:
+        info->pipeline_stage_uses.emplace(
+            semantic::PipelineStageUsage::kVertexOutput);
+        break;
+      case ast::PipelineStage::kFragment:
+        info->pipeline_stage_uses.emplace(
+            semantic::PipelineStageUsage::kFragmentOutput);
+        break;
+      case ast::PipelineStage::kCompute:
+        info->pipeline_stage_uses.emplace(
+            semantic::PipelineStageUsage::kComputeOutput);
+        break;
+      case ast::PipelineStage::kNone:
+        break;
+    }
   }
 
   if (!BlockStatement(func->body())) {
@@ -1359,7 +1407,8 @@ void Resolver::CreateSemanticNodes() const {
     builder_->Sem().Add(
         str, builder_->create<semantic::Struct>(
                  str, std::move(info->members), info->align, info->size,
-                 info->size_no_padding, info->storage_class_usage));
+                 info->size_no_padding, info->storage_class_usage,
+                 info->pipeline_stage_uses));
   }
 }
 

src/resolver/resolver.h

@@ -24,6 +24,7 @@
 #include "src/intrinsic_table.h"
 #include "src/program_builder.h"
 #include "src/scope_stack.h"
+#include "src/semantic/struct.h"
 #include "src/utils/unique_vector.h"
 
 namespace tint {
@@ -139,6 +140,7 @@ class Resolver {
     uint32_t size = 0;
     uint32_t size_no_padding = 0;
     std::unordered_set<ast::StorageClass> storage_class_usage;
+    std::unordered_set<semantic::PipelineStageUsage> pipeline_stage_uses;
   };
 
   /// Structure holding semantic information about a block (i.e. scope), such as

src/resolver/struct_pipeline_stage_use_test.cc

@@ -0,0 +1,204 @@
+// 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/resolver/resolver.h"
+
+#include "gmock/gmock.h"
+#include "src/ast/stage_decoration.h"
+#include "src/resolver/resolver_test_helper.h"
+#include "src/semantic/struct.h"
+
+using ::testing::UnorderedElementsAre;
+
+namespace tint {
+namespace resolver {
+namespace {
+
+using ResolverPipelineStageUseTest = ResolverTest;
+
+TEST_F(ResolverPipelineStageUseTest, UnusedStruct) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_TRUE(sem->PipelineStageUses().empty());
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsNonEntryPointParam) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("foo", {Const("param", s)}, ty.void_(), {}, {});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_TRUE(sem->PipelineStageUses().empty());
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsNonEntryPointReturnType) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("foo", {}, s, {Return(Construct(s, Expr(0.f)))}, {});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_TRUE(sem->PipelineStageUses().empty());
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsVertexShaderParam) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kVertex)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(sem->PipelineStageUses(),
+              UnorderedElementsAre(semantic::PipelineStageUsage::kVertexInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsVertexShaderReturnType) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {}, s, {Return(Construct(s, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kVertex)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kVertexOutput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsFragmentShaderParam) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsFragmentShaderReturnType) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("main", {}, s, {Return(Construct(s, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentOutput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsComputeShaderParam) {
+  auto* s = Structure("S", {Member("a", ty.u32(),
+                                   {create<ast::BuiltinDecoration>(
+                                       ast::Builtin::kLocalInvocationIndex)})});
+
+  Func("main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kCompute)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kComputeInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedMultipleStages) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+
+  Func("vert_main", {Const("param", s)}, s, {Return(Construct(s, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kVertex)});
+
+  Func("frag_main", {Const("param", s)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kVertexInput,
+                           semantic::PipelineStageUsage::kVertexOutput,
+                           semantic::PipelineStageUsage::kFragmentInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsShaderParamViaAlias) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+  auto* s_alias = ty.alias("S_alias", s);
+
+  Func("main", {Const("param", s_alias)}, ty.void_(), {},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentInput));
+}
+
+TEST_F(ResolverPipelineStageUseTest, StructUsedAsShaderReturnTypeViaAlias) {
+  auto* s = Structure(
+      "S", {Member("a", ty.f32(), {create<ast::LocationDecoration>(0)})});
+  auto* s_alias = ty.alias("S_alias", s);
+
+  Func("main", {}, s_alias, {Return(Construct(s_alias, Expr(0.f)))},
+       {create<ast::StageDecoration>(ast::PipelineStage::kFragment)});
+
+  ASSERT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* sem = Sem().Get(s);
+  ASSERT_NE(sem, nullptr);
+  EXPECT_THAT(
+      sem->PipelineStageUses(),
+      UnorderedElementsAre(semantic::PipelineStageUsage::kFragmentOutput));
+}
+
+}  // namespace
+}  // namespace resolver
+}  // namespace tint

src/semantic/sem_struct.cc

@@ -25,13 +25,15 @@ Struct::Struct(type::Struct* type,
                uint32_t align,
                uint32_t size,
                uint32_t size_no_padding,
-               std::unordered_set<ast::StorageClass> storage_class_usage)
+               std::unordered_set<ast::StorageClass> storage_class_usage,
+               std::unordered_set<PipelineStageUsage> pipeline_stage_uses)
     : type_(type),
       members_(std::move(members)),
       align_(align),
       size_(size),
       size_no_padding_(size_no_padding),
-      storage_class_usage_(std::move(storage_class_usage)) {}
+      storage_class_usage_(std::move(storage_class_usage)),
+      pipeline_stage_uses_(std::move(pipeline_stage_uses)) {}
 
 Struct::~Struct() = default;
 

src/semantic/struct.h

@@ -40,6 +40,16 @@ class StructMember;
 /// A vector of StructMember pointers.
 using StructMemberList = std::vector<StructMember*>;
 
+/// Metadata to capture how a structure is used in a shader module.
+enum class PipelineStageUsage {
+  kVertexInput,
+  kVertexOutput,
+  kFragmentInput,
+  kFragmentOutput,
+  kComputeInput,
+  kComputeOutput,
+};
+
 /// Struct holds the semantic information for structures.
 class Struct : public Castable<Struct, Node> {
  public:
@@ -51,12 +61,14 @@ class Struct : public Castable<Struct, Node> {
   /// @param size_no_padding size of the members without the end of structure
   /// alignment padding
   /// @param storage_class_usage a set of all the storage class usages
+  /// @param pipeline_stage_uses a set of all the pipeline stage uses
   Struct(type::Struct* type,
          StructMemberList members,
          uint32_t align,
          uint32_t size,
          uint32_t size_no_padding,
-         std::unordered_set<ast::StorageClass> storage_class_usage);
+         std::unordered_set<ast::StorageClass> storage_class_usage,
+         std::unordered_set<PipelineStageUsage> pipeline_stage_uses);
 
   /// Destructor
   ~Struct() override;
@@ -105,6 +117,11 @@ class Struct : public Castable<Struct, Node> {
     return false;
   }
 
+  /// @returns the set of entry point uses of this structure
+  const std::unordered_set<PipelineStageUsage>& PipelineStageUses() const {
+    return pipeline_stage_uses_;
+  }
+
  private:
   type::Struct* const type_;
   StructMemberList const members_;
@@ -112,6 +129,7 @@ class Struct : public Castable<Struct, Node> {
   uint32_t const size_;
   uint32_t const size_no_padding_;
   std::unordered_set<ast::StorageClass> const storage_class_usage_;
+  std::unordered_set<PipelineStageUsage> const pipeline_stage_uses_;
 };
 
 /// StructMember holds the semantic information for structure members.

test/BUILD.gn

@@ -181,6 +181,7 @@ source_set("tint_unittests_core_src") {
     "../src/resolver/resolver_test_helper.cc",
     "../src/resolver/resolver_test_helper.h",
     "../src/resolver/struct_layout_test.cc",
+    "../src/resolver/struct_pipeline_stage_use_test.cc",
     "../src/resolver/struct_storage_class_use_test.cc",
     "../src/resolver/type_constructor_validation_test.cc",
     "../src/resolver/type_validation_test.cc",