MSL: Deal with texture swizzle on arrays of images.

4 files changed, 332 insertions, 3 deletions

diff --git a/reference/opt/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag b/reference/opt/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag
new file mode 100644
index 00000000..05761604
--- /dev/null
+++ b/reference/opt/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag
@@ -0,0 +1,150 @@
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct spvAux
+{
+    uint swizzleConst[1];
+};
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+struct main0_in
+{
+    float2 vUV [[user(locn0)]];
+};
+
+enum class spvSwizzle : uint
+{
+    none = 0,
+    zero,
+    one,
+    red,
+    green,
+    blue,
+    alpha
+};
+
+template<typename T> struct spvRemoveReference { typedef T type; };
+template<typename T> struct spvRemoveReference<thread T&> { typedef T type; };
+template<typename T> struct spvRemoveReference<thread T&&> { typedef T type; };
+template<typename T> inline constexpr thread T&& spvForward(thread typename spvRemoveReference<T>::type& x)
+{
+    return static_cast<thread T&&>(x);
+}
+template<typename T> inline constexpr thread T&& spvForward(thread typename spvRemoveReference<T>::type&& x)
+{
+    return static_cast<thread T&&>(x);
+}
+
+template<typename T>
+inline T spvGetSwizzle(vec<T, 4> x, T c, spvSwizzle s)
+{
+    switch (s)
+    {
+        case spvSwizzle::none:
+            return c;
+        case spvSwizzle::zero:
+            return 0;
+        case spvSwizzle::one:
+            return 1;
+        case spvSwizzle::red:
+            return x.r;
+        case spvSwizzle::green:
+            return x.g;
+        case spvSwizzle::blue:
+            return x.b;
+        case spvSwizzle::alpha:
+            return x.a;
+    }
+}
+
+// Wrapper function that swizzles texture samples and fetches.
+template<typename T>
+inline vec<T, 4> spvTextureSwizzle(vec<T, 4> x, uint s)
+{
+    if (!s)
+        return x;
+    return vec<T, 4>(spvGetSwizzle(x, x.r, spvSwizzle((s >> 0) & 0xFF)), spvGetSwizzle(x, x.g, spvSwizzle((s >> 8) & 0xFF)), spvGetSwizzle(x, x.b, spvSwizzle((s >> 16) & 0xFF)), spvGetSwizzle(x, x.a, spvSwizzle((s >> 24) & 0xFF)));
+}
+
+template<typename T>
+inline T spvTextureSwizzle(T x, uint s)
+{
+    return spvTextureSwizzle(vec<T, 4>(x, 0, 0, 1), s).x;
+}
+
+// Wrapper function that swizzles texture gathers.
+template<typename T, typename Tex, typename... Ts>
+inline vec<T, 4> spvGatherSwizzle(sampler s, const thread Tex& t, Ts... params, component c, uint sw) METAL_CONST_ARG(c)
+{
+    if (sw)
+    {
+        switch (spvSwizzle((sw >> (uint(c) * 8)) & 0xFF))
+        {
+            case spvSwizzle::none:
+                break;
+            case spvSwizzle::zero:
+                return vec<T, 4>(0, 0, 0, 0);
+            case spvSwizzle::one:
+                return vec<T, 4>(1, 1, 1, 1);
+            case spvSwizzle::red:
+                return t.gather(s, spvForward<Ts>(params)..., component::x);
+            case spvSwizzle::green:
+                return t.gather(s, spvForward<Ts>(params)..., component::y);
+            case spvSwizzle::blue:
+                return t.gather(s, spvForward<Ts>(params)..., component::z);
+            case spvSwizzle::alpha:
+                return t.gather(s, spvForward<Ts>(params)..., component::w);
+        }
+    }
+    switch (c)
+    {
+        case component::x:
+            return t.gather(s, spvForward<Ts>(params)..., component::x);
+        case component::y:
+            return t.gather(s, spvForward<Ts>(params)..., component::y);
+        case component::z:
+            return t.gather(s, spvForward<Ts>(params)..., component::z);
+        case component::w:
+            return t.gather(s, spvForward<Ts>(params)..., component::w);
+    }
+}
+
+// Wrapper function that swizzles depth texture gathers.
+template<typename T, typename Tex, typename... Ts>
+inline vec<T, 4> spvGatherCompareSwizzle(sampler s, const thread Tex& t, Ts... params, uint sw) 
+{
+    if (sw)
+    {
+        switch (spvSwizzle(sw & 0xFF))
+        {
+            case spvSwizzle::none:
+            case spvSwizzle::red:
+                break;
+            case spvSwizzle::zero:
+            case spvSwizzle::green:
+            case spvSwizzle::blue:
+            case spvSwizzle::alpha:
+                return vec<T, 4>(0, 0, 0, 0);
+            case spvSwizzle::one:
+                return vec<T, 4>(1, 1, 1, 1);
+        }
+    }
+    return t.gather_compare(s, spvForward<Ts>(params)...);
+}
+
+fragment main0_out main0(main0_in in [[stage_in]], constant spvAux& spvAuxBuffer [[buffer(30)]], array<texture2d<float>, 4> uSampler [[texture(0)]], array<sampler, 4> uSamplerSmplr [[sampler(0)]])
+{
+    main0_out out = {};
+    constant uint32_t* uSamplerSwzl = &spvAuxBuffer.swizzleConst[0];
+    out.FragColor = spvTextureSwizzle(uSampler[2].sample(uSamplerSmplr[2], in.vUV), uSamplerSwzl[2]);
+    return out;
+}
+
diff --git a/reference/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag b/reference/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag
new file mode 100644
index 00000000..d82b741a
--- /dev/null
+++ b/reference/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag
@@ -0,0 +1,155 @@
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct spvAux
+{
+    uint swizzleConst[1];
+};
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+struct main0_in
+{
+    float2 vUV [[user(locn0)]];
+};
+
+enum class spvSwizzle : uint
+{
+    none = 0,
+    zero,
+    one,
+    red,
+    green,
+    blue,
+    alpha
+};
+
+template<typename T> struct spvRemoveReference { typedef T type; };
+template<typename T> struct spvRemoveReference<thread T&> { typedef T type; };
+template<typename T> struct spvRemoveReference<thread T&&> { typedef T type; };
+template<typename T> inline constexpr thread T&& spvForward(thread typename spvRemoveReference<T>::type& x)
+{
+    return static_cast<thread T&&>(x);
+}
+template<typename T> inline constexpr thread T&& spvForward(thread typename spvRemoveReference<T>::type&& x)
+{
+    return static_cast<thread T&&>(x);
+}
+
+template<typename T>
+inline T spvGetSwizzle(vec<T, 4> x, T c, spvSwizzle s)
+{
+    switch (s)
+    {
+        case spvSwizzle::none:
+            return c;
+        case spvSwizzle::zero:
+            return 0;
+        case spvSwizzle::one:
+            return 1;
+        case spvSwizzle::red:
+            return x.r;
+        case spvSwizzle::green:
+            return x.g;
+        case spvSwizzle::blue:
+            return x.b;
+        case spvSwizzle::alpha:
+            return x.a;
+    }
+}
+
+// Wrapper function that swizzles texture samples and fetches.
+template<typename T>
+inline vec<T, 4> spvTextureSwizzle(vec<T, 4> x, uint s)
+{
+    if (!s)
+        return x;
+    return vec<T, 4>(spvGetSwizzle(x, x.r, spvSwizzle((s >> 0) & 0xFF)), spvGetSwizzle(x, x.g, spvSwizzle((s >> 8) & 0xFF)), spvGetSwizzle(x, x.b, spvSwizzle((s >> 16) & 0xFF)), spvGetSwizzle(x, x.a, spvSwizzle((s >> 24) & 0xFF)));
+}
+
+template<typename T>
+inline T spvTextureSwizzle(T x, uint s)
+{
+    return spvTextureSwizzle(vec<T, 4>(x, 0, 0, 1), s).x;
+}
+
+// Wrapper function that swizzles texture gathers.
+template<typename T, typename Tex, typename... Ts>
+inline vec<T, 4> spvGatherSwizzle(sampler s, const thread Tex& t, Ts... params, component c, uint sw) METAL_CONST_ARG(c)
+{
+    if (sw)
+    {
+        switch (spvSwizzle((sw >> (uint(c) * 8)) & 0xFF))
+        {
+            case spvSwizzle::none:
+                break;
+            case spvSwizzle::zero:
+                return vec<T, 4>(0, 0, 0, 0);
+            case spvSwizzle::one:
+                return vec<T, 4>(1, 1, 1, 1);
+            case spvSwizzle::red:
+                return t.gather(s, spvForward<Ts>(params)..., component::x);
+            case spvSwizzle::green:
+                return t.gather(s, spvForward<Ts>(params)..., component::y);
+            case spvSwizzle::blue:
+                return t.gather(s, spvForward<Ts>(params)..., component::z);
+            case spvSwizzle::alpha:
+                return t.gather(s, spvForward<Ts>(params)..., component::w);
+        }
+    }
+    switch (c)
+    {
+        case component::x:
+            return t.gather(s, spvForward<Ts>(params)..., component::x);
+        case component::y:
+            return t.gather(s, spvForward<Ts>(params)..., component::y);
+        case component::z:
+            return t.gather(s, spvForward<Ts>(params)..., component::z);
+        case component::w:
+            return t.gather(s, spvForward<Ts>(params)..., component::w);
+    }
+}
+
+// Wrapper function that swizzles depth texture gathers.
+template<typename T, typename Tex, typename... Ts>
+inline vec<T, 4> spvGatherCompareSwizzle(sampler s, const thread Tex& t, Ts... params, uint sw) 
+{
+    if (sw)
+    {
+        switch (spvSwizzle(sw & 0xFF))
+        {
+            case spvSwizzle::none:
+            case spvSwizzle::red:
+                break;
+            case spvSwizzle::zero:
+            case spvSwizzle::green:
+            case spvSwizzle::blue:
+            case spvSwizzle::alpha:
+                return vec<T, 4>(0, 0, 0, 0);
+            case spvSwizzle::one:
+                return vec<T, 4>(1, 1, 1, 1);
+        }
+    }
+    return t.gather_compare(s, spvForward<Ts>(params)...);
+}
+
+float4 sample_in_func(thread const array<texture2d<float>, 4> uSampler, thread const array<sampler, 4> uSamplerSmplr, constant uint32_t* uSamplerSwzl, thread float2& vUV)
+{
+    return spvTextureSwizzle(uSampler[2].sample(uSamplerSmplr[2], vUV), uSamplerSwzl[2]);
+}
+
+fragment main0_out main0(main0_in in [[stage_in]], constant spvAux& spvAuxBuffer [[buffer(30)]], array<texture2d<float>, 4> uSampler [[texture(0)]], array<sampler, 4> uSamplerSmplr [[sampler(0)]])
+{
+    main0_out out = {};
+    constant uint32_t* uSamplerSwzl = &spvAuxBuffer.swizzleConst[0];
+    out.FragColor = sample_in_func(uSampler, uSamplerSmplr, uSamplerSwzl, in.vUV);
+    return out;
+}
+
diff --git a/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag b/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag
new file mode 100644
index 00000000..f7351b21
--- /dev/null
+++ b/shaders-msl/frag/array-of-texture-swizzle.msl2.swizzle.frag
@@ -0,0 +1,17 @@
+#version 450
+
+layout(set = 0, binding = 0) uniform sampler2D uSampler[4];
+layout(set = 0, binding = 1) uniform sampler2D uSamp;
+layout(location = 0) in vec2 vUV;
+
+layout(location = 0) out vec4 FragColor;
+
+vec4 sample_in_func()
+{
+	return texture(uSampler[2], vUV);
+}
+
+void main()
+{
+	FragColor = sample_in_func();
+}
diff --git a/spirv_msl.cpp b/spirv_msl.cpp
index aaa65d4b..5a891830 100644
--- a/spirv_msl.cpp
+++ b/spirv_msl.cpp
@@ -4373,7 +4373,10 @@ void CompilerMSL::emit_function_prototype(SPIRFunction &func, const Bitset &)
 
 		// Manufacture automatic swizzle arg.
 		if (msl_options.swizzle_texture_samples && has_sampled_images && is_sampled_image_type(arg_type))
-			decl += join(", constant uint32_t& ", to_swizzle_expression(arg.id));
+		{
+			bool is_array = !arg_type.array.empty();
+			decl += join(", constant uint32_t", is_array ? "* " : "& ", to_swizzle_expression(arg.id));
+		}
 
 		if (&arg != &func.arguments.back())
 			decl += ", ";
@@ -5847,9 +5850,13 @@ void CompilerMSL::fix_up_shader_inputs_outputs()
 			if (msl_options.swizzle_texture_samples && has_sampled_images && is_sampled_image_type(type))
 			{
 				auto &entry_func = this->get<SPIRFunction>(ir.default_entry_point);
-				entry_func.fixup_hooks_in.push_back([this, &var, var_id]() {
+				entry_func.fixup_hooks_in.push_back([this, &type, &var, var_id]() {
 					auto &aux_type = expression_type(aux_buffer_id);
-					statement("constant uint32_t& ", to_swizzle_expression(var_id), " = ", to_name(aux_buffer_id), ".",
+					bool is_array_type = !type.array.empty();
+
+					// If we have an array of images, we need to be able to index into it, so take a pointer instead.
+					statement("constant uint32_t", is_array_type ? "* " : "& ", to_swizzle_expression(var_id),
+					          is_array_type ? " = &" : " = ", to_name(aux_buffer_id), ".",
 					          to_member_name(aux_type, k_aux_mbr_idx_swizzle_const), "[",
 					          convert_to_string(get_metal_resource_index(var, SPIRType::Image)), "];");
 				});