Use buffer device address with push constants

3 files changed, 33 insertions, 59 deletions

diff --git a/src/vk-api.cpp b/src/vk-api.cpp
index 9ce8c05..c6dcf0b 100644
--- a/src/vk-api.cpp
+++ b/src/vk-api.cpp
@@ -61,7 +61,7 @@ Buffer vkapi::create_buffer(size_t data_size, VkBufferUsageFlags buffer_usage, V
     VkBufferCreateInfo buffer_info  = {};
     buffer_info.sType               = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
     buffer_info.size                = data_size;
-    buffer_info.usage               = buffer_usage;
+    buffer_info.usage               = buffer_usage | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT_KHR;
      
     VmaAllocationCreateInfo alloc_create_info = {};
     alloc_create_info.usage = mem_usage;
@@ -71,6 +71,13 @@ Buffer vkapi::create_buffer(size_t data_size, VkBufferUsageFlags buffer_usage, V
 
     buffer.size = data_size;
 
+    VkBufferDeviceAddressInfo buffer_device_address_info = {};
+    buffer_device_address_info.sType    = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
+    buffer_device_address_info.pNext    = VK_NULL_HANDLE;
+    buffer_device_address_info.buffer   = buffer.handle;
+
+    buffer.device_address = vkGetBufferDeviceAddress(context.device, &buffer_device_address_info);
+
     return std::move(buffer);
 }
 
@@ -440,14 +447,19 @@ Pipeline vkapi::create_compute_pipeline(const char* shader_name, std::vector<VkD
 
     VKRESULT(vkCreateDescriptorSetLayout(context.device, &descriptor_set_layout_create_info, VK_NULL_HANDLE, &pipeline.descriptor_set_layout))
 
+    VkPushConstantRange push_constant               = {};
+    push_constant.stageFlags                        = VK_SHADER_STAGE_COMPUTE_BIT;
+    push_constant.offset                            = 0;
+    push_constant.size                              = 24;
+
     VkPipelineLayoutCreateInfo layout_create_info   = {};
     layout_create_info.sType                        = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
     layout_create_info.pNext                        = VK_NULL_HANDLE;
     layout_create_info.flags                        = 0;
     layout_create_info.setLayoutCount               = 1;
     layout_create_info.pSetLayouts                  = &pipeline.descriptor_set_layout;
-    layout_create_info.pushConstantRangeCount       = 0;
-    layout_create_info.pPushConstantRanges          = VK_NULL_HANDLE;
+    layout_create_info.pushConstantRangeCount       = 1;
+    layout_create_info.pPushConstantRanges          = &push_constant;
 
     VKRESULT(vkCreatePipelineLayout(context.device, &layout_create_info, VK_NULL_HANDLE, &pipeline.layout))
 
@@ -522,6 +534,7 @@ Pipeline vkapi::create_graphics_pipeline(const char* shader_name, std::vector<Vk
 
     VKRESULT(vkCreateDescriptorSetLayout(context.device, &descriptor_set_layout_create_info, nullptr, &pipeline.descriptor_set_layout))
 
+
     VkPipelineLayoutCreateInfo layout_create_info   = {};
     layout_create_info.sType                        = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
     layout_create_info.pNext                        = nullptr;
@@ -543,24 +556,6 @@ Pipeline vkapi::create_graphics_pipeline(const char* shader_name, std::vector<Vk
     input_assembly_state_create_info.flags                                  = 0;
     input_assembly_state_create_info.topology                               = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
 
-    // VkViewport viewport = {};
-    // viewport.x          = 0.f;
-    // viewport.y          = 0.f;
-    // viewport.width      = 384.f;
-    // viewport.height     = 186.f;
-    // viewport.minDepth   = 0.f;
-    // viewport.maxDepth   = 1.f;
-
-    // VkRect2D scissor    = {};
-    // scissor.offset      = {
-    //     0,
-    //     0,
-    // };
-    // scissor.extent      = {
-    //     384,
-    //     186,
-    // };
-
     VkPipelineViewportStateCreateInfo viewport_state_create_info            = {};
     viewport_state_create_info.sType                                        = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
     viewport_state_create_info.pNext                                        = nullptr;
@@ -569,10 +564,6 @@ Pipeline vkapi::create_graphics_pipeline(const char* shader_name, std::vector<Vk
     viewport_state_create_info.pViewports                                   = VK_NULL_HANDLE;
     viewport_state_create_info.scissorCount                                 = 1;
     viewport_state_create_info.pScissors                                    = VK_NULL_HANDLE;
-    // viewport_state_create_info.viewportCount                                = 1;
-    // viewport_state_create_info.pViewports                                   = &viewport;
-    // viewport_state_create_info.scissorCount                                 = 1;
-    // viewport_state_create_info.pScissors                                    = &scissor;
 
     VkPipelineRasterizationStateCreateInfo rasterization_state_create_info  = {};
     rasterization_state_create_info.sType                                   = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
diff --git a/src/vk-context.cpp b/src/vk-context.cpp
index ce593ab..999c762 100644
--- a/src/vk-context.cpp
+++ b/src/vk-context.cpp
@@ -108,17 +108,24 @@ void vkcontext::create_device() {
     float queue_priority                        = 1.f;
     queue_create_info.pQueuePriorities          = &queue_priority;
 
-    VkPhysicalDeviceFeatures device_features    = {};
+    VkPhysicalDeviceVulkan12Features physical_device_12_features    = {};
+    physical_device_12_features.sType                               = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES;
+    physical_device_12_features.pNext                               = VK_NULL_HANDLE;
+    physical_device_12_features.bufferDeviceAddress                 = VK_TRUE;
+
+    VkPhysicalDeviceFeatures2 device_features   = {};
+    device_features.sType                       = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
+    device_features.pNext                       = &physical_device_12_features;
 
     std::vector<const char*> device_ext         = {
-        "VK_KHR_swapchain"
+        "VK_KHR_swapchain",
     };
 
     VkDeviceCreateInfo device_create_info       = {};
     device_create_info.sType                    = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
+    device_create_info.pNext                    = &device_features;
     device_create_info.pQueueCreateInfos        = &queue_create_info;
     device_create_info.queueCreateInfoCount     = 1;
-    device_create_info.pEnabledFeatures         = &device_features;
     device_create_info.ppEnabledExtensionNames  = device_ext.data();
     device_create_info.enabledExtensionCount    = device_ext.size();
     device_create_info.enabledLayerCount        = 0; // Deprecated https://www.khronos.org/registry/vulkan/specs/1.2/html/chap31.html#extendingvulkan-layers-devicelayerdeprecation
@@ -140,6 +147,7 @@ void vkcontext::create_memory_allocator() {
     allocator_info.physicalDevice = physical_device;
     allocator_info.device = device;
     allocator_info.instance = instance;
+    allocator_info.flags = VMA_ALLOCATOR_CREATE_BUFFER_DEVICE_ADDRESS_BIT;
      
     vmaCreateAllocator(&allocator_info, &allocator);
 }
diff --git a/src/vk-renderer.cpp b/src/vk-renderer.cpp
index 521bd52..0c73337 100644
--- a/src/vk-renderer.cpp
+++ b/src/vk-renderer.cpp
@@ -38,34 +38,13 @@ vkrenderer::vkrenderer(window& wnd, size_t scene_buffer_size, size_t geometry_bu
     compute_sets_bindings = {
         {
             .binding = 0,
-            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
-            .descriptorCount = 1,
-            .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
-            .pImmutableSamplers = VK_NULL_HANDLE,
-        },
-        {
-            .binding = 1,
-            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
-            .descriptorCount = 1,
-            .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
-            .pImmutableSamplers = VK_NULL_HANDLE,
-        },
-        {
-            .binding = 2,
-            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
-            .descriptorCount = 1,
-            .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
-            .pImmutableSamplers = VK_NULL_HANDLE,
-        },
-        {
-            .binding = 3,
             .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
             .descriptorCount = 1,
             .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
             .pImmutableSamplers = VK_NULL_HANDLE,
         },
         {
-            .binding = 4,
+            .binding = 1,
             .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
             .descriptorCount = 1,
             .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
@@ -111,13 +90,6 @@ vkrenderer::vkrenderer(window& wnd, size_t scene_buffer_size, size_t geometry_bu
     geometry_buffer = api.create_buffer(geometry_buffer_size, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VMA_MEMORY_USAGE_GPU_ONLY);
     bvh_buffer = api.create_buffer(bvh_buffer_size, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VMA_MEMORY_USAGE_GPU_ONLY);
 
-
-    for (size_t index = 0; index < virtual_frames_count; index++) {
-        api.update_descriptor_set_buffer(compute_shader_sets[index], compute_sets_bindings[0], scene_buffer);
-        api.update_descriptor_set_buffer(compute_shader_sets[index], compute_sets_bindings[1], geometry_buffer);
-        api.update_descriptor_set_buffer(compute_shader_sets[index], compute_sets_bindings[2], bvh_buffer);
-    }
-
     std::vector<VkFormat> attachments_format { swapchain.surface_format.format };
     render_pass = api.create_render_pass(attachments_format, VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR);
 
@@ -315,8 +287,8 @@ void vkrenderer::begin_frame() {
     auto output_image = accumulation_images[0];
     auto accumulation_image = accumulation_images[1];
 
-    api.update_descriptor_set_image(compute_shader_sets[virtual_frame_index], compute_sets_bindings[3], output_image.view, VK_NULL_HANDLE);
-    api.update_descriptor_set_image(compute_shader_sets[virtual_frame_index], compute_sets_bindings[4], accumulation_image.view, VK_NULL_HANDLE);
+    api.update_descriptor_set_image(compute_shader_sets[virtual_frame_index], compute_sets_bindings[0], output_image.view, VK_NULL_HANDLE);
+    api.update_descriptor_set_image(compute_shader_sets[virtual_frame_index], compute_sets_bindings[1], accumulation_image.view, VK_NULL_HANDLE);
 
     api.update_descriptor_set_image(tonemapping_shader_sets[virtual_frame_index], tonemapping_sets_bindings[0], swapchain.images[swapchain_image_index].view, VK_NULL_HANDLE);
     api.update_descriptor_set_image(tonemapping_shader_sets[virtual_frame_index], tonemapping_sets_bindings[1], output_image.view, VK_NULL_HANDLE);
@@ -384,6 +356,9 @@ void vkrenderer::ui() {
 void vkrenderer::compute(uint32_t width, uint32_t height) {
     auto cmd_buf = command_buffers[virtual_frame_index];
 
+    vkCmdPushConstants(cmd_buf, compute_pipeline.layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, 8, (void*)&scene_buffer.device_address);
+    vkCmdPushConstants(cmd_buf, compute_pipeline.layout, VK_SHADER_STAGE_COMPUTE_BIT, 8, 8, (void*)&geometry_buffer.device_address);
+    vkCmdPushConstants(cmd_buf, compute_pipeline.layout, VK_SHADER_STAGE_COMPUTE_BIT, 16, 8, (void*)&bvh_buffer.device_address);
     api.run_compute_pipeline(cmd_buf, compute_pipeline, compute_shader_sets[virtual_frame_index], width / 8, height / 8, 1);