MSL: Support VK_KHR_push_descriptor.

If we have argument buffers, we also need to support using plain
descriptor sets for certain cases where API wants it.

main.cpp

@@ -496,6 +496,7 @@ struct CLIArguments
 	bool msl_pad_fragment_output = false;
 	bool msl_domain_lower_left = false;
 	bool msl_argument_buffers = false;
+	vector<uint32_t> msl_push_descriptor_sets;
 	vector<PLSArg> pls_in;
 	vector<PLSArg> pls_out;
 	vector<Remap> remaps;
@@ -554,6 +555,7 @@ static void print_help()
 	                "\t[--msl-pad-fragment-output]\n"
 	                "\t[--msl-domain-lower-left]\n"
 	                "\t[--msl-argument-buffers]\n"
+	                "\t[--msl-push-descriptor-set <index>]\n"
 	                "\t[--hlsl]\n"
 	                "\t[--reflect]\n"
 	                "\t[--shader-model]\n"
@@ -726,6 +728,8 @@ static int main_inner(int argc, char *argv[])
 	cbs.add("--msl-pad-fragment-output", [&args](CLIParser &) { args.msl_pad_fragment_output = true; });
 	cbs.add("--msl-domain-lower-left", [&args](CLIParser &) { args.msl_domain_lower_left = true; });
 	cbs.add("--msl-argument-buffers", [&args](CLIParser &) { args.msl_argument_buffers = true; });
+	cbs.add("--msl-push-descriptor-set",
+	        [&args](CLIParser &parser) { args.msl_push_descriptor_sets.push_back(parser.next_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();
@@ -860,6 +864,8 @@ static int main_inner(int argc, char *argv[])
 		msl_opts.tess_domain_origin_lower_left = args.msl_domain_lower_left;
 		msl_opts.argument_buffers = args.msl_argument_buffers;
 		msl_comp->set_msl_options(msl_opts);
+		for (auto &v : args.msl_push_descriptor_sets)
+			msl_comp->add_push_descriptor_set(v);
 	}
 	else if (args.hlsl)
 		compiler.reset(new CompilerHLSL(move(spirv_parser.get_parsed_ir())));

reference/opt/shaders-msl/comp/argument-buffers-push.msl2.argument.push.comp

@@ -0,0 +1,40 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct SSBO3
+{
+    float4 v;
+};
+
+struct SSBO0
+{
+    float4 v;
+};
+
+struct SSBO1
+{
+    float4 v;
+};
+
+struct SSBO2
+{
+    float4 v;
+};
+
+struct spvDescriptorSetBuffer0
+{
+    const device SSBO0* ssbo0 [[id(0)]];
+};
+
+struct spvDescriptorSetBuffer1
+{
+    const device SSBO1* ssbo1 [[id(0)]];
+};
+
+kernel void main0(constant spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], constant spvDescriptorSetBuffer1& spvDescriptorSet1 [[buffer(1)]], const device SSBO2& ssbo2 [[buffer(5)]], device SSBO3& ssbo3 [[buffer(6)]])
+{
+    ssbo3.v = ((*spvDescriptorSet0.ssbo0).v + (*spvDescriptorSet1.ssbo1).v) + ssbo2.v;
+}
+

reference/shaders-msl/comp/argument-buffers-push.msl2.argument.push.comp

@@ -0,0 +1,40 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct SSBO3
+{
+    float4 v;
+};
+
+struct SSBO0
+{
+    float4 v;
+};
+
+struct SSBO1
+{
+    float4 v;
+};
+
+struct SSBO2
+{
+    float4 v;
+};
+
+struct spvDescriptorSetBuffer0
+{
+    const device SSBO0* ssbo0 [[id(0)]];
+};
+
+struct spvDescriptorSetBuffer1
+{
+    const device SSBO1* ssbo1 [[id(0)]];
+};
+
+kernel void main0(constant spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], constant spvDescriptorSetBuffer1& spvDescriptorSet1 [[buffer(1)]], const device SSBO2& ssbo2 [[buffer(5)]], device SSBO3& ssbo3 [[buffer(6)]])
+{
+    ssbo3.v = ((*spvDescriptorSet0.ssbo0).v + (*spvDescriptorSet1.ssbo1).v) + ssbo2.v;
+}
+

shaders-msl/comp/argument-buffers-push.msl2.argument.push.comp

@@ -0,0 +1,27 @@
+#version 450
+layout(local_size_x = 1) in;
+
+layout(set = 0, binding = 0) readonly buffer SSBO0
+{
+   vec4 v;
+} ssbo0;
+
+layout(set = 1, binding = 0) readonly buffer SSBO1
+{
+   vec4 v;
+} ssbo1;
+
+layout(set = 2, binding = 5) readonly buffer SSBO2
+{
+   vec4 v;
+} ssbo2;
+
+layout(set = 3, binding = 6) writeonly buffer SSBO3
+{
+   vec4 v;
+} ssbo3;
+
+void main()
+{
+   ssbo3.v = ssbo0.v + ssbo1.v + ssbo2.v;
+}

spirv_cross_c.cpp

@@ -703,6 +703,19 @@ spvc_result spvc_compiler_msl_add_resource_binding(spvc_compiler compiler,
 	return SPVC_SUCCESS;
 }
 
+spvc_result spvc_compiler_msl_add_push_descriptor_set(spvc_compiler compiler, unsigned desc_set)
+{
+	if (compiler->backend != SPVC_BACKEND_MSL)
+	{
+		compiler->context->report_error("MSL function used on a non-MSL backend.");
+		return SPVC_ERROR_INVALID_ARGUMENT;
+	}
+
+	auto &msl = *static_cast<CompilerMSL *>(compiler->compiler.get());
+	msl.add_push_descriptor_set(desc_set);
+	return SPVC_SUCCESS;
+}
+
 spvc_bool spvc_compiler_msl_is_vertex_attribute_used(spvc_compiler compiler, unsigned location)
 {
 	if (compiler->backend != SPVC_BACKEND_MSL)

spirv_cross_c.h

@@ -505,6 +505,7 @@ SPVC_PUBLIC_API spvc_result spvc_compiler_msl_add_vertex_attribute(spvc_compiler
                                                                    const spvc_msl_vertex_attribute *attrs);
 SPVC_PUBLIC_API spvc_result spvc_compiler_msl_add_resource_binding(spvc_compiler compiler,
                                                                    const spvc_msl_resource_binding *binding);
+SPVC_PUBLIC_API spvc_result spvc_compiler_msl_add_push_descriptor_set(spvc_compiler compiler, unsigned desc_set);
 SPVC_PUBLIC_API spvc_bool spvc_compiler_msl_is_vertex_attribute_used(spvc_compiler compiler, unsigned location);
 SPVC_PUBLIC_API spvc_bool spvc_compiler_msl_is_resource_used(spvc_compiler compiler,
                                                              SpvExecutionModel model,

spirv_msl.cpp

@@ -62,6 +62,12 @@ void CompilerMSL::add_msl_resource_binding(const MSLResourceBinding &binding)
 	resource_bindings.push_back({ binding, false });
 }
 
+void CompilerMSL::add_push_descriptor_set(uint32_t desc_set)
+{
+	if (desc_set < kMaxArgumentBuffers)
+		argument_buffer_push |= 1u << desc_set;
+}
+
 bool CompilerMSL::is_msl_vertex_attribute_used(uint32_t location)
 {
 	return vtx_attrs_in_use.count(location) != 0;
@@ -5617,25 +5623,7 @@ string CompilerMSL::entry_point_args_argument_buffer(bool append_comma)
 		next_metal_resource_index_buffer = i + 1;
 	}
 
-	// Declare the push constant block separately, it does not belong in an indirect argument buffer.
-	uint32_t push_constant_id = 0;
-	ir.for_each_typed_id<SPIRVariable>([&](uint32_t self, SPIRVariable &var) {
-		if (var.storage == StorageClassPushConstant && !is_hidden_variable(var))
-			push_constant_id = self;
-	});
-
-	if (push_constant_id)
-	{
-		add_resource_name(push_constant_id);
-		auto &var = get<SPIRVariable>(push_constant_id);
-		auto &type = get_variable_data_type(var);
-		uint32_t resource_index = get_metal_resource_index(var, type.basetype);
-		if (!ep_args.empty())
-			ep_args += ", ";
-		ep_args += get_argument_address_space(var) + " " + type_to_glsl(type) + "& " + to_name(push_constant_id);
-		ep_args += " [[buffer(" + convert_to_string(resource_index) + ")]]";
-	}
-
+	entry_point_args_push_descriptors(ep_args);
 	entry_point_args_builtin(ep_args);
 
 	if (!ep_args.empty() && append_comma)
@@ -5644,12 +5632,8 @@ string CompilerMSL::entry_point_args_argument_buffer(bool append_comma)
 	return ep_args;
 }
 
-// Returns a string containing a comma-delimited list of args for the entry point function
-// This is the "classic" method of MSL 1 when we don't have argument buffer support.
-string CompilerMSL::entry_point_args_classic(bool append_comma)
+void CompilerMSL::entry_point_args_push_descriptors(string &ep_args)
 {
-	string ep_args = entry_point_arg_stage_in();
-
 	// Output resources, sorted by resource index & type
 	// We need to sort to work around a bug on macOS 10.13 with NVidia drivers where switching between shaders
 	// with different order of buffers can result in issues with buffer assignments inside the driver.
@@ -5671,6 +5655,13 @@ string CompilerMSL::entry_point_args_classic(bool append_comma)
 			auto &type = get_variable_data_type(var);
 			uint32_t var_id = var.self;
 
+			if (var.storage != StorageClassPushConstant)
+			{
+				uint32_t desc_set = get_decoration(var_id, DecorationDescriptorSet);
+				if (descriptor_set_is_argument_buffer(desc_set))
+					return;
+			}
+
 			if (type.basetype == SPIRType::SampledImage)
 			{
 				add_resource_name(var_id);
@@ -5693,7 +5684,7 @@ string CompilerMSL::entry_point_args_classic(bool append_comma)
 		}
 	});
 
-	std::sort(resources.begin(), resources.end(), [](const Resource &lhs, const Resource &rhs) {
+	sort(resources.begin(), resources.end(), [](const Resource &lhs, const Resource &rhs) {
 		return tie(lhs.basetype, lhs.index) < tie(rhs.basetype, rhs.index);
 	});
 
@@ -5760,7 +5751,14 @@ string CompilerMSL::entry_point_args_classic(bool append_comma)
 			break;
 		}
 	}
+}
 
+// Returns a string containing a comma-delimited list of args for the entry point function
+// This is the "classic" method of MSL 1 when we don't have argument buffer support.
+string CompilerMSL::entry_point_args_classic(bool append_comma)
+{
+	string ep_args = entry_point_arg_stage_in();
+	entry_point_args_push_descriptors(ep_args);
 	entry_point_args_builtin(ep_args);
 
 	if (!ep_args.empty() && append_comma)
@@ -7527,6 +7525,16 @@ std::string CompilerMSL::to_initializer_expression(const SPIRVariable &var)
 		return CompilerGLSL::to_initializer_expression(var);
 }
 
+bool CompilerMSL::descriptor_set_is_argument_buffer(uint32_t desc_set) const
+{
+	if (!msl_options.argument_buffers)
+		return false;
+	if (desc_set >= kMaxArgumentBuffers)
+		return false;
+
+	return (argument_buffer_push & (1u << desc_set)) == 0;
+}
+
 void CompilerMSL::analyze_argument_buffers()
 {
 	// Gather all used resources and sort them out into argument buffers.
@@ -7555,6 +7563,10 @@ void CompilerMSL::analyze_argument_buffers()
 		    !is_hidden_variable(var))
 		{
 			uint32_t desc_set = get_decoration(self, DecorationDescriptorSet);
+			// Ignore if it's part of a push descriptor set.
+			if (!descriptor_set_is_argument_buffer(desc_set))
+				return;
+
 			uint32_t var_id = var.self;
 			auto &type = get_variable_data_type(var);
 
@@ -7597,6 +7609,8 @@ void CompilerMSL::analyze_argument_buffers()
 		if (resources.empty())
 			continue;
 
+		assert(descriptor_set_is_argument_buffer(desc_set));
+
 		uint32_t next_id = ir.increase_bound_by(3);
 		uint32_t type_id = next_id + 1;
 		uint32_t ptr_type_id = next_id + 2;

spirv_msl.hpp

@@ -54,10 +54,11 @@ struct MSLVertexAttr
 // Matches the binding index of a MSL resource for a binding within a descriptor set.
 // Taken together, the stage, desc_set and binding combine to form a reference to a resource
 // descriptor used in a particular shading stage.
-// If using MSL 2.0 argument buffers, the binding reference we remap to will become an [[id(N)]] attribute within
+// If using MSL 2.0 argument buffers, and the descriptor set is not marked as a push descriptor set,
+// the binding reference we remap to will become an [[id(N)]] attribute within
 // the "descriptor set" argument buffer structure.
-// For resources which are bound in the "classic" MSL 1.0 way, the remap will become a
-// [[buffer(N)]], [[texture(N)]] or [[sampler(N)]].
+// For resources which are bound in the "classic" MSL 1.0 way or push descriptors, the remap will become a
+// [[buffer(N)]], [[texture(N)]] or [[sampler(N)]] depending on the resource types used.
 struct MSLResourceBinding
 {
 	spv::ExecutionModel stage = spv::ExecutionModelMax;
@@ -285,6 +286,10 @@ public:
 	// the set/binding combination was used by the MSL code.
 	void add_msl_resource_binding(const MSLResourceBinding &resource);
 
+	// When using MSL argument buffers, we can force "classic" MSL 1.0 binding schemes for certain descriptor sets.
+	// This corresponds to VK_KHR_push_descriptor in Vulkan.
+	void add_push_descriptor_set(uint32_t desc_set);
+
 	// Query after compilation is done. This allows you to check if a location or set/binding combination was used by the shader.
 	bool is_msl_vertex_attribute_used(uint32_t location);
 	bool is_msl_resource_binding_used(spv::ExecutionModel model, uint32_t set, uint32_t binding);
@@ -427,6 +432,7 @@ protected:
 	std::string entry_point_args_argument_buffer(bool append_comma);
 	std::string entry_point_arg_stage_in();
 	void entry_point_args_builtin(std::string &args);
+	void entry_point_args_push_descriptors(std::string &args);
 	std::string to_qualified_member_name(const SPIRType &type, uint32_t index);
 	std::string ensure_valid_name(std::string name, std::string pfx);
 	std::string to_sampler_expression(uint32_t id);
@@ -527,7 +533,9 @@ protected:
 	std::vector<uint32_t> buffer_arrays;
 
 	uint32_t argument_buffer_ids[kMaxArgumentBuffers];
+	uint32_t argument_buffer_push = 0;
 	void analyze_argument_buffers();
+	bool descriptor_set_is_argument_buffer(uint32_t desc_set) const;
 
 	uint32_t get_target_components_for_fragment_location(uint32_t location) const;
 	uint32_t build_extended_vector_type(uint32_t type_id, uint32_t components);

test_shaders.py

@@ -171,6 +171,12 @@ def cross_compile_msl(shader, spirv, opt, paths):
         msl_args.append('--msl-domain-lower-left')
     if '.argument.' in shader:
         msl_args.append('--msl-argument-buffers')
+    if '.push.' in shader:
+        # Arbitrary for testing purposes.
+        msl_args.append('--msl-push-descriptor-set')
+        msl_args.append('2')
+        msl_args.append('--msl-push-descriptor-set')
+        msl_args.append('3')
 
     subprocess.check_call(msl_args)