render/vulkan: add support for output color transforms

include/render/vulkan.h

@@ -248,6 +248,8 @@ struct wlr_vk_command_buffer {
 	struct wl_list destroy_textures; // wlr_vk_texture.destroy_link
 	// Staging shared buffers to release after the command buffer completes
 	struct wl_list stage_buffers; // wlr_vk_shared_buffer.link
+	// Color transform to unref after the command buffer completes
+	struct wlr_color_transform *color_transform;
 
 	// For DMA-BUF implicit sync interop, may be NULL
 	VkSemaphore binary_semaphore;
@@ -302,6 +304,8 @@ struct wlr_vk_renderer {
 
 	struct wl_list render_buffers; // wlr_vk_render_buffer.link
 
+	struct wl_list color_transforms; // wlr_vk_color_transform.link
+
 	// Pool of command buffers
 	struct wlr_vk_command_buffer command_buffers[VULKAN_COMMAND_BUFFERS_CAP];
 
@@ -373,10 +377,11 @@ struct wlr_vk_render_pass {
 	float projection[9];
 	bool failed;
 	bool srgb_pathway; // if false, rendering via intermediate blending buffer
+	struct wlr_color_transform *color_transform;
 };
 
 struct wlr_vk_render_pass *vulkan_begin_render_pass(struct wlr_vk_renderer *renderer,
-	struct wlr_vk_render_buffer *buffer);
+	struct wlr_vk_render_buffer *buffer, const struct wlr_buffer_pass_options *options);
 
 // Suballocates a buffer span with the given size that can be mapped
 // and used as staging buffer. The allocation is implicitly released when the
@@ -487,6 +492,22 @@ struct wlr_vk_buffer_span {
 	struct wlr_vk_allocation alloc;
 };
 
+
+// Lookup table for a color transform
+struct wlr_vk_color_transform {
+	struct wlr_addon addon; // owned by: wlr_vk_renderer
+	struct wl_list link; // wlr_vk_renderer, list of all color transforms
+
+	struct {
+		VkImage image;
+		VkImageView image_view;
+		VkDeviceMemory memory;
+		VkDescriptorSet ds;
+		struct wlr_vk_descriptor_pool *ds_pool;
+	} lut_3d;
+};
+void vk_color_transform_destroy(struct wlr_addon *addon);
+
 // util
 const char *vulkan_strerror(VkResult err);
 void vulkan_change_layout(VkCommandBuffer cb, VkImage img,

include/wlr/render/pass.h

@@ -30,6 +30,9 @@ struct wlr_render_timer;
 struct wlr_buffer_pass_options {
 	/* Timer to measure the duration of the render pass */
 	struct wlr_render_timer *timer;
+	/* Color transform to apply to the output of the render pass,
+	 * leave NULL to indicate sRGB/no custom transform */
+	struct wlr_color_transform *color_transform;
 };
 
 /**

render/vulkan/pass.c

@@ -2,11 +2,14 @@
 #include <drm_fourcc.h>
 #include <stdlib.h>
 #include <wlr/util/log.h>
+#include <wlr/render/color.h>
 
+#include "render/color.h"
 #include "render/vulkan.h"
 #include "types/wlr_matrix.h"
 
 static const struct wlr_render_pass_impl render_pass_impl;
+static const struct wlr_addon_interface vk_color_transform_impl;
 
 static struct wlr_vk_render_pass *get_render_pass(struct wlr_render_pass *wlr_pass) {
 	assert(wlr_pass->impl == &render_pass_impl);
@@ -14,6 +17,16 @@ static struct wlr_vk_render_pass *get_render_pass(struct wlr_render_pass *wlr_pa
 	return pass;
 }
 
+static struct wlr_vk_color_transform *get_color_transform(
+		struct wlr_color_transform *c, struct wlr_vk_renderer *renderer) {
+	struct wlr_addon *a = wlr_addon_find(&c->addons, renderer, &vk_color_transform_impl);
+	if (!a) {
+		return NULL;
+	}
+	struct wlr_vk_color_transform *transform = wl_container_of(a, transform, addon);
+	return transform;
+}
+
 static void bind_pipeline(struct wlr_vk_render_pass *pass, VkPipeline pipeline) {
 	if (pipeline == pass->bound_pipeline) {
 		return;
@@ -103,21 +116,36 @@ static bool render_pass_submit(struct wlr_render_pass *wlr_pass) {
 			.uv_off = { 0, 0 },
 			.uv_size = { 1, 1 },
 		};
+		size_t dim = pass->color_transform ? pass->color_transform->lut3d.dim_len : 1;
 		struct wlr_vk_frag_output_pcr_data frag_pcr_data = {
-			.lut_3d_offset = 0.5f / 1,
+			.lut_3d_offset = 0.5f / dim,
-			.lut_3d_scale = (float)(1 - 1) / 1,
+			.lut_3d_scale = (float)(dim - 1) / dim,
 		};
 		mat3_to_mat4(final_matrix, vert_pcr_data.mat4);
 
+		if (pass->color_transform) {
+			bind_pipeline(pass, render_buffer->plain.render_setup->output_pipe_lut3d);
+		} else {
 			bind_pipeline(pass, render_buffer->plain.render_setup->output_pipe_srgb);
+		}
 		vkCmdPushConstants(render_cb->vk, renderer->output_pipe_layout,
 			VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(vert_pcr_data), &vert_pcr_data);
 		vkCmdPushConstants(render_cb->vk, renderer->output_pipe_layout,
 			VK_SHADER_STAGE_FRAGMENT_BIT, sizeof(vert_pcr_data),
 			sizeof(frag_pcr_data), &frag_pcr_data);
+
+		VkDescriptorSet lut_ds;
+		if (pass->color_transform && pass->color_transform->type == COLOR_TRANSFORM_LUT_3D) {
+			struct wlr_vk_color_transform *transform =
+				get_color_transform(pass->color_transform, renderer);
+			assert(transform);
+			lut_ds = transform->lut_3d.ds;
+		} else {
+			lut_ds = renderer->output_ds_lut3d_dummy;
+		}
 		VkDescriptorSet ds[] = {
 			render_buffer->plain.blend_descriptor_set, // set 0
-			renderer->output_ds_lut3d_dummy, // set 1
+			lut_ds, // set 1
 		};
 		size_t ds_len = sizeof(ds) / sizeof(ds[0]);
 		vkCmdBindDescriptorSets(render_cb->vk,
@@ -679,9 +707,247 @@ static const struct wlr_render_pass_impl render_pass_impl = {
 	.add_texture = render_pass_add_texture,
 };
 
+
+void vk_color_transform_destroy(struct wlr_addon *addon) {
+	struct wlr_vk_renderer *renderer = (struct wlr_vk_renderer *)addon->owner;
+	struct wlr_vk_color_transform *transform = wl_container_of(addon, transform, addon);
+
+	VkDevice dev = renderer->dev->dev;
+	if (transform->lut_3d.image) {
+		vkDestroyImage(dev, transform->lut_3d.image, NULL);
+		vkDestroyImageView(dev, transform->lut_3d.image_view, NULL);
+		vkFreeMemory(dev, transform->lut_3d.memory, NULL);
+		vulkan_free_ds(renderer, transform->lut_3d.ds_pool, transform->lut_3d.ds);
+	}
+
+	wl_list_remove(&transform->link);
+	wlr_addon_finish(&transform->addon);
+	free(transform);
+}
+
+static bool create_3d_lut_image(struct wlr_vk_renderer *renderer,
+		const struct wlr_color_transform_lut3d *lut_3d,
+		VkImage *image, VkImageView *image_view,
+		VkDeviceMemory *memory,	VkDescriptorSet *ds,
+		struct wlr_vk_descriptor_pool **ds_pool) {
+	VkDevice dev = renderer->dev->dev;
+	VkResult res;
+
+	*image = VK_NULL_HANDLE;
+	*memory = VK_NULL_HANDLE;
+	*image_view = VK_NULL_HANDLE;
+	*ds = VK_NULL_HANDLE;
+	*ds_pool = NULL;
+
+	// R32G32B32 is not a required Vulkan format
+	// TODO: use it when available
+	VkFormat format = VK_FORMAT_R32G32B32A32_SFLOAT;
+
+	VkImageCreateInfo img_info = {
+		.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
+		.imageType = VK_IMAGE_TYPE_3D,
+		.format = format,
+		.mipLevels = 1,
+		.arrayLayers = 1,
+		.samples = VK_SAMPLE_COUNT_1_BIT,
+		.sharingMode = VK_SHARING_MODE_EXCLUSIVE,
+		.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
+		.extent = (VkExtent3D) { lut_3d->dim_len, lut_3d->dim_len, lut_3d->dim_len },
+		.tiling = VK_IMAGE_TILING_OPTIMAL,
+		.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT,
+	};
+	res = vkCreateImage(dev, &img_info, NULL, image);
+	if (res != VK_SUCCESS) {
+		wlr_vk_error("vkCreateImage failed", res);
+		return NULL;
+	}
+
+	VkMemoryRequirements mem_reqs = {0};
+	vkGetImageMemoryRequirements(dev, *image, &mem_reqs);
+
+	int mem_type_index = vulkan_find_mem_type(renderer->dev,
+		VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, mem_reqs.memoryTypeBits);
+	if (mem_type_index == -1) {
+		wlr_log(WLR_ERROR, "Failed to find suitable memory type");
+		goto fail_image;
+	}
+
+	VkMemoryAllocateInfo mem_info = {
+		.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
+		.allocationSize = mem_reqs.size,
+		.memoryTypeIndex = mem_type_index,
+	};
+	res = vkAllocateMemory(dev, &mem_info, NULL, memory);
+	if (res != VK_SUCCESS) {
+		wlr_vk_error("vkAllocateMemory failed", res);
+		goto fail_image;
+	}
+
+	res = vkBindImageMemory(dev, *image, *memory, 0);
+	if (res != VK_SUCCESS) {
+		wlr_vk_error("vkBindMemory failed", res);
+		goto fail_memory;
+	}
+
+	VkImageViewCreateInfo view_info = {
+		.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
+		.viewType = VK_IMAGE_VIEW_TYPE_3D,
+		.format = format,
+		.components.r = VK_COMPONENT_SWIZZLE_IDENTITY,
+		.components.g = VK_COMPONENT_SWIZZLE_IDENTITY,
+		.components.b = VK_COMPONENT_SWIZZLE_IDENTITY,
+		.components.a = VK_COMPONENT_SWIZZLE_IDENTITY,
+		.subresourceRange = (VkImageSubresourceRange) {
+			.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+			.baseMipLevel = 0,
+			.levelCount = 1,
+			.baseArrayLayer = 0,
+			.layerCount = 1,
+		},
+		.image = *image,
+	};
+	res = vkCreateImageView(dev, &view_info, NULL, image_view);
+	if (res != VK_SUCCESS) {
+		wlr_vk_error("vkCreateImageView failed", res);
+		goto fail_image;
+	}
+
+	size_t bytes_per_block = 4 * sizeof(float);
+	size_t size = lut_3d->dim_len * lut_3d->dim_len * lut_3d->dim_len * bytes_per_block;
+	struct wlr_vk_buffer_span span = vulkan_get_stage_span(renderer,
+		size, bytes_per_block);
+	if (!span.buffer || span.alloc.size != size) {
+		wlr_log(WLR_ERROR, "Failed to retrieve staging buffer");
+		goto fail_imageview;
+	}
+
+	void *data;
+	res = vkMapMemory(dev, span.buffer->memory, span.alloc.start, size, 0, &data);
+	if (res != VK_SUCCESS) {
+		wlr_vk_error("vkMapMemory", res);
+		goto fail_imageview;
+	}
+
+	float *dst = data;
+	size_t dim_len = lut_3d->dim_len;
+	for (size_t b_index = 0; b_index < dim_len; b_index++) {
+		for (size_t g_index = 0; g_index < dim_len; g_index++) {
+			for (size_t r_index = 0; r_index < dim_len; r_index++) {
+				size_t sample_index = r_index + dim_len * g_index + dim_len * dim_len * b_index;
+				size_t src_offset = 3 * sample_index;
+				size_t dst_offset = 4 * sample_index;
+				dst[dst_offset] = lut_3d->lut_3d[src_offset];
+				dst[dst_offset + 1] = lut_3d->lut_3d[src_offset + 1];
+				dst[dst_offset + 2] = lut_3d->lut_3d[src_offset + 2];
+				dst[dst_offset + 3] = 1.0;
+			}
+		}
+	}
+
+	vkUnmapMemory(dev, span.buffer->memory);
+
+	VkCommandBuffer cb = vulkan_record_stage_cb(renderer);
+	vulkan_change_layout(cb, *image,
+		VK_IMAGE_LAYOUT_UNDEFINED, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, 0,
+		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT,
+		VK_ACCESS_TRANSFER_WRITE_BIT);
+	VkBufferImageCopy copy = {
+		.bufferOffset = span.alloc.start,
+		.imageExtent.width = lut_3d->dim_len,
+		.imageExtent.height = lut_3d->dim_len,
+		.imageExtent.depth = lut_3d->dim_len,
+		.imageSubresource.layerCount = 1,
+		.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+	};
+	vkCmdCopyBufferToImage(cb, span.buffer->buffer, *image,
+		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copy);
+	vulkan_change_layout(cb, *image,
+		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT,
+		VK_ACCESS_TRANSFER_WRITE_BIT,
+		VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
+		VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT, VK_ACCESS_SHADER_READ_BIT);
+
+	*ds_pool = vulkan_alloc_texture_ds(renderer,
+		renderer->output_ds_lut3d_layout, ds);
+	if (!*ds_pool) {
+		wlr_log(WLR_ERROR, "Failed to allocate descriptor");
+		goto fail_imageview;
+	}
+
+	VkDescriptorImageInfo ds_img_info = {
+		.imageView = *image_view,
+		.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
+	};
+	VkWriteDescriptorSet ds_write = {
+		.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
+		.descriptorCount = 1,
+		.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
+		.dstSet = *ds,
+		.pImageInfo = &ds_img_info,
+	};
+	vkUpdateDescriptorSets(dev, 1, &ds_write, 0, NULL);
+
+	return true;
+
+fail_imageview:
+	vkDestroyImageView(dev, *image_view, NULL);
+fail_memory:
+	vkFreeMemory(dev, *memory, NULL);
+fail_image:
+	vkDestroyImage(dev, *image, NULL);
+	return false;
+}
+
+static struct wlr_vk_color_transform *vk_color_transform_create(
+		struct wlr_vk_renderer *renderer, struct wlr_color_transform *transform) {
+	struct wlr_vk_color_transform *vk_transform =
+		calloc(1, sizeof(*vk_transform));
+	if (!vk_transform) {
+		return NULL;
+	}
+
+	if (transform->type == COLOR_TRANSFORM_LUT_3D) {
+		if (!create_3d_lut_image(renderer, &transform->lut3d,
+				&vk_transform->lut_3d.image,
+				&vk_transform->lut_3d.image_view,
+				&vk_transform->lut_3d.memory,
+				&vk_transform->lut_3d.ds,
+				&vk_transform->lut_3d.ds_pool)) {
+			free(vk_transform);
+			return NULL;
+		}
+	}
+
+	wlr_addon_init(&vk_transform->addon, &transform->addons,
+		renderer, &vk_color_transform_impl);
+	wl_list_insert(&renderer->color_transforms, &vk_transform->link);
+
+	return vk_transform;
+}
+
+
+static const struct wlr_addon_interface vk_color_transform_impl = {
+	"vk_color_transform",
+	.destroy = vk_color_transform_destroy,
+};
+
 struct wlr_vk_render_pass *vulkan_begin_render_pass(struct wlr_vk_renderer *renderer,
-		struct wlr_vk_render_buffer *buffer) {
+		struct wlr_vk_render_buffer *buffer, const struct wlr_buffer_pass_options *options) {
-	bool using_srgb_pathway = buffer->srgb.framebuffer != VK_NULL_HANDLE;
+	bool using_srgb_pathway;
+	if (options != NULL && options->color_transform != NULL) {
+		using_srgb_pathway = false;
+
+		if (!get_color_transform(options->color_transform, renderer)) {
+			/* Try to create a new color transform */
+			if (!vk_color_transform_create(renderer, options->color_transform)) {
+				wlr_log(WLR_ERROR, "Failed to create color transform");
+				return NULL;
+			}
+		}
+	} else {
+		// Use srgb pathway if it is the default/has already been set up
+		using_srgb_pathway = buffer->srgb.framebuffer != VK_NULL_HANDLE;
+	}
 
 	if (!using_srgb_pathway && !buffer->plain.image_view) {
 		struct wlr_dmabuf_attributes attribs;
@@ -700,6 +966,10 @@ struct wlr_vk_render_pass *vulkan_begin_render_pass(struct wlr_vk_renderer *rend
 	wlr_render_pass_init(&pass->base, &render_pass_impl);
 	pass->renderer = renderer;
 	pass->srgb_pathway = using_srgb_pathway;
+	if (options != NULL && options->color_transform != NULL) {
+		wlr_color_transform_ref(options->color_transform);
+		pass->color_transform = options->color_transform;
+	}
 
 	rect_union_init(&pass->updated_region);
 

render/vulkan/renderer.c

@@ -8,6 +8,7 @@
 #include <unistd.h>
 #include <drm_fourcc.h>
 #include <vulkan/vulkan.h>
+#include <wlr/render/color.h>
 #include <wlr/render/interface.h>
 #include <wlr/types/wlr_drm.h>
 #include <wlr/util/box.h>
@@ -457,6 +458,11 @@ static void release_command_buffer_resources(struct wlr_vk_command_buffer *cb,
 		wl_list_remove(&buf->link);
 		wl_list_insert(&renderer->stage.buffers, &buf->link);
 	}
+
+	if (cb->color_transform) {
+		wlr_color_transform_unref(cb->color_transform);
+		cb->color_transform = NULL;
+	}
 }
 
 static struct wlr_vk_command_buffer *get_command_buffer(
@@ -1074,6 +1080,12 @@ static void vulkan_destroy(struct wlr_renderer *wlr_renderer) {
 		destroy_render_buffer(render_buffer);
 	}
 
+	struct wlr_vk_color_transform *color_transform, *color_transform_tmp;
+	wl_list_for_each_safe(color_transform, color_transform_tmp,
+			&renderer->color_transforms, link) {
+		vk_color_transform_destroy(&color_transform->addon);
+	}
+
 	struct wlr_vk_render_format_setup *setup, *tmp_setup;
 	wl_list_for_each_safe(setup, tmp_setup,
 			&renderer->render_format_setups, link) {
@@ -1380,7 +1392,8 @@ static struct wlr_render_pass *vulkan_begin_buffer_pass(struct wlr_renderer *wlr
 		}
 	}
 
-	struct wlr_vk_render_pass *render_pass = vulkan_begin_render_pass(renderer, render_buffer);
+	struct wlr_vk_render_pass *render_pass = vulkan_begin_render_pass(
+		renderer, render_buffer, options);
 	if (render_pass == NULL) {
 		return NULL;
 	}
@@ -2397,6 +2410,7 @@ struct wlr_renderer *vulkan_renderer_create_for_device(struct wlr_vk_device *dev
 	wl_list_init(&renderer->output_descriptor_pools);
 	wl_list_init(&renderer->render_format_setups);
 	wl_list_init(&renderer->render_buffers);
+	wl_list_init(&renderer->color_transforms);
 	wl_list_init(&renderer->pipeline_layouts);
 
 	if (!init_static_render_data(renderer)) {