Add generic input/output remappers.

The shader interface between mesh shaders and pixel shaders is based
on semantics rather than packing locations, so we can use these to
capture the mesh shader output signature and forward it when compiling
the pixel shader.

4 files changed, 167 insertions, 19 deletions

diff --git a/dxil_converter.cpp b/dxil_converter.cpp
index 9d2eee1..79ec4f2 100644
--- a/dxil_converter.cpp
+++ b/dxil_converter.cpp
@@ -3529,9 +3529,24 @@ bool Converter::Impl::emit_patch_variables()
 			// The offset is deduced from the control point I/O signature.
 			// TODO: If it's possible to omit trailing CP members in domain shader, we will need to pass this offset
 			// into the compiler.
-			builder.addDecoration(variable_id, spv::DecorationLocation, start_row + patch_location_offset);
-			if (start_col != 0)
-				builder.addDecoration(variable_id, spv::DecorationComponent, start_col);
+			VulkanStageIO vk_io = { start_row + patch_location_offset, start_col, true };
+
+			if (resource_mapping_iface)
+			{
+				D3DStageIO d3d_io = { semantic_name.c_str(), semantic_index, start_row, rows };
+
+				if (execution_model == spv::ExecutionModelTessellationEvaluation)
+				{
+					if (!resource_mapping_iface->remap_stage_input(d3d_io, vk_io))
+						return false;
+				}
+				else if (!resource_mapping_iface->remap_stage_output(d3d_io, vk_io))
+					return false;
+			}
+
+			builder.addDecoration(variable_id, spv::DecorationLocation, vk_io.location);
+			if (vk_io.component != 0)
+				builder.addDecoration(variable_id, spv::DecorationComponent, vk_io.component);
 		}
 
 		builder.addDecoration(variable_id, execution_model == spv::ExecutionModelMeshEXT
@@ -3772,18 +3787,27 @@ bool Converter::Impl::emit_stage_output_variables()
 		{
 			if (execution_model == spv::ExecutionModelVertex ||
 			    execution_model == spv::ExecutionModelTessellationEvaluation ||
-			    execution_model == spv::ExecutionModelGeometry)
+			    execution_model == spv::ExecutionModelGeometry ||
+					execution_model == spv::ExecutionModelMeshEXT)
 			{
 				emit_interpolation_decorations(variable_id, interpolation);
 			}
 
-			unsigned effective_start_row = start_row;
+			VulkanStageIO vk_output = { start_row, start_col };
+
 			if (execution_model == spv::ExecutionModelGeometry && geometry_stream < 4)
-				effective_start_row += start_row_for_geometry_stream[geometry_stream];
+				vk_output.location += start_row_for_geometry_stream[geometry_stream];
 
-			builder.addDecoration(variable_id, spv::DecorationLocation, effective_start_row);
-			if (start_col != 0)
-				builder.addDecoration(variable_id, spv::DecorationComponent, start_col);
+			if (resource_mapping_iface)
+			{
+				D3DStageIO d3d_output = { semantic_name.c_str(), semantic_index, start_row, rows };
+				if (!resource_mapping_iface->remap_stage_output(d3d_output, vk_output))
+					return false;
+			}
+
+			builder.addDecoration(variable_id, spv::DecorationLocation, vk_output.location);
+			if (vk_output.component != 0)
+				builder.addDecoration(variable_id, spv::DecorationComponent, vk_output.component);
 		}
 	}
 
@@ -4381,18 +4405,29 @@ bool Converter::Impl::emit_stage_input_variables()
 			if (execution_model == spv::ExecutionModelFragment)
 				emit_interpolation_decorations(variable_id, interpolation);
 
-			VulkanVertexInput vk_input = { start_row };
-			if (execution_model == spv::ExecutionModelVertex && resource_mapping_iface)
+			VulkanStageIO vk_input = { start_row, start_col };
+
+			if (resource_mapping_iface)
 			{
-				D3DVertexInput d3d_input = { semantic_name.c_str(), semantic_index, start_row, rows };
-				if (!resource_mapping_iface->remap_vertex_input(d3d_input, vk_input))
+				D3DStageIO d3d_input = { semantic_name.c_str(), semantic_index, start_row, rows };
+
+				if (execution_model == spv::ExecutionModelVertex)
+				{
+					if (!resource_mapping_iface->remap_vertex_input(d3d_input, vk_input))
+						return false;
+				}
+
+				if (!resource_mapping_iface->remap_stage_input(d3d_input, vk_input))
 					return false;
 			}
 
 			builder.addDecoration(variable_id, spv::DecorationLocation, vk_input.location);
 
-			if (execution_model != spv::ExecutionModelVertex && start_col != 0)
-				builder.addDecoration(variable_id, spv::DecorationComponent, start_col);
+			if (execution_model != spv::ExecutionModelVertex && vk_input.component != 0)
+				builder.addDecoration(variable_id, spv::DecorationComponent, vk_input.component);
+
+			if (execution_model == spv::ExecutionModelFragment && (vk_input.flags & STAGE_IO_PER_PRIMITIVE))
+				builder.addDecoration(variable_id, spv::DecorationPerPrimitiveEXT);
 		}
 	}
 
diff --git a/dxil_converter.hpp b/dxil_converter.hpp
index 15b5cf8..8986716 100644
--- a/dxil_converter.hpp
+++ b/dxil_converter.hpp
@@ -158,7 +158,7 @@ struct VulkanCBVBinding
 	bool push_constant;
 };
 
-struct D3DVertexInput
+struct D3DStageIO
 {
 	const char *semantic;
 	unsigned semantic_index;
@@ -166,9 +166,19 @@ struct D3DVertexInput
 	unsigned rows;
 };
 
-struct VulkanVertexInput
+enum VulkanStageIoFlagBits
+{
+	STAGE_IO_NONE = 0u,
+	STAGE_IO_PER_PRIMITIVE = 0x1u,
+};
+
+using VulkanStageIoFlags = unsigned;
+
+struct VulkanStageIO
 {
 	unsigned location;
+	unsigned component;
+	VulkanStageIoFlags flags;
 };
 
 struct D3DStreamOutput
@@ -193,8 +203,10 @@ public:
 	virtual bool remap_sampler(const D3DBinding &d3d_binding, VulkanBinding &vulkan_binding) = 0;
 	virtual bool remap_uav(const D3DUAVBinding &d3d_binding, VulkanUAVBinding &vulkan_binding) = 0;
 	virtual bool remap_cbv(const D3DBinding &d3d_binding, VulkanCBVBinding &vulkan_binding) = 0;
-	virtual bool remap_vertex_input(const D3DVertexInput &d3d_input, VulkanVertexInput &vulkan_location) = 0;
+	virtual bool remap_vertex_input(const D3DStageIO &d3d_input, VulkanStageIO &vulkan_location) = 0;
 	virtual bool remap_stream_output(const D3DStreamOutput &d3d_output, VulkanStreamOutput &vulkan_output) = 0;
+	virtual bool remap_stage_input(const D3DStageIO &d3d_input, VulkanStageIO &vk_input) = 0;
+	virtual bool remap_stage_output(const D3DStageIO &d3d_output, VulkanStageIO &vk_output) = 0;
 	virtual unsigned get_root_constant_word_count() = 0;
 	virtual unsigned get_root_descriptor_count() = 0;
 };
diff --git a/dxil_spirv_c.cpp b/dxil_spirv_c.cpp
index 53ddca5..fc0e150 100644
--- a/dxil_spirv_c.cpp
+++ b/dxil_spirv_c.cpp
@@ -203,7 +203,7 @@ struct Remapper : ResourceRemappingInterface
 		}
 	}
 
-	bool remap_vertex_input(const D3DVertexInput &d3d_input, VulkanVertexInput &vk_input) override
+	bool remap_vertex_input(const D3DStageIO &d3d_input, VulkanStageIO &vk_input) override
 	{
 		dxil_spv_d3d_vertex_input c_input = { d3d_input.semantic, d3d_input.semantic_index, d3d_input.start_row,
 			                                  d3d_input.rows };
@@ -214,6 +214,7 @@ struct Remapper : ResourceRemappingInterface
 			if (input_remapper(input_userdata, &c_input, &c_vk_input) == DXIL_SPV_TRUE)
 			{
 				vk_input.location = c_vk_input.location;
+				vk_input.component = 0;
 				return true;
 			}
 			else
@@ -250,6 +251,52 @@ struct Remapper : ResourceRemappingInterface
 		}
 	}
 
+	bool remap_stage_input(const D3DStageIO &d3d_input, VulkanStageIO &vk_input) override
+	{
+		dxil_spv_d3d_shader_stage_io c_input = { d3d_input.semantic, d3d_input.semantic_index };
+		dxil_spv_vulkan_shader_stage_io c_vk_input = { vk_input.location, vk_input.component, vk_input.flags };
+
+		if (stage_input_remapper)
+		{
+			if (stage_input_remapper(stage_input_userdata, &c_input, &c_vk_input) == DXIL_SPV_TRUE)
+			{
+				vk_input.location = c_vk_input.location;
+				vk_input.component = c_vk_input.component;
+				vk_input.flags = c_vk_input.flags;
+				return true;
+			}
+			else
+				return false;
+		}
+		else
+		{
+			return true;
+		}
+	}
+
+	bool remap_stage_output(const D3DStageIO &d3d_output, VulkanStageIO &vk_output) override
+	{
+		dxil_spv_d3d_shader_stage_io c_output = { d3d_output.semantic, d3d_output.semantic_index };
+		dxil_spv_vulkan_shader_stage_io c_vk_output = { vk_output.location, vk_output.component, vk_output.flags };
+
+		if (stage_output_remapper)
+		{
+			if (stage_output_remapper(stage_output_userdata, &c_output, &c_vk_output) == DXIL_SPV_TRUE)
+			{
+				vk_output.location = c_vk_output.location;
+				vk_output.component = c_vk_output.component;
+				vk_output.flags = c_vk_output.flags;
+				return true;
+			}
+			else
+				return false;
+		}
+		else
+		{
+			return true;
+		}
+	}
+
 	unsigned get_root_constant_word_count() override
 	{
 		return root_constant_word_count;
@@ -278,6 +325,12 @@ struct Remapper : ResourceRemappingInterface
 	dxil_spv_stream_output_remapper_cb output_remapper = nullptr;
 	void *output_userdata = nullptr;
 
+	dxil_spv_shader_stage_io_remapper_cb stage_input_remapper = nullptr;
+	void *stage_input_userdata = nullptr;
+
+	dxil_spv_shader_stage_io_remapper_cb stage_output_remapper = nullptr;
+	void *stage_output_userdata = nullptr;
+
 	unsigned root_constant_word_count = 0;
 	unsigned root_descriptor_count = 0;
 };
@@ -658,6 +711,20 @@ void dxil_spv_converter_set_cbv_remapper(dxil_spv_converter converter, dxil_spv_
 	converter->remapper.cbv_userdata = userdata;
 }
 
+void dxil_spv_converter_set_stage_input_remapper(dxil_spv_converter converter,
+                                              dxil_spv_shader_stage_io_remapper_cb remapper, void *userdata)
+{
+	converter->remapper.stage_input_remapper = remapper;
+	converter->remapper.stage_input_userdata = userdata;
+}
+
+void dxil_spv_converter_set_stage_output_remapper(dxil_spv_converter converter,
+                                              dxil_spv_shader_stage_io_remapper_cb remapper, void *userdata)
+{
+	converter->remapper.stage_output_remapper = remapper;
+	converter->remapper.stage_output_userdata = userdata;
+}
+
 void dxil_spv_converter_set_vertex_input_remapper(dxil_spv_converter converter,
                                                   dxil_spv_vertex_input_remapper_cb remapper, void *userdata)
 {
diff --git a/dxil_spirv_c.h b/dxil_spirv_c.h
index 548586e..2f51c9c 100644
--- a/dxil_spirv_c.h
+++ b/dxil_spirv_c.h
@@ -126,6 +126,30 @@ typedef struct dxil_spv_vulkan_stream_output
 typedef dxil_spv_bool (*dxil_spv_stream_output_remapper_cb)(void *userdata, const dxil_spv_d3d_stream_output *d3d_output,
                                                             dxil_spv_vulkan_stream_output *vulkan_output);
 
+typedef struct dxil_spv_d3d_shader_stage_io
+{
+	const char *semantic;
+	unsigned semantic_index;
+	unsigned start_row;
+	unsigned rows;
+} dxil_spv_d3d_shader_stage_io;
+
+enum dxil_spv_vulkan_shader_stage_io_flags
+{
+	DXIL_SPV_SHADER_STAGE_IO_NONE = 0u,
+	DXIL_SPV_SHADER_STAGE_IO_PER_PRIMITIVE = 0x1u,
+};
+
+typedef struct dxil_spv_vulkan_shader_stage_io
+{
+	unsigned location;
+	unsigned component;
+	unsigned flags;
+} dxil_spv_vulkan_shader_stage_io;
+
+typedef dxil_spv_bool (*dxil_spv_shader_stage_io_remapper_cb)(void *userdata, const dxil_spv_d3d_shader_stage_io *d3d_input,
+                                                              dxil_spv_vulkan_shader_stage_io *vulkan_variable);
+
 /* Matches DXIL enum */
 typedef enum dxil_spv_resource_kind
 {
@@ -560,6 +584,16 @@ DXIL_SPV_PUBLIC_API void dxil_spv_converter_free(dxil_spv_converter converter);
 
 DXIL_SPV_PUBLIC_API void dxil_spv_converter_set_entry_point(dxil_spv_converter converter, const char *entry_point);
 
+DXIL_SPV_PUBLIC_API void dxil_spv_converter_set_stage_input_remapper(
+		dxil_spv_converter converter,
+		dxil_spv_shader_stage_io_remapper_cb remapper,
+		void *userdata);
+
+DXIL_SPV_PUBLIC_API void dxil_spv_converter_set_stage_output_remapper(
+		dxil_spv_converter converter,
+		dxil_spv_shader_stage_io_remapper_cb remapper,
+		void *userdata);
+
 DXIL_SPV_PUBLIC_API void dxil_spv_converter_set_vertex_input_remapper(
 		dxil_spv_converter converter,
 		dxil_spv_vertex_input_remapper_cb remapper,