Enabling setting a fixed sampleMask in Metal fragment shaders.

In Metal render pipelines don't have an option to set a sampleMask
parameter, the only way to get that functionality is to set the
sample_mask output of the fragment shader to this value directly.
We also need to take care to combine the fixed sample mask with the
one that the shader might possibly output.

main.cpp

@@ -560,6 +560,7 @@ struct CLIArguments
 	bool msl_enable_clip_distance_user_varying = true;
 	bool msl_multi_patch_workgroup = false;
 	bool msl_vertex_for_tessellation = false;
+	uint32_t msl_additional_fixed_sample_mask = 0xffffffff;
 	bool glsl_emit_push_constant_as_ubo = false;
 	bool glsl_emit_ubo_as_plain_uniforms = false;
 	SmallVector<pair<uint32_t, uint32_t>> glsl_ext_framebuffer_fetch;
@@ -757,7 +758,9 @@ static void print_help_msl()
 					"\t\tIn addition, this style also passes input variables in buffers directly instead of using vertex attribute processing.\n"
 					"\t\tIn a future version of SPIRV-Cross, this will become the default.\n"
 	                "\t[--msl-vertex-for-tessellation]:\n\t\tWhen handling a vertex shader, marks it as one that will be used with a new-style tessellation control shader.\n"
-					"\t\tThe vertex shader is output to MSL as a compute kernel which outputs vertices to the buffer in the order they are received, rather than in index order as with --msl-capture-output normally.\n");
+					"\t\tThe vertex shader is output to MSL as a compute kernel which outputs vertices to the buffer in the order they are received, rather than in index order as with --msl-capture-output normally.\n"
+	                "\t[--msl-additional-fixed-sample-mask <mask>]:\n"
+	                "\t\tSet an additional fixed sample mask. If the shader outputs a sample mask, then the final sample mask will be a bitwise AND of the two.\n");
 	// clang-format on
 }
 
@@ -993,6 +996,7 @@ static string compile_iteration(const CLIArguments &args, std::vector<uint32_t>
 		msl_opts.enable_clip_distance_user_varying = args.msl_enable_clip_distance_user_varying;
 		msl_opts.multi_patch_workgroup = args.msl_multi_patch_workgroup;
 		msl_opts.vertex_for_tessellation = args.msl_vertex_for_tessellation;
+		msl_opts.additional_fixed_sample_mask = args.msl_additional_fixed_sample_mask;
 		msl_comp->set_msl_options(msl_opts);
 		for (auto &v : args.msl_discrete_descriptor_sets)
 			msl_comp->add_discrete_descriptor_set(v);
@@ -1406,6 +1410,8 @@ static int main_inner(int argc, char *argv[])
 	});
 	cbs.add("--msl-multi-patch-workgroup", [&args](CLIParser &) { args.msl_multi_patch_workgroup = true; });
 	cbs.add("--msl-vertex-for-tessellation", [&args](CLIParser &) { args.msl_vertex_for_tessellation = true; });
+	cbs.add("--msl-additional-fixed-sample-mask",
+	        [&args](CLIParser &parser) { args.msl_additional_fixed_sample_mask = parser.next_hex_uint(); });
 	cbs.add("--extension", [&args](CLIParser &parser) { args.extensions.push_back(parser.next_string()); });
 	cbs.add("--rename-entry-point", [&args](CLIParser &parser) {
 		auto old_name = parser.next_string();

reference/opt/shaders-msl/frag/sample-mask-in-and-out.fixed-sample-mask.frag

@@ -0,0 +1,20 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+    uint gl_SampleMask [[sample_mask]];
+};
+
+fragment main0_out main0(uint gl_SampleMaskIn [[sample_mask]])
+{
+    main0_out out = {};
+    out.FragColor = float4(1.0);
+    out.gl_SampleMask = gl_SampleMaskIn;
+    out.gl_SampleMask &= 34;
+    return out;
+}
+

reference/opt/shaders-msl/frag/sample-mask-not-used.fixed-sample-mask.frag

@@ -0,0 +1,19 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+    uint gl_SampleMask [[sample_mask]];
+};
+
+fragment main0_out main0()
+{
+    main0_out out = {};
+    out.FragColor = float4(1.0);
+    out.gl_SampleMask = 34;
+    return out;
+}
+

reference/opt/shaders-msl/frag/sample-mask.fixed-sample-mask.frag

@@ -0,0 +1,20 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+    uint gl_SampleMask [[sample_mask]];
+};
+
+fragment main0_out main0()
+{
+    main0_out out = {};
+    out.FragColor = float4(1.0);
+    out.gl_SampleMask = 0;
+    out.gl_SampleMask &= 34;
+    return out;
+}
+

reference/shaders-msl/frag/sample-mask-in-and-out.fixed-sample-mask.frag

@@ -0,0 +1,20 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+    uint gl_SampleMask [[sample_mask]];
+};
+
+fragment main0_out main0(uint gl_SampleMaskIn [[sample_mask]])
+{
+    main0_out out = {};
+    out.FragColor = float4(1.0);
+    out.gl_SampleMask = gl_SampleMaskIn;
+    out.gl_SampleMask &= 34;
+    return out;
+}
+

reference/shaders-msl/frag/sample-mask-not-used.fixed-sample-mask.frag

@@ -0,0 +1,19 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+    uint gl_SampleMask [[sample_mask]];
+};
+
+fragment main0_out main0()
+{
+    main0_out out = {};
+    out.FragColor = float4(1.0);
+    out.gl_SampleMask = 34;
+    return out;
+}
+

reference/shaders-msl/frag/sample-mask.fixed-sample-mask.frag

@@ -0,0 +1,20 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+    uint gl_SampleMask [[sample_mask]];
+};
+
+fragment main0_out main0()
+{
+    main0_out out = {};
+    out.FragColor = float4(1.0);
+    out.gl_SampleMask = 0;
+    out.gl_SampleMask &= 34;
+    return out;
+}
+

shaders-msl/frag/sample-mask-in-and-out.fixed-sample-mask.frag

@@ -0,0 +1,10 @@
+#version 450
+
+layout(location = 0) out vec4 FragColor;
+
+void main()
+{
+	FragColor = vec4(1.0);
+	gl_SampleMask[0] = gl_SampleMaskIn[0];
+}
+

shaders-msl/frag/sample-mask-not-used.fixed-sample-mask.frag

@@ -0,0 +1,8 @@
+#version 450
+
+layout(location = 0) out vec4 FragColor;
+
+void main()
+{
+	FragColor = vec4(1.0);
+}

shaders-msl/frag/sample-mask.fixed-sample-mask.frag

@@ -0,0 +1,10 @@
+#version 450
+
+layout(location = 0) out vec4 FragColor;
+
+void main()
+{
+	FragColor = vec4(1.0);
+	gl_SampleMask[0] = 0;
+}
+

spirv_msl.cpp

@@ -156,9 +156,10 @@ void CompilerMSL::build_implicit_builtins()
 	    (active_input_builtins.get(BuiltInVertexId) || active_input_builtins.get(BuiltInVertexIndex) ||
 	     active_input_builtins.get(BuiltInBaseVertex) || active_input_builtins.get(BuiltInInstanceId) ||
 	     active_input_builtins.get(BuiltInInstanceIndex) || active_input_builtins.get(BuiltInBaseInstance));
+	bool need_sample_mask = msl_options.additional_fixed_sample_mask != 0xffffffff;
 	if (need_subpass_input || need_sample_pos || need_subgroup_mask || need_vertex_params || need_tesc_params ||
 	    need_multiview || need_dispatch_base || need_vertex_base_params || need_grid_params ||
-	    needs_subgroup_invocation_id)
+	    needs_subgroup_invocation_id || need_sample_mask)
 	{
 		bool has_frag_coord = false;
 		bool has_sample_id = false;
@@ -173,12 +174,27 @@ void CompilerMSL::build_implicit_builtins()
 		bool has_view_idx = false;
 		uint32_t workgroup_id_type = 0;
 
+		// FIXME: Investigate the fact that there are no checks for the entry point interface variables.
 		ir.for_each_typed_id<SPIRVariable>([&](uint32_t, SPIRVariable &var) {
-			if (var.storage != StorageClassInput || !ir.meta[var.self].decoration.builtin)
+			if (!ir.meta[var.self].decoration.builtin)
 				return;
 
 			// Use Metal's native frame-buffer fetch API for subpass inputs.
 			BuiltIn builtin = ir.meta[var.self].decoration.builtin_type;
+
+			if (var.storage == StorageClassOutput)
+			{
+				if (need_sample_mask && builtin == BuiltInSampleMask)
+				{
+					builtin_sample_mask_id = var.self;
+					mark_implicit_builtin(StorageClassOutput, BuiltInSampleMask, var.self);
+					does_shader_write_sample_mask = true;
+				}
+			}
+
+			if (var.storage != StorageClassInput)
+				return;
+
 			if (need_subpass_input && (!msl_options.is_ios() || !msl_options.ios_use_framebuffer_fetch_subpasses) &&
 			    builtin == BuiltInFragCoord)
 			{
@@ -561,6 +577,26 @@ void CompilerMSL::build_implicit_builtins()
 			set_name(var_id, "spvDispatchBase");
 			builtin_dispatch_base_id = var_id;
 		}
+
+		if (need_sample_mask && !does_shader_write_sample_mask)
+		{
+			uint32_t offset = ir.increase_bound_by(2);
+			uint32_t var_id = offset + 1;
+
+			// Create gl_SampleMask.
+			SPIRType uint_type_ptr_out;
+			uint_type_ptr_out = get_uint_type();
+			uint_type_ptr_out.pointer = true;
+			uint_type_ptr_out.parent_type = get_uint_type_id();
+			uint_type_ptr_out.storage = StorageClassOutput;
+			
+			auto &ptr_out_type = set<SPIRType>(offset, uint_type_ptr_out);
+			ptr_out_type.self = get_uint_type_id();
+			set<SPIRVariable>(var_id, offset, StorageClassOutput);
+			set_decoration(var_id, DecorationBuiltIn, BuiltInSampleMask);
+			builtin_sample_mask_id = var_id;
+			mark_implicit_builtin(StorageClassOutput, BuiltInSampleMask, var_id);
+		}
 	}
 
 	if (needs_swizzle_buffer_def)
@@ -1066,6 +1102,8 @@ string CompilerMSL::compile()
 		active_interface_variables.insert(builtin_layer_id);
 	if (builtin_dispatch_base_id && !msl_options.supports_msl_version(1, 2))
 		active_interface_variables.insert(builtin_dispatch_base_id);
+	if (builtin_sample_mask_id)
+		active_interface_variables.insert(builtin_sample_mask_id);
 
 	// Create structs to hold input, output and uniform variables.
 	// Do output first to ensure out. is declared at top of entry function.
@@ -10512,6 +10550,30 @@ void CompilerMSL::fix_up_shader_inputs_outputs()
 				break;
 			}
 		}
+		else if (var.storage == StorageClassOutput && is_builtin_variable(var))
+		{
+			if (bi_type == BuiltInSampleMask && get_execution_model() == ExecutionModelFragment &&
+				msl_options.additional_fixed_sample_mask != 0xffffffff)
+			{
+				// If the additional fixed sample mask was set, we need to adjust the sample_mask
+				// output to reflect that. If the shader outputs the sample_mask itself too, we need
+				// to AND the two masks to get the final one.
+				if (does_shader_write_sample_mask)
+				{
+					entry_func.fixup_hooks_out.push_back([=]() {
+						statement(to_expression(builtin_sample_mask_id), " &= ",
+							      msl_options.additional_fixed_sample_mask, ";");
+					});
+				}
+				else
+				{
+					entry_func.fixup_hooks_out.push_back([=]() {
+						statement(to_expression(builtin_sample_mask_id), " = ",
+							      msl_options.additional_fixed_sample_mask, ";");
+					});
+				}
+			}
+		}
 	});
 }
 

spirv_msl.hpp

@@ -278,6 +278,10 @@ public:
 		uint32_t shader_input_wg_index = 0;
 		uint32_t device_index = 0;
 		uint32_t enable_frag_output_mask = 0xffffffff;
+		// Metal doesn't allow setting a fixed sample mask directly in the pipeline.
+		// We can evade this restriction by ANDing the internal sample_mask output
+		// of the shader with the additional fixed sample mask.
+		uint32_t additional_fixed_sample_mask = 0xffffffff;
 		bool enable_point_size_builtin = true;
 		bool enable_frag_depth_builtin = true;
 		bool enable_frag_stencil_ref_builtin = true;
@@ -809,6 +813,7 @@ protected:
 	void emit_entry_point_declarations() override;
 	uint32_t builtin_frag_coord_id = 0;
 	uint32_t builtin_sample_id_id = 0;
+	uint32_t builtin_sample_mask_id = 0;
 	uint32_t builtin_vertex_idx_id = 0;
 	uint32_t builtin_base_vertex_id = 0;
 	uint32_t builtin_instance_idx_id = 0;
@@ -827,6 +832,8 @@ protected:
 	uint32_t dynamic_offsets_buffer_id = 0;
 	uint32_t uint_type_id = 0;
 
+	bool does_shader_write_sample_mask = false;
+
 	void bitcast_to_builtin_store(uint32_t target_id, std::string &expr, const SPIRType &expr_type) override;
 	void bitcast_from_builtin_load(uint32_t source_id, std::string &expr, const SPIRType &expr_type) override;
 	void emit_store_statement(uint32_t lhs_expression, uint32_t rhs_expression) override;

test_shaders.py

@@ -293,6 +293,9 @@ def cross_compile_msl(shader, spirv, opt, iterations, paths):
         msl_args.append('2')
     if '.for-tess.' in shader:
         msl_args.append('--msl-vertex-for-tessellation')
+    if '.fixed-sample-mask.' in shader:
+        msl_args.append('--msl-additional-fixed-sample-mask')
+        msl_args.append('0x00000022')
 
     subprocess.check_call(msl_args)