Track spirv code and shader stage info for postprocessing

Author: per-mathisen-arm

scripts/util.py

@@ -158,7 +158,8 @@ def validate_funcs(lst):
 replay_post_calls = [ 'vkCreateInstance', 'vkDestroyInstance', 'vkQueuePresentKHR', 'vkAcquireNextImageKHR', 'vkAcquireNextImage2KHR',
 	'vkGetBufferDeviceAddress', 'vkGetBufferDeviceAddressKHR', 'vkGetAccelerationStructureDeviceAddressKHR' ]
 validate_funcs(replay_post_calls)
-replay_postprocess_calls = [ 'vkCmdPushConstants', 'vkCmdPushConstants2KHR' ]
+replay_postprocess_calls = [ 'vkCmdPushConstants', 'vkCmdPushConstants2KHR', 'vkCreateRayTracingPipelinesKHR', 'vkCreateGraphicsPipelines',
+	'vkCreateComputePipelines' ]
 validate_funcs(replay_postprocess_calls)
 trace_pre_calls = [ 'vkQueueSubmit', 'vkCreateInstance', 'vkCreateDevice', 'vkFreeMemory', 'vkQueueSubmit2', 'vkQueueSubmit2KHR' ]
 validate_funcs(trace_pre_calls)
@@ -1259,6 +1260,7 @@ def save_add_tracking(name):
 			z.do('add->cache = pipelineCache;')
 			if name == 'vkCreateGraphicsPipelines': z.do('add->type = VK_PIPELINE_BIND_POINT_GRAPHICS;')
 			elif name == 'vkCreateComputePipelines': z.do('add->type = VK_PIPELINE_BIND_POINT_COMPUTE;')
+			elif name == 'PFN_vkCreateRayTracingPipelinesKHR': z.do('add->type = VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR;')
 		elif type == 'VkSwapchainKHR':
 			z.do('add->info = pCreateInfos[i];')
 		z.do('DLOG2("insert %s into %s index %%u", (unsigned)add->index);' % (type, name))
@@ -1337,6 +1339,9 @@ def load_add_tracking(name):
 				z.do('data.format = pCreateInfo->format; // as above, might be missing in json')
 			elif type == 'VkDescriptorSet':
 				z.do('data.pool = pAllocateInfo->descriptorPool;')
+			elif type == 'VkShaderModule':
+				z.do('data.code.resize(pCreateInfo->codeSize);')
+				z.do('memcpy(data.code.data(), pCreateInfo->pCode, pCreateInfo->codeSize * sizeof(uint32_t));')
 			z.do('data.enter_created();')
 		else: # multiple
 			z.do('for (unsigned i = 0; i < %s; i++)' % count)

src/hardcode_read.cpp

@@ -1061,6 +1061,63 @@ void replay_postprocess_vkCmdPushConstants2(lava_file_reader& reader, VkCommandB
 	replay_postprocess_vkCmdPushConstants2KHR(reader, commandBuffer, pPushConstantsInfo);
 }
 
+static void copy_shader_stage(shader_stage& stage, const VkPipelineShaderStageCreateInfo& info)
+{
+	stage.flags = info.flags;
+	stage.module = info.module;
+	stage.name = info.pName;
+	if (info.pSpecializationInfo)
+	{
+		stage.specialization_constants.resize(info.pSpecializationInfo->mapEntryCount);
+		memcpy(stage.specialization_constants.data(), info.pSpecializationInfo->pMapEntries, info.pSpecializationInfo->mapEntryCount * sizeof(uint32_t));
+		stage.specialization_data.resize(info.pSpecializationInfo->dataSize);
+		memcpy(stage.specialization_data.data(), info.pSpecializationInfo->pData, info.pSpecializationInfo->dataSize);
+	}
+}
+
+void replay_postprocess_vkCreateComputePipelines(lava_file_reader& reader, VkResult result, VkDevice device, VkPipelineCache pipelineCache, uint32_t createInfoCount,
+	const VkComputePipelineCreateInfo* pCreateInfos, const VkAllocationCallbacks* pAllocator, VkPipeline* pPipelines)
+{
+	for (uint32_t i = 0; i < createInfoCount; i++)
+	{
+		const uint32_t pipeline_index = index_to_VkPipeline.index(pPipelines[i]);
+		trackedpipeline& pipeline_data = VkPipeline_index.at(pipeline_index);
+		pipeline_data.shader_stages.resize(1);
+		copy_shader_stage(pipeline_data.shader_stages[0], pCreateInfos[i].stage);
+	}
+}
+
+void replay_postprocess_vkCreateGraphicsPipelines(lava_file_reader& reader, VkResult result, VkDevice device, VkPipelineCache pipelineCache, uint32_t createInfoCount,
+	const VkGraphicsPipelineCreateInfo* pCreateInfos, const VkAllocationCallbacks* pAllocator, VkPipeline* pPipelines)
+{
+	for (uint32_t i = 0; i < createInfoCount; i++)
+	{
+		const uint32_t pipeline_index = index_to_VkPipeline.index(pPipelines[i]);
+		trackedpipeline& pipeline_data = VkPipeline_index.at(pipeline_index);
+		pipeline_data.shader_stages.resize(pCreateInfos[i].stageCount);
+		for (uint32_t stage = 0; stage < pCreateInfos[i].stageCount; stage++)
+		{
+			copy_shader_stage(pipeline_data.shader_stages[0], pCreateInfos[i].pStages[stage]);
+		}
+	}
+}
+
+
+void replay_postprocess_vkCreateRayTracingPipelinesKHR(lava_file_reader& reader, VkResult result, VkDevice device, VkDeferredOperationKHR deferredOperation, VkPipelineCache pipelineCache,
+	uint32_t createInfoCount, const VkRayTracingPipelineCreateInfoKHR* pCreateInfos, const VkAllocationCallbacks* pAllocator, VkPipeline* pPipelines)
+{
+	for (uint32_t i = 0; i < createInfoCount; i++)
+	{
+		const uint32_t pipeline_index = index_to_VkPipeline.index(pPipelines[i]);
+		trackedpipeline& pipeline_data = VkPipeline_index.at(pipeline_index);
+		pipeline_data.shader_stages.resize(pCreateInfos[i].stageCount);
+		for (uint32_t stage = 0; stage < pCreateInfos[i].stageCount; stage++)
+		{
+			copy_shader_stage(pipeline_data.shader_stages[0], pCreateInfos[i].pStages[stage]);
+		}
+	}
+}
+
 void replay_pre_vkCmdPushConstants2KHR(lava_file_reader& reader, VkCommandBuffer commandBuffer, const VkPushConstantsInfoKHR* pPushConstantsInfo)
 {
 	assert(pPushConstantsInfo);

src/lavatube.h

@@ -197,6 +197,7 @@ struct trackedshadermodule : trackable
 	using trackable::trackable; // inherit constructor
 	bool enables_device_address = false;
 	size_t size = 0;
+	std::vector<uint32_t> code; // only for replayer
 };
 
 struct trackedmemoryobject : trackedobject
@@ -372,12 +373,23 @@ struct trackedfence : trackable
 	int frame_delay = -1; // delay fuse uninitialized
 };
 
+struct shader_stage
+{
+	VkPipelineShaderStageCreateFlags flags;
+	VkShaderStageFlagBits stage;
+	VkShaderModule module;
+	std::string name;
+	std::vector<VkSpecializationMapEntry> specialization_constants;
+	std::vector<char> specialization_data;
+};
+
 struct trackedpipeline : trackable
 {
 	using trackable::trackable; // inherit constructor
 	VkPipelineBindPoint type = VK_PIPELINE_BIND_POINT_MAX_ENUM;
 	VkPipelineCreateFlags flags = 0;
 	VkPipelineCache cache = VK_NULL_HANDLE;
+	std::vector<shader_stage> shader_stages; // only set for postprocessing
 
 	void self_test() const
 	{