MSL: Fix various bugs with framebuffer fetch on macOS and argument buffers.fix-1582

Introduce a helper to make it clearer if a resource can be
considered for argument buffers or not.

8 files changed, 115 insertions, 18 deletions

diff --git a/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.argument.frag b/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.argument.frag
new file mode 100644
index 00000000..8c7f67b6
--- /dev/null
+++ b/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.argument.frag
@@ -0,0 +1,23 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct spvDescriptorSetBuffer0
+{
+    sampler uSampler [[id(8)]];
+    texture2d<float> uTex [[id(9)]];
+};
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+fragment main0_out main0(constant spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], float4 uSub [[color(1)]])
+{
+    main0_out out = {};
+    out.FragColor = uSub + spvDescriptorSet0.uTex.sample(spvDescriptorSet0.uSampler, float2(0.5));
+    return out;
+}
+
diff --git a/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.frag b/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.frag
new file mode 100644
index 00000000..9108927e
--- /dev/null
+++ b/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.frag
@@ -0,0 +1,17 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+fragment main0_out main0(float4 uSub [[color(1)]], texture2d<float> uTex [[texture(9)]], sampler uSampler [[sampler(8)]])
+{
+    main0_out out = {};
+    out.FragColor = uSub + uTex.sample(uSampler, float2(0.5));
+    return out;
+}
+
diff --git a/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.ios.framebuffer-fetch.msl2.frag b/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.ios.framebuffer-fetch.msl2.frag
new file mode 100644
index 00000000..9108927e
--- /dev/null
+++ b/reference/shaders-msl-no-opt/frag/subpass-input.decoration-binding.ios.framebuffer-fetch.msl2.frag
@@ -0,0 +1,17 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+fragment main0_out main0(float4 uSub [[color(1)]], texture2d<float> uTex [[texture(9)]], sampler uSampler [[sampler(8)]])
+{
+    main0_out out = {};
+    out.FragColor = uSub + uTex.sample(uSampler, float2(0.5));
+    return out;
+}
+
diff --git a/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.argument.frag b/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.argument.frag
new file mode 100644
index 00000000..671a4d1b
--- /dev/null
+++ b/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.argument.frag
@@ -0,0 +1,12 @@
+#version 450
+
+layout(set = 0, binding = 10, input_attachment_index = 1) uniform subpassInput uSub;
+layout(location = 0) out vec4 FragColor;
+
+layout(set = 0, binding = 9) uniform texture2D uTex;
+layout(set = 0, binding = 8) uniform sampler uSampler;
+
+void main()
+{
+	FragColor = subpassLoad(uSub) + texture(sampler2D(uTex, uSampler), vec2(0.5));
+}
diff --git a/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.frag b/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.frag
new file mode 100644
index 00000000..671a4d1b
--- /dev/null
+++ b/shaders-msl-no-opt/frag/subpass-input.decoration-binding.framebuffer-fetch.msl23.frag
@@ -0,0 +1,12 @@
+#version 450
+
+layout(set = 0, binding = 10, input_attachment_index = 1) uniform subpassInput uSub;
+layout(location = 0) out vec4 FragColor;
+
+layout(set = 0, binding = 9) uniform texture2D uTex;
+layout(set = 0, binding = 8) uniform sampler uSampler;
+
+void main()
+{
+	FragColor = subpassLoad(uSub) + texture(sampler2D(uTex, uSampler), vec2(0.5));
+}
diff --git a/shaders-msl-no-opt/frag/subpass-input.decoration-binding.ios.framebuffer-fetch.msl2.frag b/shaders-msl-no-opt/frag/subpass-input.decoration-binding.ios.framebuffer-fetch.msl2.frag
new file mode 100644
index 00000000..671a4d1b
--- /dev/null
+++ b/shaders-msl-no-opt/frag/subpass-input.decoration-binding.ios.framebuffer-fetch.msl2.frag
@@ -0,0 +1,12 @@
+#version 450
+
+layout(set = 0, binding = 10, input_attachment_index = 1) uniform subpassInput uSub;
+layout(location = 0) out vec4 FragColor;
+
+layout(set = 0, binding = 9) uniform texture2D uTex;
+layout(set = 0, binding = 8) uniform sampler uSampler;
+
+void main()
+{
+	FragColor = subpassLoad(uSub) + texture(sampler2D(uTex, uSampler), vec2(0.5));
+}
diff --git a/spirv_msl.cpp b/spirv_msl.cpp
index 1060bb27..a6694570 100644
--- a/spirv_msl.cpp
+++ b/spirv_msl.cpp
@@ -10906,12 +10906,7 @@ void CompilerMSL::entry_point_args_discrete_descriptors(string &ep_args)
 		{
 			auto &type = get_variable_data_type(var);
 
-			// Very specifically, image load-store in argument buffers are disallowed on MSL on iOS.
-			// But we won't know when the argument buffer is encoded whether this image will have
-			// a NonWritable decoration. So just use discrete arguments for all storage images
-			// on iOS.
-			if (!(msl_options.is_ios() && type.basetype == SPIRType::Image && type.image.sampled == 2) &&
-			    var.storage != StorageClassPushConstant)
+			if (is_supported_argument_buffer_type(type) && var.storage != StorageClassPushConstant)
 			{
 				uint32_t desc_set = get_decoration(var_id, DecorationDescriptorSet);
 				if (descriptor_set_is_argument_buffer(desc_set))
@@ -14597,6 +14592,17 @@ bool CompilerMSL::descriptor_set_is_argument_buffer(uint32_t desc_set) const
 	return (argument_buffer_discrete_mask & (1u << desc_set)) == 0;
 }
 
+bool CompilerMSL::is_supported_argument_buffer_type(const SPIRType &type) const
+{
+	// Very specifically, image load-store in argument buffers are disallowed on MSL on iOS.
+	// But we won't know when the argument buffer is encoded whether this image will have
+	// a NonWritable decoration. So just use discrete arguments for all storage images
+	// on iOS.
+	bool is_storage_image = type.basetype == SPIRType::Image && type.image.sampled == 2;
+	bool is_supported_type = !msl_options.is_ios() || !is_storage_image;
+	return !type_is_msl_framebuffer_fetch(type) && is_supported_type;
+}
+
 void CompilerMSL::analyze_argument_buffers()
 {
 	// Gather all used resources and sort them out into argument buffers.
@@ -14679,23 +14685,20 @@ void CompilerMSL::analyze_argument_buffers()
 			{
 				inline_block_vars.push_back(var_id);
 			}
-			else if (!constexpr_sampler)
+			else if (!constexpr_sampler && is_supported_argument_buffer_type(type))
 			{
 				// 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);
+				resources_in_set[desc_set].push_back(
+					{ &var, to_name(var_id), type.basetype, get_metal_resource_index(var, type.basetype), 0 });
+
+				// Emulate texture2D atomic operations
+				if (atomic_image_vars.count(var.self))
 				{
-					add_resource_name(var_id);
+					uint32_t buffer_resource_index = get_metal_resource_index(var, SPIRType::AtomicCounter, 0);
 					resources_in_set[desc_set].push_back(
-					    { &var, to_name(var_id), type.basetype, get_metal_resource_index(var, type.basetype), 0 });
-
-					// Emulate texture2D atomic operations
-					if (atomic_image_vars.count(var.self))
-					{
-						uint32_t buffer_resource_index = get_metal_resource_index(var, SPIRType::AtomicCounter, 0);
-						resources_in_set[desc_set].push_back(
-						    { &var, to_name(var_id) + "_atomic", SPIRType::Struct, buffer_resource_index, 0 });
-					}
+						{ &var, to_name(var_id) + "_atomic", SPIRType::Struct, buffer_resource_index, 0 });
 				}
 			}
 
diff --git a/spirv_msl.hpp b/spirv_msl.hpp
index 07e322c5..8c786123 100644
--- a/spirv_msl.hpp
+++ b/spirv_msl.hpp
@@ -1026,6 +1026,7 @@ protected:
 	void activate_argument_buffer_resources();
 
 	bool type_is_msl_framebuffer_fetch(const SPIRType &type) const;
+	bool is_supported_argument_buffer_type(const SPIRType &type) const;
 
 	// OpcodeHandler that handles several MSL preprocessing operations.
 	struct OpCodePreprocessor : OpcodeHandler