Use buffer device address with push constants

7 files changed, 76 insertions, 95 deletions

diff --git a/include/scene.hpp b/include/scene.hpp
index 0002b1d..3346282 100644
--- a/include/scene.hpp
+++ b/include/scene.hpp
@@ -10,13 +10,10 @@
 

 class scene {

     struct metadata {

-        color sky_color = { 0.5, 0.7, 1.0 };

-        color ground_color = { 1.0, 1.0, 1.0 };

-

         camera cam;

 

         // uint32_t max_bounce = 50;

-        uint32_t max_bounce = 1;

+        uint32_t max_bounce = 4;

 

         // Output image resolution

         uint32_t width;

diff --git a/include/vk-api.hpp b/include/vk-api.hpp
index 2f7302e..28173f9 100644
--- a/include/vk-api.hpp
+++ b/include/vk-api.hpp
@@ -10,6 +10,7 @@ struct Buffer {
     VmaAllocationInfo   alloc_info;
     VkBuffer            handle;
     size_t              size;
+    VkDeviceAddress     device_address;
 };
 
 struct Pipeline {
diff --git a/include/vulkan-loader.hpp b/include/vulkan-loader.hpp
index b4bef68..6bdf54f 100644
--- a/include/vulkan-loader.hpp
+++ b/include/vulkan-loader.hpp
@@ -80,6 +80,7 @@
     X(vkDestroyCommandPool)           \
     X(vkBeginCommandBuffer)           \
     X(vkCmdPipelineBarrier)           \
+    X(vkCmdPushConstants)             \
     X(vkCmdClearColorImage)           \
     X(vkCmdBeginRenderPass)           \
     X(vkCmdBindPipeline)              \
@@ -113,6 +114,7 @@
     X(vkCreateDescriptorPool)         \
     X(vkAllocateDescriptorSets)       \
     X(vkUpdateDescriptorSets)         \
+    X(vkGetBufferDeviceAddress)       \
     X(vkGetBufferMemoryRequirements)  \
     X(vkGetImageMemoryRequirements)   \
     X(vkAllocateMemory)               \
diff --git a/shaders/compute.comp b/shaders/compute.comp
index 3a4a043..44e09a2 100644
--- a/shaders/compute.comp
+++ b/shaders/compute.comp
@@ -1,3 +1,5 @@
+#extension GL_EXT_buffer_reference : require
+
 layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
 
 struct ray {
@@ -67,9 +69,7 @@ struct bvh_node {
     int primitive_id;
 };
 
-layout(set = 0, binding = 0) buffer scene_metadata {
-    vec4 sky_color;
-    vec4 ground_color;
+layout(buffer_reference) readonly buffer scene_metadata {
     camera cam;
 
     uint max_bounce;
@@ -84,17 +84,23 @@ layout(set = 0, binding = 0) buffer scene_metadata {
     uint padding;
 };
 
-layout(set = 0, binding = 1) buffer geometry {
+layout(buffer_reference) readonly buffer primitives_array {
     // sphere[] spheres;
-    triangle[] triangles;
+    triangle triangles[];
 };
 
-layout(set = 0, binding = 2) buffer bvh {
+layout(buffer_reference) readonly buffer nodes_array {
     bvh_node[] nodes;
 };
 
-layout(set = 0, binding = 3, rgba32f) uniform writeonly image2D output_image;
-layout(set = 0, binding = 4, rgba32f) uniform readonly image2D accumulation_image;
+layout(set = 0, binding = 0, rgba32f) uniform writeonly image2D output_image;
+layout(set = 0, binding = 1, rgba32f) uniform readonly image2D accumulation_image;
+
+layout(push_constant) uniform buffers {
+    scene_metadata scene;
+    primitives_array geometry;
+    nodes_array bvh;
+} bufs;
 
 vec3 at(ray r, float t) {
     return r.origin + r.direction * t;
@@ -154,12 +160,12 @@ bool hit_aabb(vec3 minimum, vec3 maximum, ray r, float min_t, float max_t) {
 // uint hit_aabbs(ray r, float min_t, float max_t) {
 //     uint hit_count = 0;
 //     // for (uint idx = 0; idx <= spheres_count * 2; idx++) {
-//     for (uint idx = 0; idx <= triangles_count * 2; idx++) {
-//         hit_count += uint(hit_aabb(nodes[idx].bb, r, min_t, max_t));
+//     for (uint idx = 0; idx <= bufs.scene.triangles_count * 2; idx++) {
+//         hit_count += uint(hit_aabb(bufs.bvh.nodes[idx].bb, r, min_t, max_t));
 // 
-//         if (nodes[idx].primitive_id >= 0) {
-//             // hit_count += uint(hit_aabb(spheres[nodes[idx].primitive_id].bb, r, min_t, max_t));
-//             // hit_count += uint(hit_aabb(triangles[nodes[idx].primitive_id].bb, r, min_t, max_t));
+//         if (bufs.bvh.nodes[idx].primitive_id >= 0) {
+//             // hit_count += uint(hit_aabb(spheres[bufs.bvh.nodes[idx].primitive_id].bb, r, min_t, max_t));
+//             // hit_count += uint(hit_aabb(triangles[bufs.bvh.nodes[idx].primitive_id].bb, r, min_t, max_t));
 //         }
 //     }
 // 
@@ -218,10 +224,11 @@ bool hit_triangle(triangle tri, ray r, float min_t, float max_t, out hit_info in
     info.t = t;
     info.point = at(r, info.t);
 
-    vec3 outward_normal = normal;
+    // vec3 outward_normal = normal;
     // info.front_face = dot(outward_normal, r.direction) < 0.0;
     info.normal = normal;
-    info.mat = material(vec4(u, v, 0.0, 0.0), vec4(0.0), 1.0, 1.5, LAMBERTIAN_MATERIAL, 0);
+    // info.mat = material(vec4(u, v, 0.0, 0.0), vec4(0.0), 1.0, 1.5, LAMBERTIAN_MATERIAL, 0);
+    info.mat = material(vec4(0.55), vec4(0.0), 1.0, 1.5, LAMBERTIAN_MATERIAL, 0);
 
     return true;
 }
@@ -234,23 +241,23 @@ bool hit_node(ray r, float min_t, float max_t, out hit_info info) {
 
     while(id != -1) {
         // Leafs
-        if (nodes[id].primitive_id != -1) {
-            // if (hit_sphere(spheres[nodes[id].primitive_id], r, min_t, closest_t, temp_info)) {
+        if (bufs.bvh.nodes[id].primitive_id != -1) {
+            // if (hit_sphere(spheres[bufs.bvh.nodes[id].primitive_id], r, min_t, closest_t, temp_info)) {
             //     info = temp_info;
             //     closest_t = temp_info.t;
             //     hit = true;
             // }
-            if (hit_triangle(triangles[nodes[id].primitive_id], r, min_t, closest_t, temp_info)) {
+            if (hit_triangle(bufs.geometry.triangles[bufs.bvh.nodes[id].primitive_id], r, min_t, closest_t, temp_info)) {
                 info = temp_info;
                 closest_t = temp_info.t;
                 hit = true;
             }
-            id = nodes[id].next_id;
+            id = bufs.bvh.nodes[id].next_id;
         } else { // Internal nodes
-            if (hit_aabb(nodes[id].min, nodes[id].max, r, min_t, closest_t)) {
+            if (hit_aabb(bufs.bvh.nodes[id].min, bufs.bvh.nodes[id].max, r, min_t, closest_t)) {
                 id++;
             } else {
-                id = nodes[id].next_id;
+                id = bufs.bvh.nodes[id].next_id;
             }
         }
     }
@@ -309,7 +316,7 @@ vec4 ray_color(ray r, uint seed) {
     vec4 final_color = vec4(0.0);
     vec4 throughput = vec4(1.0);
     ray current_ray = r;
-    while (hit_node(current_ray, 0.001, 1e15, info) && bounces < max_bounce) {
+    while (hit_node(current_ray, 0.001, 1e15, info) && bounces < bufs.scene.max_bounce) {
         ray scattered;
         vec4 attenuation = vec4(0.0);
         // if ((info.mat.type == LAMBERTIAN_MATERIAL && lambertian_scattering(info.mat, current_ray, seed, info, attenuation, scattered)) ||
@@ -350,27 +357,27 @@ void main() {
     uint x = gl_GlobalInvocationID.x;
     uint y = gl_GlobalInvocationID.y;
 
-    vec2 uv = vec2(float(x) / (width - 1), 1.0 - (float(y) / (height - 1)));
+    vec2 uv = vec2(float(x) / (bufs.scene.width - 1), 1.0 - (float(y) / (bufs.scene.height - 1)));
 
-    uint seed = uint(x * uint(1973) + y * uint(9277) + sample_index * uint(26699)) | uint(1);
+    uint seed = uint(x * uint(1973) + y * uint(9277) + bufs.scene.sample_index * uint(26699)) | uint(1);
     vec2 rand_disk = disk_vec(seed);
-    vec2 jittered_uvs = uv + vec2(rand(seed) / (width - 1), rand(seed) / (height - 1));
+    vec2 jittered_uvs = uv + vec2(rand(seed) / (bufs.scene.width - 1), rand(seed) / (bufs.scene.height - 1));
     vec2 lens_disk = vec2(0.0);
-    if (enable_dof) {
-        lens_disk = rand_disk * cam.lens_radius;
+    if (bufs.scene.enable_dof) {
+        lens_disk = rand_disk * bufs.scene.cam.lens_radius;
     }
 
-    vec3 offset = lens_disk.x * cam.right.xyz + lens_disk.y * cam.up.xyz;
-    vec4 proj_plane_pos = cam.first_pixel + jittered_uvs.x * cam.horizontal + jittered_uvs.y * cam.vertical;
+    vec3 offset = lens_disk.x * bufs.scene.cam.right.xyz + lens_disk.y * bufs.scene.cam.up.xyz;
+    vec4 proj_plane_pos = bufs.scene.cam.first_pixel + jittered_uvs.x * bufs.scene.cam.horizontal + jittered_uvs.y * bufs.scene.cam.vertical;
 
-    ray r = ray(cam.position.xyz + offset, proj_plane_pos.xyz - cam.position.xyz - offset);
+    ray r = ray(bufs.scene.cam.position.xyz + offset, proj_plane_pos.xyz - bufs.scene.cam.position.xyz - offset);
     vec4 out_color = ray_color(r, seed);
 
-    // if (debug_bvh) {
+    // if (bufs.scene.debug_bvh) {
     //     out_color -= hit_aabbs(r, 0.001, 1e15) * vec4(0.0, 0.025, 0.025, 0.0);
     // }
 
-    float spp_scale = (1.0 / sample_index);
+    float spp_scale = (1.0 / bufs.scene.sample_index);
     vec4 acc_color = imageLoad(accumulation_image, ivec2(x, y));
     imageStore(output_image, ivec2(x, y), vec4(mix(acc_color, out_color, spp_scale)));
 }
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);