1 // Copyright 2024 The Amber Authors.
2 // Copyright (C) 2024 Advanced Micro Devices, Inc. All rights reserved.
3 //
4 // Licensed under the Apache License, Version 2.0 (the "License");
5 // you may not use this file except in compliance with the License.
6 // You may obtain a copy of the License at
7 //
8 // http://www.apache.org/licenses/LICENSE-2.0
9 //
10 // Unless required by applicable law or agreed to in writing, software
11 // distributed under the License is distributed on an "AS IS" BASIS,
12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 // See the License for the specific language governing permissions and
14 // limitations under the License.
15
16 #include "src/vulkan/blas.h"
17
18 #include <cstring>
19
20 #include "src/vulkan/command_buffer.h"
21
22 namespace amber {
23 namespace vulkan {
24
align(VkDeviceSize v,VkDeviceSize a)25 inline VkDeviceSize align(VkDeviceSize v, VkDeviceSize a) {
26 return (v + a - 1) & ~(a - 1);
27 }
28
BLAS(Device * device)29 BLAS::BLAS(Device* device) : device_(device) {}
30
~BLAS()31 BLAS::~BLAS() {
32 if (blas_ != VK_NULL_HANDLE) {
33 device_->GetPtrs()->vkDestroyAccelerationStructureKHR(
34 device_->GetVkDevice(), blas_, nullptr);
35 }
36 }
37
CreateBLAS(amber::BLAS * blas)38 Result BLAS::CreateBLAS(amber::BLAS* blas) {
39 if (blas_ != VK_NULL_HANDLE)
40 return Result("Cannot recreate acceleration structure");
41
42 std::vector<std::unique_ptr<Geometry>>& geometries = blas->GetGeometries();
43 std::vector<VkDeviceSize> vertexBufferOffsets;
44 VkDeviceSize vertexBufferSize = 0;
45
46 VkDeviceOrHostAddressConstKHR const_null_placeholder = {};
47 VkDeviceOrHostAddressKHR null_placeholder = {};
48
49 accelerationStructureGeometriesKHR_.resize(geometries.size());
50 accelerationStructureBuildRangeInfoKHR_.resize(geometries.size());
51 maxPrimitiveCounts_.resize(geometries.size());
52 vertexBufferOffsets.resize(geometries.size());
53
54 for (size_t geometryNdx = 0; geometryNdx < geometries.size(); ++geometryNdx) {
55 const std::unique_ptr<Geometry>& geometryData = geometries[geometryNdx];
56 VkDeviceOrHostAddressConstKHR vertexData = {};
57 VkAccelerationStructureGeometryDataKHR geometry;
58 VkGeometryTypeKHR geometryType = VK_GEOMETRY_TYPE_MAX_ENUM_KHR;
59
60 if (geometryData->IsTriangle()) {
61 VkAccelerationStructureGeometryTrianglesDataKHR
62 accelerationStructureGeometryTrianglesDataKHR = {
63 // NOLINTNEXTLINE(whitespace/line_length)
64 VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR,
65 nullptr,
66 VK_FORMAT_R32G32B32_SFLOAT,
67 vertexData,
68 3 * sizeof(float),
69 static_cast<uint32_t>(geometryData->getVertexCount()),
70 VK_INDEX_TYPE_NONE_KHR,
71 const_null_placeholder,
72 const_null_placeholder,
73 };
74
75 geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;
76 geometry.triangles = accelerationStructureGeometryTrianglesDataKHR;
77 } else if (geometryData->IsAABB()) {
78 const VkAccelerationStructureGeometryAabbsDataKHR
79 accelerationStructureGeometryAabbsDataKHR = {
80 VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR,
81 nullptr, vertexData, sizeof(VkAabbPositionsKHR)};
82
83 geometryType = VK_GEOMETRY_TYPE_AABBS_KHR;
84 geometry.aabbs = accelerationStructureGeometryAabbsDataKHR;
85 } else {
86 assert(false && "unknown geometry type");
87 }
88
89 const VkAccelerationStructureGeometryKHR accelerationStructureGeometry = {
90 VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
91 nullptr,
92 geometryType,
93 geometry,
94 VkGeometryFlagsKHR(geometryData->GetFlags())
95 };
96 const VkAccelerationStructureBuildRangeInfoKHR
97 accelerationStructureBuildRangeInfosKHR = {
98 static_cast<uint32_t>(geometryData->getPrimitiveCount()), 0, 0, 0};
99
100 accelerationStructureGeometriesKHR_[geometryNdx] =
101 accelerationStructureGeometry;
102 accelerationStructureBuildRangeInfoKHR_[geometryNdx] =
103 accelerationStructureBuildRangeInfosKHR;
104 maxPrimitiveCounts_[geometryNdx] =
105 accelerationStructureBuildRangeInfosKHR.primitiveCount;
106 vertexBufferOffsets[geometryNdx] = vertexBufferSize;
107 size_t s1 = sizeof(geometryData->GetData()[0]);
108 vertexBufferSize += align(geometryData->GetData().size() * s1, 8);
109 }
110
111 const VkAccelerationStructureGeometryKHR*
112 accelerationStructureGeometriesKHRPointer =
113 accelerationStructureGeometriesKHR_.data();
114 accelerationStructureBuildGeometryInfoKHR_ = {
115 VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
116 nullptr,
117 VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
118 0u,
119 VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR,
120 VK_NULL_HANDLE,
121 VK_NULL_HANDLE,
122 static_cast<uint32_t>(accelerationStructureGeometriesKHR_.size()),
123 accelerationStructureGeometriesKHRPointer,
124 nullptr,
125 null_placeholder,
126 };
127 VkAccelerationStructureBuildSizesInfoKHR sizeInfo = {
128 VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR, nullptr, 0,
129 0, 0};
130
131 device_->GetPtrs()->vkGetAccelerationStructureBuildSizesKHR(
132 device_->GetVkDevice(), VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
133 &accelerationStructureBuildGeometryInfoKHR_, maxPrimitiveCounts_.data(),
134 &sizeInfo);
135
136 const uint32_t accelerationStructureSize =
137 static_cast<uint32_t>(sizeInfo.accelerationStructureSize);
138
139 buffer_ =
140 MakeUnique<TransferBuffer>(device_, accelerationStructureSize, nullptr);
141 buffer_->AddUsageFlags(
142 VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR |
143 VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT);
144 buffer_->AddAllocateFlags(VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT);
145 buffer_->Initialize();
146
147 const VkAccelerationStructureCreateInfoKHR accelerationStructureCreateInfoKHR{
148 VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR,
149 nullptr,
150 0,
151 buffer_->GetVkBuffer(),
152 0,
153 accelerationStructureSize,
154 VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
155 0};
156
157 if (device_->GetPtrs()->vkCreateAccelerationStructureKHR(
158 device_->GetVkDevice(), &accelerationStructureCreateInfoKHR, nullptr,
159 &blas_) != VK_SUCCESS)
160 return Result("Vulkan::Calling vkCreateAccelerationStructureKHR failed");
161
162 accelerationStructureBuildGeometryInfoKHR_.dstAccelerationStructure = blas_;
163
164 if (sizeInfo.buildScratchSize > 0) {
165 scratch_buffer_ = MakeUnique<TransferBuffer>(
166 device_, static_cast<uint32_t>(sizeInfo.buildScratchSize), nullptr);
167 scratch_buffer_->AddUsageFlags(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
168 VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT);
169 scratch_buffer_->AddAllocateFlags(VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT);
170 scratch_buffer_->Initialize();
171
172 accelerationStructureBuildGeometryInfoKHR_.scratchData.deviceAddress =
173 scratch_buffer_->getBufferDeviceAddress();
174 }
175
176 if (vertexBufferSize > 0) {
177 vertex_buffer_ = MakeUnique<TransferBuffer>(
178 device_, static_cast<uint32_t>(vertexBufferSize), nullptr);
179 vertex_buffer_->AddUsageFlags(
180 VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR |
181 VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT);
182 vertex_buffer_->AddAllocateFlags(VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT);
183 vertex_buffer_->SetMemoryPropertiesFlags(
184 VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
185 VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
186 vertex_buffer_->Initialize();
187
188 void* memory_ptr = vertex_buffer_->HostAccessibleMemoryPtr();
189 assert(memory_ptr != nullptr);
190
191 for (size_t geometryNdx = 0; geometryNdx < geometries.size();
192 ++geometryNdx) {
193 VkDeviceOrHostAddressConstKHR p;
194
195 p.deviceAddress = vertex_buffer_.get()->getBufferDeviceAddress() +
196 vertexBufferOffsets[geometryNdx];
197
198 const auto& data = geometries[geometryNdx]->GetData();
199 std::memcpy(reinterpret_cast<char*>(memory_ptr) +
200 vertexBufferOffsets[geometryNdx],
201 data.data(), data.size() * sizeof(*data.data()));
202
203 if (geometries[geometryNdx]->IsTriangle()) {
204 accelerationStructureGeometriesKHR_[geometryNdx]
205 .geometry.triangles.vertexData = p;
206 } else if (geometries[geometryNdx]->IsAABB()) {
207 accelerationStructureGeometriesKHR_[geometryNdx].geometry.aabbs.data =
208 p;
209 } else {
210 assert(false && "unknown geometry type");
211 }
212 accelerationStructureGeometriesKHR_[geometryNdx].flags =
213 VkGeometryFlagsKHR(geometries[geometryNdx]->GetFlags());
214 }
215 }
216
217 return {};
218 }
219
BuildBLAS(CommandBuffer * command_buffer)220 Result BLAS::BuildBLAS(CommandBuffer* command_buffer) {
221 if (blas_ == VK_NULL_HANDLE)
222 return Result("Acceleration structure should be created first");
223 if (built_)
224 return {};
225
226 VkCommandBuffer cmdBuffer = command_buffer->GetVkCommandBuffer();
227
228 vertex_buffer_->CopyToDevice(command_buffer);
229
230 VkAccelerationStructureBuildRangeInfoKHR*
231 accelerationStructureBuildRangeInfoKHRPtr =
232 accelerationStructureBuildRangeInfoKHR_.data();
233
234 device_->GetPtrs()->vkCmdBuildAccelerationStructuresKHR(
235 cmdBuffer, 1, &accelerationStructureBuildGeometryInfoKHR_,
236 &accelerationStructureBuildRangeInfoKHRPtr);
237
238 const VkAccessFlags accessMasks =
239 VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR |
240 VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR;
241 const VkMemoryBarrier memBarrier{
242 VK_STRUCTURE_TYPE_MEMORY_BARRIER,
243 nullptr,
244 accessMasks,
245 accessMasks,
246 };
247
248 device_->GetPtrs()->vkCmdPipelineBarrier(
249 cmdBuffer, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
250 VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 1, &memBarrier, 0, nullptr, 0,
251 nullptr);
252
253 built_ = true;
254
255 return {};
256 }
257
getVkBLASDeviceAddress()258 VkDeviceAddress BLAS::getVkBLASDeviceAddress() {
259 VkAccelerationStructureDeviceAddressInfoKHR info = {
260 VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR, nullptr,
261 blas_};
262
263 assert(blas_ != VK_NULL_HANDLE);
264
265 return device_->GetPtrs()->vkGetAccelerationStructureDeviceAddressKHR(
266 device_->GetVkDevice(), &info);
267 }
268
269 } // namespace vulkan
270 } // namespace amber
271