Merge pull request #1239 from cdavis5e/msl-inline-uniform-blocks

MSL: Support inline uniform blocks in argument buffers.

CMakeLists.txt

@@ -315,7 +315,7 @@ if (SPIRV_CROSS_STATIC)
 endif()
 
 set(spirv-cross-abi-major 0)
-set(spirv-cross-abi-minor 23)
+set(spirv-cross-abi-minor 24)
 set(spirv-cross-abi-patch 0)
 
 if (SPIRV_CROSS_SHARED)

main.cpp

@@ -529,6 +529,7 @@ struct CLIArguments
 	SmallVector<uint32_t> msl_discrete_descriptor_sets;
 	SmallVector<uint32_t> msl_device_argument_buffers;
 	SmallVector<pair<uint32_t, uint32_t>> msl_dynamic_buffers;
+	SmallVector<pair<uint32_t, uint32_t>> msl_inline_uniform_blocks;
 	SmallVector<PLSArg> pls_in;
 	SmallVector<PLSArg> pls_out;
 	SmallVector<Remap> remaps;
@@ -612,6 +613,7 @@ static void print_help()
 	                "\t[--msl-view-index-from-device-index]\n"
 	                "\t[--msl-dispatch-base]\n"
 	                "\t[--msl-dynamic-buffer <set index> <binding>]\n"
+	                "\t[--msl-inline-uniform-block <set index> <binding>]\n"
 	                "\t[--msl-decoration-binding]\n"
 	                "\t[--msl-force-active-argument-buffer-resources]\n"
 	                "\t[--hlsl]\n"
@@ -812,6 +814,8 @@ static string compile_iteration(const CLIArguments &args, std::vector<uint32_t>
 		uint32_t i = 0;
 		for (auto &v : args.msl_dynamic_buffers)
 			msl_comp->add_dynamic_buffer(v.first, v.second, i++);
+		for (auto &v : args.msl_inline_uniform_blocks)
+			msl_comp->add_inline_uniform_block(v.first, v.second);
 	}
 	else if (args.hlsl)
 		compiler.reset(new CompilerHLSL(move(spirv_parser.get_parsed_ir())));
@@ -1153,6 +1157,13 @@ static int main_inner(int argc, char *argv[])
 	cbs.add("--msl-decoration-binding", [&args](CLIParser &) { args.msl_decoration_binding = true; });
 	cbs.add("--msl-force-active-argument-buffer-resources",
 	        [&args](CLIParser &) { args.msl_force_active_argument_buffer_resources = true; });
+	cbs.add("--msl-inline-uniform-block", [&args](CLIParser &parser) {
+		args.msl_argument_buffers = true;
+		// Make sure next_uint() is called in-order.
+		uint32_t desc_set = parser.next_uint();
+		uint32_t binding = parser.next_uint();
+		args.msl_inline_uniform_blocks.push_back(make_pair(desc_set, binding));
+	});
 	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/comp/basic.inline-block.msl2.comp

@@ -0,0 +1,53 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+typedef packed_float4 packed_float4x4[4];
+
+struct Baz
+{
+    int f;
+    int g;
+};
+
+struct X
+{
+    int x;
+    int y;
+    float z;
+};
+
+struct Foo
+{
+    int a;
+    int b;
+    packed_float4x4 c;
+    X x[2];
+};
+
+struct Bar
+{
+    int d;
+    int e;
+};
+
+constant uint3 gl_WorkGroupSize [[maybe_unused]] = uint3(3u, 3u, 2u);
+
+struct spvDescriptorSetBuffer0
+{
+    constant Bar* m_38 [[id(0)]];
+    Foo m_32 [[id(1)]];
+};
+
+struct spvDescriptorSetBuffer1
+{
+    device Baz* baz [[id(0)]][3];
+};
+
+kernel void main0(constant spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], constant spvDescriptorSetBuffer1& spvDescriptorSet1 [[buffer(1)]], uint3 gl_GlobalInvocationID [[thread_position_in_grid]])
+{
+    spvDescriptorSet1.baz[gl_GlobalInvocationID.x]->f = spvDescriptorSet0.m_32.a + (*spvDescriptorSet0.m_38).d;
+    spvDescriptorSet1.baz[gl_GlobalInvocationID.x]->g = spvDescriptorSet0.m_32.b * (*spvDescriptorSet0.m_38).e;
+}
+

reference/shaders-msl/comp/basic.inline-block.msl2.comp

@@ -0,0 +1,54 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+typedef packed_float4 packed_float4x4[4];
+
+struct Baz
+{
+    int f;
+    int g;
+};
+
+struct X
+{
+    int x;
+    int y;
+    float z;
+};
+
+struct Foo
+{
+    int a;
+    int b;
+    packed_float4x4 c;
+    X x[2];
+};
+
+struct Bar
+{
+    int d;
+    int e;
+};
+
+constant uint3 gl_WorkGroupSize [[maybe_unused]] = uint3(3u, 3u, 2u);
+
+struct spvDescriptorSetBuffer0
+{
+    constant Bar* m_38 [[id(0)]];
+    Foo m_32 [[id(1)]];
+};
+
+struct spvDescriptorSetBuffer1
+{
+    device Baz* baz [[id(0)]][3];
+};
+
+kernel void main0(constant spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], constant spvDescriptorSetBuffer1& spvDescriptorSet1 [[buffer(1)]], uint3 gl_GlobalInvocationID [[thread_position_in_grid]])
+{
+    uint3 coords = gl_GlobalInvocationID;
+    spvDescriptorSet1.baz[coords.x]->f = spvDescriptorSet0.m_32.a + (*spvDescriptorSet0.m_38).d;
+    spvDescriptorSet1.baz[coords.x]->g = spvDescriptorSet0.m_32.b * (*spvDescriptorSet0.m_38).e;
+}
+

shaders-msl/comp/basic.inline-block.msl2.comp

@@ -0,0 +1,37 @@
+#version 450
+#extension GL_EXT_scalar_block_layout : require
+layout(local_size_x = 3, local_size_y = 3, local_size_z = 2) in;
+
+struct X
+{
+	int x;
+	int y;
+	float z;
+};
+
+layout(set = 0, binding = 0, scalar) uniform Foo
+{
+	int a;
+	int b;
+	mat4 c;
+	X x[2];
+};
+
+layout(set = 0, binding = 1) uniform Bar
+{
+	int d;
+	int e;
+};
+
+layout(set = 1, binding = 2) buffer Baz
+{
+	int f;
+	int g;
+} baz[3];
+
+void main()
+{
+	uvec3 coords = gl_GlobalInvocationID;
+	baz[coords.x].f = a + d;
+	baz[coords.x].g = b * e;
+}

spirv_cross_c.cpp

@@ -1029,6 +1029,26 @@ spvc_result spvc_compiler_msl_add_dynamic_buffer(spvc_compiler compiler, unsigne
 #endif
 }
 
+spvc_result spvc_compiler_msl_add_inline_uniform_block(spvc_compiler compiler, unsigned desc_set, unsigned binding)
+{
+#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.add_inline_uniform_block(desc_set, binding);
+	return SPVC_SUCCESS;
+#else
+	(void)binding;
+	(void)desc_set;
+	compiler->context->report_error("MSL function used on a non-MSL backend.");
+	return SPVC_ERROR_INVALID_ARGUMENT;
+#endif
+}
+
 spvc_result spvc_compiler_msl_add_discrete_descriptor_set(spvc_compiler compiler, unsigned desc_set)
 {
 #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 23
+#define SPVC_C_API_VERSION_MINOR 24
 /* Bumped if internal implementation details change. */
 #define SPVC_C_API_VERSION_PATCH 0
 
@@ -691,6 +691,8 @@ SPVC_PUBLIC_API unsigned spvc_compiler_msl_get_automatic_resource_binding_second
 
 SPVC_PUBLIC_API spvc_result spvc_compiler_msl_add_dynamic_buffer(spvc_compiler compiler, unsigned desc_set, unsigned binding, unsigned index);
 
+SPVC_PUBLIC_API spvc_result spvc_compiler_msl_add_inline_uniform_block(spvc_compiler compiler, unsigned desc_set, unsigned binding);
+
 /*
  * Reflect resources.
  * Maps almost 1:1 to C++ API.

spirv_msl.cpp

@@ -68,6 +68,12 @@ void CompilerMSL::add_dynamic_buffer(uint32_t desc_set, uint32_t binding, uint32
 	buffers_requiring_dynamic_offset[pair] = { index, 0 };
 }
 
+void CompilerMSL::add_inline_uniform_block(uint32_t desc_set, uint32_t binding)
+{
+	SetBindingPair pair = { desc_set, binding };
+	inline_uniform_blocks.insert(pair);
+}
+
 void CompilerMSL::add_discrete_descriptor_set(uint32_t desc_set)
 {
 	if (desc_set < kMaxArgumentBuffers)
@@ -9912,6 +9918,11 @@ uint32_t CompilerMSL::get_metal_resource_index(SPIRVariable &var, SPIRType::Base
 	// If we did not explicitly remap, allocate bindings on demand.
 	// We cannot reliably use Binding decorations since SPIR-V and MSL's binding models are very different.
 
+	bool allocate_argument_buffer_ids = false;
+
+	if (var.storage != StorageClassPushConstant)
+		allocate_argument_buffer_ids = descriptor_set_is_argument_buffer(var_desc_set);
+
 	uint32_t binding_stride = 1;
 	auto &type = get<SPIRType>(var.basetype);
 	for (uint32_t i = 0; i < uint32_t(type.array.size()); i++)
@@ -9922,20 +9933,11 @@ uint32_t CompilerMSL::get_metal_resource_index(SPIRVariable &var, SPIRType::Base
 	// If a binding has not been specified, revert to incrementing resource indices.
 	uint32_t resource_index;
 
-	bool allocate_argument_buffer_ids = false;
-	uint32_t desc_set = 0;
-
-	if (var.storage != StorageClassPushConstant)
-	{
-		desc_set = get_decoration(var.self, DecorationDescriptorSet);
-		allocate_argument_buffer_ids = descriptor_set_is_argument_buffer(desc_set);
-	}
-
 	if (allocate_argument_buffer_ids)
 	{
 		// Allocate from a flat ID binding space.
-		resource_index = next_metal_resource_ids[desc_set];
+		resource_index = next_metal_resource_ids[var_desc_set];
-		next_metal_resource_ids[desc_set] += binding_stride;
+		next_metal_resource_ids[var_desc_set] += binding_stride;
 	}
 	else
 	{
@@ -12448,6 +12450,7 @@ void CompilerMSL::analyze_argument_buffers()
 		uint32_t plane;
 	};
 	SmallVector<Resource> resources_in_set[kMaxArgumentBuffers];
+	SmallVector<uint32_t> inline_block_vars;
 
 	bool set_needs_swizzle_buffer[kMaxArgumentBuffers] = {};
 	bool set_needs_buffer_sizes[kMaxArgumentBuffers] = {};
@@ -12480,6 +12483,7 @@ void CompilerMSL::analyze_argument_buffers()
 				}
 			}
 
+			uint32_t binding = get_decoration(var_id, DecorationBinding);
 			if (type.basetype == SPIRType::SampledImage)
 			{
 				add_resource_name(var_id);
@@ -12502,9 +12506,14 @@ void CompilerMSL::analyze_argument_buffers()
 					    { &var, to_sampler_expression(var_id), SPIRType::Sampler, sampler_resource_index, 0 });
 				}
 			}
+			else if (inline_uniform_blocks.count(SetBindingPair{ desc_set, binding }))
+			{
+				inline_block_vars.push_back(var_id);
+			}
 			else if (!constexpr_sampler)
 			{
 				// constexpr samplers are not declared as resources.
+				// Inline uniform blocks are always emitted at the end.
 				if (!msl_options.is_ios() || type.basetype != SPIRType::Image || type.image.sampled != 2)
 				{
 					add_resource_name(var_id);
@@ -12579,6 +12588,16 @@ void CompilerMSL::analyze_argument_buffers()
 		}
 	}
 
+	// Now add inline uniform blocks.
+	for (uint32_t var_id : inline_block_vars)
+	{
+		auto &var = get<SPIRVariable>(var_id);
+		uint32_t desc_set = get_decoration(var_id, DecorationDescriptorSet);
+		add_resource_name(var_id);
+		resources_in_set[desc_set].push_back(
+		    { &var, to_name(var_id), SPIRType::Struct, get_metal_resource_index(var, SPIRType::Struct), 0 });
+	}
+
 	for (uint32_t desc_set = 0; desc_set < kMaxArgumentBuffers; desc_set++)
 	{
 		auto &resources = resources_in_set[desc_set];
@@ -12676,6 +12695,12 @@ void CompilerMSL::analyze_argument_buffers()
 					buffer_type.member_types.push_back(var.basetype);
 					buffers_requiring_dynamic_offset[pair].second = var.self;
 				}
+				else if (inline_uniform_blocks.count(pair))
+				{
+					// Put the buffer block itself into the argument buffer.
+					buffer_type.member_types.push_back(get_variable_data_type_id(var));
+					set_qualified_name(var.self, join(to_name(buffer_variable_id), ".", mbr_name));
+				}
 				else
 				{
 					// Resources will be declared as pointers not references, so automatically dereference as appropriate.

spirv_msl.hpp

@@ -432,6 +432,13 @@ public:
 	// an offset taken from the dynamic offset buffer.
 	void add_dynamic_buffer(uint32_t desc_set, uint32_t binding, uint32_t index);
 
+	// desc_set and binding are the SPIR-V descriptor set and binding of a buffer resource
+	// in this shader. This function marks that resource as an inline uniform block
+	// (VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT). This function only has any effect if argument buffers
+	// are enabled. If so, the buffer block will be directly embedded into the argument
+	// buffer, instead of being referenced indirectly via pointer.
+	void add_inline_uniform_block(uint32_t desc_set, uint32_t binding);
+
 	// 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_discrete_descriptor_set(uint32_t desc_set);
@@ -854,6 +861,8 @@ protected:
 	// Must be ordered since array is in a specific order.
 	std::map<SetBindingPair, std::pair<uint32_t, uint32_t>> buffers_requiring_dynamic_offset;
 
+	std::unordered_set<SetBindingPair, InternalHasher> inline_uniform_blocks;
+
 	uint32_t argument_buffer_ids[kMaxArgumentBuffers];
 	uint32_t argument_buffer_discrete_mask = 0;
 	uint32_t argument_buffer_device_storage_mask = 0;

test_shaders.py

@@ -244,6 +244,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 '.inline-block.' in shader:
+        # Arbitrary for testing purposes.
+        msl_args.append('--msl-inline-uniform-block')
+        msl_args.append('0')
+        msl_args.append('0')
     if '.device-argument-buffer.' in shader:
         msl_args.append('--msl-device-argument-buffer')
         msl_args.append('0')