From 96e2c56f116d1045e0707096cc46545dd9de168f Mon Sep 17 00:00:00 2001
From: Mikulas Florek <mikulas.florek@gamedev.sk>
Date: Thu, 21 Sep 2023 21:47:47 +0200
Subject: [PATCH] fsr2 support

---
 data/pipelines/main.pln   | 24 +++++++++++++++++-----
 src/renderer/pipeline.cpp | 43 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/data/pipelines/main.pln b/data/pipelines/main.pln
index a4a20b0d9..94dd3c5f0 100644
--- a/data/pipelines/main.pln
+++ b/data/pipelines/main.pln
@@ -31,6 +31,7 @@ local taa_history = -1
 local render_grass = true
 local render_impostors = true
 local render_terrain = true
+local fsr2_enable = false
 
 type GBuffer = {
 	A : RenderBuffer,
@@ -463,8 +464,8 @@ local function render_preview()
 end
 
 local function TAA(res, gbuffer)
-	PIXEL_JITTER = taa_enabled
 	if taa_enabled then
+		PIXEL_JITTER = true
 		beginBlock("taa")
 		if taa_history == -1 then
 			taa_history = createRenderbuffer(taa_history_desc)
@@ -504,6 +505,7 @@ local function TAA(res, gbuffer)
 end
 
 main = function()
+	PIXEL_JITTER = false
 	if PREVIEW then
 		render_preview()
 		return
@@ -526,14 +528,21 @@ main = function()
 	sss:postprocess(getfenv(1), nil, gbuffer, shadowmap)
 
 	local hdr_buffer = lightPass(view_params, gbuffer, shadowmap)
-	
-	local res = hdr_buffer
 	cubemap_sky:postprocess(getfenv(1), hdr_buffer, gbuffer, shadowmap)
 	atmo:postprocess(getfenv(1), hdr_buffer, gbuffer, shadowmap)
 	
+	local res 
+	if fsr2Dispatch and fsr2_enable then
+		PIXEL_JITTER = true
+		res = createRenderbuffer(taa_desc)	
+		fsr2Dispatch(hdr_buffer, gbuffer.DS, gbuffer.D, res)
+	else
+		res = hdr_buffer
+	end
+
 	waterPass(view_params, entities, res, gbuffer, shadowmap)
 	transparentPass(view_params, entities, res, gbuffer, shadowmap)
-	setRenderTargetsReadonlyDS(hdr_buffer, gbuffer.DS)
+	setRenderTargetsReadonlyDS(res, gbuffer.DS)
 
 	res = dof:postprocess(getfenv(1), res, gbuffer, shadowmap)
 	bloom:postprocess(getfenv(1), res, gbuffer, shadowmap)
@@ -576,7 +585,9 @@ main = function()
 		end
 	end
 
-	TAA(res, gbuffer)
+	if not fsr2_enable then
+		TAA(res, gbuffer)
+	end
 
 	render2D()
 
@@ -650,6 +661,9 @@ onGUI = function()
 		_, debug_shadow_buf = ImGui.Checkbox("GBuffer shadow", debug_shadow_buf)
 		_, debug_clusters = ImGui.Checkbox("Clusters", debug_clusters)
 		_, taa_enabled = ImGui.Checkbox("TAA", taa_enabled)
+		if fsr2Dispatch then
+			_, fsr2_enable = ImGui.Checkbox("FSR2", fsr2_enable)
+		end
 		_, render_grass = ImGui.Checkbox("Grass", render_grass)
 		_, render_impostors = ImGui.Checkbox("Impostors", render_impostors)
 		_, render_terrain = ImGui.Checkbox("Terrain", render_terrain)
diff --git a/src/renderer/pipeline.cpp b/src/renderer/pipeline.cpp
index 7e0fdbbf7..c9a6fbfd5 100644
--- a/src/renderer/pipeline.cpp
+++ b/src/renderer/pipeline.cpp
@@ -34,6 +34,10 @@
 #include "terrain.h"
 #include "texture.h"
 
+#ifdef LUMIX_FSR2
+	#include "../plugins/dx/src/fsr2.h"
+#endif
+
 namespace Lumix
 {
 
@@ -3912,6 +3916,38 @@ struct PipelineImpl final : Pipeline
 		stream.freeMemory(mem.data, m_renderer.getAllocator());
 	}
 
+	#ifdef LUMIX_FSR2
+		void fsr2Dispatch(PipelineTexture color, PipelineTexture depth, PipelineTexture motion_vectors, PipelineTexture output) {
+			DrawStream& stream = m_renderer.getDrawStream();	
+			gpu::FSR2DispatchParams params;
+			params.jitter_x = m_viewport.pixel_offset.x;
+			params.jitter_y = m_viewport.pixel_offset.y;
+			params.time_delta = m_renderer.getEngine().getLastTimeDelta() * 1000;
+			params.render_width = m_viewport.w;
+			params.render_height = m_viewport.h;
+			params.near_plane = m_viewport.near;
+			params.fov = m_viewport.fov;
+			params.color = toHandle(color);
+			params.depth = toHandle(depth);
+			params.motion_vectors = toHandle(motion_vectors);
+			params.output = toHandle(output);
+			beginBlock("FSR2");
+			stream.pushLambda([this, params](){
+				IVec2 s(m_viewport.w, m_viewport.h);
+				if (m_fsr2 && m_fsr2_size != s) {
+					gpu::fsr2Shutdown(*m_fsr2);
+					m_fsr2 = nullptr;
+				}
+				if (!m_fsr2) {
+					m_fsr2 = gpu::fsr2Init(m_viewport.w, m_viewport.h, m_allocator);
+					m_fsr2_size = s;
+				}
+				gpu::fsr2Dispatch(*m_fsr2, params);
+			});
+			endBlock();
+		}
+	#endif
+
 	void registerLuaAPI(lua_State* L)
 	{
 		lua_rawgeti(m_lua_state, LUA_REGISTRYINDEX, m_lua_env);
@@ -3951,6 +3987,9 @@ struct PipelineImpl final : Pipeline
 		REGISTER_FUNCTION(endBlock);
 		REGISTER_FUNCTION(environmentCastShadows);
 		REGISTER_FUNCTION(executeCustomCommand);
+		#ifdef LUMIX_FSR2
+			REGISTER_FUNCTION(fsr2Dispatch);
+		#endif		
 		REGISTER_FUNCTION(getCameraParams);
 		REGISTER_FUNCTION(getShadowCameraParams);
 		REGISTER_FUNCTION(keepRenderbufferAlive);
@@ -4084,6 +4123,10 @@ struct PipelineImpl final : Pipeline
 		Buffer refl_probes;
 	} m_cluster_buffers;
 	Viewport m_shadow_camera_viewports[4];
+	#ifdef LUMIX_FSR2
+		Lumix::gpu::FSR2Context* m_fsr2 = nullptr;
+		IVec2 m_fsr2_size = IVec2(0);
+	#endif
 };