Track whether shaders use buffer device addresses

scripts/util.py

@@ -148,7 +148,8 @@
 		'vkGetPhysicalDeviceFormatProperties', 'vkGetPhysicalDeviceFormatProperties2', 'vkCmdPushDescriptorSetKHR', 'vkCreateSwapchainKHR',
 		'vkGetBufferMemoryRequirements2KHR', 'vkGetDeviceBufferMemoryRequirements', 'vkGetDeviceBufferMemoryRequirementsKHR',
 		'vkGetDeviceImageMemoryRequirements', 'vkGetDeviceImageMemoryRequirementsKHR', 'vkGetPhysicalDeviceFeatures2', 'vkGetPhysicalDeviceFeatures2KHR',
-		'vkGetPhysicalDeviceMemoryProperties2', 'vkGetDeviceImageSparseMemoryRequirementsKHR', 'vkGetDeviceImageSparseMemoryRequirements' ]
+		'vkGetPhysicalDeviceMemoryProperties2', 'vkGetDeviceImageSparseMemoryRequirementsKHR', 'vkGetDeviceImageSparseMemoryRequirements',
+		'vkCreateShaderModule' ]
 skip_post_calls = [ 'vkGetQueryPoolResults', 'vkGetPhysicalDeviceXcbPresentationSupportKHR' ]
 # Workaround to be able to rewrite parameter inputs while tracing: These input variables are copied and replaced to not be const anymore.
 deconstify = {
@@ -168,7 +169,8 @@
 # Subclassing of trackable
 trackable_type_map_general = { 'VkBuffer': 'trackedbuffer', 'VkImage': 'trackedimage', 'VkCommandBuffer': 'trackedcmdbuffer', 'VkDescriptorSet': 'trackeddescriptorset',
 	'VkDeviceMemory': 'trackedmemory', 'VkFence': 'trackedfence', 'VkPipeline': 'trackedpipeline', 'VkImageView': 'trackedimageview', 'VkBufferView': 'trackedbufferview',
-	'VkDevice': 'trackeddevice', 'VkFramebuffer': 'trackedframebuffer', 'VkRenderPass': 'trackedrenderpass', 'VkQueue': 'trackedqueue', 'VkPhysicalDevice': 'trackedphysicaldevice' }
+	'VkDevice': 'trackeddevice', 'VkFramebuffer': 'trackedframebuffer', 'VkRenderPass': 'trackedrenderpass', 'VkQueue': 'trackedqueue', 'VkPhysicalDevice': 'trackedphysicaldevice',
+	'VkShaderModule': 'trackedshadermodule' }
 trackable_type_map_trace = trackable_type_map_general.copy()
 trackable_type_map_trace.update({ 'VkCommandBuffer': 'trackedcmdbuffer_trace', 'VkSwapchainKHR': 'trackedswapchain_trace', 'VkDescriptorSet': 'trackeddescriptorset_trace',
 	'VkEvent': 'trackedevent_trace', 'VkDescriptorPool': 'trackeddescriptorpool_trace', 'VkCommandPool': 'trackedcommandpool_trace' })
@@ -1150,6 +1152,9 @@ def save_add_tracking(name):
 			z.do('add->mipLevels = pCreateInfo->mipLevels;')
 			z.do('add->arrayLayers = pCreateInfo->arrayLayers;')
 			z.do('if (pCreateInfo->flags & VK_IMAGE_CREATE_ALIAS_BIT) ELOG("Image aliasing detected! We need to implement support for this!");')
+		elif type == 'VkShaderModule':
+			z.do('add->type = VK_OBJECT_TYPE_SHADER_MODULE;')
+			z.do('add->size = pCreateInfo->codeSize;')
 		elif type == 'VkRenderPass' and name == 'vkCreateRenderPass':
 			z.do('add->attachments.resize(pCreateInfo->attachmentCount);')
 			z.do('for (unsigned ii = 0; ii < pCreateInfo->attachmentCount; ii++) add->attachments[ii] = pCreateInfo->pAttachments[ii]; // struct copy')

src/hardcode_read.cpp

@@ -1782,6 +1782,16 @@ static trackedframebuffer trackedframebuffer_json(const Json::Value& v)
 	return t;
 }
 
+static trackedshadermodule trackedshadermodule_json(const Json::Value& v)
+{
+	trackedshadermodule t(v["frame_created"].asInt());
+	t.frame_destroyed = v["frame_destroyed"].asInt();
+	if (v.isMember("name")) t.name = v["name"].asString();
+	if (v.isMember("size")) t.name = v["size"].asInt();
+	if (v.isMember("enables_buffer_device_address")) t.enables_buffer_device_address = v["enables_buffer_device_address"].asBool();
+	return t;
+}
+
 static trackedrenderpass trackedrenderpass_json(const Json::Value& v)
 {
 	trackedrenderpass t(v["frame_created"].asInt());

src/hardcode_write.cpp

@@ -139,6 +139,17 @@ static trackable* object_trackable(const trace_records& r, VkObjectType type, ui
 	return nullptr;
 }
 
+static void trace_post_vkCreateShaderModule(lava_file_writer& writer, VkResult result, VkDevice device, const VkShaderModuleCreateInfo* pCreateInfo, const VkAllocationCallbacks* pAllocator, VkShaderModule* pShaderModule)
+{
+	if (result != VK_SUCCESS) return; // ignore rest on failure
+	if (shader_has_buffer_devices_addresses(pCreateInfo->pCode, pCreateInfo->codeSize))
+	{
+		auto* shader = writer.parent->records.VkShaderModule_index.at(*pShaderModule);
+		shader->enables_buffer_device_address = true;
+		ILOG("Shader %u enables buffer references!", shader->index); // remove this later
+	}
+}
+
 static void trace_post_vkAcquireNextImageKHR(lava_file_writer& writer, VkResult result, VkDevice device, VkSwapchainKHR swapchain, uint64_t timeout, VkSemaphore semaphore, VkFence fence, uint32_t* pImageIndex)
 {
 	DLOG("Acquired swapchain image index=%u", *pImageIndex);
@@ -2380,6 +2391,14 @@ static Json::Value trackeddevice_json(const trackeddevice* t)
 	return v;
 }
 
+static Json::Value trackedshadermodule_json(const trackedshadermodule* t)
+{
+	Json::Value v = trackable_json(t);
+	if (t->enables_buffer_device_address) v["enables_buffer_device_address"] = true;
+	v["size"] = (unsigned)t->size;
+	return v;
+}
+
 static Json::Value trackedphysicaldevice_json(const trackedphysicaldevice* t)
 {
 	Json::Value v = trackable_json(t);

src/lavatube.h

@@ -138,6 +138,13 @@ struct trackedobject : trackable
 	}
 };
 
+struct trackedshadermodule : trackedobject
+{
+	using trackedobject::trackedobject; // inherit constructor
+	bool enables_buffer_device_address = false;
+	size_t size = 0;
+};
+
 struct trackedbuffer : trackedobject
 {
 	using trackedobject::trackedobject; // inherit constructor

src/util.cpp

@@ -1,6 +1,7 @@
 #include "util.h"
 
 #include <string.h>
+#include <spirv/unified1/spirv.h>
 
 #if defined(_GNU_SOURCE) || defined(__BIONIC__)
 #include <pthread.h>
@@ -347,6 +348,22 @@ const void* find_extension(const void* sptr, VkStructureType sType)
 	return ptr;
 }
 
+bool shader_has_buffer_devices_addresses(const uint32_t* code, uint32_t code_size)
+{
+	uint16_t opcode;
+	uint16_t word_count;
+	const uint32_t* insn = code + 5;
+	code_size /= 4; // bytes to words
+	do {
+		opcode = uint16_t(insn[0]);
+		word_count = uint16_t(insn[0] >> 16);
+		if (opcode == SpvOpExtension && strcmp((char*)&insn[2], "KHR_physical_storage_buffer") == 0) return true;
+		insn += word_count;
+	}
+	while (insn != code + code_size && opcode != SpvOpMemoryModel);
+	return false;
+}
+
 const char* pretty_print_VkObjectType(VkObjectType val)
 {
 	switch (val)

src/util.h

@@ -178,6 +178,9 @@ static __attribute__((pure)) inline uint64_t gettime()
 
 std::string get_vulkan_lib_path();
 
+/// Whether a SPIRV shader has enabled support for buffer device addresses
+bool shader_has_buffer_devices_addresses(const uint32_t* code, uint32_t codeSize);
+
 /// Faster than std::vector but with much the same interface. The performance improvement mostly
 /// comes from not filling memory beforehand.
 struct buffer