msl: Generate a fixed sample mask if requested

The CanonicalizeEntryPointIO transform now takes an optional config
option for a fixed sample mask. If there was no sample mask in the
authored shader, add one and return the fixed mask, otherwise AND the
fixed mask with the authored value.

Add a config option to the MSL sanitizer to receive a fixed sample
mask and pass it through to CanonicalizeEntryPointIO.

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

src/program_builder.h

@@ -1490,6 +1490,16 @@ class ProgramBuilder {
                                          Expr(std::forward<RHS>(rhs)));
   }
 
+  /// @param lhs the left hand argument to the and operation
+  /// @param rhs the right hand argument to the and operation
+  /// @returns a `ast::BinaryExpression` bitwise anding `lhs` and `rhs`
+  template <typename LHS, typename RHS>
+  ast::BinaryExpression* And(LHS&& lhs, RHS&& rhs) {
+    return create<ast::BinaryExpression>(ast::BinaryOp::kAnd,
+                                         Expr(std::forward<LHS>(lhs)),
+                                         Expr(std::forward<RHS>(rhs)));
+  }
+
   /// @param lhs the left hand argument to the subtraction operation
   /// @param rhs the right hand argument to the subtraction operation
   /// @returns a `ast::BinaryExpression` subtracting `rhs` from `lhs`

src/transform/canonicalize_entry_point_io.cc

@@ -99,6 +99,14 @@ Output CanonicalizeEntryPointIO::Run(const Program* in, const DataMap& data) {
     }
   }
 
+  // Returns true if `decos` contains a `sample_mask` builtin.
+  auto has_sample_mask_builtin = [](const ast::DecorationList& decos) {
+    if (auto* builtin = ast::GetDecoration<ast::BuiltinDecoration>(decos)) {
+      return builtin->value() == ast::Builtin::kSampleMask;
+    }
+    return false;
+  };
+
   for (auto* func_ast : ctx.src->AST().Functions()) {
     if (!func_ast->IsEntryPoint()) {
       continue;
@@ -106,6 +114,10 @@ Output CanonicalizeEntryPointIO::Run(const Program* in, const DataMap& data) {
 
     auto* func = ctx.src->Sem().Get(func_ast);
 
+    bool needs_fixed_sample_mask =
+        func_ast->pipeline_stage() == ast::PipelineStage::kFragment &&
+        cfg->fixed_sample_mask != 0xFFFFFFFF;
+
     ast::VariableList new_parameters;
 
     if (!func->Parameters().empty()) {
@@ -216,11 +228,13 @@ Output CanonicalizeEntryPointIO::Run(const Program* in, const DataMap& data) {
     // Handle return type.
     auto* ret_type = func->ReturnType();
     std::function<ast::Type*()> new_ret_type;
-    if (ret_type->Is<sem::Void>()) {
+    if (ret_type->Is<sem::Void>() && !needs_fixed_sample_mask) {
       new_ret_type = [&ctx] { return ctx.dst->ty.void_(); };
     } else {
       ast::StructMemberList new_struct_members;
 
+      bool has_authored_sample_mask = false;
+
       if (auto* str = ret_type->As<sem::Struct>()) {
         // Rebuild struct with only the entry point IO attributes.
         for (auto* member : str->Members()) {
@@ -238,12 +252,28 @@ Output CanonicalizeEntryPointIO::Run(const Program* in, const DataMap& data) {
           auto* member_ty = ctx.Clone(member->Declaration()->type());
           new_struct_members.push_back(
               ctx.dst->Member(symbol, member_ty, new_decorations));
+
+          if (has_sample_mask_builtin(new_decorations)) {
+            has_authored_sample_mask = true;
+          }
         }
-      } else {
+      } else if (!ret_type->Is<sem::Void>()) {
         auto* member_ty = ctx.Clone(func->Declaration()->return_type());
         auto decos = ctx.Clone(func_ast->return_type_decorations());
         new_struct_members.push_back(
             ctx.dst->Member("value", member_ty, std::move(decos)));
+
+        if (has_sample_mask_builtin(func_ast->return_type_decorations())) {
+          has_authored_sample_mask = true;
+        }
+      }
+
+      // If a sample mask builtin is required and the shader source did not
+      // contain one, create one now.
+      if (needs_fixed_sample_mask && !has_authored_sample_mask) {
+        new_struct_members.push_back(
+            ctx.dst->Member(ctx.dst->Sym(), ctx.dst->ty.u32(),
+                            {ctx.dst->Builtin(ast::Builtin::kSampleMask)}));
       }
 
       // Sort struct members to satisfy HLSL interfacing matching rules.
@@ -283,17 +313,54 @@ Output CanonicalizeEntryPointIO::Run(const Program* in, const DataMap& data) {
           }
 
           for (auto* member : new_struct_members) {
-            ret_values.push_back(
+            ast::Expression* expr = nullptr;
-                ctx.dst->MemberAccessor(new_ret_value(), member->symbol()));
+
+            if (needs_fixed_sample_mask &&
+                has_sample_mask_builtin(member->decorations())) {
+              // Use the fixed sample mask, combining it with the authored value
+              // if there is one.
+              expr = ctx.dst->Expr(cfg->fixed_sample_mask);
+              if (has_authored_sample_mask) {
+                expr = ctx.dst->And(
+                    ctx.dst->MemberAccessor(new_ret_value(), member->symbol()),
+                    expr);
+              }
+            } else {
+              expr = ctx.dst->MemberAccessor(new_ret_value(), member->symbol());
+            }
+            ret_values.push_back(expr);
           }
         } else {
-          ret_values.push_back(new_ret_value());
+          if (!ret_type->Is<sem::Void>()) {
+            ret_values.push_back(new_ret_value());
+          }
+
+          if (needs_fixed_sample_mask) {
+            // If the original return value was a sample mask, `and` it with the
+            // fixed mask and return the result.
+            // Otherwise, append the fixed mask to the list of return values,
+            // since it will be the last element of the output struct.
+            if (has_authored_sample_mask) {
+              ret_values[0] =
+                  ctx.dst->And(ret_values[0], cfg->fixed_sample_mask);
+            } else {
+              ret_values.push_back(ctx.dst->Expr(cfg->fixed_sample_mask));
+            }
+          }
         }
 
         auto* new_ret =
             ctx.dst->Return(ctx.dst->Construct(new_ret_type(), ret_values));
         ctx.Replace(ret, new_ret);
       }
+
+      if (needs_fixed_sample_mask && func->ReturnStatements().empty()) {
+        // There we no return statements but we need to return a fixed sample
+        // mask, so add a return statement that does this.
+        ctx.InsertBack(func_ast->body()->statements(),
+                       ctx.dst->Return(ctx.dst->Construct(
+                           new_ret_type(), cfg->fixed_sample_mask)));
+      }
     }
 
     // Rewrite the function header with the new parameters.
@@ -308,13 +375,12 @@ Output CanonicalizeEntryPointIO::Run(const Program* in, const DataMap& data) {
   return Output(Program(std::move(out)));
 }
 
-CanonicalizeEntryPointIO::Config::Config(BuiltinStyle builtins)
+CanonicalizeEntryPointIO::Config::Config(BuiltinStyle builtins,
-    : builtin_style(builtins) {}
+                                         uint32_t sample_mask)
+    : builtin_style(builtins), fixed_sample_mask(sample_mask) {}
 
 CanonicalizeEntryPointIO::Config::Config(const Config&) = default;
 CanonicalizeEntryPointIO::Config::~Config() = default;
-CanonicalizeEntryPointIO::Config& CanonicalizeEntryPointIO::Config::operator=(
-    const Config&) = default;
 
 }  // namespace transform
 }  // namespace tint

src/transform/canonicalize_entry_point_io.h

@@ -82,7 +82,8 @@ class CanonicalizeEntryPointIO : public Transform {
   struct Config : public Castable<Config, Data> {
     /// Constructor
     /// @param builtins the approach to use for emitting builtins.
-    explicit Config(BuiltinStyle builtins);
+    /// @param sample_mask an optional sample mask to combine with shader masks
+    explicit Config(BuiltinStyle builtins, uint32_t sample_mask = 0xFFFFFFFF);
 
     /// Copy constructor
     Config(const Config&);
@@ -90,12 +91,11 @@ class CanonicalizeEntryPointIO : public Transform {
     /// Destructor
     ~Config() override;
 
-    /// Assignment operator
-    /// @returns this Config
-    Config& operator=(const Config&);
-
     /// The approach to use for emitting builtins.
-    BuiltinStyle builtin_style;
+    BuiltinStyle const builtin_style;
+
+    /// A fixed sample mask to combine into masks produced by fragment shaders.
+    uint32_t const fixed_sample_mask;
   };
 
   /// Constructor

src/transform/canonicalize_entry_point_io_test.cc

@@ -790,6 +790,276 @@ fn tint_symbol_1(tint_symbol : tint_symbol_2) {
   EXPECT_EQ(expect, str(got));
 }
 
+TEST_F(CanonicalizeEntryPointIOTest, FixedSampleMask_VoidNoReturn) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main() {
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol_1 {
+  [[builtin(sample_mask)]]
+  tint_symbol : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> tint_symbol_1 {
+  return tint_symbol_1(3u);
+}
+)";
+
+  DataMap data;
+  data.Add<CanonicalizeEntryPointIO::Config>(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, 0x03u);
+  auto got = Run<CanonicalizeEntryPointIO>(src, data);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(CanonicalizeEntryPointIOTest, FixedSampleMask_VoidWithReturn) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main() {
+  return;
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol_1 {
+  [[builtin(sample_mask)]]
+  tint_symbol : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> tint_symbol_1 {
+  return tint_symbol_1(3u);
+}
+)";
+
+  DataMap data;
+  data.Add<CanonicalizeEntryPointIO::Config>(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, 0x03u);
+  auto got = Run<CanonicalizeEntryPointIO>(src, data);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(CanonicalizeEntryPointIOTest, FixedSampleMask_WithAuthoredMask) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main() -> [[builtin(sample_mask)]] u32 {
+  return 7u;
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol {
+  [[builtin(sample_mask)]]
+  value : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> tint_symbol {
+  return tint_symbol((7u & 3u));
+}
+)";
+
+  DataMap data;
+  data.Add<CanonicalizeEntryPointIO::Config>(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, 0x03u);
+  auto got = Run<CanonicalizeEntryPointIO>(src, data);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(CanonicalizeEntryPointIOTest, FixedSampleMask_WithoutAuthoredMask) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main() -> [[location(0)]] f32 {
+  return 1.0;
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol_1 {
+  [[location(0)]]
+  value : f32;
+  [[builtin(sample_mask)]]
+  tint_symbol : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> tint_symbol_1 {
+  return tint_symbol_1(1.0, 3u);
+}
+)";
+
+  DataMap data;
+  data.Add<CanonicalizeEntryPointIO::Config>(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, 0x03u);
+  auto got = Run<CanonicalizeEntryPointIO>(src, data);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(CanonicalizeEntryPointIOTest, FixedSampleMask_StructWithAuthoredMask) {
+  auto* src = R"(
+struct Output {
+  [[builtin(frag_depth)]] depth : f32;
+  [[builtin(sample_mask)]] mask : u32;
+  [[location(0)]] value : f32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> Output {
+  return Output(0.5, 7u, 1.0);
+}
+)";
+
+  auto* expect = R"(
+struct Output {
+  depth : f32;
+  mask : u32;
+  value : f32;
+};
+
+struct tint_symbol {
+  [[location(0)]]
+  value : f32;
+  [[builtin(frag_depth)]]
+  depth : f32;
+  [[builtin(sample_mask)]]
+  mask : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> tint_symbol {
+  let tint_symbol_1 : Output = Output(0.5, 7u, 1.0);
+  return tint_symbol(tint_symbol_1.value, tint_symbol_1.depth, (tint_symbol_1.mask & 3u));
+}
+)";
+
+  DataMap data;
+  data.Add<CanonicalizeEntryPointIO::Config>(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, 0x03u);
+  auto got = Run<CanonicalizeEntryPointIO>(src, data);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(CanonicalizeEntryPointIOTest,
+       FixedSampleMask_StructWithoutAuthoredMask) {
+  auto* src = R"(
+struct Output {
+  [[builtin(frag_depth)]] depth : f32;
+  [[location(0)]] value : f32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> Output {
+  return Output(0.5, 1.0);
+}
+)";
+
+  auto* expect = R"(
+struct Output {
+  depth : f32;
+  value : f32;
+};
+
+struct tint_symbol_1 {
+  [[location(0)]]
+  value : f32;
+  [[builtin(frag_depth)]]
+  depth : f32;
+  [[builtin(sample_mask)]]
+  tint_symbol : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main() -> tint_symbol_1 {
+  let tint_symbol_2 : Output = Output(0.5, 1.0);
+  return tint_symbol_1(tint_symbol_2.value, tint_symbol_2.depth, 3u);
+}
+)";
+
+  DataMap data;
+  data.Add<CanonicalizeEntryPointIO::Config>(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, 0x03u);
+  auto got = Run<CanonicalizeEntryPointIO>(src, data);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(CanonicalizeEntryPointIOTest, FixedSampleMask_MultipleShaders) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main1() -> [[builtin(sample_mask)]] u32 {
+  return 7u;
+}
+
+[[stage(fragment)]]
+fn frag_main2() -> [[location(0)]] f32 {
+  return 1.0;
+}
+
+[[stage(vertex)]]
+fn vert_main1() -> [[builtin(position)]] vec4<f32> {
+  return vec4<f32>();
+}
+
+[[stage(compute)]]
+fn comp_main1() {
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol {
+  [[builtin(sample_mask)]]
+  value : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main1() -> tint_symbol {
+  return tint_symbol((7u & 3u));
+}
+
+struct tint_symbol_2 {
+  [[location(0)]]
+  value : f32;
+  [[builtin(sample_mask)]]
+  tint_symbol_1 : u32;
+};
+
+[[stage(fragment)]]
+fn frag_main2() -> tint_symbol_2 {
+  return tint_symbol_2(1.0, 3u);
+}
+
+struct tint_symbol_3 {
+  [[builtin(position)]]
+  value : vec4<f32>;
+};
+
+[[stage(vertex)]]
+fn vert_main1() -> tint_symbol_3 {
+  return tint_symbol_3(vec4<f32>());
+}
+
+[[stage(compute)]]
+fn comp_main1() {
+}
+)";
+
+  DataMap data;
+  data.Add<CanonicalizeEntryPointIO::Config>(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, 0x03u);
+  auto got = Run<CanonicalizeEntryPointIO>(src, data);
+
+  EXPECT_EQ(expect, str(got));
+}
+
 }  // namespace
 }  // namespace transform
 }  // namespace tint

src/transform/msl.cc

@@ -51,13 +51,17 @@ Output Msl::Run(const Program* in, const DataMap& inputs) {
 
   auto* cfg = inputs.Get<Config>();
 
-  // Build the config for the array length transform.
+  // Build the configs for the internal transforms.
   uint32_t buffer_size_ubo_index = kDefaultBufferSizeUniformIndex;
+  uint32_t fixed_sample_mask = 0xFFFFFFFF;
   if (cfg) {
     buffer_size_ubo_index = cfg->buffer_size_ubo_index;
+    fixed_sample_mask = cfg->fixed_sample_mask;
   }
   auto array_length_from_uniform_cfg = ArrayLengthFromUniform::Config(
       sem::BindingPoint{0, buffer_size_ubo_index});
+  auto entry_point_io_cfg = CanonicalizeEntryPointIO::Config(
+      CanonicalizeEntryPointIO::BuiltinStyle::kParameter, fixed_sample_mask);
 
   // Use the SSBO binding numbers as the indices for the buffer size lookups.
   for (auto* var : in->AST().GlobalVariables()) {
@@ -84,7 +88,7 @@ Output Msl::Run(const Program* in, const DataMap& inputs) {
   internal_inputs.Add<ArrayLengthFromUniform::Config>(
       std::move(array_length_from_uniform_cfg));
   internal_inputs.Add<CanonicalizeEntryPointIO::Config>(
-      CanonicalizeEntryPointIO::BuiltinStyle::kParameter);
+      std::move(entry_point_io_cfg));
   auto out = manager.Run(in, internal_inputs);
   if (!out.program.IsValid()) {
     return out;
@@ -274,8 +278,9 @@ void Msl::HandleModuleScopeVariables(CloneContext& ctx) const {
   }
 }
 
-Msl::Config::Config(uint32_t buffer_size_ubo_idx)
+Msl::Config::Config(uint32_t buffer_size_ubo_idx, uint32_t sample_mask)
-    : buffer_size_ubo_index(buffer_size_ubo_idx) {}
+    : buffer_size_ubo_index(buffer_size_ubo_idx),
+      fixed_sample_mask(sample_mask) {}
 Msl::Config::Config(const Config&) = default;
 Msl::Config::~Config() = default;
 

src/transform/msl.h

@@ -32,7 +32,8 @@ class Msl : public Transform {
   struct Config : public Castable<Data, transform::Data> {
     /// Constructor
     /// @param buffer_size_ubo_idx the index to use for the buffer size UBO
-    explicit Config(uint32_t buffer_size_ubo_idx);
+    /// @param sample_mask the fixed sample mask to use for fragment shaders
+    explicit Config(uint32_t buffer_size_ubo_idx, uint32_t sample_mask);
 
     /// Copy constructor
     Config(const Config&);
@@ -42,6 +43,9 @@ class Msl : public Transform {
 
     /// The index to use when generating a UBO to receive storage buffer sizes.
     uint32_t buffer_size_ubo_index;
+
+    /// The fixed sample mask to combine with fragment shader outputs.
+    uint32_t fixed_sample_mask;
   };
 
   /// Information produced by the sanitizer that users may need to act on.