Merge pull request #1366 from KhronosGroup/fix-13652020-05-19

HLSL: Implement image queries for UAV images.

8 files changed, 355 insertions, 69 deletions

diff --git a/reference/opt/shaders-hlsl/frag/image-query-uav.frag b/reference/opt/shaders-hlsl/frag/image-query-uav.frag
new file mode 100644
index 00000000..3b50282f
--- /dev/null
+++ b/reference/opt/shaders-hlsl/frag/image-query-uav.frag
@@ -0,0 +1,8 @@
+void frag_main()
+{
+}
+
+void main()
+{
+    frag_main();
+}
diff --git a/reference/opt/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag b/reference/opt/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag
new file mode 100644
index 00000000..3b50282f
--- /dev/null
+++ b/reference/opt/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag
@@ -0,0 +1,8 @@
+void frag_main()
+{
+}
+
+void main()
+{
+    frag_main();
+}
diff --git a/reference/shaders-hlsl/frag/image-query-uav.frag b/reference/shaders-hlsl/frag/image-query-uav.frag
new file mode 100644
index 00000000..381a9557
--- /dev/null
+++ b/reference/shaders-hlsl/frag/image-query-uav.frag
@@ -0,0 +1,64 @@
+RWTexture1D<float4> uImage1D : register(u0);
+RWTexture2D<float2> uImage2D : register(u1);
+RWTexture2DArray<float> uImage2DArray : register(u2);
+RWTexture3D<unorm float4> uImage3D : register(u3);
+RWBuffer<snorm float4> uImageBuffer : register(u6);
+
+uint3 SPIRV_Cross_imageSize(RWTexture2DArray<float> Tex, out uint Param)
+{
+    uint3 ret;
+    Tex.GetDimensions(ret.x, ret.y, ret.z);
+    Param = 0u;
+    return ret;
+}
+
+uint2 SPIRV_Cross_imageSize(RWTexture2D<float2> Tex, out uint Param)
+{
+    uint2 ret;
+    Tex.GetDimensions(ret.x, ret.y);
+    Param = 0u;
+    return ret;
+}
+
+uint SPIRV_Cross_imageSize(RWTexture1D<float4> Tex, out uint Param)
+{
+    uint ret;
+    Tex.GetDimensions(ret.x);
+    Param = 0u;
+    return ret;
+}
+
+uint3 SPIRV_Cross_imageSize(RWTexture3D<unorm float4> Tex, out uint Param)
+{
+    uint3 ret;
+    Tex.GetDimensions(ret.x, ret.y, ret.z);
+    Param = 0u;
+    return ret;
+}
+
+uint SPIRV_Cross_imageSize(RWBuffer<snorm float4> Tex, out uint Param)
+{
+    uint ret;
+    Tex.GetDimensions(ret.x);
+    Param = 0u;
+    return ret;
+}
+
+void frag_main()
+{
+    uint _14_dummy_parameter;
+    int a = int(SPIRV_Cross_imageSize(uImage1D, _14_dummy_parameter));
+    uint _22_dummy_parameter;
+    int2 b = int2(SPIRV_Cross_imageSize(uImage2D, _22_dummy_parameter));
+    uint _30_dummy_parameter;
+    int3 c = int3(SPIRV_Cross_imageSize(uImage2DArray, _30_dummy_parameter));
+    uint _36_dummy_parameter;
+    int3 d = int3(SPIRV_Cross_imageSize(uImage3D, _36_dummy_parameter));
+    uint _42_dummy_parameter;
+    int e = int(SPIRV_Cross_imageSize(uImageBuffer, _42_dummy_parameter));
+}
+
+void main()
+{
+    frag_main();
+}
diff --git a/reference/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag b/reference/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag
new file mode 100644
index 00000000..5db5d359
--- /dev/null
+++ b/reference/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag
@@ -0,0 +1,63 @@
+RWTexture1D<float4> uImage1D : register(u0);
+RWTexture2D<float2> uImage2D : register(u1);
+Texture2DArray<float4> uImage2DArray : register(t2);
+RWTexture3D<unorm float4> uImage3D : register(u3);
+RWBuffer<snorm float4> uImageBuffer : register(u6);
+
+uint3 SPIRV_Cross_textureSize(Texture2DArray<float4> Tex, uint Level, out uint Param)
+{
+    uint3 ret;
+    Tex.GetDimensions(Level, ret.x, ret.y, ret.z, Param);
+    return ret;
+}
+
+uint2 SPIRV_Cross_imageSize(RWTexture2D<float2> Tex, out uint Param)
+{
+    uint2 ret;
+    Tex.GetDimensions(ret.x, ret.y);
+    Param = 0u;
+    return ret;
+}
+
+uint SPIRV_Cross_imageSize(RWTexture1D<float4> Tex, out uint Param)
+{
+    uint ret;
+    Tex.GetDimensions(ret.x);
+    Param = 0u;
+    return ret;
+}
+
+uint3 SPIRV_Cross_imageSize(RWTexture3D<unorm float4> Tex, out uint Param)
+{
+    uint3 ret;
+    Tex.GetDimensions(ret.x, ret.y, ret.z);
+    Param = 0u;
+    return ret;
+}
+
+uint SPIRV_Cross_imageSize(RWBuffer<snorm float4> Tex, out uint Param)
+{
+    uint ret;
+    Tex.GetDimensions(ret.x);
+    Param = 0u;
+    return ret;
+}
+
+void frag_main()
+{
+    uint _14_dummy_parameter;
+    int a = int(SPIRV_Cross_imageSize(uImage1D, _14_dummy_parameter));
+    uint _22_dummy_parameter;
+    int2 b = int2(SPIRV_Cross_imageSize(uImage2D, _22_dummy_parameter));
+    uint _30_dummy_parameter;
+    int3 c = int3(SPIRV_Cross_textureSize(uImage2DArray, 0u, _30_dummy_parameter));
+    uint _36_dummy_parameter;
+    int3 d = int3(SPIRV_Cross_imageSize(uImage3D, _36_dummy_parameter));
+    uint _42_dummy_parameter;
+    int e = int(SPIRV_Cross_imageSize(uImageBuffer, _42_dummy_parameter));
+}
+
+void main()
+{
+    frag_main();
+}
diff --git a/shaders-hlsl/frag/image-query-uav.frag b/shaders-hlsl/frag/image-query-uav.frag
new file mode 100644
index 00000000..25103e6e
--- /dev/null
+++ b/shaders-hlsl/frag/image-query-uav.frag
@@ -0,0 +1,18 @@
+#version 450
+
+layout(rgba32f, binding = 0) uniform writeonly image1D uImage1D;
+layout(rg32f, binding = 1) uniform writeonly image2D uImage2D;
+layout(r32f, binding = 2) uniform readonly image2DArray uImage2DArray;
+layout(rgba8, binding = 3) uniform writeonly image3D uImage3D;
+layout(rgba8_snorm, binding = 6) uniform writeonly imageBuffer uImageBuffer;
+
+// There is no RWTexture2DMS.
+
+void main()
+{
+	int a = imageSize(uImage1D);
+	ivec2 b = imageSize(uImage2D);
+	ivec3 c = imageSize(uImage2DArray);
+	ivec3 d = imageSize(uImage3D);
+	int e = imageSize(uImageBuffer);
+}
diff --git a/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag b/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag
new file mode 100644
index 00000000..25103e6e
--- /dev/null
+++ b/shaders-hlsl/frag/image-query-uav.nonwritable-uav-texture.frag
@@ -0,0 +1,18 @@
+#version 450
+
+layout(rgba32f, binding = 0) uniform writeonly image1D uImage1D;
+layout(rg32f, binding = 1) uniform writeonly image2D uImage2D;
+layout(r32f, binding = 2) uniform readonly image2DArray uImage2DArray;
+layout(rgba8, binding = 3) uniform writeonly image3D uImage3D;
+layout(rgba8_snorm, binding = 6) uniform writeonly imageBuffer uImageBuffer;
+
+// There is no RWTexture2DMS.
+
+void main()
+{
+	int a = imageSize(uImage1D);
+	ivec2 b = imageSize(uImage2D);
+	ivec3 c = imageSize(uImage2DArray);
+	ivec3 d = imageSize(uImage3D);
+	int e = imageSize(uImageBuffer);
+}
diff --git a/spirv_hlsl.cpp b/spirv_hlsl.cpp
index f27163c2..680d3d79 100644
--- a/spirv_hlsl.cpp
+++ b/spirv_hlsl.cpp
@@ -23,6 +23,41 @@ using namespace spv;
 using namespace SPIRV_CROSS_NAMESPACE;
 using namespace std;
 
+enum class ImageFormatNormalizedState
+{
+	None = 0,
+	Unorm = 1,
+	Snorm = 2
+};
+
+static ImageFormatNormalizedState image_format_to_normalized_state(ImageFormat fmt)
+{
+	switch (fmt)
+	{
+	case ImageFormatR8:
+	case ImageFormatR16:
+	case ImageFormatRg8:
+	case ImageFormatRg16:
+	case ImageFormatRgba8:
+	case ImageFormatRgba16:
+	case ImageFormatRgb10A2:
+		return ImageFormatNormalizedState::Unorm;
+
+	case ImageFormatR8Snorm:
+	case ImageFormatR16Snorm:
+	case ImageFormatRg8Snorm:
+	case ImageFormatRg16Snorm:
+	case ImageFormatRgba8Snorm:
+	case ImageFormatRgba16Snorm:
+		return ImageFormatNormalizedState::Snorm;
+
+	default:
+		break;
+	}
+
+	return ImageFormatNormalizedState::None;
+}
+
 static unsigned image_format_to_components(ImageFormat fmt)
 {
 	switch (fmt)
@@ -1423,66 +1458,14 @@ void CompilerHLSL::emit_resources()
 		}
 	}
 
-	if (required_textureSizeVariants != 0)
+	emit_texture_size_variants(required_texture_size_variants.srv, "4", false, "");
+	for (uint32_t norm = 0; norm < 3; norm++)
 	{
-		static const char *types[QueryTypeCount] = { "float4", "int4", "uint4" };
-		static const char *dims[QueryDimCount] = { "Texture1D",   "Texture1DArray",  "Texture2D",   "Texture2DArray",
-			                                       "Texture3D",   "Buffer",          "TextureCube", "TextureCubeArray",
-			                                       "Texture2DMS", "Texture2DMSArray" };
-
-		static const bool has_lod[QueryDimCount] = { true, true, true, true, true, false, true, true, false, false };
-
-		static const char *ret_types[QueryDimCount] = {
-			"uint", "uint2", "uint2", "uint3", "uint3", "uint", "uint2", "uint3", "uint2", "uint3",
-		};
-
-		static const uint32_t return_arguments[QueryDimCount] = {
-			1, 2, 2, 3, 3, 1, 2, 3, 2, 3,
-		};
-
-		for (uint32_t index = 0; index < QueryDimCount; index++)
+		for (uint32_t comp = 0; comp < 4; comp++)
 		{
-			for (uint32_t type_index = 0; type_index < QueryTypeCount; type_index++)
-			{
-				uint32_t bit = 16 * type_index + index;
-				uint64_t mask = 1ull << bit;
-
-				if ((required_textureSizeVariants & mask) == 0)
-					continue;
-
-				statement(ret_types[index], " SPIRV_Cross_textureSize(", dims[index], "<", types[type_index],
-				          "> Tex, uint Level, out uint Param)");
-				begin_scope();
-				statement(ret_types[index], " ret;");
-				switch (return_arguments[index])
-				{
-				case 1:
-					if (has_lod[index])
-						statement("Tex.GetDimensions(Level, ret.x, Param);");
-					else
-					{
-						statement("Tex.GetDimensions(ret.x);");
-						statement("Param = 0u;");
-					}
-					break;
-				case 2:
-					if (has_lod[index])
-						statement("Tex.GetDimensions(Level, ret.x, ret.y, Param);");
-					else
-						statement("Tex.GetDimensions(ret.x, ret.y, Param);");
-					break;
-				case 3:
-					if (has_lod[index])
-						statement("Tex.GetDimensions(Level, ret.x, ret.y, ret.z, Param);");
-					else
-						statement("Tex.GetDimensions(ret.x, ret.y, ret.z, Param);");
-					break;
-				}
-
-				statement("return ret;");
-				end_scope();
-				statement("");
-			}
+			static const char *qualifiers[] = { "", "unorm ", "snorm " };
+			static const char *vecsizes[] = { "", "2", "3", "4" };
+			emit_texture_size_variants(required_texture_size_variants.uav[norm][comp], vecsizes[comp], true, qualifiers[norm]);
 		}
 	}
 
@@ -1845,6 +1828,83 @@ void CompilerHLSL::emit_resources()
 	}
 }
 
+void CompilerHLSL::emit_texture_size_variants(uint64_t variant_mask, const char *vecsize_qualifier, bool uav, const char *type_qualifier)
+{
+	if (variant_mask == 0)
+		return;
+
+	static const char *types[QueryTypeCount] = { "float", "int", "uint" };
+	static const char *dims[QueryDimCount] = { "Texture1D",   "Texture1DArray",  "Texture2D",   "Texture2DArray",
+	                                           "Texture3D",   "Buffer",          "TextureCube", "TextureCubeArray",
+	                                           "Texture2DMS", "Texture2DMSArray" };
+
+	static const bool has_lod[QueryDimCount] = { true, true, true, true, true, false, true, true, false, false };
+
+	static const char *ret_types[QueryDimCount] = {
+		"uint", "uint2", "uint2", "uint3", "uint3", "uint", "uint2", "uint3", "uint2", "uint3",
+	};
+
+	static const uint32_t return_arguments[QueryDimCount] = {
+		1, 2, 2, 3, 3, 1, 2, 3, 2, 3,
+	};
+
+	for (uint32_t index = 0; index < QueryDimCount; index++)
+	{
+		for (uint32_t type_index = 0; type_index < QueryTypeCount; type_index++)
+		{
+			uint32_t bit = 16 * type_index + index;
+			uint64_t mask = 1ull << bit;
+
+			if ((variant_mask & mask) == 0)
+				continue;
+
+			statement(ret_types[index], " SPIRV_Cross_", (uav ? "image" : "texture"), "Size(", (uav ? "RW" : ""),
+			          dims[index], "<", type_qualifier, types[type_index], vecsize_qualifier,
+					  "> Tex, ", (uav ? "" : "uint Level, "), "out uint Param)");
+			begin_scope();
+			statement(ret_types[index], " ret;");
+			switch (return_arguments[index])
+			{
+			case 1:
+				if (has_lod[index] && !uav)
+					statement("Tex.GetDimensions(Level, ret.x, Param);");
+				else
+				{
+					statement("Tex.GetDimensions(ret.x);");
+					statement("Param = 0u;");
+				}
+				break;
+			case 2:
+				if (has_lod[index] && !uav)
+					statement("Tex.GetDimensions(Level, ret.x, ret.y, Param);");
+				else if (!uav)
+					statement("Tex.GetDimensions(ret.x, ret.y, Param);");
+				else
+				{
+					statement("Tex.GetDimensions(ret.x, ret.y);");
+					statement("Param = 0u;");
+				}
+				break;
+			case 3:
+				if (has_lod[index] && !uav)
+					statement("Tex.GetDimensions(Level, ret.x, ret.y, ret.z, Param);");
+				else if (!uav)
+					statement("Tex.GetDimensions(ret.x, ret.y, ret.z, Param);");
+				else
+				{
+					statement("Tex.GetDimensions(ret.x, ret.y, ret.z);");
+					statement("Param = 0u;");
+				}
+				break;
+			}
+
+			statement("return ret;");
+			end_scope();
+			statement("");
+		}
+	}
+}
+
 string CompilerHLSL::layout_for_member(const SPIRType &type, uint32_t index)
 {
 	auto &flags = get_member_decoration_bitset(type.self, index);
@@ -4834,8 +4894,7 @@ void CompilerHLSL::emit_instruction(const Instruction &instruction)
 		auto result_type = ops[0];
 		auto id = ops[1];
 
-		require_texture_query_variant(expression_type(ops[2]));
-
+		require_texture_query_variant(ops[2]);
 		auto dummy_samples_levels = join(get_fallback_name(id), "_dummy_parameter");
 		statement("uint ", dummy_samples_levels, ";");
 
@@ -4853,12 +4912,22 @@ void CompilerHLSL::emit_instruction(const Instruction &instruction)
 		auto result_type = ops[0];
 		auto id = ops[1];
 
-		require_texture_query_variant(expression_type(ops[2]));
+		require_texture_query_variant(ops[2]);
+		bool uav = expression_type(ops[2]).image.sampled == 2;
+
+		if (const auto *var = maybe_get_backing_variable(ops[2]))
+			if (hlsl_options.nonwritable_uav_texture_as_srv && has_decoration(var->self, DecorationNonWritable))
+				uav = false;
 
 		auto dummy_samples_levels = join(get_fallback_name(id), "_dummy_parameter");
 		statement("uint ", dummy_samples_levels, ";");
 
-		auto expr = join("SPIRV_Cross_textureSize(", to_expression(ops[2]), ", 0u, ", dummy_samples_levels, ")");
+		string expr;
+		if (uav)
+			expr = join("SPIRV_Cross_imageSize(", to_expression(ops[2]), ", ", dummy_samples_levels, ")");
+		else
+			expr = join("SPIRV_Cross_textureSize(", to_expression(ops[2]), ", 0u, ", dummy_samples_levels, ")");
+
 		auto &restype = get<SPIRType>(ops[0]);
 		expr = bitcast_expression(restype, SPIRType::UInt, expr);
 		emit_op(result_type, id, expr, true);
@@ -4871,14 +4940,25 @@ void CompilerHLSL::emit_instruction(const Instruction &instruction)
 		auto result_type = ops[0];
 		auto id = ops[1];
 
-		require_texture_query_variant(expression_type(ops[2]));
+		require_texture_query_variant(ops[2]);
+		bool uav = expression_type(ops[2]).image.sampled == 2;
+		if (opcode == OpImageQueryLevels && uav)
+			SPIRV_CROSS_THROW("Cannot query levels for UAV images.");
+
+		if (const auto *var = maybe_get_backing_variable(ops[2]))
+			if (hlsl_options.nonwritable_uav_texture_as_srv && has_decoration(var->self, DecorationNonWritable))
+				uav = false;
 
 		// Keep it simple and do not emit special variants to make this look nicer ...
 		// This stuff is barely, if ever, used.
 		forced_temporaries.insert(id);
 		auto &type = get<SPIRType>(result_type);
 		statement(variable_decl(type, to_name(id)), ";");
-		statement("SPIRV_Cross_textureSize(", to_expression(ops[2]), ", 0u, ", to_name(id), ");");
+
+		if (uav)
+			statement("SPIRV_Cross_imageSize(", to_expression(ops[2]), ", ", to_name(id), ");");
+		else
+			statement("SPIRV_Cross_textureSize(", to_expression(ops[2]), ", 0u, ", to_name(id), ");");
 
 		auto &restype = get<SPIRType>(ops[0]);
 		auto expr = bitcast_expression(restype, SPIRType::UInt, to_name(id));
@@ -5192,8 +5272,16 @@ void CompilerHLSL::emit_instruction(const Instruction &instruction)
 	}
 }
 
-void CompilerHLSL::require_texture_query_variant(const SPIRType &type)
+void CompilerHLSL::require_texture_query_variant(uint32_t var_id)
 {
+	if (const auto *var = maybe_get_backing_variable(var_id))
+		var_id = var->self;
+
+	auto &type = expression_type(var_id);
+	bool uav = type.image.sampled == 2;
+	if (hlsl_options.nonwritable_uav_texture_as_srv && has_decoration(var_id, DecorationNonWritable))
+		uav = false;
+
 	uint32_t bit = 0;
 	switch (type.image.dim)
 	{
@@ -5242,11 +5330,15 @@ void CompilerHLSL::require_texture_query_variant(const SPIRType &type)
 		SPIRV_CROSS_THROW("Unsupported query type.");
 	}
 
+	auto norm_state = image_format_to_normalized_state(type.image.format);
+	auto &variant = uav ? required_texture_size_variants.uav[uint32_t(norm_state)][image_format_to_components(type.image.format) - 1] :
+	                required_texture_size_variants.srv;
+
 	uint64_t mask = 1ull << bit;
-	if ((required_textureSizeVariants & mask) == 0)
+	if ((variant & mask) == 0)
 	{
 		force_recompile();
-		required_textureSizeVariants |= mask;
+		variant |= mask;
 	}
 }
 
diff --git a/spirv_hlsl.hpp b/spirv_hlsl.hpp
index e98c6443..fa3b656f 100644
--- a/spirv_hlsl.hpp
+++ b/spirv_hlsl.hpp
@@ -264,8 +264,23 @@ private:
 	bool requires_scalar_reflect = false;
 	bool requires_scalar_refract = false;
 	bool requires_scalar_faceforward = false;
-	uint64_t required_textureSizeVariants = 0;
-	void require_texture_query_variant(const SPIRType &type);
+
+	struct TextureSizeVariants
+	{
+		// MSVC 2013 workaround.
+		TextureSizeVariants()
+		{
+			srv = 0;
+			for (auto &unorm : uav)
+				for (auto &u : unorm)
+					u = 0;
+		}
+		uint64_t srv;
+		uint64_t uav[3][4];
+	} required_texture_size_variants;
+
+	void require_texture_query_variant(uint32_t var_id);
+	void emit_texture_size_variants(uint64_t variant_mask, const char *vecsize_qualifier, bool uav, const char *type_qualifier);
 
 	enum TextureQueryVariantDim
 	{