GLSL: Emit nonuniformEXT in correct place for late-combined samplers.

Need to emit nonuniformEXT(sampler2D()) since constructor expressions in
Vulkan GLSL do not propgate the nonuniform qualifier.

9 files changed, 133 insertions, 2 deletions

diff --git a/reference/shaders-hlsl-no-opt/frag/nonuniform-constructor.sm51.fxconly.frag b/reference/shaders-hlsl-no-opt/frag/nonuniform-constructor.sm51.fxconly.frag
new file mode 100644
index 00000000..ca9a116f
--- /dev/null
+++ b/reference/shaders-hlsl-no-opt/frag/nonuniform-constructor.sm51.fxconly.frag
@@ -0,0 +1,32 @@
+Texture2D<float4> uTex[] : register(t0, space0);
+SamplerState Immut : register(s0, space1);
+
+static float4 FragColor;
+static int vIndex;
+static float2 vUV;
+
+struct SPIRV_Cross_Input
+{
+    float2 vUV : TEXCOORD0;
+    nointerpolation int vIndex : TEXCOORD1;
+};
+
+struct SPIRV_Cross_Output
+{
+    float4 FragColor : SV_Target0;
+};
+
+void frag_main()
+{
+    FragColor = uTex[NonUniformResourceIndex(vIndex)].Sample(Immut, vUV);
+}
+
+SPIRV_Cross_Output main(SPIRV_Cross_Input stage_input)
+{
+    vIndex = stage_input.vIndex;
+    vUV = stage_input.vUV;
+    frag_main();
+    SPIRV_Cross_Output stage_output;
+    stage_output.FragColor = FragColor;
+    return stage_output;
+}
diff --git a/reference/shaders-msl-no-opt/frag/nonuniform-constructor.msl2.frag b/reference/shaders-msl-no-opt/frag/nonuniform-constructor.msl2.frag
new file mode 100644
index 00000000..f1ad5c5f
--- /dev/null
+++ b/reference/shaders-msl-no-opt/frag/nonuniform-constructor.msl2.frag
@@ -0,0 +1,23 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using namespace metal;
+
+struct main0_out
+{
+    float4 FragColor [[color(0)]];
+};
+
+struct main0_in
+{
+    float2 vUV [[user(locn0)]];
+    int vIndex [[user(locn1)]];
+};
+
+fragment main0_out main0(main0_in in [[stage_in]], array<texture2d<float>, 10> uTex [[texture(0)]], sampler Immut [[sampler(0)]])
+{
+    main0_out out = {};
+    out.FragColor = uTex[in.vIndex].sample(Immut, in.vUV);
+    return out;
+}
+
diff --git a/reference/shaders-no-opt/asm/frag/nonuniform-qualifier-propagation.vk.nocompat.asm.frag.vk b/reference/shaders-no-opt/asm/frag/nonuniform-qualifier-propagation.vk.nocompat.asm.frag.vk
index 5f7ddeee..289f576f 100644
--- a/reference/shaders-no-opt/asm/frag/nonuniform-qualifier-propagation.vk.nocompat.asm.frag.vk
+++ b/reference/shaders-no-opt/asm/frag/nonuniform-qualifier-propagation.vk.nocompat.asm.frag.vk
@@ -24,7 +24,7 @@ void main()
     int i = vIndex;
     int _59 = i + 10;
     int _64 = i + 40;
-    FragColor = texture(sampler2D(uSamplers[nonuniformEXT(_59)], uSamps[nonuniformEXT(_64)]), vUV);
+    FragColor = texture(nonuniformEXT(sampler2D(uSamplers[_59], uSamps[_64])), vUV);
     int _71 = i + 10;
     FragColor = texture(uCombinedSamplers[nonuniformEXT(_71)], vUV);
     int _77 = i + 20;
diff --git a/reference/shaders-no-opt/frag/nonuniform-constructor.vk.nocompat.frag.vk b/reference/shaders-no-opt/frag/nonuniform-constructor.vk.nocompat.frag.vk
new file mode 100644
index 00000000..ab58862f
--- /dev/null
+++ b/reference/shaders-no-opt/frag/nonuniform-constructor.vk.nocompat.frag.vk
@@ -0,0 +1,15 @@
+#version 450
+#extension GL_EXT_nonuniform_qualifier : require
+
+layout(set = 0, binding = 0) uniform texture2D uTex[];
+layout(set = 1, binding = 0) uniform sampler Immut;
+
+layout(location = 0) out vec4 FragColor;
+layout(location = 1) flat in int vIndex;
+layout(location = 0) in vec2 vUV;
+
+void main()
+{
+    FragColor = texture(nonuniformEXT(sampler2D(uTex[vIndex], Immut)), vUV);
+}
+
diff --git a/shaders-hlsl-no-opt/frag/nonuniform-constructor.sm51.fxconly.frag b/shaders-hlsl-no-opt/frag/nonuniform-constructor.sm51.fxconly.frag
new file mode 100644
index 00000000..452aa953
--- /dev/null
+++ b/shaders-hlsl-no-opt/frag/nonuniform-constructor.sm51.fxconly.frag
@@ -0,0 +1,14 @@
+#version 450
+#extension GL_EXT_nonuniform_qualifier : require
+
+layout(location = 0) out vec4 FragColor;
+layout(location = 0) in vec2 vUV;
+layout(location = 1) flat in int vIndex;
+
+layout(set = 0, binding = 0) uniform texture2D uTex[];
+layout(set = 1, binding = 0) uniform sampler Immut;
+
+void main()
+{
+	FragColor = texture(nonuniformEXT(sampler2D(uTex[vIndex], Immut)), vUV);
+}
diff --git a/shaders-msl-no-opt/frag/nonuniform-constructor.msl2.frag b/shaders-msl-no-opt/frag/nonuniform-constructor.msl2.frag
new file mode 100644
index 00000000..4e0460af
--- /dev/null
+++ b/shaders-msl-no-opt/frag/nonuniform-constructor.msl2.frag
@@ -0,0 +1,14 @@
+#version 450
+#extension GL_EXT_nonuniform_qualifier : require
+
+layout(location = 0) out vec4 FragColor;
+layout(location = 0) in vec2 vUV;
+layout(location = 1) flat in int vIndex;
+
+layout(set = 0, binding = 0) uniform texture2D uTex[10];
+layout(set = 1, binding = 0) uniform sampler Immut;
+
+void main()
+{
+	FragColor = texture(nonuniformEXT(sampler2D(uTex[vIndex], Immut)), vUV);
+}
diff --git a/shaders-no-opt/frag/nonuniform-constructor.vk.nocompat.frag b/shaders-no-opt/frag/nonuniform-constructor.vk.nocompat.frag
new file mode 100644
index 00000000..452aa953
--- /dev/null
+++ b/shaders-no-opt/frag/nonuniform-constructor.vk.nocompat.frag
@@ -0,0 +1,14 @@
+#version 450
+#extension GL_EXT_nonuniform_qualifier : require
+
+layout(location = 0) out vec4 FragColor;
+layout(location = 0) in vec2 vUV;
+layout(location = 1) flat in int vIndex;
+
+layout(set = 0, binding = 0) uniform texture2D uTex[];
+layout(set = 1, binding = 0) uniform sampler Immut;
+
+void main()
+{
+	FragColor = texture(nonuniformEXT(sampler2D(uTex[vIndex], Immut)), vUV);
+}
diff --git a/spirv_glsl.cpp b/spirv_glsl.cpp
index f1653e22..cfa6bab4 100644
--- a/spirv_glsl.cpp
+++ b/spirv_glsl.cpp
@@ -6169,6 +6169,7 @@ std::string CompilerGLSL::to_texture_op(const Instruction &i, bool sparse, bool
 	bool gather = false;
 	bool proj = false;
 	bool fetch = false;
+	bool nonuniform_expression = false;
 	const uint32_t *opt = nullptr;
 
 	auto &result_type = get<SPIRType>(result_type_id);
@@ -6177,7 +6178,17 @@ std::string CompilerGLSL::to_texture_op(const Instruction &i, bool sparse, bool
 
 	// Make sure non-uniform decoration is back-propagated to where it needs to be.
 	if (has_decoration(img, DecorationNonUniformEXT))
-		propagate_nonuniform_qualifier(img);
+	{
+		// In Vulkan GLSL, we cannot back-propgate nonuniform qualifiers if we
+		// use a combined image sampler constructor.
+		// We're only interested in back-propagating if we can trace back through access chains.
+		// If not, we will apply nonuniform to the sampled image expression itself.
+		auto *backing = maybe_get_backing_variable(img);
+		if (backing)
+			propagate_nonuniform_qualifier(img);
+		else
+			nonuniform_expression = true;
+	}
 
 	switch (op)
 	{
@@ -6362,6 +6373,7 @@ std::string CompilerGLSL::to_texture_op(const Instruction &i, bool sparse, bool
 	args.sample = sample;
 	args.sparse_texel = sparse_texel_id;
 	args.min_lod = minlod;
+	args.nonuniform_expression = nonuniform_expression;
 	expr += to_function_args(args, forward);
 	expr += ")";
 
@@ -6550,6 +6562,12 @@ string CompilerGLSL::to_function_args(const TextureFunctionArguments &args, bool
 	else
 		farg_str = to_expression(img);
 
+	if (args.nonuniform_expression && farg_str.find_first_of('[') != string::npos)
+	{
+		// Only emit nonuniformEXT() wrapper if the underlying expression is arrayed in some way.
+		farg_str = join(backend.nonuniform_qualifier, "(", farg_str, ")");
+	}
+
 	bool swizz_func = backend.swizzle_is_function;
 	auto swizzle = [swizz_func](uint32_t comps, uint32_t in_comps) -> const char * {
 		if (comps == in_comps)
diff --git a/spirv_glsl.hpp b/spirv_glsl.hpp
index 3c41efef..add549aa 100644
--- a/spirv_glsl.hpp
+++ b/spirv_glsl.hpp
@@ -405,6 +405,7 @@ protected:
 		uint32_t coord = 0, coord_components = 0, dref = 0;
 		uint32_t grad_x = 0, grad_y = 0, lod = 0, coffset = 0, offset = 0;
 		uint32_t bias = 0, component = 0, sample = 0, sparse_texel = 0, min_lod = 0;
+		bool nonuniform_expression = false;
 	};
 	virtual std::string to_function_args(const TextureFunctionArguments &args, bool *p_forward);