From 95d5343c9ad0da15a1100228a3ac7ed6faba4f78 Mon Sep 17 00:00:00 2001
From: Jakob Bornecrantz <jakob@collabora.com>
Date: Mon, 18 Apr 2022 15:41:29 +0100
Subject: [PATCH] c/render: Make distortion shader optionally do timewarp

---
 src/xrt/compositor/render/render_resources.c | 22 ++++++++----
 src/xrt/compositor/shaders/distortion.comp   | 36 +++++++++++++++++++-
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/src/xrt/compositor/render/render_resources.c b/src/xrt/compositor/render/render_resources.c
index 51540492e..ee99a96e7 100644
--- a/src/xrt/compositor/render/render_resources.c
+++ b/src/xrt/compositor/render/render_resources.c
@@ -242,6 +242,7 @@ create_compute_descriptor_set_layout(struct vk_bundle *vk,
 struct compute_distortion_params
 {
 	uint32_t distortion_texel_count;
+	VkBool32 do_timewarp;
 };
 
 static VkResult
@@ -252,13 +253,18 @@ create_compute_distortion_pipeline(struct vk_bundle *vk,
                                    const struct compute_distortion_params *params,
                                    VkPipeline *out_compute_pipeline)
 {
-	VkSpecializationMapEntry entries[1] = {
-	    {
-	        .constantID = 0,
-	        .offset = offsetof(struct compute_distortion_params, distortion_texel_count),
-	        .size = sizeof(params->distortion_texel_count),
-	    },
+#define ENTRY(ID, FIELD)                                                                                               \
+	{                                                                                                              \
+	    .constantID = ID,                                                                                          \
+	    .offset = offsetof(struct compute_distortion_params, FIELD),                                               \
+	    sizeof(params->FIELD),                                                                                     \
+	}
+
+	VkSpecializationMapEntry entries[2] = {
+	    ENTRY(0, distortion_texel_count),
+	    ENTRY(1, do_timewarp),
 	};
+#undef ENTRY
 
 	VkSpecializationInfo specialization_info = {
 	    .mapEntryCount = ARRAY_SIZE(entries),
@@ -693,6 +699,7 @@ render_resources_init(struct render_resources *r,
 
 	struct compute_distortion_params distortion_params = {
 	    .distortion_texel_count = COMP_DISTORTION_IMAGE_DIMENSIONS,
+	    .do_timewarp = false,
 	};
 
 	C(create_compute_distortion_pipeline(  //
@@ -705,12 +712,13 @@ render_resources_init(struct render_resources *r,
 
 	struct compute_distortion_params distortion_timewarp_params = {
 	    .distortion_texel_count = COMP_DISTORTION_IMAGE_DIMENSIONS,
+	    .do_timewarp = true,
 	};
 
 	C(create_compute_distortion_pipeline(           //
 	    vk,                                         // vk_bundle
 	    r->pipeline_cache,                          // pipeline_cache
-	    r->shaders->distortion_timewarp_comp,       // shader
+	    r->shaders->distortion_comp,                // shader
 	    r->compute.pipeline_layout,                 // pipeline_layout
 	    &distortion_timewarp_params,                // params
 	    &r->compute.distortion_timewarp_pipeline)); // out_compute_pipeline
diff --git a/src/xrt/compositor/shaders/distortion.comp b/src/xrt/compositor/shaders/distortion.comp
index 3481eefa3..4be6bbbdb 100644
--- a/src/xrt/compositor/shaders/distortion.comp
+++ b/src/xrt/compositor/shaders/distortion.comp
@@ -11,6 +11,9 @@
 // The size of the distortion texture dimensions in texels.
 layout(constant_id = 0) const int distortion_texel_count = 2;
 
+// Should we do timewarp.
+layout(constant_id = 1) const bool do_timewarp = false;
+
 layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
 
 layout(set = 0, binding = 0) uniform sampler2D source[2];
@@ -57,7 +60,7 @@ vec2 position_to_uv(ivec2 extent, uint ix, uint iy)
 	return dist_uv;
 }
 
-vec2 transform_uv(vec2 uv, uint iz)
+vec2 transform_uv_subimage(vec2 uv, uint iz)
 {
 	vec2 values = uv;
 
@@ -68,6 +71,37 @@ vec2 transform_uv(vec2 uv, uint iz)
 	return values.xy;
 }
 
+vec2 transform_uv_timewarp(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;
+	values.y = -values.y; // Flip to OpenXR coordinate system.
+
+	// 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;
+}
+
+vec2 transform_uv(vec2 uv, uint iz)
+{
+	if (do_timewarp) {
+		return transform_uv_timewarp(uv, iz);
+	} else {
+		return transform_uv_subimage(uv, iz);
+	}
+}
+
 void main()
 {
 	uint ix = gl_GlobalInvocationID.x;