seperate descriptor sets from pipelines to allow swapping out shader resources without recreating pipelines

1 files changed, 193 insertions, 109 deletions

diff --git a/video.cpp b/video.cpp
index 1ce0cba..57b5c73 100644
--- a/video.cpp
+++ b/video.cpp
@@ -7,6 +7,7 @@
 #define max_vertex_formats 64
 #define max_rpos 64
 #define max_pipelines 64
+#define max_descriptor_sets 64
 #define max_shaders 32
 #define max_samplers 16
 
@@ -515,24 +516,27 @@ struct Renderpass_Vk {
 };
 
 struct Pso_Key {
-	Pipeline pip;
+	Pipeline pso;
 	Render_Pass rpo;
 
 	bool operator==(const Pso_Key& other) const {
-		int size = sizeof *this, i;
-		uint8_t* bba = (uint8_t*)this;
-		uint8_t* bbb = (uint8_t*)&other;
-		for (i = 0; i < size; i++, bba++, bbb++)
-			if (*bba != *bbb) return false;
-		return true;
+		return
+			rpo == other.rpo &&
+			pso.desc_layout_eq(other.pso);
+	}
+};
+
+struct Dso_Key {
+	Pipeline pip;
+
+	bool operator==(const Dso_Key& other) const {
+		return pip.desc_resources_eq(other.pip);
 	}
 };
 
 struct Pipeline_Vk {
 	VkPipeline pip;
 	VkPipelineLayout lay;
-	VkDescriptorPool dp;
-	VkDescriptorSet dset;
 	VkDescriptorSetLayout dlay;
 	int age;
 
@@ -593,9 +597,25 @@ struct Pipeline_Vk {
 	);
 	void init_descriptors(
 		Device_Vk* dev,
-		const Pipeline& desc,
-		int count
+		const Pipeline& desc
+	);
+
+	void on_submit() {
+		age = 0;
+	}
+};
+
+struct Descriptor_Set_Vk {
+	VkDescriptorPool dp;
+	VkDescriptorSet dset;
+	int age;
+
+	void init(
+		Device_Vk* dev,
+		const Pipeline_Vk& pip,
+		const Pipeline& desc
 	);
+	void destroy(Device_Vk* dev);
 
 	void on_submit() {
 		age = 0;
@@ -636,8 +656,20 @@ struct Hash_Function<Render_Pass>
 template<>
 struct Hash_Function<Pso_Key>
 {
-	size_t operator()(const Pso_Key& rp) const {
-		return (size_t)fnv1a64((uint8_t*)&rp, sizeof rp);
+	size_t operator()(const Pso_Key& k) const {
+		return fnv1a64_2(
+			k.pso.pipeline_hash,
+			(uint8_t*)&k.rpo,
+			sizeof k.rpo
+		);
+	}
+};
+
+template<>
+struct Hash_Function<Dso_Key>
+{
+	size_t operator()(const Dso_Key& k) const {
+		return k.pip.descriptor_resource_hash;
 	}
 };
 
@@ -764,6 +796,7 @@ struct Device_Vk : public Device {
 
 	Hash_Map<Render_Pass, Renderpass_Vk, max_rpos> rpo_cache;
 	Hash_Map<Pso_Key, Pipeline_Vk, max_pipelines> pso_cache;
+	Hash_Map<Dso_Key, Descriptor_Set_Vk, max_descriptor_sets> dso_cache;
 
 	Terminator* terminators;
 	uint32_t terminator_index;
@@ -793,6 +826,14 @@ struct Device_Vk : public Device {
 	Renderpass_Vk& get_rpo(const Render_Pass& rp);
 	Pipeline_Vk& create_pso(const Pso_Key& pip);
 	Pipeline_Vk& get_pso(const Pso_Key& pop);
+	Descriptor_Set_Vk& create_dso(
+		const Pipeline_Vk& pip,
+		const Dso_Key& k
+	);
+	Descriptor_Set_Vk& get_dso(
+		const Pipeline_Vk& pip,
+		const Dso_Key& k
+	);
 
 	void collect_garbage();
 	void queue_destroy(Late_Terminated* obj);
@@ -1113,6 +1154,7 @@ void Device_Vk::init_internal() {
 	sampler_count = 1;
 	rpo_cache.init();
 	pso_cache.init();
+	dso_cache.init();
 	shader_loader.init(this);
 	texture_loader.init(this);
 	register_asset_loader("CSH2", &shader_loader);
@@ -1162,6 +1204,8 @@ void Device_Vk::deinit_internal() {
 		i.second.destroy(this);
 	for (auto i : pso_cache)
 		i.second.destroy(this);
+	for (auto i : dso_cache)
+		i.second.destroy(this);
 	for (i = 0; i < max_contexts; i++) {
 		auto& context = contexts[i];
 		if (context.state & context_state_init)
@@ -1226,6 +1270,26 @@ Pipeline_Vk& Device_Vk::get_pso(const Pso_Key& pip) {
 	return *pso;
 }
 
+Descriptor_Set_Vk& Device_Vk::create_dso(
+	const Pipeline_Vk& pip,
+	const Dso_Key& k
+) {
+	Descriptor_Set_Vk dso;
+	dso.age = 0;
+	dso.init(this, pip, k.pip);
+	return dso_cache.set(k, dso);
+}
+
+Descriptor_Set_Vk& Device_Vk::get_dso(
+	const Pipeline_Vk& pip,
+	const Dso_Key& k
+) {
+	Descriptor_Set_Vk* dso = dso_cache.get(k);
+	if (!dso)
+		return create_dso(pip, k);
+	return *dso;
+}
+
 void Renderpass_Vk::destroy(Device_Vk* dev) {
 	vkDestroyRenderPass(dev->dev, rpo, &dev->ac);
 	vkDestroyFramebuffer(dev->dev, fbo, &dev->ac);
@@ -1249,6 +1313,14 @@ void Device_Vk::collect_garbage() {
 			pso_cache.remove(i.first);
 		}
 	}
+	for (const auto& i: dso_cache) {
+		auto& dso = i.second;
+		dso.age++;
+		if (dso.age > max_age) {
+			dso.destroy(this);
+			dso_cache.remove(i.first);
+		}
+	}
 }
 
 void Device_Vk::queue_destroy(Late_Terminated* obj) {
@@ -1665,6 +1737,7 @@ void Context::submit(
 	Vertex_Buffer_Binding* binding;
 	Pso_Key pso_key = { p, rp };
 	Pipeline_Vk& pso = dev->get_pso(pso_key);
+	Descriptor_Set_Vk& dso = dev->get_dso(pso, *(Dso_Key*)&p);
 	auto& rpo = ctx->begin_rp(rp);
 	vkCmdBindPipeline(
 		ctx->cb,
@@ -1677,7 +1750,7 @@ void Context::submit(
 		pso.lay,
 		0,
 		1,
-		&pso.dset,
+		&dso.dset,
 		0,
 		0
 	);
@@ -1695,6 +1768,7 @@ void Context::submit(
 	);
 	ctx->end_rp(rpo);
 	pso.on_submit();
+	dso.on_submit();
 }
 
 void Context::submit(
@@ -2169,51 +2243,13 @@ void Pipeline_Vk::init_blending(
 
 void Pipeline_Vk::init_descriptors(
 	Device_Vk* dev,
-	const Pipeline& desc,
-	int count
+	const Pipeline& desc
 ) {
-	int sampler_count = 0, cbuffer_count = 0;
-	const Descriptor* d = desc.descriptors;
+	const Descriptor* sdescs = desc.descriptors;
 	Shader_Vk& shader = *(Shader_Vk*)&dev->get_shader(desc.shader);
-	VkDescriptorPoolSize sizes[4];
 	VkResult r;
-	int size_count = 0, i;
-	for (; d; d = d->next) {
-		switch (d->type) {
-			case Descriptor::Type::texture:
-				sampler_count++;
-				break;
-			default:
-				assert(0);
-				break;
-		}
-	}
-	if (sampler_count) {
-		int idx = size_count++;
-		sizes[idx] = {
-			.type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-			.descriptorCount = (uint32_t)sampler_count
-		};
-	}
-	if (cbuffer_count) {
-		int idx = size_count++;
-		sizes[idx] = {
-			.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-			.descriptorCount = (uint32_t)cbuffer_count
-		};
-	}
-	{
-		VkDescriptorPoolCreateInfo di{};
-		di.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
-		di.poolSizeCount = (uint32_t)size_count;
-		di.pPoolSizes = sizes;
-		di.maxSets = (uint32_t)count;
-		r = vkCreateDescriptorPool(dev->dev, &di, &dev->ac, &dp);
-		if (r != VK_SUCCESS) {
-			print_err("Failed to create a descriptor pool.\n");
-			pbreak(r);
-		}
-	}
+	int count = desc.descriptor_count;
+	int i;
 	{
 		VkDescriptorSetLayoutBinding* descs =
 			(VkDescriptorSetLayoutBinding*)heap_alloc(
@@ -2221,13 +2257,12 @@ void Pipeline_Vk::init_descriptors(
 				count * sizeof *descs
 			);
 		VkDescriptorSetLayoutCreateInfo di{};
-		VkDescriptorSetAllocateInfo da{};
 		memset(descs, 0, count * sizeof *descs);
-		Descriptor* src = desc.descriptors;
-		for (i = count - 1; i >= 0; i--) {
+		for (i = 0; i < count; i++) {
 			int j, stage;
 			auto& dst = descs[i];
-			switch (src->type) {
+			auto& src = sdescs[i];
+			switch (src.type) {
 				case Descriptor::Type::texture:
 					dst.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
 					break;
@@ -2235,16 +2270,15 @@ void Pipeline_Vk::init_descriptors(
 					assert(0);
 					break;
 			}
-			dst.binding = src->slot;
+			dst.binding = src.slot;
 			dst.descriptorCount = 1;
 			dst.stageFlags = 0;
-			stage = shader.descriptor_stage(src->slot);
+			stage = shader.descriptor_stage(src.slot);
 			for (j = 0; j < shader_type_count; j++) {
 				if (stage & (1 << j)) {
 					dst.stageFlags |= Shader_Vk::stage((Shader_Type)j);
 				}
 			}
-			src = src->next;
 		}
 		di.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
 		di.bindingCount = (uint32_t)count;
@@ -2259,47 +2293,6 @@ void Pipeline_Vk::init_descriptors(
 			print_err("Failed to create descriptor set layout.\n");
 			pbreak(r);
 		}
-		da.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
-		da.descriptorPool = dp;
-		da.descriptorSetCount = 1;
-		da.pSetLayouts = &dlay;
-		r = vkAllocateDescriptorSets(
-			dev->dev,
-			&da,
-			&dset
-		);
-		/* todo this should be refactored to allow
-		 * swapping the resources without recreating pipelines */
-		src = desc.descriptors;
-		for (i = 0; i < count; i++) {
-			VkDescriptorImageInfo img{};
-			VkWriteDescriptorSet wd{};
-			wd.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
-			wd.dstSet = dset;
-			wd.dstBinding = src->slot;
-			wd.dstArrayElement = 0;
-			wd.descriptorCount = 1;
-
-			switch (src->type) {
-				case Descriptor::Type::texture: {
-					Texture_Descriptor* td = (Texture_Descriptor*)src;
-					assert(td->texture);
-					assert(td->sampler);
-					Texture_Vk& t = *(Texture_Vk*)&dev->get_texture(td->texture);
-					Sampler_Vk& s = *(Sampler_Vk*)&dev->samplers[td->sampler];
-					img.imageView = t.view;
-					img.sampler = s.sampler;
-					img.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
-					wd.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
-					wd.pImageInfo = &img;
-				} break;
-				default:
-					assert(0);
-					break;
-			}
-			vkUpdateDescriptorSets(dev->dev, 1, &wd, 0, 0);
-			src = src->next;
-		}
 		heap_free(dev->heap, descs);
 	}
 }
@@ -2311,10 +2304,9 @@ void Pipeline_Vk::init_layout(
 	VkResult r;
 	VkPipelineLayoutCreateInfo li{};
 	(void)desc;
-	int desc_count = desc.count_descriptors();
-	int set_count = desc_count? 1: 0;
+	int set_count = desc.descriptor_count? 1: 0;
 	if (set_count) {
-		init_descriptors(dev, desc, desc_count);
+		init_descriptors(dev, desc);
 	}
 	li.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
 	li.setLayoutCount = set_count;
@@ -2336,7 +2328,7 @@ void Pipeline_Vk::init(Device_Vk* dev, const Pso_Key& key) {
 	char buffer[1024];
 	Arena scope;
 	VkResult r;
-	const auto& desc = key.pip;
+	const auto& desc = key.pso;
 	VkGraphicsPipelineCreateInfo info{};
 	init_arena(&scope, buffer, sizeof buffer);
 	init_layout(dev, desc);
@@ -2368,12 +2360,104 @@ void Pipeline_Vk::init(Device_Vk* dev, const Pso_Key& key) {
 }
 
 void Pipeline_Vk::destroy(Device_Vk* dev) {
-	vkDestroyDescriptorPool(dev->dev, dp, &dev->ac);
 	vkDestroyDescriptorSetLayout(dev->dev, dlay, &dev->ac);
 	vkDestroyPipelineLayout(dev->dev, lay, &dev->ac);
 	vkDestroyPipeline(dev->dev, pip, &dev->ac);
 }
 
+void Descriptor_Set_Vk::init(
+	Device_Vk* dev,
+	const Pipeline_Vk& pip,
+	const Pipeline& desc
+) {
+	int count = desc.descriptor_count, i;
+	int sampler_count = 0, cbuffer_count = 0;
+	int size_count = 0;
+	VkDescriptorSetAllocateInfo da{};
+	VkDescriptorPoolSize sizes[4];
+	VkResult r;
+	for (i = 0; i < count; i++) {
+		auto& src = desc.descriptors[i];
+		switch (src.type) {
+			case Descriptor::Type::texture:
+				sampler_count++;
+				break;
+			default:
+				assert(0);
+				break;
+		}
+	}
+	if (sampler_count) {
+		int idx = size_count++;
+		sizes[idx] = {
+			.type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
+			.descriptorCount = (uint32_t)sampler_count
+		};
+	}
+	if (cbuffer_count) {
+		int idx = size_count++;
+		sizes[idx] = {
+			.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
+			.descriptorCount = (uint32_t)cbuffer_count
+		};
+	}
+	{
+		VkDescriptorPoolCreateInfo di{};
+		di.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
+		di.poolSizeCount = (uint32_t)size_count;
+		di.pPoolSizes = sizes;
+		di.maxSets = (uint32_t)count;
+		r = vkCreateDescriptorPool(dev->dev, &di, &dev->ac, &dp);
+		if (r != VK_SUCCESS) {
+			print_err("Failed to create a descriptor pool.\n");
+			pbreak(r);
+		}
+	}
+	da.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
+	da.descriptorPool = dp;
+	da.descriptorSetCount = 1;
+	da.pSetLayouts = &pip.dlay;
+	r = vkAllocateDescriptorSets(
+		dev->dev,
+		&da,
+		&dset
+	);
+	if (r != VK_SUCCESS) {
+		print_err("Failed to allocate descriptor set.\n");
+		pbreak(r);
+	}
+	for (i = 0; i < count; i++) {
+		VkDescriptorImageInfo img{};
+		VkWriteDescriptorSet wd{};
+		auto& src = desc.descriptors[i];
+		wd.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
+		wd.dstSet = dset;
+		wd.dstBinding = src.slot;
+		wd.dstArrayElement = 0;
+		wd.descriptorCount = 1;
+
+		switch (src.type) {
+			case Descriptor::Type::texture: {
+				Texture_Descriptor* td = (Texture_Descriptor*)src.payload;
+				assert(td->texture);
+				assert(td->sampler);
+				Texture_Vk& t = *(Texture_Vk*)&dev->get_texture(td->texture);
+				Sampler_Vk& s = *(Sampler_Vk*)&dev->samplers[td->sampler];
+				img.imageView = t.view;
+				img.sampler = s.sampler;
+				img.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
+				wd.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
+				wd.pImageInfo = &img;
+			} break;
+		}
+		vkUpdateDescriptorSets(dev->dev, 1, &wd, 0, 0);
+	}
+}
+
+void Descriptor_Set_Vk::destroy(Device_Vk* dev) {
+	vkDestroyDescriptorPool(dev->dev, dp, &dev->ac);
+}
+
 int Vertex_Format_Vk::svariable_type_size(SVariable_Type type) {
 	switch (type) {
 		case svariable_type_float: return 4;