Clean up suballocator and meta data storage for Vulkan objects

Fixing one of lavatube's earliest design mistakes that tried
to avoid relying on information stored outside of the
command stream.

Improve suballocator interface and make it mix linear and
optimal allocations if bufferImageGranularity is one.

Author: per-mathisen-arm

scripts/util.py

@@ -716,15 +716,17 @@ def print_load(self, name, owner): # called for each parameter
 			z.do('%s = reader.read_%s();' % (varname, self.type))
 
 		if self.funcname in ['vkBindImageMemory', 'VkBindImageMemoryInfoKHR', 'VkBindImageMemoryInfo'] and self.name == 'image':
-			z.do('const VkMemoryPropertyFlags special_flags = static_cast<VkMemoryPropertyFlags>(reader.read_uint32_t()); // fetch memory flags especially added')
-			z.do('const VkImageTiling tiling = static_cast<VkImageTiling>(reader.read_uint32_t()); // fetch tiling property especially added')
-			z.do('const VkDeviceSize min_size = static_cast<VkDeviceSize>(reader.read_uint64_t()); // fetch padded memory size')
 			z.do('trackedimage& image_data = VkImage_index.at(image_index);')
-			z.do('suballoc_location loc = reader.parent->allocator.add_image(reader.thread_index(), reader.device, %s, image_index, special_flags, tiling, min_size);' % varname)
+			z.do('image_data.memory_flags = static_cast<VkMemoryPropertyFlags>(reader.read_uint32_t()); // fetch memory flags especially added') # TBD remove me
+			z.do('const VkImageTiling tiling = static_cast<VkImageTiling>(reader.read_uint32_t()); // fetch tiling property especially added') # TBD remove me
+			z.do('assert(tiling == image_data.tiling);')
+			z.do('const VkDeviceSize min_size = static_cast<VkDeviceSize>(reader.read_uint64_t()); // fetch padded memory size') # TBD remove me
+			z.do('assert(min_size == image_data.size);')
+			z.do('suballoc_location loc = reader.parent->allocator.add_image(reader.thread_index(), reader.device, %s, image_data);' % varname)
 		elif self.funcname in ['vkBindBufferMemory', 'VkBindBufferMemoryInfo', 'VkBindBufferMemoryInfoKHR'] and self.name == 'buffer':
-			z.do('const VkMemoryPropertyFlags special_flags = static_cast<VkMemoryPropertyFlags>(reader.read_uint32_t()); // fetch memory flags especially added')
 			z.do('trackedbuffer& buffer_data = VkBuffer_index.at(buffer_index);')
-			z.do('suballoc_location loc = reader.parent->allocator.add_buffer(reader.thread_index(), reader.device, %s, special_flags, buffer_data);' % varname)
+			z.do('buffer_data.memory_flags = static_cast<VkMemoryPropertyFlags>(reader.read_uint32_t()); // fetch memory flags especially added') # TBD remove me
+			z.do('suballoc_location loc = reader.parent->allocator.add_buffer(reader.thread_index(), reader.device, %s, buffer_data);' % varname)
 
 		if self.funcname in ['vkBindImageMemory', 'vkBindBufferMemory', 'VkBindBufferMemoryInfo', 'VkBindBufferMemoryInfoKHR', 'VkBindImageMemoryInfoKHR', 'VkBindImageMemoryInfo']:
 			if self.name == 'memory':
@@ -1007,12 +1009,13 @@ def print_save(self, name, owner, postprocess = False): # called for each parame
 			z.do('writer.commandBuffer = %s;' % varname) # always earlier in the parameter list than images and buffers, fortunately
 			z.do('writer.device = commandbuffer_data->device;')
 			z.do('writer.physicalDevice = commandbuffer_data->physicalDevice;')
-		if self.funcname in ['vkBindImageMemory', 'vkBindBufferMemory', 'VkBindImageMemoryInfo', 'VkBindImageMemoryInfoKHR', 'VkBindBufferMemoryInfo', 'VkBindBufferMemoryInfoKHR'] and self.name in ['image', 'buffer']:
+		if self.funcname in ['vkBindImageMemory', 'vkBindBufferMemory', 'VkBindImageMemoryInfo', 'VkBindImageMemoryInfoKHR', 'VkBindBufferMemoryInfo', 'VkBindBufferMemoryInfoKHR', 'VkBindTensorMemoryInfoARM'] and self.name in ['image', 'buffer', 'tensor']:
 			z.do('const auto* meminfo = writer.parent->records.VkDeviceMemory_index.at(%s);' % (owner + 'memory'))
-			z.do('writer.write_uint32_t(static_cast<uint32_t>(meminfo->propertyFlags)); // save memory flags')
-		if self.funcname in ['vkBindImageMemory', 'VkBindImageMemoryInfo', 'VkBindImageMemoryInfoKHR'] and self.name == 'image':
-			z.do('writer.write_uint32_t(static_cast<uint32_t>(image_data->tiling)); // save tiling info')
-			z.do('writer.write_uint64_t(static_cast<uint64_t>(image_data->size)); // save padded image size')
+			z.do('%s->memory_flags = meminfo->propertyFlags;' % totrackable(self.type))
+			z.do('writer.write_uint32_t(static_cast<uint32_t>(meminfo->propertyFlags)); // save memory flags') # TBD remove
+		if self.funcname in ['vkBindImageMemory', 'VkBindImageMemoryInfo', 'VkBindImageMemoryInfoKHR'] and self.name == 'image': # TBD remove
+			z.do('writer.write_uint32_t(static_cast<uint32_t>(image_data->tiling)); // save tiling info') # TBD remove
+			z.do('writer.write_uint64_t(static_cast<uint64_t>(image_data->size)); // save padded image size') # TBD remove
 		if self.funcname == 'vkAllocateMemory' and self.name == 'pAllocateInfo' and not postprocess:
 			z.do('frame_mutex.lock();')
 			z.do('assert(real_memory_properties.memoryTypeCount > 0);')

src/hardcode_write.cpp

@@ -400,7 +400,8 @@ static void trace_post_vkBindBufferMemory(lava_file_writer& writer, VkResult res
 	{
 		wrap_vkGetBufferMemoryRequirements(device, buffer, &buffer_data->req);
 	}
-	// we pass in size recorded from vkCreateBuffer, which is the actually used size, rather than required allocation size
+	// we store size recorded from vkCreateBuffer, which is the actually used size, rather than required allocation size
+	assert(buffer_data->size >= buffer_data->req.size);
 	buffer_data->accessible = (memory_data->propertyFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT);
 	buffer_data->enter_bound();
 	writer.parent->memory_mutex.unlock();

src/json_helpers.cpp

@@ -36,6 +36,9 @@ static Json::Value trackedobject_json(const trackedobject *t)
 		v["alias_index"] = t->alias_index;
 		v["alias_type"] = (unsigned)t->alias_type;
 	}
+	v["req_size"] = (Json::Value::UInt64)t->req.size; // info only, do not read
+	v["req_alignment"] = (unsigned)t->req.alignment; // info only, do not read
+	v["memory_flags"] = (unsigned)t->memory_flags;
 	return v;
 }
 
@@ -55,8 +58,6 @@ Json::Value trackedbuffer_json(const trackedbuffer* t)
 	v["sharingMode"] = (unsigned)t->sharingMode;
 	v["usage"] = (unsigned)t->usage;
 	v["usage2"] = (Json::Value::UInt64)t->usage2;
-	v["req_size"] = (Json::Value::UInt64)t->req.size;
-	v["req_alignment"] = (unsigned)t->req.alignment;
 	v["written"] = (Json::Value::UInt64)t->written;
 	v["updates"] = (unsigned)t->updates;
 	return v;
@@ -66,8 +67,6 @@ Json::Value trackedtensor_json(const trackedtensor* t)
 {
 	Json::Value v = trackedobject_json(t);
 	v["sharingMode"] = (unsigned)t->sharingMode;
-	v["req_size"] = (Json::Value::UInt64)t->req.size;
-	v["req_alignment"] = (unsigned)t->req.alignment;
 	v["tiling"] = (unsigned)t->tiling;
 	v["format"] = (unsigned)t->format;
 	v["usage"] = (unsigned)t->usage;
@@ -92,8 +91,6 @@ Json::Value trackedimage_json(const trackedimage* t)
 	v["sharingMode"] = (unsigned)t->sharingMode;
 	v["usage"] = (unsigned)t->usage;
 	v["imageType"] = (unsigned)t->imageType;
-	v["req_size"] = (Json::Value::UInt64)t->req.size;
-	v["req_alignment"] = (unsigned)t->req.alignment;
 	v["format"] = (unsigned)t->format;
 	v["written"] = (Json::Value::UInt64)t->written;
 	v["updates"] = (unsigned)t->updates;
@@ -305,6 +302,7 @@ trackedaccelerationstructure trackedaccelerationstructure_json(const Json::Value
 		t.alias_index = v["alias_index"].asUInt();
 		t.alias_type = (VkObjectType)v["alias_type"].asUInt();
 	}
+	if (v.isMember("memory_flags")) t.memory_flags = (VkMemoryPropertyFlags)v["memory_flags"].asUInt();
 	t.enter_initialized();
 	return t;
 }
@@ -318,15 +316,13 @@ trackedbuffer trackedbuffer_json(const Json::Value& v)
 	t.sharingMode = (VkSharingMode)v["sharingMode"].asUInt();
 	t.usage = (VkBufferUsageFlags)v["usage"].asUInt();
 	if (v.isMember("usage2")) t.usage2 = (VkBufferUsageFlags2)v["usage2"].asUInt64();
-	t.req.size = v["req_size"].asUInt64();
-	t.req.alignment = v["req_alignment"].asUInt();
-	t.req.memoryTypeBits = 0;
 	t.object_type = VK_OBJECT_TYPE_BUFFER;
 	if (v.isMember("alias_index"))
 	{
 		t.alias_index = v["alias_index"].asUInt();
 		t.alias_type = (VkObjectType)v["alias_type"].asUInt();
 	}
+	if (v.isMember("memory_flags")) t.memory_flags = (VkMemoryPropertyFlags)v["memory_flags"].asUInt();
 	t.enter_initialized();
 	return t;
 }
@@ -340,9 +336,6 @@ trackedimage trackedimage_json(const Json::Value& v)
 	t.sharingMode = (VkSharingMode)v["sharingMode"].asUInt();
 	t.usage = (VkImageUsageFlags)v["usage"].asUInt();
 	t.imageType = (VkImageType)v["imageType"].asUInt();
-	t.req.size = v["req_size"].asUInt64();
-	t.req.alignment = v["req_alignment"].asUInt();
-	t.req.memoryTypeBits = 0;
 	t.initialLayout = (VkImageLayout)(v.get("initialLayout", 0).asUInt());
 	t.currentLayout = t.initialLayout;
 	t.samples = (VkSampleCountFlagBits)(v.get("samples", 0).asUInt());
@@ -360,6 +353,7 @@ trackedimage trackedimage_json(const Json::Value& v)
 		t.alias_index = v["alias_index"].asUInt();
 		t.alias_type = (VkObjectType)v["alias_type"].asUInt();
 	}
+	if (v.isMember("memory_flags")) t.memory_flags = (VkMemoryPropertyFlags)v["memory_flags"].asUInt();
 	t.object_type = VK_OBJECT_TYPE_IMAGE;
 	t.enter_initialized();
 	return t;
@@ -370,9 +364,6 @@ trackedtensor trackedtensor_json(const Json::Value& v)
 	trackedtensor t;
 	trackable_helper(t, v);
 	t.sharingMode = (VkSharingMode)v["sharingMode"].asUInt();
-	t.req.size = v["req_size"].asUInt64();
-	t.req.alignment = v["req_alignment"].asUInt();
-	t.req.memoryTypeBits = 0;
 	t.object_type = VK_OBJECT_TYPE_TENSOR_ARM;
 	t.tiling = (VkTensorTilingARM)v["tiling"].asUInt();
 	t.format = (VkFormat)v["format"].asUInt();
@@ -387,6 +378,7 @@ trackedtensor trackedtensor_json(const Json::Value& v)
 		t.alias_index = v["alias_index"].asUInt();
 		t.alias_type = (VkObjectType)v["alias_type"].asUInt();
 	}
+	if (v.isMember("memory_flags")) t.memory_flags = (VkMemoryPropertyFlags)v["memory_flags"].asUInt();
 	t.enter_initialized();
 	return t;
 }

src/lavatube.h

@@ -175,6 +175,7 @@ struct trackedobject : trackable
 	VkObjectType alias_type = VK_OBJECT_TYPE_UNKNOWN;
 	uint32_t alias_index = UINT32_MAX;
 	VkDeviceAddress device_address = 0;
+	VkMemoryPropertyFlags memory_flags = 0;
 
 	bool is_state(states s) const { return (uint8_t)s == state; }
 	void set_state(states s) { state = (uint8_t)s; }

src/suballocator.cpp

@@ -73,6 +73,8 @@ struct suballocator_private
 	std::vector<lookup> image_lookup;
 	std::vector<lookup> buffer_lookup;
 	std::vector<lookup> tensor_lookup;
+	/// Does this device have the an annoying optimal-to-linear padding requirement? If so, put optimal and linear objects in different memory heaps
+	bool allow_mixed_tiling = true;
 
 	void print_memory_usage();
 	uint32_t get_device_memory_type(uint32_t type_filter, VkMemoryPropertyFlags& properties);
@@ -90,6 +92,15 @@ struct suballocator_private
 
 // helpers
 
+static VkMemoryPropertyFlags prune_memory_flags(VkMemoryPropertyFlags flags)
+{
+	if ((flags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT) || (flags & VK_MEMORY_PROPERTY_HOST_CACHED_BIT))
+	{
+		flags &= ~VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT; // do not require this bit in these cases
+	}
+	return flags;
+}
+
 void suballocator_private::print_memory_usage()
 {
 	printf("Suballocator memory usage:\n");
@@ -166,6 +177,9 @@ void suballocator::setup(VkPhysicalDevice physicaldevice)
 	if (priv->run)
 	{
 		wrap_vkGetPhysicalDeviceMemoryProperties(physicaldevice, &priv->memory_properties);
+		VkPhysicalDeviceProperties pdprops = {};
+		wrap_vkGetPhysicalDeviceProperties(physicaldevice, &pdprops);
+		priv->allow_mixed_tiling = (pdprops.limits.bufferImageGranularity == 1);
 	}
 	else
 	{
@@ -175,6 +189,7 @@ void suballocator::setup(VkPhysicalDevice physicaldevice)
 		priv->memory_properties.memoryHeapCount = 1;
 		priv->memory_properties.memoryHeaps[0].size = UINT64_MAX;
 		priv->memory_properties.memoryHeaps[0].flags = VK_MEMORY_HEAP_DEVICE_LOCAL_BIT;
+		priv->allow_mixed_tiling = true;
 	}
 	assert(priv->memory_properties.memoryTypeCount > 0);
 }
@@ -262,12 +277,11 @@ suballoc_location suballocator_private::add_object(VkDevice device, uint16_t tid
 			h.deletes.clear();
 		}
 		// find suballocation
-		if (h.tid == tid && (flags & f) == flags && h.free >= s.size && h.memoryTypeIndex == memoryTypeIndex && h.tiling == tiling)
+		if (h.tid == tid && (flags & f) == flags && h.free >= s.size && h.memoryTypeIndex == memoryTypeIndex && (h.tiling == tiling || allow_mixed_tiling))
 		{
 			// First case: nothing allocated in heap. In this case, we do not care about alignment, because according to the spec:
 			// "Allocations returned by vkAllocateMemory are guaranteed to meet any alignment requirement of the implementation."
-			// Also second case: some memory is available before the first allocation. Like above, we are guaranteed not to have
-			// to care about alignment in this case.
+			// Also second case: We place our allocation first. Like above, we are guaranteed not to have to care about alignment.
 			if (h.subs.empty() || (h.subs.front().offset >= s.size))
 			{
 				s.offset = 0;
@@ -354,30 +368,28 @@ void suballocator_private::bind(heap& h, const suballocation& s)
 	assert(s.offset + s.size <= h.total);
 }
 
-suballoc_location suballocator::add_image(uint16_t tid, VkDevice device, VkImage image, uint32_t image_index, VkMemoryPropertyFlags flags, VkImageTiling tiling, VkDeviceSize min_size)
+suballoc_location suballocator::add_image(uint16_t tid, VkDevice device, VkImage image, const trackedimage& image_data)
 {
-	if ((flags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT) || (flags & VK_MEMORY_PROPERTY_HOST_CACHED_BIT))
-	{
-		flags &= ~VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT; // do not require this bit in these cases
-	}
 	VkMemoryRequirements2 req = {};
-	const bool dedicated = priv->fill_image_memreq(device, image, req, min_size);
-	const uint32_t memoryTypeIndex = priv->get_device_memory_type(req.memoryRequirements.memoryTypeBits, flags);
+	VkMemoryPropertyFlags memory_flags = prune_memory_flags(image_data.memory_flags);
+	const bool dedicated = priv->fill_image_memreq(device, image, req, image_data.size);
+	const uint32_t memoryTypeIndex = priv->get_device_memory_type(req.memoryRequirements.memoryTypeBits, memory_flags);
 	suballocation s;
 	s.type = VK_OBJECT_TYPE_IMAGE;
 	s.handle.image = image;
-	s.size = std::max(req.memoryRequirements.size, min_size);
+	s.size = std::max(req.memoryRequirements.size, image_data.size);
 	s.offset = 0;
-	s.index = image_index;
+	s.index = image_data.index;
 	s.alignment = req.memoryRequirements.alignment;
-	auto r = priv->add_object(device, tid, memoryTypeIndex, s, flags, tiling, dedicated, 0);
+	auto r = priv->add_object(device, tid, memoryTypeIndex, s, memory_flags, image_data.tiling, dedicated, 0);
 	assert(r.offset == s.offset);
 	return r;
 }
 
-void suballocator::virtualswap_images(VkDevice device, const std::vector<VkImage>& images, VkMemoryPropertyFlags flags)
+void suballocator::virtualswap_images(VkDevice device, const std::vector<VkImage>& images, VkMemoryPropertyFlags memory_flags)
 {
 	assert(priv->run);
+	VkMemoryPropertyFlags flags = prune_memory_flags(memory_flags);
 	VkMemoryRequirements2 req = {};
 	const bool dedicated = priv->fill_image_memreq(device, images.at(0), req, 0);
 	const uint32_t memoryTypeIndex = priv->get_device_memory_type(req.memoryRequirements.memoryTypeBits, flags);
@@ -407,13 +419,10 @@ void suballocator::virtualswap_images(VkDevice device, const std::vector<VkImage
 	}
 }
 
-suballoc_location suballocator::add_buffer(uint16_t tid, VkDevice device, VkBuffer buffer, VkMemoryPropertyFlags mempropflags, const trackedbuffer& buffer_data)
+suballoc_location suballocator::add_buffer(uint16_t tid, VkDevice device, VkBuffer buffer, const trackedbuffer& buffer_data)
 {
+	VkMemoryPropertyFlags memory_flags = prune_memory_flags(buffer_data.memory_flags);
 	const VkBufferUsageFlags buffer_flags = buffer_data.usage;
-	if ((mempropflags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT) || (mempropflags & VK_MEMORY_PROPERTY_HOST_CACHED_BIT))
-	{
-		mempropflags &= ~VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT; // do not require this bit in these cases
-	}
 	VkMemoryRequirements2 req = { VK_STRUCTURE_TYPE_MEMORY_REQUIREMENTS_2, nullptr };
 	VkMemoryDedicatedRequirements dedicated = { VK_STRUCTURE_TYPE_MEMORY_DEDICATED_REQUIREMENTS, nullptr };
 	if (use_dedicated_allocation() && priv->run)
@@ -433,7 +442,7 @@ suballoc_location suballocator::add_buffer(uint16_t tid, VkDevice device, VkBuff
 		req.memoryRequirements.alignment = 1;
 		req.memoryRequirements.memoryTypeBits = 1;
 	}
-	uint32_t memoryTypeIndex = priv->get_device_memory_type(req.memoryRequirements.memoryTypeBits, mempropflags);
+	uint32_t memoryTypeIndex = priv->get_device_memory_type(req.memoryRequirements.memoryTypeBits, memory_flags);
 	suballocation s;
 	s.type = VK_OBJECT_TYPE_BUFFER;
 	s.handle.buffer = buffer;
@@ -443,7 +452,7 @@ suballoc_location suballocator::add_buffer(uint16_t tid, VkDevice device, VkBuff
 	s.alignment = req.memoryRequirements.alignment;
 	VkMemoryAllocateFlags allocflags = 0;
 	if (buffer_flags & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT) { dedicated.prefersDedicatedAllocation = VK_TRUE; allocflags |= VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT_KHR; }
-	auto r = priv->add_object(device, tid, memoryTypeIndex, s, mempropflags, VK_IMAGE_TILING_LINEAR, dedicated.prefersDedicatedAllocation, allocflags);
+	auto r = priv->add_object(device, tid, memoryTypeIndex, s, memory_flags, VK_IMAGE_TILING_LINEAR, dedicated.prefersDedicatedAllocation, allocflags);
 	assert(r.offset == s.offset);
 	return r;
 }

src/suballocator.h

@@ -29,10 +29,10 @@ struct suballocator
 
 	/// Add an image to our memory pools. Thread safe because each thread gets its own set of memory pools that only they
 	/// can modify. Other threads may access the objects stored inside subject to Vulkan external synchronization rules.
-	suballoc_location add_image(uint16_t tid, VkDevice device, VkImage image, uint32_t image_index, VkMemoryPropertyFlags flags, VkImageTiling tiling, VkDeviceSize min_size);
+	suballoc_location add_image(uint16_t tid, VkDevice device, VkImage image, const trackedimage& image_data);
 
 	/// Add a buffer to our memory pools. See above.
-	suballoc_location add_buffer(uint16_t tid, VkDevice device, VkBuffer buffer, VkMemoryPropertyFlags memory_flags, const trackedbuffer& buffer_data);
+	suballoc_location add_buffer(uint16_t tid, VkDevice device, VkBuffer buffer, const trackedbuffer& buffer_data);
 
 	/// Delete an image from our memory pools. Thread safe because the internal data structure is preallocated and never resized,
 	/// and deleted entries are never reused.