[spirv-writer] Add support for sample index/mask builtins

Use a sanitizing transform to convert scalar `sample_mask_{in,out}`
variables to single element arrays.

Add the `SampleRateShading` capability if the `sample_index` builtin
is used.

Bug: tint:372
Change-Id: Id7280e3ddb21e0a098d83587d123c97e3c34fa1b
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/41662
Reviewed-by: dan sinclair <dsinclair@chromium.org>
Commit-Queue: James Price <jrprice@google.com>
Auto-Submit: James Price <jrprice@google.com>

src/transform/spirv.cc

@@ -14,9 +14,13 @@
 
 #include "src/transform/spirv.h"
 
+#include <string>
 #include <utility>
 
+#include "src/ast/variable.h"
 #include "src/program_builder.h"
+#include "src/semantic/expression.h"
+#include "src/semantic/variable.h"
 
 namespace tint {
 namespace transform {
@@ -26,9 +30,60 @@ Spirv::~Spirv() = default;
 
 Transform::Output Spirv::Run(const Program* in) {
   ProgramBuilder out;
-  CloneContext(&out, in).Clone();
+  CloneContext ctx(&out, in);
+  HandleSampleMaskBuiltins(ctx);
+  ctx.Clone();
   return Output{Program(std::move(out))};
 }
 
+void Spirv::HandleSampleMaskBuiltins(CloneContext& ctx) const {
+  // Find global variables decorated with [[builtin(sample_mask_{in,out})]] and
+  // change their type from `u32` to `array<u32, 1>`, as required by Vulkan.
+  //
+  // Before:
+  // ```
+  // [[builtin(sample_mask_out)]] var<out> mask_out : u32;
+  // fn main() -> void {
+  //   mask_out = 1u;
+  // }
+  // ```
+  // After:
+  // ```
+  // [[builtin(sample_mask_out)]] var<out> mask_out : array<u32, 1>;
+  // fn main() -> void {
+  //   mask_out[0] = 1u;
+  // }
+  // ```
+
+  for (auto* var : ctx.src->AST().GlobalVariables()) {
+    for (auto* deco : var->decorations()) {
+      if (auto* builtin = deco->As<ast::BuiltinDecoration>()) {
+        if (builtin->value() != ast::Builtin::kSampleMaskIn &&
+            builtin->value() != ast::Builtin::kSampleMaskOut) {
+          continue;
+        }
+
+        // Use the same name as the old variable.
+        std::string var_name = ctx.src->Symbols().NameFor(var->symbol());
+        // Use `array<u32, 1>` for the new variable.
+        auto* type = ctx.dst->ty.array(ctx.dst->ty.u32(), 1u);
+        // Create the new variable.
+        auto* var_arr =
+            ctx.dst->Var(var->source(), var_name, var->declared_storage_class(),
+                         type, nullptr, ctx.Clone(var->decorations()));
+        // Replace the variable with the arrayed version.
+        ctx.Replace(var, var_arr);
+
+        // Replace all uses of the old variable with `var_arr[0]`.
+        for (auto* user : ctx.src->Sem().Get(var)->Users()) {
+          auto* new_ident = ctx.dst->IndexAccessor(
+              ctx.dst->Expr(var_arr->symbol()), ctx.dst->Expr(0));
+          ctx.Replace<ast::Expression>(user->Declaration(), new_ident);
+        }
+      }
+    }
+  }
+}
+
 }  // namespace transform
 }  // namespace tint

src/transform/spirv.h

@@ -18,6 +18,10 @@
 #include "src/transform/transform.h"
 
 namespace tint {
+
+// Forward declarations
+class CloneContext;
+
 namespace transform {
 
 /// Spirv is a transform used to sanitize a Program for use with the Spirv
@@ -36,6 +40,10 @@ class Spirv : public Transform {
   /// @param program the source program to transform
   /// @returns the transformation result
   Output Run(const Program* program) override;
+
+ private:
+  /// Change type of sample mask builtin variables to single element arrays.
+  void HandleSampleMaskBuiltins(CloneContext& ctx) const;
 };
 
 }  // namespace transform

src/transform/spirv_test.cc

@@ -0,0 +1,107 @@
+// 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/spirv.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using SpirvTest = TransformTest;
+
+TEST_F(SpirvTest, HandleSampleMaskBuiltins_Basic) {
+  auto* src = R"(
+[[builtin(sample_index)]] var<in> sample_index : u32;
+
+[[builtin(sample_mask_in)]] var<in> mask_in : u32;
+
+[[builtin(sample_mask_out)]] var<out> mask_out : u32;
+
+[[stage(fragment)]]
+fn main() -> void {
+  mask_out = mask_in;
+}
+)";
+
+  auto* expect = R"(
+[[builtin(sample_index)]] var<in> sample_index : u32;
+
+[[builtin(sample_mask_in)]] var<in> mask_in : array<u32, 1>;
+
+[[builtin(sample_mask_out)]] var<out> mask_out : array<u32, 1>;
+
+[[stage(fragment)]]
+fn main() -> void {
+  mask_out[0] = mask_in[0];
+}
+)";
+
+  auto got = Transform<Spirv>(src);
+
+  EXPECT_EQ(expect, got);
+}
+
+TEST_F(SpirvTest, HandleSampleMaskBuiltins_FunctionArg) {
+  auto* src = R"(
+[[builtin(sample_mask_in)]] var<in> mask_in : u32;
+
+[[builtin(sample_mask_out)]] var<out> mask_out : u32;
+
+fn filter(mask: u32) -> u32 {
+  return (mask & 3u);
+}
+
+fn set_mask(input : u32) -> void {
+  mask_out = input;
+}
+
+[[stage(fragment)]]
+fn main() -> void {
+  set_mask(filter(mask_in));
+}
+)";
+
+  auto* expect = R"(
+[[builtin(sample_mask_in)]] var<in> mask_in : array<u32, 1>;
+
+[[builtin(sample_mask_out)]] var<out> mask_out : array<u32, 1>;
+
+fn filter(mask : u32) -> u32 {
+  return (mask & 3u);
+}
+
+fn set_mask(input : u32) -> void {
+  mask_out[0] = input;
+}
+
+[[stage(fragment)]]
+fn main() -> void {
+  set_mask(filter(mask_in[0]));
+}
+)";
+
+  auto got = Transform<Spirv>(src);
+
+  EXPECT_EQ(expect, got);
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint