diff --git a/samples/CMakeLists.txt b/samples/CMakeLists.txt index 30e8d1e3e..c81ab18cb 100644 --- a/samples/CMakeLists.txt +++ b/samples/CMakeLists.txt @@ -62,6 +62,7 @@ set(ORDER_LIST "fragment_shading_rate" "push_descriptors" "raytracing_basic" + "raytracing_reflection" "timeline_semaphore" "synchronization_2" "buffer_device_address" diff --git a/samples/README.md b/samples/README.md index 74787da8b..1970494b5 100644 --- a/samples/README.md +++ b/samples/README.md @@ -209,3 +209,9 @@ Demonstrates how to use descriptor indexing to enable update-after-bind and non- ### [Fragment shading rate](./extensions/fragment_shading_rate)
**Extension**: [```VK_KHR_fragment_shading_rate```](https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_fragment_shading_rate.html)
Uses a special framebuffer attachment to control fragment shading rates for different framebuffer regions. This allows explicit control over the number of fragment shader invocations for each pixel covered by a fragment, which is e.g. useful for foveated rendering. + +### [Ray tracing: reflection, shadow rays](./extensions/ray_tracing_reflection)
+**Extensions**: [```VK_KHR_ray_tracing_pipeline```](https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#VK_KHR_ray_tracing_pipeline), [```VK_KHR_acceleration_structure```](https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#VK_KHR_acceleration_structure), [```VK_EXT_descriptor_indexing```](https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_EXT_descriptor_indexing.html), [```VK_EXT_scalar_block_layout```](https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_EXT_scalar_block_layout.html) +
+Render a simple scene showing the basics of ray tracing, including reflection and shadow rays. The sample creates some geometries and create a bottom acceleration structure for each, then make instances of those, using different materials and placing them at different locations. +
diff --git a/samples/extensions/ray_tracing_reflection/CMakeLists.txt b/samples/extensions/ray_tracing_reflection/CMakeLists.txt new file mode 100644 index 000000000..2d3b004fc --- /dev/null +++ b/samples/extensions/ray_tracing_reflection/CMakeLists.txt @@ -0,0 +1,32 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION +# SPDX-License-Identifier: Apache-2.0 + + +get_filename_component(FOLDER_NAME ${CMAKE_CURRENT_LIST_DIR} NAME) +get_filename_component(PARENT_DIR ${CMAKE_CURRENT_LIST_DIR} PATH) +get_filename_component(CATEGORY_NAME ${PARENT_DIR} NAME) + +file(GLOB GLSL_FILES "${CMAKE_SOURCE_DIR}/shaders/${FOLDER_NAME}/*.*") + +add_sample_with_tags( + ID ${FOLDER_NAME} + FILES ${GLSL_FILES} + CATEGORY ${CATEGORY_NAME} + AUTHOR "Martin-Karl Lefrançois" + NAME "Ray tracing reflection" + DESCRIPTION "Example for hardware accelerated ray tracing") + diff --git a/samples/extensions/ray_tracing_reflection/README.md b/samples/extensions/ray_tracing_reflection/README.md new file mode 100644 index 000000000..96b172de2 --- /dev/null +++ b/samples/extensions/ray_tracing_reflection/README.md @@ -0,0 +1,155 @@ + + +# Ray tracing - reflection + +![](img1.png) + +## Overview + +This sample is a extended version of the [ray tracing basic](../raytracing_basic) with the addition of multiple geometries, +instances and materials. + +In addition, this sample is showing how to cast shadow rays and reflections. + +## Geometries, bottom, and top-level acceleration structures + +The structures of the geometry are like those described in the [OBJ format](https://en.wikipedia.org/wiki/Wavefront_.obj_file). For our geometry there is a list of vertices (position and normal) and a triplet of indices for each triangle. + +Each triangle also has an index for a material. In this example, each object has its own list of materials, but we could have made it so that the materials are shared by all objects in the scene. + +The example scene has two geometries: a plane and a cube. + +You can see the creation of the scene under `RaytracingReflection::create_scene()`. + +### `create_model()` + +This function allocates and upload the geometry to the GPU. There are four buffers per geometry. + +* Vertices: the position and normal +* Indices: index of vertex to form a triangle +* Material index: material id per triangle +* Materials: a list of materials (albedo, specular, reflection) + +### `create_buffer_references()` + +In this example, buffer references are used. Instead of having a descriptor set with multiple arrays to access the buffers, we create a buffer that contains the addresses of the scene models. With this method, we can easily access the data of the model we hit in the shader. + +In the shader, `VkDeviceAddress` are `uint64_t`, therefore we will access a buffer of an array of structure `ObjBuffers`. + +````cpp +struct ObjBuffers +{ + uint64_t vertices; + uint64_t indices; + uint64_t materials; + uint64_t materialIndices; +}; +layout(set = 0, binding = 3) buffer _scene_desc { ObjBuffers i[]; } scene_desc; +```` + +The addresses correspond to buffers, so we will declare them like this: + +````cpp +layout(buffer_reference, scalar) buffer Vertices {Vertex v[]; }; // Positions of an object +layout(buffer_reference, scalar) buffer Indices {uvec3 i[]; }; // Triangle indices +layout(buffer_reference, scalar) buffer Materials {WaveFrontMaterial m[]; }; // Array of all materials on an object +layout(buffer_reference, scalar) buffer MatIndices {int i[]; }; // Material ID for each triangle +```` + +The in the shader, to access the data of an object, we will do the following: + +````cpp + ObjBuffers objResource = scene_desc.i[gl_InstanceCustomIndexEXT]; + MatIndices matIndices = MatIndices(objResource.materialIndices); + Materials materials = Materials(objResource.materials); + Indices indices = Indices(objResource.indices); + Vertices vertices = Vertices(objResource.vertices); +```` + +Note that `gl_InstanceCustomIndexEXT` was set with one of the three scene objects. See `RaytracingReflection::create_blas_instance`. + +### `create_bottom_level_acceleration_structure()` + +We build a lower-level acceleration structure (BLAS) for each geometry: a cube with one material on each face (0), a plane (1), and a mirror cube (2). + +These BLAS are instantiated by the top-level acceleration structure (TLAS) with a transformation matrix. + +In this example w are calling separately the construction of all BLAS, allocating scratch buffer each time. A [better way](https://nvpro-samples.github.io/vk_raytracing_tutorial_KHR/#accelerationstructure/bottom-levelaccelerationstructure/helperdetails:raytracingbuilder::buildblas()) would be to build them, knowing the size the biggest scratch buffer and doing all at once. Also provide in [the helper](https://nvpro-samples.github.io/vk_raytracing_tutorial_KHR/#accelerationstructure/bottom-levelaccelerationstructure/helperdetails:raytracingbuilder::buildblas()) , is the ability to compact the memory used to store the BLAS when using the `VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_COMPACTION_BIT_KHR` flag. + +### `create_top_level_acceleration_structure()` + +The top-level acceleration structure (TLAS) embeds multiple BLAS. It is possible to reuse the same BLAS and give it a different transformation matrix to place it in a different position in the scene. We see this in `create_scene()`, the same BLAS id is reused with different matrices. + +Note, the BLAS id will be identified by the `gl_InstanceCustomIndexEXT` in the shader. + +![](img2.png) + +## Ray tracing pipeline + +The difference with [ray tracing basic](../raytracing_basic), is the addition of the second miss-shade module. This is called from the closest-hit shader to detect if there is an object between the hit point and the light. + +More on ray tracing pipeline [here](https://nvpro-samples.github.io/vk_raytracing_tutorial_KHR/#raytracingpipeline). + +### Shadows + +In the closest-hit shader, we trace a ray to determine if there is an object between the hit point and the light. For this trace ray, we use the shadow-miss shader (index 1) and a different payload (index 1) that contains only one boolean value, "`isShadowed`". We assume that an object is blocking the light, so we initialize the value to `true`. Then, if the ray hits nothing, we [set this value to false](missShadow.rmiss). + +The origin of the ray is the hit position and the direction of the ray is toward the light. Note, we are using an hardcoded infinite light **L**. + +This method for shooting shadow rays is fast because we set the trace flag to skip execution of the closest-hit and terminate on the first hit, then only execute the shadow-miss-shader and set a small payload. + +````cpp +uint flags = gl_RayFlagsTerminateOnFirstHitEXT | gl_RayFlagsOpaqueEXT | gl_RayFlagsSkipClosestHitShaderEXT; +```` + +### Reflection + +For reflection, we do not change `maxPipelineRayRecursionDepth` and leave the value at 2. Instead of recursively looping from the closest-hit shader, we store the information of the next ray in the payload and send new rays from the ray generation shader. + +When we call `traceRayEXT` from the closest-hit shader, it must store the state of all variables needed after execution. Recursively calling `traceRayEXT` requires storing a lot of data per ray call, and that is typically slow. + +Instead, we store in the payload the ray origin and the ray direction that the ray generation shader will use. This method also removes the pipeline ray recursion depth constraint. + +Here is how the payload is defined: + +````cpp +struct hitPayload +{ + vec3 radiance; + vec3 attenuation; + int done; + vec3 rayOrigin; + vec3 rayDir; +}; +```` + +The radiance is the value at the point of impact multiplied by the attenuation. The first time the attenuation is vec3(1) (no attenuation), but the shininess of the material reduces the attenuation in the following hits. After a few passes, the radiance will be close to vec3(0). The `done` is an indication that the ray did not hit anything. The miss shader sets it to true and the loop can be terminated. The origin of the ray starts at the camera, and is replaced by the position of the target. The direction starts at the camera direction, and then is reflected purely at the surface of the object. + +The recursion limit is set in the ray-generation shader. Currently it is set to **64**, changing its value will change the number of times the ray bounces off. + +Note: we could add a test on the attenuation and exist the loop if the value is below a certain threshold. + +## Diagram of the ray pipeline + +![](img3.png) + +## Other Tutorial + +The [following tutorial](https://github.com/nvpro-samples/vk_raytracing_tutorial_KHR/tree/master/ray_tracing_reflections) is showing the limitation of hitting the recursion limits. Note, the spec does not guarantee a recursion check at runtime. If you exceed either the recursion depth you reported in the raytrace pipeline create info, or the physical device recursion limit, undefined behavior results. diff --git a/samples/extensions/ray_tracing_reflection/img1.png b/samples/extensions/ray_tracing_reflection/img1.png new file mode 100644 index 000000000..7670af1e9 Binary files /dev/null and b/samples/extensions/ray_tracing_reflection/img1.png differ diff --git a/samples/extensions/ray_tracing_reflection/img2.png b/samples/extensions/ray_tracing_reflection/img2.png new file mode 100644 index 000000000..e06b7f7aa Binary files /dev/null and b/samples/extensions/ray_tracing_reflection/img2.png differ diff --git a/samples/extensions/ray_tracing_reflection/img3.png b/samples/extensions/ray_tracing_reflection/img3.png new file mode 100644 index 000000000..50120117e Binary files /dev/null and b/samples/extensions/ray_tracing_reflection/img3.png differ diff --git a/samples/extensions/ray_tracing_reflection/ray_tracing_reflection.cpp b/samples/extensions/ray_tracing_reflection/ray_tracing_reflection.cpp new file mode 100644 index 000000000..b65ca5fc2 --- /dev/null +++ b/samples/extensions/ray_tracing_reflection/ray_tracing_reflection.cpp @@ -0,0 +1,978 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION + * SPDX-License-Identifier: Apache-2.0 + */ + +/* + * More complex example for hardware accelerated ray tracing using VK_KHR_ray_tracing_pipeline and VK_KHR_acceleration_structure + */ + +#define TINYOBJLOADER_IMPLEMENTATION + +#include "ray_tracing_reflection.h" +#include + +struct ObjPlane : ObjModelCpu +{ + ObjPlane() + { + vertices = { + {{+1, 0, +1}, {0, 1, 0}}, + {{-1, 0, +1}, {0, 1, 0}}, + {{+1, 0, -1}, {0, 1, 0}}, + {{-1, 0, -1}, {0, 1, 0}}, + }; + indices = {0, 1, 2, 1, 2, 3}; + mat_index = {0, 0}; + } +}; + +struct ObjCube : ObjModelCpu +{ + ObjCube() + { + vertices = { + {{+0.5f, +0.5f, +0.5f}, {+0.f, +1.f, +0.f}}, // Top + {{-0.5f, +0.5f, +0.5f}, {+0.f, +1.f, +0.f}}, + {{+0.5f, +0.5f, -0.5f}, {+0.f, +1.f, +0.f}}, + {{-0.5f, +0.5f, -0.5f}, {+0.f, +1.f, +0.f}}, + {{+0.5f, -0.5f, +0.5f}, {+0.f, -1.f, +0.f}}, // Bottom + {{-0.5f, -0.5f, +0.5f}, {+0.f, -1.f, +0.f}}, + {{+0.5f, -0.5f, -0.5f}, {+0.f, -1.f, +0.f}}, + {{-0.5f, -0.5f, -0.5f}, {+0.f, -1.f, +0.f}}, + {{+0.5f, +0.5f, +0.5f}, {+1.f, +0.f, +0.f}}, // Right + {{+0.5f, +0.5f, -0.5f}, {+1.f, +0.f, +0.f}}, + {{+0.5f, -0.5f, -0.5f}, {+1.f, +0.f, +0.f}}, + {{+0.5f, -0.5f, +0.5f}, {+1.f, +0.f, +0.f}}, + {{-0.5f, +0.5f, +0.5f}, {-1.f, +0.f, +0.f}}, // left + {{-0.5f, +0.5f, -0.5f}, {-1.f, +0.f, +0.f}}, + {{-0.5f, -0.5f, -0.5f}, {-1.f, +0.f, +0.f}}, + {{-0.5f, -0.5f, +0.5f}, {-1.f, +0.f, +0.f}}, + {{-0.5f, +0.5f, +0.5f}, {+0.f, +0.f, +1.f}}, // front + {{+0.5f, +0.5f, +0.5f}, {+0.f, +0.f, +1.f}}, + {{+0.5f, -0.5f, +0.5f}, {+0.f, +0.f, +1.f}}, + {{-0.5f, -0.5f, +0.5f}, {+0.f, +0.f, +1.f}}, + {{-0.5f, +0.5f, -0.5f}, {+0.f, +0.f, -1.f}}, // back + {{+0.5f, +0.5f, -0.5f}, {+0.f, +0.f, -1.f}}, + {{+0.5f, -0.5f, -0.5f}, {+0.f, +0.f, -1.f}}, + {{-0.5f, -0.5f, -0.5f}, {+0.f, +0.f, -1.f}}, + + }; + indices = { + 0, 1, 2, 1, 2, 3, /*top*/ + 4, 5, 6, 5, 6, 7, /*bottom*/ + 8, 9, 10, 8, 10, 11, /*right*/ + 12, 13, 14, 12, 14, 15, /*left*/ + 16, 17, 18, 16, 18, 19, /*front*/ + 20, 21, 22, 20, 22, 23, /*back*/ + }; + mat_index = {0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5}; + } +}; + +RaytracingReflection::RaytracingReflection() +{ + title = "Hardware accelerated ray tracing"; + + set_api_version(VK_API_VERSION_1_2); + + // Ray tracing related extensions required by this sample + add_device_extension(VK_KHR_ACCELERATION_STRUCTURE_EXTENSION_NAME); + add_device_extension(VK_KHR_RAY_TRACING_PIPELINE_EXTENSION_NAME); + + // Required by VK_KHR_acceleration_structure + add_device_extension(VK_KHR_BUFFER_DEVICE_ADDRESS_EXTENSION_NAME); + add_device_extension(VK_KHR_DEFERRED_HOST_OPERATIONS_EXTENSION_NAME); +} + +RaytracingReflection::~RaytracingReflection() +{ + if (device) + { + vkDestroyPipeline(get_device().get_handle(), pipeline, nullptr); + vkDestroyPipelineLayout(get_device().get_handle(), pipeline_layout, nullptr); + vkDestroyDescriptorSetLayout(get_device().get_handle(), descriptor_set_layout, nullptr); + vkDestroyImageView(get_device().get_handle(), storage_image.view, nullptr); + vkDestroyImage(get_device().get_handle(), storage_image.image, nullptr); + vkFreeMemory(get_device().get_handle(), storage_image.memory, nullptr); + delete_acceleration_structure(top_level_acceleration_structure); + for (auto &b : bottom_level_acceleration_structure) + { + delete_acceleration_structure(b); + } + + for (auto &obj : obj_models) + { + obj.vertex_buffer.reset(); + obj.index_buffer.reset(); + obj.mat_color_buffer.reset(); + obj.mat_index_buffer.reset(); + } + + ubo.reset(); + } +} + +/* + Enable extension features required by this sample + These are passed to device creation via a pNext structure chain +*/ +void RaytracingReflection::request_gpu_features(vkb::PhysicalDevice &gpu) +{ + // The request is filling with the capabilities (all on by default) + auto &vulkan12_features = gpu.request_extension_features(VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES); + auto &vulkan11_features = gpu.request_extension_features(VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES); + + auto &ray_tracing_features = gpu.request_extension_features(VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_PIPELINE_FEATURES_KHR); + auto &acceleration_structure_features = gpu.request_extension_features(VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_FEATURES_KHR); + + // Enabling all Vulkan features (Int64) + gpu.get_mutable_requested_features() = gpu.get_features(); +} + +/* + Set up a storage image that the ray generation shader will be writing to +*/ +void RaytracingReflection::create_storage_image() +{ + storage_image.width = width; + storage_image.height = height; + + VkImageCreateInfo image{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; + image.imageType = VK_IMAGE_TYPE_2D; + image.format = VK_FORMAT_B8G8R8A8_UNORM; + image.extent.width = storage_image.width; + image.extent.height = storage_image.height; + image.extent.depth = 1; + image.mipLevels = 1; + image.arrayLayers = 1; + image.samples = VK_SAMPLE_COUNT_1_BIT; + image.tiling = VK_IMAGE_TILING_OPTIMAL; + image.usage = VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_STORAGE_BIT; + image.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; + VK_CHECK(vkCreateImage(get_device().get_handle(), &image, nullptr, &storage_image.image)); + + VkMemoryRequirements memory_requirements; + vkGetImageMemoryRequirements(get_device().get_handle(), storage_image.image, &memory_requirements); + VkMemoryAllocateInfo memory_allocate_info{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; + memory_allocate_info.allocationSize = memory_requirements.size; + memory_allocate_info.memoryTypeIndex = get_device().get_memory_type(memory_requirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + VK_CHECK(vkAllocateMemory(get_device().get_handle(), &memory_allocate_info, nullptr, &storage_image.memory)); + VK_CHECK(vkBindImageMemory(get_device().get_handle(), storage_image.image, storage_image.memory, 0)); + + VkImageViewCreateInfo color_image_view{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; + color_image_view.viewType = VK_IMAGE_VIEW_TYPE_2D; + color_image_view.format = VK_FORMAT_B8G8R8A8_UNORM; + color_image_view.subresourceRange = {}; + color_image_view.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT; + color_image_view.subresourceRange.baseMipLevel = 0; + color_image_view.subresourceRange.levelCount = 1; + color_image_view.subresourceRange.baseArrayLayer = 0; + color_image_view.subresourceRange.layerCount = 1; + color_image_view.image = storage_image.image; + VK_CHECK(vkCreateImageView(get_device().get_handle(), &color_image_view, nullptr, &storage_image.view)); + + VkCommandBuffer command_buffer = get_device().create_command_buffer(VK_COMMAND_BUFFER_LEVEL_PRIMARY, true); + vkb::set_image_layout(command_buffer, storage_image.image, + VK_IMAGE_LAYOUT_UNDEFINED, + VK_IMAGE_LAYOUT_GENERAL, + {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}); + get_device().flush_command_buffer(command_buffer, queue); +} + +/* + Create the bottom level acceleration structure that contains the scene's geometry (triangles) +*/ +void RaytracingReflection::create_bottom_level_acceleration_structure(ObjModelGpu &obj_model) +{ + // Note that the buffer usage flags for buffers consumed by the bottom level acceleration structure require special flags + const VkBufferUsageFlags buffer_usage_flags = VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + + // Setup a single transformation matrix that can be used to transform the whole geometry for a single bottom level acceleration structure + VkTransformMatrixKHR transform_matrix = { + 1.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 1.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 1.0f, 0.0f}; + std::unique_ptr transform_matrix_buffer = std::make_unique(get_device(), sizeof(transform_matrix), buffer_usage_flags, VMA_MEMORY_USAGE_CPU_TO_GPU); + transform_matrix_buffer->update(&transform_matrix, sizeof(transform_matrix)); + + VkDeviceOrHostAddressConstKHR vertex_data_device_address{}; + VkDeviceOrHostAddressConstKHR index_data_device_address{}; + VkDeviceOrHostAddressConstKHR transform_matrix_device_address{}; + + vertex_data_device_address.deviceAddress = obj_model.vertex_buffer->get_device_address(); + index_data_device_address.deviceAddress = obj_model.index_buffer->get_device_address(); + transform_matrix_device_address.deviceAddress = transform_matrix_buffer->get_device_address(); + + VkAccelerationStructureGeometryTrianglesDataKHR triangles{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR}; + triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR; + triangles.vertexFormat = VK_FORMAT_R32G32B32_SFLOAT; + triangles.vertexData = vertex_data_device_address; + triangles.maxVertex = obj_model.nb_vertices; + triangles.vertexStride = sizeof(ObjVertex); + triangles.indexType = VK_INDEX_TYPE_UINT32; + triangles.indexData = index_data_device_address; + triangles.transformData = transform_matrix_device_address; + + // The bottom level acceleration structure contains one set of triangles as the input geometry + VkAccelerationStructureGeometryKHR acceleration_structure_geometry{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR}; + acceleration_structure_geometry.geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR; + acceleration_structure_geometry.flags = VK_GEOMETRY_OPAQUE_BIT_KHR; + acceleration_structure_geometry.geometry.triangles = triangles; + + // Get the size requirements for buffers involved in the acceleration structure build process + VkAccelerationStructureBuildGeometryInfoKHR acceleration_structure_build_geometry_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR}; + acceleration_structure_build_geometry_info.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR; + acceleration_structure_build_geometry_info.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR; + acceleration_structure_build_geometry_info.geometryCount = 1; + acceleration_structure_build_geometry_info.pGeometries = &acceleration_structure_geometry; + + const uint32_t triangle_count = obj_model.nb_indices / 3; + + VkAccelerationStructureBuildSizesInfoKHR acceleration_structure_build_sizes_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR}; + vkGetAccelerationStructureBuildSizesKHR(device->get_handle(), + VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR, + &acceleration_structure_build_geometry_info, + &triangle_count, + &acceleration_structure_build_sizes_info); + + // Create a buffer to hold the acceleration structure + AccelerationStructure blas; + blas.buffer = std::make_unique(get_device(), acceleration_structure_build_sizes_info.accelerationStructureSize, VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR, VMA_MEMORY_USAGE_GPU_ONLY); + + // Create the acceleration structure + VkAccelerationStructureCreateInfoKHR acceleration_structure_create_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR}; + acceleration_structure_create_info.buffer = blas.buffer->get_handle(); + acceleration_structure_create_info.size = acceleration_structure_build_sizes_info.accelerationStructureSize; + acceleration_structure_create_info.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR; + vkCreateAccelerationStructureKHR(device->get_handle(), &acceleration_structure_create_info, nullptr, &blas.handle); + + // The actual build process starts here + + // Create a scratch buffer as a temporary storage for the acceleration structure build + std::unique_ptr sc_buffer; + sc_buffer = std::make_unique(get_device(), acceleration_structure_build_sizes_info.buildScratchSize, + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, + VMA_MEMORY_USAGE_CPU_TO_GPU); + + VkAccelerationStructureBuildGeometryInfoKHR acceleration_build_geometry_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR}; + acceleration_build_geometry_info.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR; + acceleration_build_geometry_info.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR; + acceleration_build_geometry_info.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR; + acceleration_build_geometry_info.dstAccelerationStructure = blas.handle; + acceleration_build_geometry_info.geometryCount = 1; + acceleration_build_geometry_info.pGeometries = &acceleration_structure_geometry; + acceleration_build_geometry_info.scratchData.deviceAddress = sc_buffer->get_device_address(); + + VkAccelerationStructureBuildRangeInfoKHR acceleration_structure_build_range_info; + acceleration_structure_build_range_info.primitiveCount = triangle_count; + acceleration_structure_build_range_info.primitiveOffset = 0; + acceleration_structure_build_range_info.firstVertex = 0; + acceleration_structure_build_range_info.transformOffset = 0; + std::vector acceleration_build_structure_range_infos = {&acceleration_structure_build_range_info}; + + // Build the acceleration structure on the device via a one-time command buffer submission + // Some implementations may support acceleration structure building on the host (VkPhysicalDeviceAccelerationStructureFeaturesKHR->accelerationStructureHostCommands), but we prefer device builds + VkCommandBuffer command_buffer = get_device().create_command_buffer(VK_COMMAND_BUFFER_LEVEL_PRIMARY, true); + vkCmdBuildAccelerationStructuresKHR(command_buffer, + 1, + &acceleration_build_geometry_info, + acceleration_build_structure_range_infos.data()); + get_device().flush_command_buffer(command_buffer, queue); + + //delete_scratch_buffer(scratch_buffer); + sc_buffer.reset(); + + // Store the blas to be re-used as instance + bottom_level_acceleration_structure.push_back(std::move(blas)); +} + +/* + Create the top level acceleration structure containing geometry instances of the bottom level acceleration structure(s) +*/ +void RaytracingReflection::create_top_level_acceleration_structure(std::vector &blas_instances) +{ + std::unique_ptr instances_buffer = std::make_unique(get_device(), + sizeof(VkAccelerationStructureInstanceKHR) * blas_instances.size(), + VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, + VMA_MEMORY_USAGE_CPU_TO_GPU); + instances_buffer->update(blas_instances.data(), sizeof(VkAccelerationStructureInstanceKHR) * blas_instances.size()); + + VkDeviceOrHostAddressConstKHR instance_data_device_address{}; + instance_data_device_address.deviceAddress = instances_buffer->get_device_address(); + + // The top level acceleration structure contains (bottom level) instance as the input geometry + VkAccelerationStructureGeometryKHR acceleration_structure_geometry{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR}; + acceleration_structure_geometry.geometryType = VK_GEOMETRY_TYPE_INSTANCES_KHR; + acceleration_structure_geometry.flags = VK_GEOMETRY_OPAQUE_BIT_KHR; + acceleration_structure_geometry.geometry.instances.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_INSTANCES_DATA_KHR; + acceleration_structure_geometry.geometry.instances.arrayOfPointers = VK_FALSE; + acceleration_structure_geometry.geometry.instances.data = instance_data_device_address; + + // Get the size requirements for buffers involved in the acceleration structure build process + VkAccelerationStructureBuildGeometryInfoKHR acceleration_structure_build_geometry_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR}; + acceleration_structure_build_geometry_info.type = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR; + acceleration_structure_build_geometry_info.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR; + acceleration_structure_build_geometry_info.geometryCount = 1; + acceleration_structure_build_geometry_info.pGeometries = &acceleration_structure_geometry; + + const auto primitive_count = static_cast(blas_instances.size()); + + VkAccelerationStructureBuildSizesInfoKHR acceleration_structure_build_sizes_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR}; + vkGetAccelerationStructureBuildSizesKHR( + device->get_handle(), VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR, + &acceleration_structure_build_geometry_info, + &primitive_count, + &acceleration_structure_build_sizes_info); + + // Create a buffer to hold the acceleration structure + top_level_acceleration_structure.buffer = std::make_unique( + get_device(), + acceleration_structure_build_sizes_info.accelerationStructureSize, + VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR, + VMA_MEMORY_USAGE_GPU_ONLY); + + // Create the acceleration structure + VkAccelerationStructureCreateInfoKHR acceleration_structure_create_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR}; + acceleration_structure_create_info.buffer = top_level_acceleration_structure.buffer->get_handle(); + acceleration_structure_create_info.size = acceleration_structure_build_sizes_info.accelerationStructureSize; + acceleration_structure_create_info.type = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR; + vkCreateAccelerationStructureKHR(device->get_handle(), &acceleration_structure_create_info, nullptr, &top_level_acceleration_structure.handle); + + // The actual build process starts here + + // Create a scratch buffer as a temporary storage for the acceleration structure build + std::unique_ptr sc_buffer; + sc_buffer = std::make_unique(get_device(), acceleration_structure_build_sizes_info.buildScratchSize, + VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, + VMA_MEMORY_USAGE_CPU_TO_GPU); + + VkAccelerationStructureBuildGeometryInfoKHR acceleration_build_geometry_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR}; + acceleration_build_geometry_info.type = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR; + acceleration_build_geometry_info.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR; + acceleration_build_geometry_info.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR; + acceleration_build_geometry_info.dstAccelerationStructure = top_level_acceleration_structure.handle; + acceleration_build_geometry_info.geometryCount = 1; + acceleration_build_geometry_info.pGeometries = &acceleration_structure_geometry; + acceleration_build_geometry_info.scratchData.deviceAddress = sc_buffer->get_device_address(); + + VkAccelerationStructureBuildRangeInfoKHR acceleration_structure_build_range_info; + acceleration_structure_build_range_info.primitiveCount = primitive_count; + acceleration_structure_build_range_info.primitiveOffset = 0; + acceleration_structure_build_range_info.firstVertex = 0; + acceleration_structure_build_range_info.transformOffset = 0; + std::vector acceleration_build_structure_range_infos = {&acceleration_structure_build_range_info}; + + // Build the acceleration structure on the device via a one-time command buffer submission + // Some implementations may support acceleration structure building on the host (VkPhysicalDeviceAccelerationStructureFeaturesKHR->accelerationStructureHostCommands), but we prefer device builds + VkCommandBuffer command_buffer = get_device().create_command_buffer(VK_COMMAND_BUFFER_LEVEL_PRIMARY, true); + vkCmdBuildAccelerationStructuresKHR( + command_buffer, + 1, + &acceleration_build_geometry_info, + acceleration_build_structure_range_infos.data()); + get_device().flush_command_buffer(command_buffer, queue); + + //delete_scratch_buffer(scratch_buffer); + sc_buffer.reset(); +} + +inline uint32_t aligned_size(uint32_t value, uint32_t alignment) +{ + return (value + alignment - 1) & ~(alignment - 1); +} + +/* +Create the GPU representation of the model +*/ +void RaytracingReflection::create_model(ObjModelCpu &obj, const std::vector &materials) +{ + ObjModelGpu model; + model.nb_indices = static_cast(obj.indices.size()); + model.nb_vertices = static_cast(obj.vertices.size()); + + auto vertex_buffer_size = obj.vertices.size() * sizeof(ObjVertex); + auto index_buffer_size = obj.indices.size() * sizeof(uint32_t); + auto mat_index_buffer_size = obj.mat_index.size() * sizeof(int32_t); + auto mat_buffer_size = materials.size() * sizeof(ObjMaterial); + + // Making sure the material triangle index don't exceed the number of materials + auto max_index = static_cast(materials.size() - 1); + std::vector mat_index(obj.mat_index.size()); + for (auto i = 0; i < obj.mat_index.size(); i++) + mat_index[i] = std::min(max_index, obj.mat_index[i]); + + // Note that the buffer usage flags for buffers consumed by the bottom level acceleration structure require special flags + VkBufferUsageFlags buffer_usage_flags = VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + + model.vertex_buffer = std::make_unique(get_device(), vertex_buffer_size, buffer_usage_flags, VMA_MEMORY_USAGE_CPU_TO_GPU); + model.vertex_buffer->update(obj.vertices.data(), vertex_buffer_size); + + // Acceleration structure flag is not needed for the rest + buffer_usage_flags = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + + model.index_buffer = std::make_unique(get_device(), index_buffer_size, buffer_usage_flags, VMA_MEMORY_USAGE_CPU_TO_GPU); + model.index_buffer->update(obj.indices.data(), index_buffer_size); + + model.mat_index_buffer = std::make_unique(get_device(), mat_index_buffer_size, buffer_usage_flags, VMA_MEMORY_USAGE_CPU_TO_GPU); + model.mat_index_buffer->update(mat_index.data(), mat_index_buffer_size); + + model.mat_color_buffer = std::make_unique(get_device(), mat_buffer_size, buffer_usage_flags, VMA_MEMORY_USAGE_CPU_TO_GPU); + model.mat_color_buffer->update(reinterpret_cast(materials.data()), mat_buffer_size); + + obj_models.push_back(std::move(model)); +} + +auto RaytracingReflection::create_blas_instance(uint32_t blas_id, glm::mat4 &mat) +{ + VkTransformMatrixKHR transform_matrix; + glm::mat3x4 rtxT = glm::transpose(mat); + memcpy(&transform_matrix, glm::value_ptr(rtxT), sizeof(VkTransformMatrixKHR)); + + AccelerationStructure &blas = bottom_level_acceleration_structure[blas_id]; + + // Get the bottom acceleration structure's handle, which will be used during the top level acceleration build + VkAccelerationStructureDeviceAddressInfoKHR acceleration_device_address_info{VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR}; + acceleration_device_address_info.accelerationStructure = blas.handle; + auto device_address = vkGetAccelerationStructureDeviceAddressKHR(device->get_handle(), &acceleration_device_address_info); + + VkAccelerationStructureInstanceKHR blas_instance{}; + blas_instance.transform = transform_matrix; + blas_instance.instanceCustomIndex = blas_id; + blas_instance.mask = 0xFF; + blas_instance.instanceShaderBindingTableRecordOffset = 0; + blas_instance.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR; + blas_instance.accelerationStructureReference = device_address; + + return blas_instance; +} + +/* + Create a buffer holding the address of model buffers (buffer reference) +*/ +void RaytracingReflection::create_buffer_references() +{ + // For each model that was created, we retrieved the address of buffers + // used by them. So in the shader, we have direct access to the data + std::vector obj_data; + auto nbObj = static_cast(obj_models.size()); + for (uint32_t i = 0; i < nbObj; ++i) + { + ObjBuffers data; + data.vertices = obj_models[i].vertex_buffer->get_device_address(); + data.indices = obj_models[i].index_buffer->get_device_address(); + data.materials = obj_models[i].mat_color_buffer->get_device_address(); + data.materialIndices = obj_models[i].mat_index_buffer->get_device_address(); + obj_data.emplace_back(data); + } + VkBufferUsageFlags buffer_usage_flags = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + scene_desc = std::make_unique(get_device(), nbObj * sizeof(ObjBuffers), buffer_usage_flags, VMA_MEMORY_USAGE_CPU_TO_GPU); + scene_desc->update(obj_data.data(), nbObj * sizeof(ObjBuffers)); +} + +/* + Create scene geometry and ray tracing acceleration structures +*/ +void RaytracingReflection::create_scene() +{ + // Materials + ObjMaterial mat_red = {{1, 0, 0}, {1, 1, 1}, {0}}; + ObjMaterial mat_green = {{0, 1, 0}, {1, 1, 1}, {0}}; + ObjMaterial mat_blue = {{0, 0, 1}, {1, 1, 1}, {0}}; + ObjMaterial mat_yellow = {{1, 1, 0}, {1, 1, 1}, {0}}; + ObjMaterial mat_cyan = {{0, 1, 1}, {1, 1, 1}, {0}}; + ObjMaterial mat_magenta = {{1, 0, 1}, {1, 1, 1}, {0}}; + ObjMaterial mat_grey = {{0.7f, 0.7f, 0.7f}, {0.9f, 0.9f, 0.9f}, {0.1f}}; // Slightly reflective + ObjMaterial mat_mirror = {{0.3f, 0.9f, 1.0f}, {0.9f, 0.9f, 0.9f}, {0.9f}}; // Mirror Slightly blue + + // Geometries + auto cube = ObjCube(); + auto plane = ObjPlane(); + + // Upload geometries to GPU + create_model(cube, {mat_red, mat_green, mat_blue, mat_yellow, mat_cyan, mat_magenta}); // 6 color faces + create_model(plane, {mat_grey}); + create_model(cube, {mat_mirror}); + + // Create a buffer holding the address of model buffers (buffer reference) + create_buffer_references(); + + // Create as many bottom acceleration structures (blas) as there are geometries/models + create_bottom_level_acceleration_structure(obj_models[0]); + create_bottom_level_acceleration_structure(obj_models[1]); + create_bottom_level_acceleration_structure(obj_models[2]); + + // Matrices to position the instances + glm::mat4 m_mirror_back = glm::scale(glm::translate(glm::mat4(1.f), glm::vec3(0.0f, 0.0f, -7.0f)), glm::vec3(5.0f, 5.0f, 0.1f)); + glm::mat4 m_mirror_front = glm::scale(glm::translate(glm::mat4(1.f), glm::vec3(0.0f, 0.0f, 7.0f)), glm::vec3(5.0f, 5.0f, 0.1f)); + glm::mat4 m_plane = glm::scale(glm::translate(glm::mat4(1.f), glm::vec3(0.0f, -1.0f, 0.0f)), glm::vec3(15.0f, 15.0f, 15.0f)); + glm::mat4 m_cube_left = glm::translate(glm::mat4(1.f), glm::vec3(-1.0f, 0.0f, 0.0f)); + glm::mat4 m_cube_right = glm::translate(glm::mat4(1.f), glm::vec3(1.0f, 0.0f, 0.0f)); + + // Creating instances of the blas to the top level acceleration structure + std::vector blas_instances; + blas_instances.push_back(create_blas_instance(0, m_cube_left)); + blas_instances.push_back(create_blas_instance(0, m_cube_right)); + blas_instances.push_back(create_blas_instance(1, m_plane)); + blas_instances.push_back(create_blas_instance(2, m_mirror_back)); + blas_instances.push_back(create_blas_instance(2, m_mirror_front)); + + // Building the TLAS + create_top_level_acceleration_structure(blas_instances); +} + +/* + Create the Shader Binding Tables that connects the ray tracing pipelines' programs and the top-level acceleration structure + + SBT Layout used in this sample: + + /-------------\ + | raygen | + |-------------| + | miss | + |-------------| + | miss shadow | + |-------------| + | hit | + \-------------/ +*/ + +void RaytracingReflection::create_shader_binding_tables() +{ + // Index position of the groups in the generated ray tracing pipeline + // To be generic, this should be pass in parameters + std::vector rgen_index{0}; + std::vector miss_index{1, 2}; + std::vector hit_index{3}; + + const uint32_t handle_size = ray_tracing_pipeline_properties.shaderGroupHandleSize; + const uint32_t handle_alignment = ray_tracing_pipeline_properties.shaderGroupHandleAlignment; + const uint32_t handle_size_aligned = aligned_size(handle_size, handle_alignment); + + const VkBufferUsageFlags sbt_buffer_usage_flags = VK_BUFFER_USAGE_SHADER_BINDING_TABLE_BIT_KHR | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + const VmaMemoryUsage sbt_memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU; + + // Create binding table buffers for each shader type + raygen_shader_binding_table = std::make_unique(get_device(), handle_size_aligned * rgen_index.size(), sbt_buffer_usage_flags, sbt_memory_usage, 0); + miss_shader_binding_table = std::make_unique(get_device(), handle_size_aligned * miss_index.size(), sbt_buffer_usage_flags, sbt_memory_usage, 0); + hit_shader_binding_table = std::make_unique(get_device(), handle_size_aligned * hit_index.size(), sbt_buffer_usage_flags, sbt_memory_usage, 0); + + // Copy the pipeline's shader handles into a host buffer + const auto group_count = static_cast(rgen_index.size() + miss_index.size() + hit_index.size()); + const auto sbt_size = group_count * handle_size_aligned; + std::vector shader_handle_storage(sbt_size); + VK_CHECK(vkGetRayTracingShaderGroupHandlesKHR(get_device().get_handle(), pipeline, 0, group_count, sbt_size, shader_handle_storage.data())); + + // Write the handles in the SBT buffer + auto copyHandles = [&](auto &buffer, std::vector &indices, uint32_t stride) { + auto *pBuffer = static_cast(buffer->map()); + for (uint32_t index = 0; index < static_cast(indices.size()); index++) + { + auto *pStart = pBuffer; + // Copy the handle + memcpy(pBuffer, shader_handle_storage.data() + (indices[index] * handle_size), handle_size); + pBuffer = pStart + stride; // Jumping to next group + } + buffer->unmap(); + }; + + copyHandles(raygen_shader_binding_table, rgen_index, handle_size_aligned); + copyHandles(miss_shader_binding_table, miss_index, handle_size_aligned); + copyHandles(hit_shader_binding_table, hit_index, handle_size_aligned); +} + +/* + Create the descriptor sets used for the ray tracing dispatch +*/ +void RaytracingReflection::create_descriptor_sets() +{ + uint32_t nbObj = static_cast(obj_models.size()); + + std::vector pool_sizes = { + {VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, 1}, + {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1}, + {VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1}, + {VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1}, + }; + VkDescriptorPoolCreateInfo descriptor_pool_create_info = vkb::initializers::descriptor_pool_create_info(pool_sizes, 1); + VK_CHECK(vkCreateDescriptorPool(get_device().get_handle(), &descriptor_pool_create_info, nullptr, &descriptor_pool)); + + VkDescriptorSetAllocateInfo descriptor_set_allocate_info = vkb::initializers::descriptor_set_allocate_info(descriptor_pool, &descriptor_set_layout, 1); + VK_CHECK(vkAllocateDescriptorSets(get_device().get_handle(), &descriptor_set_allocate_info, &descriptor_set)); + + // Setup the descriptor for binding our top level acceleration structure to the ray tracing shaders + VkWriteDescriptorSetAccelerationStructureKHR descriptor_acceleration_structure_info{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR}; + descriptor_acceleration_structure_info.accelerationStructureCount = 1; + descriptor_acceleration_structure_info.pAccelerationStructures = &top_level_acceleration_structure.handle; + + VkWriteDescriptorSet acceleration_structure_write{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET}; + acceleration_structure_write.dstSet = descriptor_set; + acceleration_structure_write.dstBinding = 0; + acceleration_structure_write.descriptorCount = 1; + acceleration_structure_write.descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR; + // The acceleration structure descriptor has to be chained via pNext + acceleration_structure_write.pNext = &descriptor_acceleration_structure_info; + + VkDescriptorImageInfo image_descriptor{}; + image_descriptor.imageView = storage_image.view; + image_descriptor.imageLayout = VK_IMAGE_LAYOUT_GENERAL; + + VkDescriptorBufferInfo uniform_descriptor = create_descriptor(*ubo); + VkDescriptorBufferInfo scene_descriptor = create_descriptor(*scene_desc); + + VkWriteDescriptorSet result_image_write = vkb::initializers::write_descriptor_set(descriptor_set, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1, &image_descriptor); + VkWriteDescriptorSet uniform_buffer_write = vkb::initializers::write_descriptor_set(descriptor_set, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 2, &uniform_descriptor); + VkWriteDescriptorSet scene_buffer_write = vkb::initializers::write_descriptor_set(descriptor_set, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 3, &scene_descriptor); + + std::vector write_descriptor_sets = { + acceleration_structure_write, + result_image_write, + uniform_buffer_write, + scene_buffer_write, + }; + vkUpdateDescriptorSets(get_device().get_handle(), static_cast(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, VK_NULL_HANDLE); +} + +/* + Create our ray tracing pipeline +*/ +void RaytracingReflection::create_ray_tracing_pipeline() +{ + // Slot for binding top level acceleration structures to the ray generation shader + VkDescriptorSetLayoutBinding acceleration_structure_layout_binding{}; + acceleration_structure_layout_binding.binding = 0; + acceleration_structure_layout_binding.descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR; + acceleration_structure_layout_binding.descriptorCount = 1; + acceleration_structure_layout_binding.stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR | VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR; + + VkDescriptorSetLayoutBinding result_image_layout_binding{}; + result_image_layout_binding.binding = 1; + result_image_layout_binding.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + result_image_layout_binding.descriptorCount = 1; + result_image_layout_binding.stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR; + + VkDescriptorSetLayoutBinding uniform_buffer_binding{}; + uniform_buffer_binding.binding = 2; + uniform_buffer_binding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER; + uniform_buffer_binding.descriptorCount = 1; + uniform_buffer_binding.stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR; + + // Scene description + VkDescriptorSetLayoutBinding scene_buffer_binding{}; + scene_buffer_binding.binding = 3; + scene_buffer_binding.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; + scene_buffer_binding.descriptorCount = 1; + scene_buffer_binding.stageFlags = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR; + + std::vector bindings = { + acceleration_structure_layout_binding, + result_image_layout_binding, + uniform_buffer_binding, + scene_buffer_binding, + }; + + VkDescriptorSetLayoutCreateInfo layout_info{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO}; + layout_info.bindingCount = static_cast(bindings.size()); + layout_info.pBindings = bindings.data(); + VK_CHECK(vkCreateDescriptorSetLayout(get_device().get_handle(), &layout_info, nullptr, &descriptor_set_layout)); + + VkPipelineLayoutCreateInfo pipeline_layout_create_info{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; + pipeline_layout_create_info.setLayoutCount = 1; + pipeline_layout_create_info.pSetLayouts = &descriptor_set_layout; + + VK_CHECK(vkCreatePipelineLayout(get_device().get_handle(), &pipeline_layout_create_info, nullptr, &pipeline_layout)); + + // Ray tracing shaders + buffer reference require SPIR-V 1.5, so we need to set the appropriate target environment for the glslang compiler + vkb::GLSLCompiler::set_target_environment(glslang::EShTargetSpv, glslang::EShTargetSpv_1_5); + + /* + Setup ray tracing shader groups + Each shader group points at the corresponding shader in the pipeline + */ + std::vector shader_stages; + + // Ray generation group + { + shader_stages.push_back(load_shader("ray_tracing_reflection/raygen.rgen", VK_SHADER_STAGE_RAYGEN_BIT_KHR)); + VkRayTracingShaderGroupCreateInfoKHR raygen_group_ci{VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR}; + raygen_group_ci.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR; + raygen_group_ci.generalShader = static_cast(shader_stages.size()) - 1; + raygen_group_ci.closestHitShader = VK_SHADER_UNUSED_KHR; + raygen_group_ci.anyHitShader = VK_SHADER_UNUSED_KHR; + raygen_group_ci.intersectionShader = VK_SHADER_UNUSED_KHR; + shader_groups.push_back(raygen_group_ci); + } + + // Ray miss group + { + shader_stages.push_back(load_shader("ray_tracing_reflection/miss.rmiss", VK_SHADER_STAGE_MISS_BIT_KHR)); + VkRayTracingShaderGroupCreateInfoKHR miss_group_ci{VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR}; + miss_group_ci.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR; + miss_group_ci.generalShader = static_cast(shader_stages.size()) - 1; + miss_group_ci.closestHitShader = VK_SHADER_UNUSED_KHR; + miss_group_ci.anyHitShader = VK_SHADER_UNUSED_KHR; + miss_group_ci.intersectionShader = VK_SHADER_UNUSED_KHR; + shader_groups.push_back(miss_group_ci); + } + + // Ray miss (shadow) group + { + shader_stages.push_back(load_shader("ray_tracing_reflection/missShadow.rmiss", VK_SHADER_STAGE_MISS_BIT_KHR)); + VkRayTracingShaderGroupCreateInfoKHR miss_group_ci{VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR}; + miss_group_ci.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR; + miss_group_ci.generalShader = static_cast(shader_stages.size()) - 1; + miss_group_ci.closestHitShader = VK_SHADER_UNUSED_KHR; + miss_group_ci.anyHitShader = VK_SHADER_UNUSED_KHR; + miss_group_ci.intersectionShader = VK_SHADER_UNUSED_KHR; + shader_groups.push_back(miss_group_ci); + } + + // Ray closest hit group + { + shader_stages.push_back(load_shader("ray_tracing_reflection/closesthit.rchit", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR)); + VkRayTracingShaderGroupCreateInfoKHR closes_hit_group_ci{VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR}; + closes_hit_group_ci.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR; + closes_hit_group_ci.generalShader = VK_SHADER_UNUSED_KHR; + closes_hit_group_ci.closestHitShader = static_cast(shader_stages.size()) - 1; + closes_hit_group_ci.anyHitShader = VK_SHADER_UNUSED_KHR; + closes_hit_group_ci.intersectionShader = VK_SHADER_UNUSED_KHR; + shader_groups.push_back(closes_hit_group_ci); + } + + /* + Create the ray tracing pipeline + */ + VkRayTracingPipelineCreateInfoKHR raytracing_pipeline_create_info{VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR}; + raytracing_pipeline_create_info.stageCount = static_cast(shader_stages.size()); + raytracing_pipeline_create_info.pStages = shader_stages.data(); + raytracing_pipeline_create_info.groupCount = static_cast(shader_groups.size()); + raytracing_pipeline_create_info.pGroups = shader_groups.data(); + raytracing_pipeline_create_info.maxPipelineRayRecursionDepth = 2; + raytracing_pipeline_create_info.layout = pipeline_layout; + VK_CHECK(vkCreateRayTracingPipelinesKHR(get_device().get_handle(), VK_NULL_HANDLE, VK_NULL_HANDLE, 1, &raytracing_pipeline_create_info, nullptr, &pipeline)); +} + +/* + Deletes all resources acquired by an acceleration structure +*/ +void RaytracingReflection::delete_acceleration_structure(AccelerationStructure &acceleration_structure) +{ + if (acceleration_structure.buffer) + { + acceleration_structure.buffer.reset(); + } + + if (acceleration_structure.handle) + { + vkDestroyAccelerationStructureKHR(device->get_handle(), acceleration_structure.handle, nullptr); + } +} + +/* + Create the uniform buffer used to pass matrices to the ray tracing ray generation shader +*/ +void RaytracingReflection::create_uniform_buffer() +{ + ubo = std::make_unique(get_device(), sizeof(uniform_data), VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, VMA_MEMORY_USAGE_CPU_TO_GPU); + ubo->convert_and_update(uniform_data); + update_uniform_buffers(); +} + +/* + Command buffer generation +*/ +void RaytracingReflection::build_command_buffers() +{ + if (width != storage_image.width || height != storage_image.height) + { + // If the view port size has changed, we need to recreate the storage image + vkDestroyImageView(get_device().get_handle(), storage_image.view, nullptr); + vkDestroyImage(get_device().get_handle(), storage_image.image, nullptr); + vkFreeMemory(get_device().get_handle(), storage_image.memory, nullptr); + create_storage_image(); + + // The descriptor also needs to be updated to reference the new image + VkDescriptorImageInfo image_descriptor{}; + image_descriptor.imageView = storage_image.view; + image_descriptor.imageLayout = VK_IMAGE_LAYOUT_GENERAL; + VkWriteDescriptorSet result_image_write = vkb::initializers::write_descriptor_set(descriptor_set, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1, &image_descriptor); + vkUpdateDescriptorSets(get_device().get_handle(), 1, &result_image_write, 0, VK_NULL_HANDLE); + build_command_buffers(); + } + + VkCommandBufferBeginInfo command_buffer_begin_info{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; + + VkImageSubresourceRange subresource_range = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + + for (int32_t i = 0; i < draw_cmd_buffers.size(); ++i) + { + VK_CHECK(vkBeginCommandBuffer(draw_cmd_buffers[i], &command_buffer_begin_info)); + + /* + Setup the strided device address regions pointing at the shader identifiers in the shader binding table + */ + + const uint32_t handle_size_aligned = aligned_size(ray_tracing_pipeline_properties.shaderGroupHandleSize, ray_tracing_pipeline_properties.shaderGroupHandleAlignment); + + VkStridedDeviceAddressRegionKHR raygen_shader_sbt_entry{}; + raygen_shader_sbt_entry.deviceAddress = raygen_shader_binding_table->get_device_address(); + raygen_shader_sbt_entry.stride = handle_size_aligned; + raygen_shader_sbt_entry.size = handle_size_aligned; + + VkStridedDeviceAddressRegionKHR miss_shader_sbt_entry{}; + miss_shader_sbt_entry.deviceAddress = miss_shader_binding_table->get_device_address(); + miss_shader_sbt_entry.stride = handle_size_aligned; + miss_shader_sbt_entry.size = handle_size_aligned * 2; + + VkStridedDeviceAddressRegionKHR hit_shader_sbt_entry{}; + hit_shader_sbt_entry.deviceAddress = hit_shader_binding_table->get_device_address(); + hit_shader_sbt_entry.stride = handle_size_aligned; + hit_shader_sbt_entry.size = handle_size_aligned; + + VkStridedDeviceAddressRegionKHR callable_shader_sbt_entry{}; + + /* + Dispatch the ray tracing commands + */ + vkCmdBindPipeline(draw_cmd_buffers[i], VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, pipeline); + vkCmdBindDescriptorSets(draw_cmd_buffers[i], VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, pipeline_layout, 0, 1, &descriptor_set, 0, 0); + + vkCmdTraceRaysKHR( + draw_cmd_buffers[i], + &raygen_shader_sbt_entry, + &miss_shader_sbt_entry, + &hit_shader_sbt_entry, + &callable_shader_sbt_entry, + width, + height, + 1); + + /* + Copy ray tracing output to swap chain image + */ + + // Prepare current swap chain image as transfer destination + vkb::set_image_layout( + draw_cmd_buffers[i], + get_render_context().get_swapchain().get_images()[i], + VK_IMAGE_LAYOUT_UNDEFINED, + VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, + subresource_range); + + // Prepare ray tracing output image as transfer source + vkb::set_image_layout( + draw_cmd_buffers[i], + storage_image.image, + VK_IMAGE_LAYOUT_GENERAL, + VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, + subresource_range); + + VkImageCopy copy_region{}; + copy_region.srcSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; + copy_region.srcOffset = {0, 0, 0}; + copy_region.dstSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; + copy_region.dstOffset = {0, 0, 0}; + copy_region.extent = {width, height, 1}; + vkCmdCopyImage(draw_cmd_buffers[i], storage_image.image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, + get_render_context().get_swapchain().get_images()[i], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, ©_region); + + // Transition swap chain image back for presentation + vkb::set_image_layout(draw_cmd_buffers[i], + get_render_context().get_swapchain().get_images()[i], + VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, + VK_IMAGE_LAYOUT_PRESENT_SRC_KHR, + subresource_range); + + // Transition ray tracing output image back to general layout + vkb::set_image_layout(draw_cmd_buffers[i], + storage_image.image, + VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, + VK_IMAGE_LAYOUT_GENERAL, + subresource_range); + VK_CHECK(vkEndCommandBuffer(draw_cmd_buffers[i])); + } +} + +void RaytracingReflection::update_uniform_buffers() +{ + auto mat = camera.matrices.perspective; + mat[1][1] *= -1; // Flipping Y axis + + uniform_data.proj_inverse = glm::inverse(mat); + uniform_data.view_inverse = glm::inverse(camera.matrices.view); + ubo->convert_and_update(uniform_data); +} + +bool RaytracingReflection::prepare(vkb::Platform &platform) +{ + if (!ApiVulkanSample::prepare(platform)) + { + return false; + } + + // This sample copies the ray traced output to the swap chain image, so we need to enable the required image usage flags + std::set image_usage_flags = {VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT, VK_IMAGE_USAGE_TRANSFER_DST_BIT}; + get_render_context().update_swapchain(image_usage_flags); + + // Get the ray tracing pipeline properties, which we'll need later on in the sample + VkPhysicalDeviceProperties2 device_properties{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2}; + device_properties.pNext = &ray_tracing_pipeline_properties; + vkGetPhysicalDeviceProperties2(get_device().get_gpu().get_handle(), &device_properties); + + // Get the acceleration structure features, which we'll need later on in the sample + VkPhysicalDeviceFeatures2 device_features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; + device_features.pNext = &acceleration_structure_features; + vkGetPhysicalDeviceFeatures2(get_device().get_gpu().get_handle(), &device_features); + + camera.type = vkb::CameraType::LookAt; + camera.set_perspective(60.0f, (float) width / (float) height, 0.1f, 512.0f); + camera.set_rotation(glm::vec3(0.0f, 0.0f, 0.0f)); + camera.set_translation(glm::vec3(0.0f, 0.0f, -2.5f)); + + create_storage_image(); + create_scene(); + create_uniform_buffer(); + create_ray_tracing_pipeline(); + create_shader_binding_tables(); + create_descriptor_sets(); + build_command_buffers(); + prepared = true; + return true; +} + +void RaytracingReflection::draw() +{ + ApiVulkanSample::prepare_frame(); + submit_info.commandBufferCount = 1; + submit_info.pCommandBuffers = &draw_cmd_buffers[current_buffer]; + VK_CHECK(vkQueueSubmit(queue, 1, &submit_info, VK_NULL_HANDLE)); + ApiVulkanSample::submit_frame(); +} + +void RaytracingReflection::render(float delta_time) +{ + if (!prepared) + return; + draw(); + if (camera.updated) + update_uniform_buffers(); +} + +std::unique_ptr create_ray_tracing_reflection() +{ + return std::make_unique(); +} diff --git a/samples/extensions/ray_tracing_reflection/ray_tracing_reflection.h b/samples/extensions/ray_tracing_reflection/ray_tracing_reflection.h new file mode 100644 index 000000000..ccadcac82 --- /dev/null +++ b/samples/extensions/ray_tracing_reflection/ray_tracing_reflection.h @@ -0,0 +1,139 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION + * SPDX-License-Identifier: Apache-2.0 + */ + +/* + * More complex example for hardware accelerated ray tracing using VK_KHR_ray_tracing_pipeline and VK_KHR_acceleration_structure + */ + +#pragma once + +#include "api_vulkan_sample.h" +#include "glsl_compiler.h" + +struct ObjMaterial +{ + glm::vec3 diffuse{0.7f, 0.7f, 0.7f}; + glm::vec3 specular{0.7f, 0.7f, 0.7f}; + float shininess{0.f}; +}; + +struct ObjVertex +{ + glm::vec3 pos; + glm::vec3 nrm; +}; + +struct ObjModelCpu +{ + std::vector vertices; + std::vector indices; + std::vector mat_index; +}; + +struct ObjModelGpu +{ + uint32_t nb_indices{0}; + uint32_t nb_vertices{0}; + std::unique_ptr vertex_buffer; // Device buffer of all 'Vertex' + std::unique_ptr index_buffer; // Device buffer of the indices forming triangles + std::unique_ptr mat_color_buffer; // Device buffer of array of 'Wavefront material' + std::unique_ptr mat_index_buffer; // Device buffer of array of 'Wavefront material' +}; + +class RaytracingReflection : public ApiVulkanSample +{ + struct AccelerationStructure + { + VkAccelerationStructureKHR handle; + std::unique_ptr buffer; + }; + + public: + VkPhysicalDeviceRayTracingPipelinePropertiesKHR ray_tracing_pipeline_properties{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_PIPELINE_PROPERTIES_KHR}; + VkPhysicalDeviceAccelerationStructureFeaturesKHR acceleration_structure_features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_FEATURES_KHR}; + + std::vector bottom_level_acceleration_structure; + AccelerationStructure top_level_acceleration_structure; + + std::vector obj_models; // Array of objects and instances in the scene + std::vector shader_groups; // Shader groups + + // Shading Binding Table + std::unique_ptr raygen_shader_binding_table; + std::unique_ptr miss_shader_binding_table; + std::unique_ptr hit_shader_binding_table; + + struct StorageImage + { + VkDeviceMemory memory{nullptr}; + VkImage image{nullptr}; + VkImageView view{nullptr}; + VkFormat format{VK_FORMAT_UNDEFINED}; + uint32_t width{0}; + uint32_t height{0}; + } storage_image; + + struct UniformData + { + glm::mat4 view_inverse; + glm::mat4 proj_inverse; + } uniform_data; + std::unique_ptr ubo; + + struct ObjBuffers + { + VkDeviceAddress vertices; + VkDeviceAddress indices; + VkDeviceAddress materials; + VkDeviceAddress materialIndices; + } obj_buffers; + std::unique_ptr scene_desc; + + VkPipeline pipeline{nullptr}; + VkPipelineLayout pipeline_layout{nullptr}; + VkDescriptorSet descriptor_set{nullptr}; + VkDescriptorSetLayout descriptor_set_layout{nullptr}; + + RaytracingReflection(); + ~RaytracingReflection(); + + void create_storage_image(); + void create_bottom_level_acceleration_structure(ObjModelGpu &obj_model); + void create_top_level_acceleration_structure(std::vector &blas_instances); + void create_model(ObjModelCpu &obj, const std::vector &materials); + auto create_blas_instance(uint32_t blas_id, glm::mat4 &mat); + void delete_acceleration_structure(AccelerationStructure &acceleration_structure); + void create_scene(); + + void create_buffer_references(); + + void create_shader_binding_tables(); + void create_descriptor_sets(); + void create_ray_tracing_pipeline(); + void create_uniform_buffer(); + void update_uniform_buffers(); + void draw(); + + void build_command_buffers() override; + void request_gpu_features(vkb::PhysicalDevice &gpu) override; + bool prepare(vkb::Platform &platform) override; + void render(float delta_time) override; +}; + +std::unique_ptr create_ray_tracing_reflection(); diff --git a/shaders/ray_tracing_reflection/closesthit.rchit b/shaders/ray_tracing_reflection/closesthit.rchit new file mode 100644 index 000000000..7f1a6903b --- /dev/null +++ b/shaders/ray_tracing_reflection/closesthit.rchit @@ -0,0 +1,166 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION + * SPDX-License-Identifier: Apache-2.0 + */ + +#version 460 +#extension GL_EXT_ray_tracing : enable +#extension GL_EXT_scalar_block_layout : enable +#extension GL_EXT_nonuniform_qualifier : enable + +#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require +#extension GL_EXT_buffer_reference2 : require +struct hitPayload +{ + vec3 radiance; + vec3 attenuation; + int done; + vec3 rayOrigin; + vec3 rayDir; +}; + +layout(location = 0) rayPayloadInEXT hitPayload prd; +layout(location = 1) rayPayloadEXT bool isShadowed; + +hitAttributeEXT vec3 attribs; + +struct WaveFrontMaterial +{ + vec3 diffuse; + vec3 specular; + float shininess; +}; + +struct Vertex +{ + vec3 pos; + vec3 nrm; +}; + +struct ObjBuffers +{ + uint64_t vertices; + uint64_t indices; + uint64_t materials; + uint64_t materialIndices; +}; + +// clang-format off +layout(buffer_reference, scalar) buffer Vertices {Vertex v[]; }; // Positions of an object +layout(buffer_reference, scalar) buffer Indices {uvec3 i[]; }; // Triangle indices +layout(buffer_reference, scalar) buffer Materials {WaveFrontMaterial m[]; }; // Array of all materials on an object +layout(buffer_reference, scalar) buffer MatIndices {int i[]; }; // Material ID for each triangle + +layout(set = 0, binding = 0) uniform accelerationStructureEXT topLevelAS; +layout(set = 0, binding = 3) buffer _scene_desc { ObjBuffers i[]; } scene_desc; +// clang-format on + +vec3 computeSpecular(WaveFrontMaterial mat, vec3 V, vec3 L, vec3 N) +{ + const float kPi = 3.14159265; + const float kShininess = max(mat.shininess, 4.0); + + // Specular + const float kEnergyConservation = (2.0 + kShininess) / (2.0 * kPi); + V = normalize(-V); + vec3 R = reflect(-L, N); + float specular = kEnergyConservation * pow(max(dot(V, R), 0.0), kShininess); + + return vec3(mat.specular * specular); +} + +void main() +{ + // When contructing the TLAS, we stored the model id in InstanceCustomIndexEXT, so the + // the instance can quickly have access to the data + + // Object data + ObjBuffers objResource = scene_desc.i[gl_InstanceCustomIndexEXT]; + MatIndices matIndices = MatIndices(objResource.materialIndices); + Materials materials = Materials(objResource.materials); + Indices indices = Indices(objResource.indices); + Vertices vertices = Vertices(objResource.vertices); + + // Retrieve the material used on this triangle 'PrimitiveID' + int mat_idx = matIndices.i[gl_PrimitiveID]; + WaveFrontMaterial mat = materials.m[mat_idx]; // Material for this triangle + + // Indices of the triangle + uvec3 ind = indices.i[gl_PrimitiveID]; + + // Vertex of the triangle + Vertex v0 = vertices.v[ind.x]; + Vertex v1 = vertices.v[ind.y]; + Vertex v2 = vertices.v[ind.z]; + + // Barycentric coordinates of the triangle + const vec3 barycentrics = vec3(1.0f - attribs.x - attribs.y, attribs.x, attribs.y); + + // Computing the normal at hit position + vec3 N = v0.nrm.xyz * barycentrics.x + v1.nrm.xyz * barycentrics.y + v2.nrm.xyz * barycentrics.z; + N = normalize(vec3(N.xyz * gl_WorldToObjectEXT)); // Transforming the normal to world space + + // Computing the coordinates of the hit position + vec3 P = v0.pos.xyz * barycentrics.x + v1.pos.xyz * barycentrics.y + v2.pos.xyz * barycentrics.z; + P = vec3(gl_ObjectToWorldEXT * vec4(P, 1.0)); // Transforming the position to world space + + // Hardocded (to) light direction + vec3 L = normalize(vec3(1, 1, 1)); + + float NdotL = dot(N, L); + + // Fake Lambertian to avoid black + vec3 diffuse = mat.diffuse * max(NdotL, 0.3); + vec3 specular = vec3(0); + + // Tracing shadow ray only if the light is visible from the surface + if (NdotL > 0) + { + float tMin = 0.001; + float tMax = 1e32; // infinite + vec3 origin = P; + vec3 rayDir = L; + uint flags = gl_RayFlagsTerminateOnFirstHitEXT | gl_RayFlagsOpaqueEXT | gl_RayFlagsSkipClosestHitShaderEXT; + isShadowed = true; + + traceRayEXT(topLevelAS, // acceleration structure + flags, // rayFlags + 0xFF, // cullMask + 0, // sbtRecordOffset + 0, // sbtRecordStride + 1, // missIndex + origin, // ray origin + tMin, // ray min range + rayDir, // ray direction + tMax, // ray max range + 1 // payload (location = 1) + ); + + if (isShadowed) + diffuse *= 0.3; + else + // Add specular only if not in shadow + specular = computeSpecular(mat, gl_WorldRayDirectionEXT, L, N); + } + prd.radiance = (diffuse + specular) * (1 - mat.shininess) * prd.attenuation; + + // Reflect + vec3 rayDir = reflect(gl_WorldRayDirectionEXT, N); + prd.attenuation *= vec3(mat.shininess); + prd.rayOrigin = P; + prd.rayDir = rayDir; +} diff --git a/shaders/ray_tracing_reflection/miss.rmiss b/shaders/ray_tracing_reflection/miss.rmiss new file mode 100644 index 000000000..12848e819 --- /dev/null +++ b/shaders/ray_tracing_reflection/miss.rmiss @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION + * SPDX-License-Identifier: Apache-2.0 + */ + +#version 460 +#extension GL_EXT_ray_tracing : enable + +struct hitPayload +{ + vec3 radiance; + vec3 attenuation; + int done; + vec3 rayOrigin; + vec3 rayDir; +}; + +layout(location = 0) rayPayloadInEXT hitPayload prd; + +void main() +{ + prd.radiance = vec3(0.3) * prd.attenuation; + prd.done = 1; +} diff --git a/shaders/ray_tracing_reflection/missShadow.rmiss b/shaders/ray_tracing_reflection/missShadow.rmiss new file mode 100644 index 000000000..8825d317f --- /dev/null +++ b/shaders/ray_tracing_reflection/missShadow.rmiss @@ -0,0 +1,28 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION + * SPDX-License-Identifier: Apache-2.0 + */ + +#version 460 +#extension GL_EXT_ray_tracing : require + +layout(location = 1) rayPayloadInEXT bool isShadowed; + +void main() +{ + isShadowed = false; +} diff --git a/shaders/ray_tracing_reflection/raygen.rgen b/shaders/ray_tracing_reflection/raygen.rgen new file mode 100644 index 000000000..b3368fe41 --- /dev/null +++ b/shaders/ray_tracing_reflection/raygen.rgen @@ -0,0 +1,73 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-FileCopyrightText: Copyright (c) 2014-2021 NVIDIA CORPORATION + * SPDX-License-Identifier: Apache-2.0 + */ + +#version 460 +#extension GL_EXT_ray_tracing : enable + +layout(set = 0, binding = 0) uniform accelerationStructureEXT topLevelAS; +layout(set = 0, binding = 1, rgba8) uniform image2D image; +layout(set = 0, binding = 2) uniform CameraProperties +{ + mat4 viewInverse; + mat4 projInverse; +} +cam; + +struct hitPayload +{ + vec3 radiance; + vec3 attenuation; + int done; + vec3 rayOrigin; + vec3 rayDir; +}; + +layout(location = 0) rayPayloadEXT hitPayload prd; + +void main() +{ + const vec2 pixelCenter = vec2(gl_LaunchIDEXT.xy) + vec2(0.5); + const vec2 inUV = pixelCenter / vec2(gl_LaunchSizeEXT.xy); + vec2 d = inUV * 2.0 - 1.0; + + vec4 origin = cam.viewInverse * vec4(0, 0, 0, 1); + vec4 target = cam.projInverse * vec4(d.x, d.y, 1, 1); + vec4 direction = cam.viewInverse * vec4(normalize(target.xyz), 0); + + float tmin = 0.001; + float tmax = 1e32; + + prd.rayOrigin = origin.xyz; + prd.rayDir = direction.xyz; + prd.radiance = vec3(0.0); + prd.attenuation = vec3(1.0); + prd.done = 0; + + vec3 hitValue = vec3(0); + + for (int depth = 0; depth < 64; depth++) + { + traceRayEXT(topLevelAS, gl_RayFlagsOpaqueEXT, 0xff, 0, 0, 0, prd.rayOrigin, tmin, prd.rayDir, tmax, 0); + hitValue += prd.radiance; + if (prd.done == 1 || length(prd.attenuation) < 0.1) + break; + } + + imageStore(image, ivec2(gl_LaunchIDEXT.xy), vec4(hitValue, 0.0)); +}