Suballocator refactor: Move it to the Vulkan device metadata

We now have one memory suballocator for each Vulkan device,
allowing better support for esoteric scenarios with multiple
devices. More importantly, we now have all the data we need
for the bigger rewrite.

We currently preallocate some space for all the traced objects
on every device. This is incorrect, but we will fix this later.

Author: per-mathisen-arm

scripts/lava.py

@@ -77,6 +77,7 @@ def out(lst, str=''):
 out(targets_read, '#include "vulkan/vulkan.h"')
 out(targets_read, '#include "window.h"')
 out(targets_read, '#include "suballocator.h"')
+out(targets_read, '#include "memory.h"')
 out([uh,wrh], '#include <unordered_map>')
 out([uh,wrh], '#include "vulkan_feature_detect.h"')
 out([r], '#include "read_auto.h"')
@@ -383,11 +384,8 @@ def out(lst, str=''):
 out(targets_read_headers, 'trackable& object_trackable(VkObjectType type, uint64_t handle);')
 out(targets_read_headers)
 
-out(targets_read, 'void retrace_init(lava_reader& replayer, const Json::Value& v, int heap_size, bool run)')
+out(targets_read, 'void retrace_init(lava_reader& replayer, const Json::Value& v)')
 out(targets_read, '{')
-out(targets_read, '\tint images = 0;')
-out(targets_read, '\tint buffers = 0;')
-out(targets_read, '\tint tensors = 0;')
 out(targets_read, '\tfor (const auto& p : v.getMemberNames())')
 out(targets_read, '\t{')
 for v in spec.root.findall('types/type'):
@@ -403,17 +401,10 @@ def out(lst, str=''):
 			out(targets_read, '\t\t\tif (v[p].asInt() == 0) ELOG("No Vulkan instances recorded. Broken trace file!");')
 		if v.find('name').text == 'VkSurfaceKHR':
 			out(targets_read, '\t\t\twindow_preallocate(v[p].asInt());')
-		if v.find('name').text == 'VkImage':
-			out(targets_read, '\t\t\timages = v[p].asInt();')
-		if v.find('name').text == 'VkBuffer':
-			out(targets_read, '\t\t\tbuffers = v[p].asInt();')
-		if v.find('name').text == 'VkTensorARM':
-			out(targets_read, '\t\t\ttensors = v[p].asInt();')
 		out(targets_read, '\t\t}')
 out(targets_read, '\t}')
-out(targets_read, '\treplayer.allocator.init(images, buffers, tensors, heap_size, !run);')
 out(targets_read, '}')
-out(targets_read_headers, 'void retrace_init(lava_reader& replayer, const Json::Value& v, int heap_size = -1, bool run = true);')
+out(targets_read_headers, 'void retrace_init(lava_reader& replayer, const Json::Value& v);')
 
 out(targets_write)
 out(targets_write, 'Json::Value trace_limits(const lava_writer* instance)')

scripts/util.py

@@ -573,6 +573,9 @@ def print_load(self, name, owner): # called for each parameter
 
 		if self.funcname in ['vkBindImageMemory', 'vkBindBufferMemory', 'VkBindBufferMemoryInfo', 'VkBindBufferMemoryInfoKHR', 'VkBindImageMemoryInfoKHR', 'VkBindImageMemoryInfo', 'VkBindTensorMemoryInfoARM']:
 			if self.name in ['image', 'buffer', 'tensor']:
+				if self.funcname in ['VkBindBufferMemoryInfo', 'VkBindBufferMemoryInfoKHR', 'VkBindImageMemoryInfoKHR', 'VkBindImageMemoryInfo', 'VkBindTensorMemoryInfoARM']:
+					z.do('const uint32_t device_index = index_to_VkDevice.index(reader.device);')
+				z.do('%s& %s = VkDevice_index.at(%s);' % (vk.trackable_type_map_replay['VkDevice'], totrackable('VkDevice'), toindex('VkDevice')))
 				z.do('%s& %s = %s_index.at(%s);' % (vk.trackable_type_map_replay[self.type], totrackable(self.type), self.type, toindex(self.type)))
 				z.do('%s.memory_flags = static_cast<VkMemoryPropertyFlags>(reader.read_uint32_t()); // fetch memory flags especially added' % totrackable(self.type)) # TBD remove
 				if self.funcname in ['vkBindImageMemory', 'VkBindImageMemoryInfoKHR', 'VkBindImageMemoryInfo'] and self.name == 'image': # TBD remove me
@@ -583,17 +586,22 @@ def print_load(self, name, owner): # called for each parameter
 				z.do('memory_requirements reqs;')
 				z.do('if (reader.run) reqs = get_%s_memory_requirements(reader.device, %s);' % (vk.trackable_type_map_replay[self.type], totrackable(self.type)))
 				z.do('else reqs = get_fake_memory_requirements(reader.device, %s);' % totrackable(self.type))
-				z.do('suballoc_location loc = reader.parent->allocator.add_trackedobject(reader.thread_index(), reader.device, reqs, (uint64_t)%s, %s);' % (varname, totrackable(self.type)))
+				z.do('suballoc_location loc = device_data.allocator->add_trackedobject(reader.thread_index(), reqs, (uint64_t)%s, %s);' % (varname, totrackable(self.type)))
 			if self.name == 'memory':
 				z.do('assert(loc.memory != VK_NULL_HANDLE);')
 				z.do('%s = loc.memory;' % varname) # relying on the order of arguments here; see case above
 			elif self.name == 'memoryOffset':
 				z.do('%s = loc.offset;' % varname) # relying on the order of arguments here; see case above
 
+		if self.funcname in ['vkDestroyBuffer', 'vkDestroyImage', 'vkDestroyTensorARM', 'vkDestroyDevice', 'VkCreateDevice'] and self.name == 'device':  #['image', 'buffer', 'tensor', 'device']:
+			z.do('%s& %s = %s_index.at(%s);' % (vk.trackable_type_map_replay['VkDevice'], totrackable('VkDevice'), self.type, toindex('VkDevice')))
+
 		if self.funcname == 'vkDestroyBuffer' and self.name == 'buffer':
-			z.do('if (buffer_index != CONTAINER_INVALID_INDEX) reader.parent->allocator.free_buffer(buffer_index);')
+			z.do('if (buffer_index != CONTAINER_INVALID_INDEX) device_data.allocator->free_buffer(buffer_index);')
 		elif self.funcname == 'vkDestroyImage' and self.name == 'image':
-			z.do('reader.parent->allocator.free_image(image_index);')
+			z.do('device_data.allocator->free_image(image_index);')
+		elif self.funcname == 'vkDestroyTensorARM' and self.name == 'tensor':
+			z.do('device_data.allocator->free_tensor(tensorarm_index);')
 		elif self.name == 'sType':
 			orig = z.struct_last()
 			stype = spec.type2sType[orig]
@@ -1205,7 +1213,13 @@ def load_add_pre(name):
 		count = spec.functions_destroy[name][1]
 		type = spec.functions_destroy[name][2]
 		if type == 'VkDevice':
-			z.do('reader.parent->allocator.destroy(device);')
+			z.do('device_data.allocator->self_test();')
+			z.do('device_data.allocator->destroy();')
+			z.do('suballoc_metrics sm = device_data.allocator->performance();')
+			z.do('ILOG("Suballocator used=%lu allocated=%lu heaps=%u objects=%u efficiency=%g", (unsigned long)sm.used, (unsigned long)sm.allocated, (unsigned)sm.heaps, (unsigned)sm.objects, sm.efficiency);')
+			z.do('assert(device_data.allocator->self_test() == 0);')
+			z.do('delete device_data.allocator;')
+			z.do('device_data.allocator = nullptr;')
 
 	if name == 'vkCreatePipelineCache': # cannot be autogenerated, need the index parameter
 		z.do('if (reader.run) replay_pre_vkCreatePipelineCache(reader, device, device_index, pCreateInfo, pipelinecache_index);')
@@ -1234,7 +1248,8 @@ def load_add_tracking(name):
 				z.do('data.device = device;')
 			elif type == 'VkDevice':
 				z.do('data.physicalDevice = physicalDevice; // track parentage')
-				z.do('reader.parent->allocator.setup(selected_physical_device);')
+				z.do('data.allocator = new suballocator();')
+				z.do('data.allocator->create(selected_physical_device, pDevice, VkImage_index, VkBuffer_index, VkTensorARM_index, reader.run);')
 			elif type == 'VkImage':
 				z.do('data.tiling = (lava_tiling)pCreateInfo->tiling;')
 				z.do('data.usage = pCreateInfo->usage;')

src/execute_commands.cpp

@@ -1,7 +1,6 @@
-static bool run_spirv(lava_file_reader& reader, const trackedpipeline& pipeline_data, const shader_stage& stage, const std::vector<std::byte>& push_constants,
+static bool run_spirv(const trackeddevice& device_data, const trackedpipeline& pipeline_data, const shader_stage& stage, const std::vector<std::byte>& push_constants,
 	const std::unordered_map<uint32_t, std::unordered_map<uint32_t, buffer_access>>& descriptorsets)
 {
-	assert(!reader.run); // this code is only run when post-processing
 	assert(stage.module != VK_NULL_HANDLE);
 	const uint32_t shader_index = index_to_VkShaderModule.index(stage.module);
 	trackedshadermodule& shader_data = VkShaderModule_index.at(shader_index);
@@ -22,7 +21,7 @@ static bool run_spirv(lava_file_reader& reader, const trackedpipeline& pipeline_
 			const buffer_access& access = binding_pair.second;
 			if (!access.buffer_data) continue;
 			const uint32_t buffer_index = access.buffer_data->index;
-			suballoc_location loc = reader.parent->allocator.find_buffer_memory(buffer_index);
+			suballoc_location loc = device_data.allocator->find_buffer_memory(buffer_index);
 			std::byte* base = (std::byte*)loc.memory;
 			std::byte* binding_ptr = base + access.offset;
 			set_bindings[binding_pair.first] = binding_ptr;
@@ -73,8 +72,9 @@ static bool run_spirv(lava_file_reader& reader, const trackedpipeline& pipeline_
 	return true;
 }
 
-static bool execute_commands(lava_file_reader& reader, VkCommandBuffer commandBuffer)
+static bool execute_commands(lava_file_reader& reader, const trackeddevice& device_data, VkCommandBuffer commandBuffer)
 {
+	assert(!reader.run); // this code is only run when post-processing
 	std::vector<std::byte> push_constants; // current state of the push constants
 	uint32_t compute_pipeline_bound = CONTAINER_INVALID_INDEX; // currently bound pipeline
 	uint32_t graphics_pipeline_bound = CONTAINER_INVALID_INDEX; // currently bound pipeline
@@ -117,8 +117,8 @@ static bool execute_commands(lava_file_reader& reader, VkCommandBuffer commandBu
 			break;
 		case VKCMDCOPYBUFFER:
 			{
-				suballoc_location src = reader.parent->allocator.find_buffer_memory(c.data.copy_buffer.src_buffer_index);
-				suballoc_location dst = reader.parent->allocator.find_buffer_memory(c.data.copy_buffer.dst_buffer_index);
+				suballoc_location src = device_data.allocator->find_buffer_memory(c.data.copy_buffer.src_buffer_index);
+				suballoc_location dst = device_data.allocator->find_buffer_memory(c.data.copy_buffer.dst_buffer_index);
 				for (uint32_t i = 0; i < c.data.copy_buffer.regionCount; i++)
 				{
 					VkBufferCopy& r = c.data.copy_buffer.pRegions[i];
@@ -129,7 +129,7 @@ static bool execute_commands(lava_file_reader& reader, VkCommandBuffer commandBu
 			break;
 		case VKCMDUPDATEBUFFER:
 			{
-				suballoc_location sub = reader.parent->allocator.find_buffer_memory(c.data.update_buffer.buffer_index);
+				suballoc_location sub = device_data.allocator->find_buffer_memory(c.data.update_buffer.buffer_index);
 				memcpy((char*)sub.memory + c.data.update_buffer.offset, c.data.update_buffer.values, c.data.update_buffer.size);
 			}
 			free(c.data.update_buffer.values);
@@ -153,15 +153,15 @@ static bool execute_commands(lava_file_reader& reader, VkCommandBuffer commandBu
 				const auto& pipeline_data = VkPipeline_index.at(compute_pipeline_bound);
 				assert(pipeline_data.shader_stages.size() == 1);
 				assert(pipeline_data.shader_stages[0].stage == VK_SHADER_STAGE_COMPUTE_BIT);
-				run_spirv(reader, pipeline_data, pipeline_data.shader_stages[0], push_constants, descriptorsets);
+				run_spirv(device_data, pipeline_data, pipeline_data.shader_stages[0], push_constants, descriptorsets);
 			}
 			break;
 		case VKCMDDRAW: // proxy for all draw commands
 			{
 				const auto& pipeline_data = VkPipeline_index.at(graphics_pipeline_bound);
 				for (const auto& stage : pipeline_data.shader_stages)
 				{
-					run_spirv(reader, pipeline_data, stage, push_constants, descriptorsets);
+					run_spirv(device_data, pipeline_data, stage, push_constants, descriptorsets);
 				}
 			}
 			break;

src/hardcode_read.cpp

@@ -1138,14 +1138,15 @@ void retrace_vkSyncBufferTRACETOOLTEST(lava_file_reader& reader)
 void retrace_vkAssertBufferARM(lava_file_reader& reader)
 {
 	const uint32_t device_index = reader.read_handle(DEBUGPARAM("VkDevice"));
+	const auto& device_data = VkDevice_index.at(device_index);
 	const uint32_t buffer_index = reader.read_handle(DEBUGPARAM("VkBuffer"));
 	const VkDeviceSize offset = reader.read_uint64_t();
 	VkDeviceSize size = reader.read_uint64_t();
 	const char* comment = reader.read_string();
 	const uint32_t checksum = reader.read_uint32_t();
 	trackedobject& tbuf = VkBuffer_index.at(buffer_index);
 	VkDevice device = index_to_VkDevice.at(device_index);
-	suballoc_location loc = reader.parent->allocator.find_buffer_memory(buffer_index);
+	suballoc_location loc = device_data.allocator->find_buffer_memory(buffer_index);
 	if (size == VK_WHOLE_SIZE)
 	{
 		size = tbuf.size - offset; // set to remaining size
@@ -1436,11 +1437,16 @@ void replay_postprocess_vkCmdBindDescriptorSets2(lava_file_reader& reader, VkCom
 
 void replay_postprocess_vkQueueSubmit2(lava_file_reader& reader, VkResult result, VkQueue queue, uint32_t submitCount, const VkSubmitInfo2* pSubmits, VkFence fence)
 {
+	const uint32_t queue_index = index_to_VkQueue.index(queue);
+	auto& queue_data = VkQueue_index.at(queue_index);
+	const uint32_t device_index = queue_data.device_index;
+	assert(device_index != UINT32_MAX);
+	auto& device_data = VkDevice_index.at(device_index);
 	for (uint32_t i = 0; i < submitCount; i++)
 	{
 		for (uint32_t j = 0; j < pSubmits[i].commandBufferInfoCount; j++)
 		{
-			execute_commands(reader, pSubmits[i].pCommandBufferInfos[j].commandBuffer);
+			execute_commands(reader, device_data, pSubmits[i].pCommandBufferInfos[j].commandBuffer);
 		}
 	}
 }
@@ -1452,11 +1458,16 @@ void replay_postprocess_vkQueueSubmit2KHR(lava_file_reader& reader, VkResult res
 
 void replay_postprocess_vkQueueSubmit(lava_file_reader& reader, VkResult result, VkQueue queue, uint32_t submitCount, const VkSubmitInfo* pSubmits, VkFence fence)
 {
+	const uint32_t queue_index = index_to_VkQueue.index(queue);
+	auto& queue_data = VkQueue_index.at(queue_index);
+	const uint32_t device_index = queue_data.device_index;
+	assert(device_index != UINT32_MAX);
+	auto& device_data = VkDevice_index.at(device_index);
 	for (uint32_t i = 0; i < submitCount; i++)
 	{
 		for (uint32_t j = 0; j < pSubmits[i].commandBufferCount; j++)
 		{
-			execute_commands(reader, pSubmits[i].pCommandBuffers[j]);
+			execute_commands(reader, device_data, pSubmits[i].pCommandBuffers[j]);
 		}
 	}
 }
@@ -1908,6 +1919,7 @@ void retrace_vkGetSwapchainImagesKHR(lava_file_reader& reader)
 	const uint32_t device_index = reader.read_handle(DEBUGPARAM("VkDevice"));
 	const uint32_t swapchain_index = reader.read_handle(DEBUGPARAM("VkSwapchainKHR"));
 	VkDevice device = index_to_VkDevice.at(device_index);
+	auto& device_data = VkDevice_index.at(device_index);
 	const uint8_t do_call = reader.read_uint8_t();
 	const VkResult stored_retval = (VkResult)reader.read_uint32_t();
 	const uint32_t stored_image_count = reader.read_uint32_t();
@@ -1953,7 +1965,7 @@ void retrace_vkGetSwapchainImagesKHR(lava_file_reader& reader)
 			result = wrap_vkCreateImage(device, &pinfo, nullptr, &data.virtual_images[i]);
 			assert(result == VK_SUCCESS);
 		}
-		reader.parent->allocator.virtualswap_images(device, data.virtual_images);
+		device_data.allocator->virtualswap_images(data.virtual_images);
 
 		if (is_noscreen())
 		{
@@ -2458,18 +2470,20 @@ void retrace_vkGetPhysicalDeviceXlibPresentationSupportKHR(lava_file_reader& rea
 
 void image_update(lava_file_reader& reader, uint32_t device_index, uint32_t image_index, uint64_t size, const VkBaseOutStructure* sptr)
 {
-	suballoc_location loc = reader.parent->allocator.find_image_memory(image_index);
+	VkDevice device = index_to_VkDevice.at(device_index);
+	const auto& device_data = VkDevice_index.at(device_index);
+	suballoc_location loc = device_data.allocator->find_image_memory(image_index);
 	DLOG2("image update idx=%u flush=%s init=%s size=%lu", image_index, loc.needs_flush ? "yes" : "no", loc.needs_init ? "yes" : "no", (unsigned long)loc.size);
 	assert(sptr == nullptr);
-	VkDevice device = index_to_VkDevice.at(device_index);
 	char* ptr = mem_map(reader, device, loc);
 	int32_t changed = reader.read_patch(ptr, loc.size);
 	mem_unmap(reader, device, loc, nullptr, ptr);
 }
 
 void buffer_update(lava_file_reader& reader, uint32_t device_index, uint32_t buffer_index, uint64_t size, const VkBaseOutStructure* sptr)
 {
-	suballoc_location loc = reader.parent->allocator.find_buffer_memory(buffer_index);
+	const auto& device_data = VkDevice_index.at(device_index);
+	suballoc_location loc = device_data.allocator->find_buffer_memory(buffer_index);
 	DLOG2("buffer update idx=%u flush=%s init=%s size=%lu", buffer_index, loc.needs_flush ? "yes" : "no", loc.needs_init ? "yes" : "no", (unsigned long)loc.size);
 	assert(sptr == nullptr || sptr->sType == VK_STRUCTURE_TYPE_MARKED_OFFSETS_ARM);
 	VkDevice device = index_to_VkDevice.at(device_index);
@@ -2486,7 +2500,8 @@ void buffer_update(lava_file_reader& reader, uint32_t device_index, uint32_t buf
 
 void tensor_update(lava_file_reader& reader, uint32_t device_index, uint32_t tensor_index, uint64_t size, const VkBaseOutStructure* sptr)
 {
-	suballoc_location loc = reader.parent->allocator.find_tensor_memory(tensor_index);
+	const auto& device_data = VkDevice_index.at(device_index);
+	suballoc_location loc = device_data.allocator->find_tensor_memory(tensor_index);
 	DLOG2("tensor update idx=%u flush=%s init=%s size=%lu", tensor_index, loc.needs_flush ? "yes" : "no", loc.needs_init ? "yes" : "no", (unsigned long)loc.size);
 	assert(sptr == nullptr);
 	VkDevice device = index_to_VkDevice.at(device_index);

src/hardcode_write.cpp

@@ -2187,6 +2187,7 @@ VKAPI_ATTR void VKAPI_CALL trace_vkGetDeviceQueue2(VkDevice device, const VkDevi
 		queue_data->queueFamily = pQueueInfo->queueFamilyIndex;
 		queue_data->queueFlags = pQueueInfo->flags;
 		queue_data->device = device;
+		queue_data->device_index = device_data->index;
 		queue_data->realIndex = realIndex;
 		queue_data->realFamily = realFamily;
 		queue_data->realQueue = *pQueue;
@@ -2236,6 +2237,7 @@ VKAPI_ATTR void VKAPI_CALL trace_vkGetDeviceQueue(VkDevice device, uint32_t queu
 		queue_data->queueFamily = queueFamilyIndex;
 		queue_data->queueFlags = physicaldevice_data->queueFamilyProperties.at(queueFamilyIndex).queueFlags;
 		queue_data->device = device;
+		queue_data->device_index = device_data->index;
 		queue_data->realIndex = realIndex;
 		queue_data->realFamily = realFamily;
 		queue_data->realQueue = *pQueue;

src/json_helpers.cpp

@@ -106,6 +106,7 @@ Json::Value trackedimage_json(const trackedimage* t)
 	v["samples"] = (unsigned)t->samples;
 	v["mipLevels"] = (unsigned)t->mipLevels;
 	v["arrayLayers"] = (unsigned)t->arrayLayers;
+	if (t->is_swapchain_image) v["swapchain_image"] = true;
 	return v;
 }
 
@@ -154,6 +155,7 @@ Json::Value trackedqueue_json(const trackedqueue* t)
 	v["queueFamily"] = (unsigned)t->queueFamily;
 	v["queueIndex"] = (unsigned)t->queueIndex;
 	v["queueFlags"] = (unsigned)t->queueFlags;
+	v["parent_device_index"] = (unsigned)t->device_index;
 	return v;
 }
 
@@ -351,7 +353,7 @@ trackedimage trackedimage_json(const Json::Value& v)
 	t.currentLayout = t.initialLayout;
 	t.samples = (VkSampleCountFlagBits)(v.get("samples", 0).asUInt());
 	t.mipLevels = (unsigned)v.get("mipLevels", 0).asUInt();
-	t.arrayLayers = (unsigned)v.get("arrayLevels", 0).asUInt();
+	t.arrayLayers = (unsigned)v.get("arrayLayers", 0).asUInt();
 	t.format = (VkFormat)v.get("format", VK_FORMAT_MAX_ENUM).asUInt();
 	if (v.isMember("extent"))
 	{
@@ -365,6 +367,7 @@ trackedimage trackedimage_json(const Json::Value& v)
 		t.alias_type = (VkObjectType)v["alias_type"].asUInt();
 	}
 	if (v.isMember("memory_flags")) t.memory_flags = (VkMemoryPropertyFlags)v["memory_flags"].asUInt();
+	if (v.isMember("swapchain_image")) t.is_swapchain_image = v["swapchain_image"].asBool();
 	t.object_type = VK_OBJECT_TYPE_IMAGE;
 
 	if (v.isMember("parent_device_index")) t.parent_device_index = v["parent_device_index"].asUInt();
@@ -456,6 +459,7 @@ trackedqueue trackedqueue_json(const Json::Value& v)
 {
 	trackedqueue t;
 	trackable_helper(t, v);
+	t.device_index = v["parent_device_index"].asUInt();
 	t.enter_initialized();
 	return t;
 }