Use two images for accumulation (TAA style)

4 files changed, 50 insertions, 29 deletions

diff --git a/src/main.cpp b/src/main.cpp
index c87e57a..9e3634f 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -46,7 +46,7 @@ int main() {
     const float aspect_ratio = 16.0 / 9.0;
     uint32_t width = 400;
     uint32_t height = width / aspect_ratio;
-    const uint32_t samples_per_pixel = 10;
+    const uint32_t samples_per_pixel = 25;
 
     window wnd { width , height };
 
@@ -80,6 +80,7 @@ int main() {
                         canRender = true;
                         renderer.recreate_swapchain();
                         ((input_data*) renderer.compute_shader_buffer.mapped_ptr)->cam.set_aspect_ratio((float)event.width / (float)event.height);
+                        ((input_data*) renderer.compute_shader_buffer.mapped_ptr)->sample_index = 0;
                     }
                     break;
                 }
diff --git a/src/utils.cpp b/src/utils.cpp
index a2cf8f9..c05cc77 100644
--- a/src/utils.cpp
+++ b/src/utils.cpp
@@ -93,7 +93,7 @@ struct input_data create_inputs(uint32_t width, uint32_t height, uint32_t sample
     // Image dimensions
     const float aspect_ratio = (float)width / (float) height;
 
-    const uint32_t max_depth = 5;
+    const uint32_t max_depth = 20;
 
     const vec3 camera_position{ 13.0, 2.0, 3.0 };
     const vec3 camera_target{ 0.0, 0.0, 0.0 };
diff --git a/src/vk-api.cpp b/src/vk-api.cpp
index 92e9a80..cd401a1 100644
--- a/src/vk-api.cpp
+++ b/src/vk-api.cpp
@@ -57,7 +57,7 @@ Buffer vkapi::create_buffer(size_t data_size, VkBufferUsageFlags buffer_usage, V
     VmaAllocationCreateInfo alloc_create_info = {};
     alloc_create_info.usage = mem_usage;
      
-    vmaCreateBuffer(context.allocator, &buffer_info, &alloc_create_info, &buffer.handle, &buffer.alloc, nullptr);
+    VKRESULT(vmaCreateBuffer(context.allocator, &buffer_info, &alloc_create_info, &buffer.handle, &buffer.alloc, nullptr))
 
     vmaMapMemory(context.allocator, buffer.alloc, (void**) &buffer.mapped_ptr);
 
@@ -101,7 +101,7 @@ Image vkapi::create_image(VkExtent3D size, VkFormat format, VkImageUsageFlags us
     VmaAllocationCreateInfo alloc_create_info = {};
     alloc_create_info.usage = VMA_MEMORY_USAGE_GPU_ONLY;
 
-    vmaCreateImage(context.allocator, &img_create_info, &alloc_create_info, &image.handle, &image.alloc, nullptr);
+    VKRESULT(vmaCreateImage(context.allocator, &img_create_info, &alloc_create_info, &image.handle, &image.alloc, nullptr))
 
     VkImageViewCreateInfo image_view_create_info    = {};
     image_view_create_info.sType                    = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
@@ -475,7 +475,7 @@ void vkapi::destroy_swapchain(Swapchain swapchain) {
 }
 
 
-void vkapi::update_descriptor_set_buffer(VkDescriptorSet set, VkDescriptorSetLayoutBinding binding,Buffer& buffer) {
+void vkapi::update_descriptor_set_buffer(VkDescriptorSet set, VkDescriptorSetLayoutBinding binding, Buffer& buffer) {
     VkDescriptorBufferInfo descriptor_buffer_info = {};
     descriptor_buffer_info.buffer           = buffer.handle;
     descriptor_buffer_info.offset           = 0;
@@ -485,7 +485,7 @@ void vkapi::update_descriptor_set_buffer(VkDescriptorSet set, VkDescriptorSetLay
     write_descriptor.sType                  = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
     write_descriptor.pNext                  = nullptr;
     write_descriptor.dstSet                 = set;
-    write_descriptor.dstBinding             = 0;
+    write_descriptor.dstBinding             = binding.binding;
     write_descriptor.dstArrayElement        = 0;
     write_descriptor.descriptorCount        = 1;
     write_descriptor.descriptorType         = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
@@ -506,7 +506,7 @@ void vkapi::update_descriptor_set_image(VkDescriptorSet set, VkDescriptorSetLayo
     write_descriptor.sType                      = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
     write_descriptor.pNext                      = nullptr;
     write_descriptor.dstSet                     = set;
-    write_descriptor.dstBinding                 = 1;
+    write_descriptor.dstBinding                 = binding.binding;
     write_descriptor.dstArrayElement            = 0;
     write_descriptor.descriptorCount            = binding.descriptorCount;
     write_descriptor.descriptorType             = binding.descriptorType;
diff --git a/src/vk-renderer.cpp b/src/vk-renderer.cpp
index 32bee48..4e5971a 100644
--- a/src/vk-renderer.cpp
+++ b/src/vk-renderer.cpp
@@ -17,7 +17,6 @@ vkrenderer::vkrenderer(window& wnd, const input_data& inputs) {
         platform_surface,
         min_swapchain_image_count,
         VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT |
-        VK_IMAGE_USAGE_STORAGE_BIT |
         VK_IMAGE_USAGE_TRANSFER_DST_BIT,
         std::nullopt
     );
@@ -36,6 +35,13 @@ vkrenderer::vkrenderer(window& wnd, const input_data& inputs) {
             .descriptorCount = 1,
             .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
             .pImmutableSamplers = VK_NULL_HANDLE,
+        },
+        {
+            .binding = 2,
+            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
+            .descriptorCount = 1,
+            .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
+            .pImmutableSamplers = VK_NULL_HANDLE,
         }
     };
 
@@ -50,7 +56,7 @@ vkrenderer::vkrenderer(window& wnd, const input_data& inputs) {
         .height = swapchain.extent.height,
         .depth = 1,
     };
-    accumulation_image = api.create_image(image_size, swapchain.surface_format.format, VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT);
+    accumulation_images = api.create_images(image_size, swapchain.surface_format.format, VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT, 2);
 
     submission_fences = api.create_fences(swapchain.image_count);
     execution_semaphores = api.create_semaphores(swapchain.image_count);
@@ -61,7 +67,6 @@ vkrenderer::vkrenderer(window& wnd, const input_data& inputs) {
 
     for (size_t index = 0; index < swapchain.image_count; index++) {
         api.update_descriptor_set_buffer(compute_shader_sets[index], bindings[0], compute_shader_buffer);
-        api.update_descriptor_set_image(compute_shader_sets[index], bindings[1], accumulation_image.view);
     }
 }
 
@@ -74,7 +79,7 @@ vkrenderer::~vkrenderer() {
     api.destroy_semaphores(acquire_semaphores);
 
     api.destroy_descriptor_sets(compute_shader_sets);
-    api.destroy_image(accumulation_image);
+    api.destroy_images(accumulation_images);
 
     api.destroy_command_buffers(command_buffers);
     api.destroy_compute_pipeline(compute_pipeline);
@@ -92,6 +97,28 @@ void vkrenderer::compute(uint32_t width, uint32_t height, bool clear) {
     auto acquire_result = vkAcquireNextImageKHR(api.context.device, swapchain.handle, UINT64_MAX, acquire_semaphores[frame_index], VK_NULL_HANDLE, &swapchain_image_index);
     handle_swapchain_result(acquire_result);
 
+    auto accumulation_image_index = frame_index % 2;
+    auto output_image = accumulation_images[accumulation_image_index];
+    auto accumulation_image = accumulation_images[(accumulation_image_index + 1) % 2];
+
+    VkDescriptorSetLayoutBinding images[2] {
+        {
+            .binding = 1,
+            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
+            .descriptorCount = 1,
+            .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
+            .pImmutableSamplers = VK_NULL_HANDLE,
+        }, {
+            .binding = 2,
+            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
+            .descriptorCount = 1,
+            .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
+            .pImmutableSamplers = VK_NULL_HANDLE,
+        }
+    };
+    api.update_descriptor_set_image(compute_shader_sets[swapchain_image_index], images[0], output_image.view);
+    api.update_descriptor_set_image(compute_shader_sets[swapchain_image_index], images[1], accumulation_image.view);
+
     auto cmd_buf = command_buffers[frame_index];
 
     api.start_record(cmd_buf);
@@ -102,17 +129,21 @@ void vkrenderer::compute(uint32_t width, uint32_t height, bool clear) {
         VkClearColorValue color = { 0.f, 0.f, 0.f, 1.f };
         vkCmdClearColorImage(cmd_buf, accumulation_image.handle, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, &color, 1, &accumulation_image.subresource_range);
 
-        api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_WRITE_BIT, accumulation_image);
+        api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT, accumulation_image);
+
+        api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, VK_ACCESS_SHADER_WRITE_BIT, output_image);
     } else {
-        api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, VK_ACCESS_SHADER_WRITE_BIT, accumulation_image);
+        api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, VK_ACCESS_SHADER_READ_BIT, accumulation_image);
+        api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, VK_ACCESS_SHADER_WRITE_BIT, output_image);
     }
+
     api.run_compute_pipeline(cmd_buf, compute_pipeline, compute_shader_sets[swapchain_image_index], width / 8, height / 8, 1);
 
-    api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, accumulation_image);
+    api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, output_image);
 
     api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, VK_ACCESS_TRANSFER_WRITE_BIT, swapchain.images[swapchain_image_index]);
 
-    api.blit_full(cmd_buf, accumulation_image, swapchain.images[swapchain_image_index]);
+    api.blit_full(cmd_buf, output_image, swapchain.images[swapchain_image_index]);
 
     api.image_barrier(cmd_buf, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_COLOR_ATTACHMENT_READ_BIT, swapchain.images[swapchain_image_index]);
 
@@ -132,28 +163,17 @@ void vkrenderer::recreate_swapchain() {
     swapchain = api.create_swapchain(
         platform_surface,
         min_swapchain_image_count,
-        VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_STORAGE_BIT,
+        VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT,
         std::optional<Swapchain> { swapchain }
     );
 
+    api.destroy_images(accumulation_images);
     VkExtent3D image_size = {
         .width = swapchain.extent.width,
         .height = swapchain.extent.height,
         .depth = 1,
     };
-    accumulation_image = api.create_image(image_size, swapchain.surface_format.format, VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT);
-
-    VkDescriptorSetLayoutBinding binding {
-        .binding = 1,
-        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
-        .descriptorCount = 1,
-        .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
-        .pImmutableSamplers = VK_NULL_HANDLE,
-    };
-
-    for (size_t index = 0; index < swapchain.image_count; index++) {
-        api.update_descriptor_set_image(compute_shader_sets[index], binding, accumulation_image.view);
-    }
+    accumulation_images = api.create_images(image_size, swapchain.surface_format.format, VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT, 2);
 }
 
 void vkrenderer::handle_swapchain_result(VkResult function_result) {