MSL: Add opt-in support for huge IABs.

If there are enough members in an IAB, we cannot use the constant
address space as MSL compiler complains about there being too many
members. Support emitting the device address space instead.

CMakeLists.txt

@@ -308,7 +308,7 @@ if (SPIRV_CROSS_STATIC)
 endif()
 
 set(spirv-cross-abi-major 0)
-set(spirv-cross-abi-minor 18)
+set(spirv-cross-abi-minor 19)
 set(spirv-cross-abi-patch 0)
 
 if (SPIRV_CROSS_SHARED)

main.cpp

@@ -522,6 +522,7 @@ struct CLIArguments
 	bool vulkan_glsl_disable_ext_samplerless_texture_functions = false;
 	bool emit_line_directives = false;
 	SmallVector<uint32_t> msl_discrete_descriptor_sets;
+	SmallVector<uint32_t> msl_device_argument_buffers;
 	SmallVector<pair<uint32_t, uint32_t>> msl_dynamic_buffers;
 	SmallVector<PLSArg> pls_in;
 	SmallVector<PLSArg> pls_out;
@@ -598,6 +599,7 @@ static void print_help()
 	                "\t[--msl-argument-buffers]\n"
 	                "\t[--msl-texture-buffer-native]\n"
 	                "\t[--msl-discrete-descriptor-set <index>]\n"
+	                "\t[--msl-device-argument-buffer <index>]\n"
 	                "\t[--msl-multiview]\n"
 	                "\t[--msl-view-index-from-device-index]\n"
 	                "\t[--msl-dispatch-base]\n"
@@ -766,6 +768,8 @@ static string compile_iteration(const CLIArguments &args, std::vector<uint32_t>
 		msl_comp->set_msl_options(msl_opts);
 		for (auto &v : args.msl_discrete_descriptor_sets)
 			msl_comp->add_discrete_descriptor_set(v);
+		for (auto &v : args.msl_device_argument_buffers)
+			msl_comp->set_argument_buffer_device_address_space(v, true);
 		uint32_t i = 0;
 		for (auto &v : args.msl_dynamic_buffers)
 			msl_comp->add_dynamic_buffer(v.first, v.second, i++);
@@ -1086,6 +1090,8 @@ static int main_inner(int argc, char *argv[])
 	cbs.add("--msl-argument-buffers", [&args](CLIParser &) { args.msl_argument_buffers = true; });
 	cbs.add("--msl-discrete-descriptor-set",
 	        [&args](CLIParser &parser) { args.msl_discrete_descriptor_sets.push_back(parser.next_uint()); });
+	cbs.add("--msl-device-argument-buffer",
+	        [&args](CLIParser &parser) { args.msl_device_argument_buffers.push_back(parser.next_uint()); });
 	cbs.add("--msl-texture-buffer-native", [&args](CLIParser &) { args.msl_texture_buffer_native = true; });
 	cbs.add("--msl-multiview", [&args](CLIParser &) { args.msl_multiview = true; });
 	cbs.add("--msl-view-index-from-device-index",

reference/opt/shaders-msl/frag/huge-argument-buffer.device-argument-buffer.argument.msl2.frag

@@ -0,0 +1,44 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct UBO
+{
+    float4 v;
+};
+
+struct spvDescriptorSetBuffer0
+{
+    array<texture2d<float>, 10000> uSamplers [[id(0)]];
+    array<sampler, 10000> uSamplersSmplr [[id(10000)]];
+};
+
+struct spvDescriptorSetBuffer1
+{
+    constant UBO* vs [[id(0)]][10000];
+};
+
+struct spvDescriptorSetBuffer2
+{
+    texture2d<float> uSampler [[id(0)]];
+    sampler uSamplerSmplr [[id(1)]];
+};
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+struct main0_in
+{
+    float2 vUV [[user(locn0)]];
+};
+
+fragment main0_out main0(main0_in in [[stage_in]], const device spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], const device spvDescriptorSetBuffer1& spvDescriptorSet1 [[buffer(1)]], constant spvDescriptorSetBuffer2& spvDescriptorSet2 [[buffer(2)]])
+{
+    main0_out out = {};
+    out.FragColor = (spvDescriptorSet0.uSamplers[9999].sample(spvDescriptorSet0.uSamplersSmplr[9999], in.vUV) + spvDescriptorSet1.vs[5000]->v) + spvDescriptorSet2.uSampler.sample(spvDescriptorSet2.uSamplerSmplr, in.vUV);
+    return out;
+}
+

reference/shaders-msl/frag/huge-argument-buffer.device-argument-buffer.argument.msl2.frag

@@ -0,0 +1,56 @@
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct UBO
+{
+    float4 v;
+};
+
+struct spvDescriptorSetBuffer0
+{
+    array<texture2d<float>, 10000> uSamplers [[id(0)]];
+    array<sampler, 10000> uSamplersSmplr [[id(10000)]];
+};
+
+struct spvDescriptorSetBuffer1
+{
+    constant UBO* vs [[id(0)]][10000];
+};
+
+struct spvDescriptorSetBuffer2
+{
+    texture2d<float> uSampler [[id(0)]];
+    sampler uSamplerSmplr [[id(1)]];
+};
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+struct main0_in
+{
+    float2 vUV [[user(locn0)]];
+};
+
+inline float4 samp_array(thread const array<texture2d<float>, 10000> uSamplers, thread const array<sampler, 10000> uSamplersSmplr, thread float2& vUV, constant UBO* const device (&vs)[10000])
+{
+    return uSamplers[9999].sample(uSamplersSmplr[9999], vUV) + vs[5000]->v;
+}
+
+inline float4 samp_single(thread float2& vUV, thread texture2d<float> uSampler, thread const sampler uSamplerSmplr)
+{
+    return uSampler.sample(uSamplerSmplr, vUV);
+}
+
+fragment main0_out main0(main0_in in [[stage_in]], const device spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], const device spvDescriptorSetBuffer1& spvDescriptorSet1 [[buffer(1)]], constant spvDescriptorSetBuffer2& spvDescriptorSet2 [[buffer(2)]])
+{
+    main0_out out = {};
+    out.FragColor = samp_array(spvDescriptorSet0.uSamplers, spvDescriptorSet0.uSamplersSmplr, in.vUV, spvDescriptorSet1.vs) + samp_single(in.vUV, spvDescriptorSet2.uSampler, spvDescriptorSet2.uSamplerSmplr);
+    return out;
+}
+

shaders-msl/frag/huge-argument-buffer.device-argument-buffer.argument.msl2.frag

@@ -0,0 +1,26 @@
+#version 450
+
+layout(location = 0) out vec4 FragColor;
+layout(location = 0) in vec2 vUV;
+layout(set = 0, binding = 0) uniform sampler2D uSamplers[10000];
+layout(set = 2, binding = 0) uniform sampler2D uSampler;
+
+layout(set = 1, binding = 0) uniform UBO
+{
+	vec4 v;
+} vs[10000];
+
+vec4 samp_array()
+{
+	return texture(uSamplers[9999], vUV) + vs[5000].v;
+}
+
+vec4 samp_single()
+{
+	return texture(uSampler, vUV);
+}
+
+void main()
+{
+	FragColor = samp_array() + samp_single();
+}

spirv_cross_c.cpp

@@ -946,6 +946,26 @@ spvc_result spvc_compiler_msl_add_discrete_descriptor_set(spvc_compiler compiler
 #endif
 }
 
+spvc_result spvc_compiler_msl_set_argument_buffer_device_address_space(spvc_compiler compiler, unsigned desc_set, spvc_bool device_address)
+{
+#if SPIRV_CROSS_C_API_MSL
+	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.set_argument_buffer_device_address_space(desc_set, bool(device_address));
+	return SPVC_SUCCESS;
+#else
+	(void)desc_set;
+	(void)device_address;
+	compiler->context->report_error("MSL function used on a non-MSL backend.");
+	return SPVC_ERROR_INVALID_ARGUMENT;
+#endif
+}
+
 spvc_bool spvc_compiler_msl_is_vertex_attribute_used(spvc_compiler compiler, unsigned location)
 {
 #if SPIRV_CROSS_C_API_MSL

spirv_cross_c.h

@@ -33,7 +33,7 @@ extern "C" {
 /* Bumped if ABI or API breaks backwards compatibility. */
 #define SPVC_C_API_VERSION_MAJOR 0
 /* Bumped if APIs or enumerations are added in a backwards compatible way. */
-#define SPVC_C_API_VERSION_MINOR 18
+#define SPVC_C_API_VERSION_MINOR 19
 /* Bumped if internal implementation details change. */
 #define SPVC_C_API_VERSION_PATCH 0
 
@@ -619,6 +619,7 @@ SPVC_PUBLIC_API spvc_result spvc_compiler_msl_add_vertex_attribute(spvc_compiler
 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_discrete_descriptor_set(spvc_compiler compiler, unsigned desc_set);
+SPVC_PUBLIC_API spvc_result spvc_compiler_msl_set_argument_buffer_device_address_space(spvc_compiler compiler, unsigned desc_set, spvc_bool device_address);
 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

@@ -73,6 +73,17 @@ void CompilerMSL::add_discrete_descriptor_set(uint32_t desc_set)
 		argument_buffer_discrete_mask |= 1u << desc_set;
 }
 
+void CompilerMSL::set_argument_buffer_device_address_space(uint32_t desc_set, bool device_storage)
+{
+	if (desc_set < kMaxArgumentBuffers)
+	{
+		if (device_storage)
+			argument_buffer_device_storage_mask |= 1u << desc_set;
+		else
+			argument_buffer_device_storage_mask &= ~(1u << desc_set);
+	}
+}
+
 bool CompilerMSL::is_msl_vertex_attribute_used(uint32_t location)
 {
 	return vtx_attrs_in_use.count(location) != 0;
@@ -8950,6 +8961,12 @@ string CompilerMSL::argument_decl(const SPIRFunction::Parameter &arg)
 				// constant SSBO * constant (&array)[N].
 				// However, this only matters for argument buffers, since for MSL 1.0 style codegen,
 				// we emit the buffer array on stack instead, and that seems to work just fine apparently.
+
+				// If the argument was marked as being in device address space, any pointer to member would
+				// be const device, not constant.
+				if (argument_buffer_device_storage_mask & (1u << desc_set))
+					decl += " const device";
+				else
 					decl += " constant";
 			}
 		}
@@ -11213,8 +11230,20 @@ void CompilerMSL::analyze_argument_buffers()
 		argument_buffer_ids[desc_set] = next_id;
 
 		auto &buffer_type = set<SPIRType>(type_id);
-		buffer_type.storage = StorageClassUniform;
+
 		buffer_type.basetype = SPIRType::Struct;
+
+		if ((argument_buffer_device_storage_mask & (1u << desc_set)) != 0)
+		{
+			buffer_type.storage = StorageClassStorageBuffer;
+			// Make sure the argument buffer gets marked as const device.
+			set_decoration(next_id, DecorationNonWritable);
+			// Need to mark the type as a Block to enable this.
+			set_decoration(type_id, DecorationBlock);
+		}
+		else
+			buffer_type.storage = StorageClassUniform;
+
 		set_name(type_id, join("spvDescriptorSetBuffer", desc_set));
 
 		auto &ptr_type = set<SPIRType>(ptr_type_id);

spirv_msl.hpp

@@ -410,6 +410,10 @@ public:
 	// This corresponds to VK_KHR_push_descriptor in Vulkan.
 	void add_discrete_descriptor_set(uint32_t desc_set);
 
+	// If an argument buffer is large enough, it may need to be in the device storage space rather than
+	// constant. Opt-in to this behavior here on a per set basis.
+	void set_argument_buffer_device_address_space(uint32_t desc_set, bool device_storage);
+
 	// 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);
 
@@ -795,6 +799,8 @@ protected:
 
 	uint32_t argument_buffer_ids[kMaxArgumentBuffers];
 	uint32_t argument_buffer_discrete_mask = 0;
+	uint32_t argument_buffer_device_storage_mask = 0;
+
 	void analyze_argument_buffers();
 	bool descriptor_set_is_argument_buffer(uint32_t desc_set) const;
 

test_shaders.py

@@ -217,6 +217,11 @@ def cross_compile_msl(shader, spirv, opt, iterations, paths):
         msl_args.append('--msl-dynamic-buffer')
         msl_args.append('1')
         msl_args.append('2')
+    if '.device-argument-buffer.' in shader:
+        msl_args.append('--msl-device-argument-buffer')
+        msl_args.append('0')
+        msl_args.append('--msl-device-argument-buffer')
+        msl_args.append('1')
 
     subprocess.check_call(msl_args)