[spirv-writer] Handle entry point parameters

Add a sanitizing transform to hoist entry point parameters out as
global variables.

Bug: tint:509
Change-Id: Ic18f69386a58d82ee11571fa9ec0c54cb5bdf2cf
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/44083
Commit-Queue: James Price <jrprice@google.com>
Auto-Submit: James Price <jrprice@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/transform/spirv.cc

@@ -29,9 +29,118 @@ Spirv::~Spirv() = default;
 Transform::Output Spirv::Run(const Program* in) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
-  HandleSampleMaskBuiltins(ctx);
+  HandleEntryPointIOTypes(ctx);
   ctx.Clone();
-  return Output{Program(std::move(out))};
+
+  // TODO(jrprice): Look into combining these transforms into a single clone.
+  Program tmp(std::move(out));
+
+  ProgramBuilder out2;
+  CloneContext ctx2(&out2, &tmp);
+  HandleSampleMaskBuiltins(ctx2);
+  ctx2.Clone();
+
+  return Output{Program(std::move(out2))};
+}
+
+void Spirv::HandleEntryPointIOTypes(CloneContext& ctx) const {
+  // Hoist entry point parameters, return values, and struct members out to
+  // global variables. Declare and construct struct parameters in the function
+  // body. Replace entry point return statements with variable assignments.
+  //
+  // Before:
+  // ```
+  // struct FragmentInput {
+  //   [[builtin(sample_index)]] sample_index : u32;
+  //   [[builtin(sample_mask_in)]] sample_mask_in : u32;
+  // };
+  // struct FragmentOutput {
+  //   [[builtin(frag_depth)]] depth: f32;
+  //   [[builtin(sample_mask_out)]] mask_out : u32;
+  // };
+  //
+  // [[stage(fragment)]]
+  // fn fs_main(
+  //   [[builtin(frag_coord)]] coord : vec4<f32>,
+  //   samples : FragmentInput
+  // ) -> FragmentOutput {
+  //   return FragmentOutput(1.0, samples.sample_mask_in);
+  // }
+  // ```
+  //
+  // After:
+  // ```
+  // struct FragmentInput {
+  //   sample_index : u32;
+  //   sample_mask_in : u32;
+  // };
+  // struct FragmentOutput {
+  //   depth: f32;
+  //   mask_out : u32;
+  // };
+  //
+  // [[builtin(frag_coord)]] var<in> coord : vec4<f32>,
+  // [[builtin(sample_index)]] var<in> sample_index : u32,
+  // [[builtin(sample_mask_in)]] var<in> sample_mask_in : u32,
+  // [[builtin(frag_depth)]] var<out> depth: f32;
+  // [[builtin(sample_mask_out)]] var<out> mask_out : u32;
+  //
+  // [[stage(fragment)]]
+  // fn fs_main() -> void {
+  //   const samples : FragmentInput(sample_index, sample_mask_in);
+  //   depth = 1.0;
+  //   mask_out = samples.sample_mask_in;
+  //   return;
+  // }
+  // ```
+
+  // TODO(jrprice): Hoist struct members decorated as entry point IO types out
+  // of struct declarations, and redeclare the structs without the decorations.
+
+  for (auto* func : ctx.src->AST().Functions()) {
+    if (!func->IsEntryPoint()) {
+      continue;
+    }
+
+    for (auto* param : func->params()) {
+      // TODO(jrprice): Handle structures by moving the declaration and
+      // construction to the function body.
+      if (param->type()->Is<type::Struct>()) {
+        TINT_ICE(ctx.dst->Diagnostics())
+            << "structures as entry point parameters are not yet supported";
+        continue;
+      }
+
+      // Create a new symbol for the global variable.
+      auto var_symbol = ctx.dst->Symbols().New();
+      // Create the global variable.
+      ctx.dst->Global(var_symbol, ctx.Clone(param->type()),
+                      ast::StorageClass::kInput, nullptr,
+                      ctx.Clone(param->decorations()));
+
+      // Replace all uses of the function parameter with the global variable.
+      for (auto* user : ctx.src->Sem().Get(param)->Users()) {
+        ctx.Replace<ast::Expression>(user->Declaration(),
+                                     ctx.dst->Expr(var_symbol));
+      }
+    }
+
+    // TODO(jrprice): Hoist the return type out to a global variable, and
+    // replace return statements with variable assignments.
+    if (!func->return_type()->Is<type::Void>()) {
+      TINT_ICE(ctx.dst->Diagnostics())
+          << "entry point return values are not yet supported";
+      continue;
+    }
+
+    // Rewrite the function header to remove the parameters.
+    // TODO(jrprice): Change return type to void when return values are handled.
+    auto* new_func = ctx.dst->create<ast::Function>(
+        func->source(), ctx.Clone(func->symbol()), ast::VariableList{},
+        ctx.Clone(func->return_type()), ctx.Clone(func->body()),
+        ctx.Clone(func->decorations()));
+    ctx.Replace(func, new_func);
+  }
 }
 
 void Spirv::HandleSampleMaskBuiltins(CloneContext& ctx) const {

src/transform/spirv.h

@@ -39,6 +39,9 @@ class Spirv : public Transform {
   Output Run(const Program* program) override;
 
  private:
+  /// Hoist entry point parameters, return values, and struct members out to
+  /// global variables.
+  void HandleEntryPointIOTypes(CloneContext& ctx) const;
   /// Change type of sample mask builtin variables to single element arrays.
   void HandleSampleMaskBuiltins(CloneContext& ctx) const;
 };

src/transform/spirv_test.cc

@@ -22,6 +22,46 @@ namespace {
 
 using SpirvTest = TransformTest;
 
+TEST_F(SpirvTest, HandleEntryPointIOTypes_Parameters) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>,
+             [[location(1)]] loc1 : f32) -> void {
+  var col : f32 = (coord.x * loc1);
+}
+
+[[stage(compute)]]
+fn compute_main([[builtin(local_invocation_id)]] local_id : vec3<u32>,
+                [[builtin(local_invocation_index)]] local_index : u32) -> void {
+  var id_x : u32 = local_id.x;
+}
+)";
+
+  auto* expect = R"(
+[[builtin(frag_coord)]] var<in> tint_symbol_1 : vec4<f32>;
+
+[[location(1)]] var<in> tint_symbol_2 : f32;
+
+[[builtin(local_invocation_id)]] var<in> tint_symbol_6 : vec3<u32>;
+
+[[builtin(local_invocation_index)]] var<in> tint_symbol_7 : u32;
+
+[[stage(fragment)]]
+fn frag_main() -> void {
+  var col : f32 = (tint_symbol_1.x * tint_symbol_2);
+}
+
+[[stage(compute)]]
+fn compute_main() -> void {
+  var id_x : u32 = tint_symbol_6.x;
+}
+)";
+
+  auto got = Transform<Spirv>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
 TEST_F(SpirvTest, HandleSampleMaskBuiltins_Basic) {
   auto* src = R"(
 [[builtin(sample_index)]] var<in> sample_index : u32;
@@ -98,6 +138,37 @@ fn main() -> void {
   EXPECT_EQ(expect, str(got));
 }
 
+// Test that different transforms within the sanitizer interact correctly.
+TEST_F(SpirvTest, MultipleTransforms) {
+  // TODO(jrprice): Make `mask_out` a return value when supported.
+  auto* src = R"(
+[[builtin(sample_mask_out)]] var<out> mask_out : u32;
+
+[[stage(fragment)]]
+fn main([[builtin(sample_index)]] sample_index : u32,
+        [[builtin(sample_mask_in)]] mask_in : u32) -> void {
+  mask_out = mask_in;
+}
+)";
+
+  auto* expect = R"(
+[[builtin(sample_index)]] var<in> tint_symbol_1 : u32;
+
+[[builtin(sample_mask_in)]] var<in> tint_symbol_2 : array<u32, 1>;
+
+[[builtin(sample_mask_out)]] var<out> mask_out : array<u32, 1>;
+
+[[stage(fragment)]]
+fn main() -> void {
+  mask_out[0] = tint_symbol_2[0];
+}
+)";
+
+  auto got = Transform<Spirv>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
 }  // namespace
 }  // namespace transform
 }  // namespace tint