c/render: Add ATW support for compute rendering

src/xrt/compositor/CMakeLists.txt

@@ -4,6 +4,7 @@
 spirv_shaders(SHADER_HEADERS
 	shaders/clear.comp
 	shaders/distortion.comp
+	shaders/distortion_timewarp.comp
 	shaders/mesh.frag
 	shaders/mesh.vert
 	shaders/layer.frag

src/xrt/compositor/main/comp_compositor.c

@@ -1509,6 +1509,7 @@ xrt_gfx_provider_create_system(struct xrt_device *xdev, struct xrt_system_compos
 
 	u_var_add_root(c, "Compositor", true);
 	u_var_add_ro_f32(c, &c->compositor_frame_times.fps, "FPS (Compositor)");
+	u_var_add_bool(c, &c->debug.atw_off, "Debug: ATW OFF");
 
 	struct u_var_timing *ft = U_TYPED_CALLOC(struct u_var_timing);
 

src/xrt/compositor/main/comp_compositor.h

@@ -141,6 +141,7 @@ struct comp_shaders
 {
 	VkShaderModule clear_comp;
 	VkShaderModule distortion_comp;
+	VkShaderModule distortion_timewarp_comp;
 
 	VkShaderModule mesh_vert;
 	VkShaderModule mesh_frag;
@@ -245,6 +246,12 @@ struct comp_compositor
 	} threading;
 
 
+	struct
+	{
+		//! Temporarily disable ATW
+		bool atw_off;
+	} debug;
+
 	struct comp_resources nr;
 };
 

src/xrt/compositor/main/comp_renderer.c

@@ -21,6 +21,8 @@
 
 #include "main/comp_layer_renderer.h"
 #include "math/m_api.h"
+#include "math/m_vec3.h"
+#include "math/m_matrix_4x4_f64.h"
 
 #include <string.h>
 #include <stdlib.h>
@@ -728,6 +730,43 @@ dispatch_graphics(struct comp_renderer *r)
  *
  */
 
+static void
+get_view_poses(struct comp_renderer *r, struct xrt_pose out_results[2])
+{
+	COMP_TRACE_MARKER();
+
+	struct xrt_space_relation relation;
+
+	xrt_device_get_tracked_pose(                         //
+	    r->c->xdev,                                      //
+	    XRT_INPUT_GENERIC_HEAD_POSE,                     //
+	    r->c->frame.rendering.predicted_display_time_ns, //
+	    &relation);                                      //
+
+	struct xrt_vec3 eye_relation = {
+	    0.063000f, /* TODO: get actual ipd_meters */
+	    0.0f,
+	    0.0f,
+	};
+
+	for (uint32_t i = 0; i < 2; i++) {
+		struct xrt_fov fov = r->c->xdev->hmd->views[i].fov;
+
+		comp_layer_renderer_set_fov(r->lr, &fov, i);
+
+		struct xrt_pose eye_pose = XRT_POSE_IDENTITY;
+		xrt_device_get_view_pose(r->c->xdev, &eye_relation, i, &eye_pose);
+
+		struct xrt_space_relation result = {0};
+		struct xrt_space_graph xsg = {0};
+		m_space_graph_add_pose_if_not_identity(&xsg, &eye_pose);
+		m_space_graph_add_relation(&xsg, &relation);
+		m_space_graph_resolve(&xsg, &result);
+
+		out_results[i] = result.pose;
+	}
+}
+
 static void
 do_projection_layers(struct comp_renderer *r,
                      struct comp_rendering_compute *crc,
@@ -747,6 +786,9 @@ do_projection_layers(struct comp_renderer *r,
 	VkImage target_image = r->c->target->images[r->acquired_buffer].handle;
 	VkImageView target_image_view = r->c->target->images[r->acquired_buffer].view;
 
+	struct xrt_pose new_view_poses[2];
+	get_view_poses(r, new_view_poses);
+
 	VkSampler src_samplers[2] = {
 	    left->sampler,
 	    right->sampler,
@@ -765,14 +807,38 @@ do_projection_layers(struct comp_renderer *r,
 		src_norm_rects[1].y = 1 + src_norm_rects[1].y;
 	}
 
-	comp_rendering_compute_projection( //
+	struct xrt_pose src_poses[2] = {
-	    crc,                           //
+	    lvd->pose,
-	    src_samplers,                  //
+	    rvd->pose,
-	    src_image_views,               //
+	};
-	    src_norm_rects,                //
+
-	    target_image,                  //
+	struct xrt_fov src_fovs[2] = {
-	    target_image_view,             //
+	    lvd->fov,
-	    views);                        //
+	    rvd->fov,
+	};
+
+	if (crc->c->debug.atw_off) {
+		comp_rendering_compute_projection( //
+		    crc,                           //
+		    src_samplers,                  //
+		    src_image_views,               //
+		    src_norm_rects,                //
+		    target_image,                  //
+		    target_image_view,             //
+		    views);                        //
+	} else {
+		comp_rendering_compute_projection_timewarp( //
+		    crc,                                    //
+		    src_samplers,                           //
+		    src_image_views,                        //
+		    src_norm_rects,                         //
+		    src_poses,                              //
+		    src_fovs,                               //
+		    new_view_poses,                         //
+		    target_image,                           //
+		    target_image_view,                      //
+		    views);                                 //
+	}
 }
 
 static void

src/xrt/compositor/main/comp_shaders.c

@@ -22,6 +22,7 @@
 
 #include "shaders/clear.comp.h"
 #include "shaders/distortion.comp.h"
+#include "shaders/distortion_timewarp.comp.h"
 #include "shaders/layer.frag.h"
 #include "shaders/layer.vert.h"
 #include "shaders/equirect1.frag.h"
@@ -90,6 +91,11 @@ comp_shaders_load(struct vk_bundle *vk, struct comp_shaders *s)
 	              sizeof(shaders_distortion_comp), // size
 	              &s->distortion_comp));           // out
 
+	C(shader_load(vk,                                       // vk_bundle
+	              shaders_distortion_timewarp_comp,         // data
+	              sizeof(shaders_distortion_timewarp_comp), // size
+	              &s->distortion_timewarp_comp));           // out
+
 	C(shader_load(vk,                        // vk_bundle
 	              shaders_mesh_vert,         // data
 	              sizeof(shaders_mesh_vert), // size
@@ -142,6 +148,7 @@ comp_shaders_close(struct vk_bundle *vk, struct comp_shaders *s)
 {
 	D(clear_comp);
 	D(distortion_comp);
+	D(distortion_timewarp_comp);
 	D(mesh_vert);
 	D(mesh_frag);
 	D(equirect1_vert);

src/xrt/compositor/render/comp_compute.c

@@ -8,6 +8,8 @@
  */
 
 #include "math/m_api.h"
+#include "math/m_matrix_4x4_f64.h"
+
 #include "main/comp_compositor.h"
 #include "render/comp_render.h"
 
@@ -68,6 +70,109 @@ calc_dispatch_dims(const struct comp_viewport_data views[2], uint32_t *out_w, ui
 	*out_h = h;
 }
 
+/*!
+ * Create a simplified projection matrix for timewarp.
+ */
+static void
+calc_projection(const struct xrt_fov *fov, struct xrt_matrix_4x4_f64 *result)
+{
+	const double tan_left = tan(fov->angle_left);
+	const double tan_right = tan(fov->angle_right);
+
+	const double tan_down = tan(fov->angle_down);
+	const double tan_up = tan(fov->angle_up);
+
+	const double tan_width = tan_right - tan_left;
+	const double tan_height = tan_up - tan_down;
+
+	const double near = 0.5;
+	const double far = 1.5;
+
+	const double a11 = 2 / tan_width;
+	const double a22 = 2 / tan_height;
+
+	const double a31 = (tan_right + tan_left) / tan_width;
+	const double a32 = (tan_up + tan_down) / tan_height;
+
+	const float a33 = -far / (far - near);
+	const float a43 = -(far * near) / (far - near);
+
+
+#if 1
+	// We skip a33 & a43 because we don't have depth.
+	(void)a33;
+	(void)a43;
+
+	// clang-format off
+	*result = (struct xrt_matrix_4x4_f64){
+		{
+			      a11,         0,  0,  0,
+			        0,       a22,  0,  0,
+			      a31,       a32, -1,  0,
+			        0,         0,  0,  1,
+		}
+	};
+	// clang-format on
+#else
+	// clang-format off
+	*result = (struct xrt_matrix_4x4_f64) {
+		.v = {
+			a11, 0, 0, 0,
+			0, a22, 0, 0,
+			a31, a32, a33, -1,
+			0, 0, a43, 0,
+		}
+	};
+	// clang-format on
+#endif
+}
+
+static void
+calc_time_warp_matrix(struct comp_rendering_compute *crc,
+                      const struct xrt_pose *src_pose,
+                      const struct xrt_fov *src_fov,
+                      const struct xrt_pose *new_pose,
+                      struct xrt_matrix_4x4 *matrix)
+{
+	// Src projection matrix.
+	struct xrt_matrix_4x4_f64 src_proj;
+	calc_projection(src_fov, &src_proj);
+
+	// Src rotation matrix.
+	struct xrt_matrix_4x4_f64 src_rot_inv;
+	struct xrt_quat src_q = src_pose->orientation;
+	src_q.x = -src_q.x;                           // I don't know why we need to do this.
+	src_q.z = -src_q.z;                           // I don't know why we need to do this.
+	m_mat4_f64_orientation(&src_q, &src_rot_inv); // This is a model matrix, a inverted view matrix.
+
+	// New rotation matrix.
+	struct xrt_matrix_4x4_f64 new_rot, new_rot_inv;
+	struct xrt_quat new_q = new_pose->orientation;
+	new_q.x = -new_q.x;                           // I don't know why we need to do this.
+	new_q.z = -new_q.z;                           // I don't know why we need to do this.
+	m_mat4_f64_orientation(&new_q, &new_rot_inv); // This is a model matrix, a inverted view matrix.
+	m_mat4_f64_invert(&new_rot_inv, &new_rot);    // Invert to make it a view matrix.
+
+	// Combine both rotation matricies to get difference.
+	struct xrt_matrix_4x4_f64 delta_rot, delta_rot_inv;
+	m_mat4_f64_multiply(&new_rot, &src_rot_inv, &delta_rot);
+	m_mat4_f64_invert(&delta_rot, &delta_rot_inv);
+
+	// Combine the source projection matrix and
+	struct xrt_matrix_4x4_f64 result;
+	m_mat4_f64_multiply(&src_proj, &delta_rot_inv, &result);
+
+	// Reset if timewarp is off.
+	if (crc->c->debug.atw_off) {
+		result = src_proj;
+	}
+
+	// Convert from f64 to f32.
+	for (int i = 0; i < 16; i++) {
+		matrix->v[i] = result.v[i];
+	}
+}
+
 
 /*
  *
@@ -417,6 +522,147 @@ comp_rendering_compute_close(struct comp_rendering_compute *crc)
 	crc->r = NULL;
 }
 
+void
+comp_rendering_compute_projection_timewarp(struct comp_rendering_compute *crc,
+                                           VkSampler src_samplers[2],
+                                           VkImageView src_image_views[2],
+                                           const struct xrt_normalized_rect src_norm_rects[2],
+                                           const struct xrt_pose src_poses[2],
+                                           const struct xrt_fov src_fovs[2],
+                                           const struct xrt_pose new_poses[2],
+                                           VkImage target_image,
+                                           VkImageView target_image_view,
+                                           const struct comp_viewport_data views[2])
+{
+	assert(crc->c != NULL);
+	assert(crc->r != NULL);
+
+	struct vk_bundle *vk = &crc->c->vk;
+	struct comp_resources *r = crc->r;
+
+
+	/*
+	 * UBO
+	 */
+
+	struct xrt_matrix_4x4 time_warp_matrix[2];
+	calc_time_warp_matrix(     //
+	    crc,                   //
+	    &src_poses[0],         //
+	    &src_fovs[0],          //
+	    &new_poses[0],         //
+	    &time_warp_matrix[0]); //
+	calc_time_warp_matrix(     //
+	    crc,                   //
+	    &src_poses[1],         //
+	    &src_fovs[1],          //
+	    &new_poses[1],         //
+	    &time_warp_matrix[1]); //
+
+	struct comp_ubo_compute_data *data = (struct comp_ubo_compute_data *)r->compute.ubo.mapped;
+	data->views[0] = views[0];
+	data->views[1] = views[1];
+	data->pre_transforms[0] = r->distortion.uv_to_tanangle[0];
+	data->pre_transforms[1] = r->distortion.uv_to_tanangle[1];
+	data->transforms[0] = time_warp_matrix[0];
+	data->transforms[1] = time_warp_matrix[1];
+	data->post_transforms[0] = src_norm_rects[0];
+	data->post_transforms[1] = src_norm_rects[1];
+
+
+	/*
+	 * Source, target and distortion images.
+	 */
+
+	VkImageSubresourceRange subresource_range = {
+	    .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+	    .baseMipLevel = 0,
+	    .levelCount = VK_REMAINING_MIP_LEVELS,
+	    .baseArrayLayer = 0,
+	    .layerCount = VK_REMAINING_ARRAY_LAYERS,
+	};
+
+	vk_set_image_layout(            //
+	    vk,                         //
+	    crc->cmd,                   //
+	    target_image,               //
+	    0,                          //
+	    VK_ACCESS_SHADER_WRITE_BIT, //
+	    VK_IMAGE_LAYOUT_UNDEFINED,  //
+	    VK_IMAGE_LAYOUT_GENERAL,    //
+	    subresource_range);         //
+
+	VkSampler sampler = r->compute.default_sampler;
+	VkSampler distortion_samplers[6] = {
+	    sampler, sampler, sampler, sampler, sampler, sampler,
+	};
+
+	update_compute_discriptor_set(     //
+	    vk,                            //
+	    r->compute.src_binding,        //
+	    src_samplers,                  //
+	    src_image_views,               //
+	    r->compute.distortion_binding, //
+	    distortion_samplers,           //
+	    r->distortion.image_views,     //
+	    r->compute.target_binding,     //
+	    target_image_view,             //
+	    r->compute.ubo_binding,        //
+	    r->compute.ubo.buffer,         //
+	    VK_WHOLE_SIZE,                 //
+	    crc->clear_descriptor_set);    //
+
+	vk->vkCmdBindPipeline(                        //
+	    crc->cmd,                                 // commandBuffer
+	    VK_PIPELINE_BIND_POINT_COMPUTE,           // pipelineBindPoint
+	    r->compute.distortion_timewarp_pipeline); // pipeline
+
+	vk->vkCmdBindDescriptorSets(        //
+	    crc->cmd,                       // commandBuffer
+	    VK_PIPELINE_BIND_POINT_COMPUTE, // pipelineBindPoint
+	    r->compute.pipeline_layout,     // layout
+	    0,                              // firstSet
+	    1,                              // descriptorSetCount
+	    &crc->clear_descriptor_set,     // pDescriptorSets
+	    0,                              // dynamicOffsetCount
+	    NULL);                          // pDynamicOffsets
+
+
+	uint32_t w = 0, h = 0;
+	calc_dispatch_dims(views, &w, &h);
+	assert(w != 0 && h != 0);
+
+	vk->vkCmdDispatch( //
+	    crc->cmd,      // commandBuffer
+	    w,             // groupCountX
+	    h,             // groupCountY
+	    2);            // groupCountZ
+
+	VkImageMemoryBarrier memoryBarrier = {
+	    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
+	    .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
+	    .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT,
+	    .oldLayout = VK_IMAGE_LAYOUT_GENERAL,
+	    .newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
+	    .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+	    .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+	    .image = target_image,
+	    .subresourceRange = subresource_range,
+	};
+
+	vk->vkCmdPipelineBarrier(                 //
+	    crc->cmd,                             //
+	    VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, //
+	    VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,    //
+	    0,                                    //
+	    0,                                    //
+	    NULL,                                 //
+	    0,                                    //
+	    NULL,                                 //
+	    1,                                    //
+	    &memoryBarrier);                      //
+}
+
 void
 comp_rendering_compute_projection(struct comp_rendering_compute *crc,
                                   VkSampler src_samplers[2],

src/xrt/compositor/render/comp_render.h

@@ -194,12 +194,18 @@ struct comp_resources
 		//! Doesn't depend on target so is static.
 		VkPipeline distortion_pipeline;
 
+		//! Doesn't depend on target so is static.
+		VkPipeline distortion_timewarp_pipeline;
+
 		//! Target info.
 		struct comp_buffer ubo;
 	} compute;
 
 	struct
 	{
+		//! Transform to go from UV to tangle angles.
+		struct xrt_normalized_rect uv_to_tanangle[2];
+
 		//! Backing memory to distortion images.
 		VkDeviceMemory device_memories[COMP_DISTORTION_NUM_IMAGES];
 
@@ -481,6 +487,18 @@ comp_rendering_compute_close(struct comp_rendering_compute *crc);
 bool
 comp_rendering_compute_begin(struct comp_rendering_compute *crc);
 
+void
+comp_rendering_compute_projection_timewarp(struct comp_rendering_compute *crc,
+                                           VkSampler src_samplers[2],
+                                           VkImageView src_image_views[2],
+                                           const struct xrt_normalized_rect src_rects[2],
+                                           const struct xrt_pose src_poses[2],
+                                           const struct xrt_fov src_fovs[2],
+                                           const struct xrt_pose new_poses[2],
+                                           VkImage target_image,
+                                           VkImageView target_image_view,
+                                           const struct comp_viewport_data views[2]);
+
 void
 comp_rendering_compute_projection(struct comp_rendering_compute *crc,            //
                                   VkSampler src_samplers[2],                     //

src/xrt/compositor/render/comp_resources.c

@@ -11,6 +11,7 @@
 
 #include "main/comp_compositor.h"
 #include "render/comp_render.h"
+#include "math/m_vec2.h"
 
 #include <stdio.h>
 
@@ -514,6 +515,38 @@ struct texture
 	struct xrt_vec2 pixels[COMP_DISTORTION_IMAGE_DIMENSIONS][COMP_DISTORTION_IMAGE_DIMENSIONS];
 };
 
+struct tan_angles_transforms
+{
+	struct xrt_vec2 offset;
+	struct xrt_vec2 scale;
+};
+
+static void
+calc_uv_to_tanangle(struct xrt_device *xdev, uint32_t view, struct xrt_normalized_rect *out_rect)
+{
+	const struct xrt_fov fov = xdev->hmd->views[view].fov;
+	const double tan_left = tan(fov.angle_left);
+	const double tan_right = tan(fov.angle_right);
+
+	const double tan_down = tan(fov.angle_down);
+	const double tan_up = tan(fov.angle_up);
+
+	const double tan_width = tan_right - tan_left;
+	const double tan_height = tan_up - tan_down;
+
+	const double tan_offset_x = (tan_right + tan_left) - tan_width / 2;
+	const double tan_offset_y = (tan_up + tan_down) - tan_height / 2;
+
+	struct xrt_normalized_rect transform = {
+	    .x = tan_offset_x,
+	    .y = tan_offset_y,
+	    .w = tan_width,
+	    .h = tan_height,
+	};
+
+	*out_rect = transform;
+}
+
 static XRT_MAYBE_UNUSED VkResult
 create_and_file_in_distortion_buffer_for_view(struct vk_bundle *vk,
                                               struct xrt_device *xdev,
@@ -683,6 +716,14 @@ comp_resources_init(struct comp_compositor *c, struct comp_resources *r)
 	    r->compute.pipeline_layout,        // pipeline_layout
 	    &r->compute.distortion_pipeline)); // out_compute_pipeline
 
+	C(create_compute_pipeline(                      //
+	    vk,                                         // vk_bundle
+	    r->pipeline_cache,                          // pipeline_cache
+	    c->shaders.distortion_timewarp_comp,        // shader
+	    r->compute.pipeline_layout,                 // pipeline_layout
+	    &r->compute.distortion_timewarp_pipeline)); // out_compute_pipeline
+
+
 	VkBufferUsageFlags ubo_usage_flags = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
 	VkMemoryPropertyFlags memory_property_flags = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT |
 	                                              VK_MEMORY_PROPERTY_HOST_COHERENT_BIT |
@@ -702,6 +743,9 @@ comp_resources_init(struct comp_compositor *c, struct comp_resources *r)
 
 	struct comp_buffer buffers[COMP_DISTORTION_NUM_IMAGES];
 
+	calc_uv_to_tanangle(c->xdev, 0, &r->distortion.uv_to_tanangle[0]);
+	calc_uv_to_tanangle(c->xdev, 1, &r->distortion.uv_to_tanangle[1]);
+
 	create_and_file_in_distortion_buffer_for_view(vk, c->xdev, &buffers[0], &buffers[2], &buffers[4], 0);
 	create_and_file_in_distortion_buffer_for_view(vk, c->xdev, &buffers[1], &buffers[3], &buffers[5], 1);
 
@@ -754,6 +798,7 @@ comp_resources_close(struct comp_compositor *c, struct comp_resources *r)
 	D(DescriptorSetLayout, r->compute.descriptor_set_layout);
 	D(Pipeline, r->compute.clear_pipeline);
 	D(Pipeline, r->compute.distortion_pipeline);
+	D(Pipeline, r->compute.distortion_timewarp_pipeline);
 	D(PipelineLayout, r->compute.pipeline_layout);
 	D(Sampler, r->compute.default_sampler);
 	for (uint32_t i = 0; i < ARRAY_SIZE(r->distortion.image_views); i++) {

src/xrt/compositor/shaders/distortion_timewarp.comp

@@ -0,0 +1,93 @@
+// Copyright 2021, Collabora Ltd.
+// Author: Jakob Bornecrantz <jakob@collabora.com>
+// SPDX-License-Identifier: BSL-1.0
+
+#version 460
+#extension GL_GOOGLE_include_directive : require
+
+#include "srgb.inc.glsl"
+
+layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
+
+layout(set = 0, binding = 0) uniform sampler2D source[2];
+layout(set = 0, binding = 1) uniform sampler2D distortion[6];
+layout(set = 0, binding = 2) uniform writeonly restrict image2D target;
+layout(set = 0, binding = 3, std140) uniform restrict Config
+{
+	ivec4 views[2];
+	vec4 pre_transform[2];
+	vec4 post_transform[2];
+	mat4 transform[2];
+} ubo;
+
+
+vec2 position_to_uv(ivec2 extent, uint ix, uint iy)
+{
+	float x = float(ix) / float(extent.x);
+	float y = float(iy) / float(extent.y);
+
+	vec2 dist_uv = vec2(x, y);
+
+#define DIM (128.0)
+#define STRETCH ((DIM - 1.0) / DIM)
+#define OFFSET (1.0 / (DIM * 2.0))
+
+	dist_uv = (dist_uv * STRETCH) + OFFSET;
+
+	return dist_uv;
+}
+
+vec2 transform_uv(vec2 uv, uint iz)
+{
+	vec4 values = vec4(uv, -1, 1);
+
+	// From uv to tan angle (tanget space).
+	values.xy = values.xy * ubo.pre_transform[iz].zw + ubo.pre_transform[iz].xy;
+
+	// Timewarp.
+	values = ubo.transform[iz] * values;
+	values.xy = values.xy * (1.0 / max(values.w, 0.00001));
+
+	// From [-1, 1] to [0, 1]
+	values.xy = values.xy * 0.5 + 0.5;
+
+	// To deal with OpenGL flip and sub image view.
+	values.xy = values.xy * ubo.post_transform[iz].zw + ubo.post_transform[iz].xy;
+
+	// Done.
+	return values.xy;
+}
+
+void main()
+{
+	uint ix = gl_GlobalInvocationID.x;
+	uint iy = gl_GlobalInvocationID.y;
+	uint iz = gl_GlobalInvocationID.z;
+
+	ivec2 offset = ivec2(ubo.views[iz].xy);
+	ivec2 extent = ivec2(ubo.views[iz].zw);
+
+	if (ix >= extent.x || iy >= extent.y) {
+		return;
+	}
+
+	vec2 dist_uv = position_to_uv(extent, ix, iy);
+
+	vec2 r_uv = texture(distortion[iz + 0], dist_uv).xy;
+	vec2 g_uv = texture(distortion[iz + 2], dist_uv).xy;
+	vec2 b_uv = texture(distortion[iz + 4], dist_uv).xy;
+
+	r_uv = transform_uv(r_uv, iz);
+	g_uv = transform_uv(g_uv, iz);
+	b_uv = transform_uv(b_uv, iz);
+
+	vec4 colour = vec4(
+		texture(source[iz], r_uv).r,
+		texture(source[iz], g_uv).g,
+		texture(source[iz], b_uv).b,
+		1);
+
+	colour = vec4(from_linear_to_srgb(colour.rgb), 1);
+
+	imageStore(target, ivec2(offset.x + ix, offset.y + iy), colour);
+}

src/xrt/compositor/shaders/meson.build

@@ -4,6 +4,7 @@
 shader_srcs = [
 	'clear.comp',
 	'distortion.comp',
+	'distortion_timewarp.comp',
 	'mesh.frag',
 	'mesh.vert',
 	'layer.vert',