Skip to content

Commit 28a788f

Browse files
JoseEmilio-ARMJoseEmilio-ARM
committed
Create Acceleration Structures
1 parent 18c3a0e commit 28a788f

File tree

2 files changed

+263
-9
lines changed

2 files changed

+263
-9
lines changed

‎attachments/38_ray_tracing.cpp‎

Lines changed: 260 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,20 @@ class HelloTriangleApplication {
143143
vk::raii::Buffer indexBuffer = nullptr;
144144
vk::raii::DeviceMemory indexBufferMemory = nullptr;
145145

146+
std::vector<vk::raii::Buffer> blasBuffers;
147+
std::vector<vk::raii::DeviceMemory> blasMemories;
148+
std::vector<vk::raii::AccelerationStructureKHR> blasHandles;
149+
150+
std::vector<vk::AccelerationStructureInstanceKHR> instances;
151+
vk::raii::Buffer instanceBuffer = nullptr;
152+
vk::raii::DeviceMemory instanceMemory = nullptr;
153+
154+
vk::raii::Buffer tlasBuffer = nullptr;
155+
vk::raii::DeviceMemory tlasMemory = nullptr;
156+
vk::raii::Buffer tlasScratchBuffer = nullptr;
157+
vk::raii::DeviceMemory tlasScratchMemory = nullptr;
158+
vk::raii::AccelerationStructureKHR tlas = nullptr;
159+
146160
std::vector<vk::raii::Buffer> uniformBuffers;
147161
std::vector<vk::raii::DeviceMemory> uniformBuffersMemory;
148162
std::vector<void*> uniformBuffersMapped;
@@ -178,7 +192,10 @@ class HelloTriangleApplication {
178192
vk::KHRSwapchainExtensionName,
179193
vk::KHRSpirv14ExtensionName,
180194
vk::KHRSynchronization2ExtensionName,
181-
vk::KHRCreateRenderpass2ExtensionName
195+
vk::KHRCreateRenderpass2ExtensionName,
196+
vk::KHRAccelerationStructureExtensionName,
197+
vk::KHRBufferDeviceAddressExtensionName,
198+
vk::KHRDeferredHostOperationsExtensionName
182199
};
183200

184201
void initWindow() {
@@ -212,6 +229,7 @@ class HelloTriangleApplication {
212229
createTextureSampler();
213230
createVertexBuffer();
214231
createIndexBuffer();
232+
createAccelerationStructures();
215233
createUniformBuffers();
216234
createDescriptorPool();
217235
createDescriptorSets();
@@ -353,15 +371,18 @@ class HelloTriangleApplication {
353371
auto features = device.template getFeatures2<vk::PhysicalDeviceFeatures2,
354372
vk::PhysicalDeviceVulkan12Features,
355373
vk::PhysicalDeviceVulkan13Features,
356-
vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT>();
374+
vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT,
375+
vk::PhysicalDeviceAccelerationStructureFeaturesKHR>();
357376
bool supportsRequiredFeatures = features.template get<vk::PhysicalDeviceFeatures2>().features.samplerAnisotropy &&
358377
features.template get<vk::PhysicalDeviceVulkan13Features>().dynamicRendering &&
359378
features.template get<vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT>().extendedDynamicState &&
360379
features.template get<vk::PhysicalDeviceVulkan12Features>().descriptorBindingSampledImageUpdateAfterBind &&
361380
features.template get<vk::PhysicalDeviceVulkan12Features>().descriptorBindingPartiallyBound &&
362381
features.template get<vk::PhysicalDeviceVulkan12Features>().descriptorBindingVariableDescriptorCount &&
363382
features.template get<vk::PhysicalDeviceVulkan12Features>().runtimeDescriptorArray &&
364-
features.template get<vk::PhysicalDeviceVulkan12Features>().shaderSampledImageArrayNonUniformIndexing;
383+
features.template get<vk::PhysicalDeviceVulkan12Features>().shaderSampledImageArrayNonUniformIndexing &&
384+
features.template get<vk::PhysicalDeviceVulkan12Features>().bufferDeviceAddress &&
385+
features.template get<vk::PhysicalDeviceAccelerationStructureFeaturesKHR>().accelerationStructure;
365386

366387
return supportsVulkan1_3 && supportsGraphics && supportsAllRequiredExtensions && supportsRequiredFeatures;
367388
} );
@@ -425,13 +446,15 @@ class HelloTriangleApplication {
425446

426447
// query for Vulkan 1.3 features
427448
vk::StructureChain<vk::PhysicalDeviceFeatures2, vk::PhysicalDeviceVulkan12Features,
428-
vk::PhysicalDeviceVulkan13Features, vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT> featureChain = {
449+
vk::PhysicalDeviceVulkan13Features, vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT,
450+
vk::PhysicalDeviceAccelerationStructureFeaturesKHR> featureChain = {
429451
{.features = {.samplerAnisotropy = true } }, // vk::PhysicalDeviceFeatures2
430452
{.shaderSampledImageArrayNonUniformIndexing = true, .descriptorBindingSampledImageUpdateAfterBind = true,
431453
.descriptorBindingPartiallyBound = true, .descriptorBindingVariableDescriptorCount = true,
432-
.runtimeDescriptorArray = true }, // vk::PhysicalDeviceVulkan12Features
454+
.runtimeDescriptorArray = true, .bufferDeviceAddress = true }, // vk::PhysicalDeviceVulkan12Features
433455
{.synchronization2 = true, .dynamicRendering = true }, // vk::PhysicalDeviceVulkan13Features
434-
{.extendedDynamicState = true } // vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT
456+
{.extendedDynamicState = true }, // vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT
457+
{.accelerationStructure = true }, // vk::PhysicalDeviceAccelerationStructureFeaturesKHR
435458
};
436459

437460
// create a Device
@@ -484,6 +507,7 @@ class HelloTriangleApplication {
484507
// Use descriptor set 0 for global data
485508
std::array global_bindings = {
486509
vk::DescriptorSetLayoutBinding( 0, vk::DescriptorType::eUniformBuffer, 1, vk::ShaderStageFlagBits::eVertex | vk::ShaderStageFlagBits::eFragment, nullptr),
510+
vk::DescriptorSetLayoutBinding( 1, vk::DescriptorType::eAccelerationStructureKHR, 1, vk::ShaderStageFlagBits::eFragment, nullptr),
487511
};
488512

489513
vk::DescriptorSetLayoutCreateInfo globalLayoutInfo{ .bindingCount = static_cast<uint32_t>(global_bindings.size()), .pBindings = global_bindings.data() };
@@ -904,7 +928,8 @@ class HelloTriangleApplication {
904928
memcpy(dataStaging, vertices.data(), bufferSize);
905929
stagingBufferMemory.unmapMemory();
906930

907-
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eVertexBuffer, vk::MemoryPropertyFlagBits::eDeviceLocal, vertexBuffer, vertexBufferMemory);
931+
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eVertexBuffer | vk::BufferUsageFlagBits::eShaderDeviceAddress |
932+
vk::BufferUsageFlagBits::eAccelerationStructureBuildInputReadOnlyKHR, vk::MemoryPropertyFlagBits::eDeviceLocal, vertexBuffer, vertexBufferMemory);
908933

909934
copyBuffer(stagingBuffer, vertexBuffer, bufferSize);
910935
}
@@ -920,7 +945,8 @@ class HelloTriangleApplication {
920945
memcpy(data, indices.data(), bufferSize);
921946
stagingBufferMemory.unmapMemory();
922947

923-
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eIndexBuffer, vk::MemoryPropertyFlagBits::eDeviceLocal, indexBuffer, indexBufferMemory);
948+
createBuffer(bufferSize, vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eIndexBuffer | vk::BufferUsageFlagBits::eShaderDeviceAddress |
949+
vk::BufferUsageFlagBits::eAccelerationStructureBuildInputReadOnlyKHR | vk::BufferUsageFlagBits::eStorageBuffer, vk::MemoryPropertyFlagBits::eDeviceLocal, indexBuffer, indexBufferMemory);
924950

925951
copyBuffer(stagingBuffer, indexBuffer, bufferSize);
926952
}
@@ -941,9 +967,214 @@ class HelloTriangleApplication {
941967
}
942968
}
943969

970+
void createAccelerationStructures() {
971+
vk::BufferDeviceAddressInfo vai{ .buffer = *vertexBuffer };
972+
vk::DeviceAddress vertexAddr = device.getBufferAddressKHR(vai);
973+
vk::BufferDeviceAddressInfo iai{ .buffer = *indexBuffer };
974+
vk::DeviceAddress indexAddr = device.getBufferAddressKHR(iai);
975+
976+
instances.reserve(submeshes.size());
977+
blasBuffers.reserve(submeshes.size());
978+
blasMemories.reserve(submeshes.size());
979+
blasHandles.reserve(submeshes.size());
980+
981+
vk::TransformMatrixKHR tm{};
982+
tm.matrix = std::array<std::array<float,4>,3>{{
983+
std::array<float,4>{1.f, 0.f, 0.f, 0.f},
984+
std::array<float,4>{0.f, 1.f, 0.f, 0.f},
985+
std::array<float,4>{0.f, 0.f, 1.f, 0.f}
986+
}};
987+
988+
// Build a bottom level acceleration structure for each submesh
989+
for (size_t i = 0; i < submeshes.size(); ++i) {
990+
const auto& submesh = submeshes[i];
991+
992+
// Prepare geometry data
993+
auto trianglesData = vk::AccelerationStructureGeometryTrianglesDataKHR{
994+
.vertexFormat = vk::Format::eR32G32B32Sfloat,
995+
.vertexData = vertexAddr,
996+
.vertexStride = sizeof(Vertex),
997+
.maxVertex = submesh.maxVertex,
998+
.indexType = vk::IndexType::eUint32,
999+
.indexData = indexAddr + submesh.indexOffset * sizeof(uint32_t)
1000+
};
1001+
1002+
vk::AccelerationStructureGeometryDataKHR geomData(trianglesData);
1003+
vk::AccelerationStructureGeometryKHR blasGeometry{
1004+
.geometryType = vk::GeometryTypeKHR::eTriangles,
1005+
.geometry = geomData
1006+
};
1007+
1008+
if (!submesh.alphaCut)
1009+
{
1010+
blasGeometry.flags = vk::GeometryFlagBitsKHR::eOpaque;
1011+
}
1012+
1013+
vk::AccelerationStructureBuildRangeInfoKHR blasRangeInfo{
1014+
.primitiveCount = static_cast<uint32_t>(submesh.indexCount / 3),
1015+
.primitiveOffset = 0,
1016+
.firstVertex = submesh.firstVertex,
1017+
.transformOffset = 0
1018+
};
1019+
1020+
// Prepare BLAS handle and buffer
1021+
vk::AccelerationStructureBuildGeometryInfoKHR blasBuildInfo{
1022+
.type = vk::AccelerationStructureTypeKHR::eBottomLevel,
1023+
.mode = vk::BuildAccelerationStructureModeKHR::eBuild,
1024+
.geometryCount = 1,
1025+
.pGeometries = &blasGeometry,
1026+
};
1027+
1028+
vk::AccelerationStructureBuildSizesInfoKHR blasSizeInfo =
1029+
device.getAccelerationStructureBuildSizesKHR(
1030+
vk::AccelerationStructureBuildTypeKHR::eDevice,
1031+
blasBuildInfo,
1032+
{ blasRangeInfo.primitiveCount }
1033+
);
1034+
1035+
vk::raii::Buffer blasBuffer = nullptr;
1036+
vk::raii::DeviceMemory blasMemory = nullptr;
1037+
blasBuffers.emplace_back(std::move(blasBuffer));
1038+
blasMemories.emplace_back(std::move(blasMemory));
1039+
createBuffer(blasSizeInfo.accelerationStructureSize,
1040+
vk::BufferUsageFlagBits::eAccelerationStructureStorageKHR |
1041+
vk::BufferUsageFlagBits::eShaderDeviceAddress |
1042+
vk::BufferUsageFlagBits::eAccelerationStructureBuildInputReadOnlyKHR,
1043+
vk::MemoryPropertyFlagBits::eDeviceLocal,
1044+
blasBuffers[i], blasMemories[i]);
1045+
1046+
vk::AccelerationStructureCreateInfoKHR blasCreateInfo{
1047+
.buffer = blasBuffers[i],
1048+
.offset = 0,
1049+
.size = blasSizeInfo.accelerationStructureSize,
1050+
.type = vk::AccelerationStructureTypeKHR::eBottomLevel,
1051+
};
1052+
1053+
blasHandles.emplace_back(device.createAccelerationStructureKHR(blasCreateInfo));
1054+
1055+
// Allocate a scratch buffer for BLAS and build it
1056+
vk::raii::Buffer scratchBuffer = nullptr;
1057+
vk::raii::DeviceMemory scratchMemory = nullptr;
1058+
createBuffer(blasSizeInfo.buildScratchSize,
1059+
vk::BufferUsageFlagBits::eStorageBuffer |
1060+
vk::BufferUsageFlagBits::eShaderDeviceAddress,
1061+
vk::MemoryPropertyFlagBits::eDeviceLocal,
1062+
scratchBuffer, scratchMemory);
1063+
1064+
vk::BufferDeviceAddressInfo scratchAddressInfo{ .buffer = *scratchBuffer };
1065+
vk::DeviceAddress scratchAddr = device.getBufferAddressKHR(scratchAddressInfo);
1066+
1067+
auto cmd = beginSingleTimeCommands();
1068+
blasBuildInfo.dstAccelerationStructure = blasHandles[i];
1069+
blasBuildInfo.scratchData.deviceAddress = scratchAddr;
1070+
cmd->buildAccelerationStructuresKHR({ blasBuildInfo }, { &blasRangeInfo });
1071+
endSingleTimeCommands(*cmd);
1072+
1073+
// Create an instance for the TLAS
1074+
vk::AccelerationStructureDeviceAddressInfoKHR addrInfo{
1075+
.accelerationStructure = *blasHandles[i]
1076+
};
1077+
vk::DeviceAddress blasDeviceAddr = device.getAccelerationStructureAddressKHR(addrInfo);
1078+
1079+
vk::AccelerationStructureInstanceKHR instance{};
1080+
instance.setTransform(tm)
1081+
.setMask(0xFF)
1082+
.setAccelerationStructureReference(blasDeviceAddr)
1083+
.setFlags(vk::GeometryInstanceFlagBitsKHR::eTriangleFacingCullDisable);
1084+
1085+
instances.push_back(instance);
1086+
}
1087+
1088+
// Prepare instance data for the TLAS
1089+
vk::DeviceSize instBufferSize = sizeof(instances[0]) * instances.size();
1090+
createBuffer(instBufferSize,
1091+
vk::BufferUsageFlagBits::eShaderDeviceAddress |
1092+
vk::BufferUsageFlagBits::eTransferDst |
1093+
vk::BufferUsageFlagBits::eAccelerationStructureBuildInputReadOnlyKHR,
1094+
vk::MemoryPropertyFlagBits::eHostVisible |
1095+
vk::MemoryPropertyFlagBits::eHostCoherent,
1096+
instanceBuffer, instanceMemory);
1097+
1098+
void *ptr = instanceMemory.mapMemory(0, instBufferSize);
1099+
memcpy(ptr, instances.data(), instBufferSize);
1100+
instanceMemory.unmapMemory();
1101+
1102+
vk::BufferDeviceAddressInfo instanceAddrInfo{ .buffer = instanceBuffer };
1103+
vk::DeviceAddress instanceAddr = device.getBufferAddressKHR(instanceAddrInfo);
1104+
1105+
vk::AccelerationStructureGeometryKHR tlasGeometry{
1106+
.geometryType = vk::GeometryTypeKHR::eInstances,
1107+
.geometry = vk::AccelerationStructureGeometryDataKHR{
1108+
vk::AccelerationStructureGeometryInstancesDataKHR{
1109+
.arrayOfPointers = vk::False,
1110+
.data = instanceAddr
1111+
}
1112+
}
1113+
};
1114+
1115+
vk::AccelerationStructureBuildRangeInfoKHR tlasRangeInfo{
1116+
.primitiveCount = static_cast<uint32_t>(instances.size()),
1117+
.primitiveOffset = 0,
1118+
.firstVertex = 0,
1119+
.transformOffset = 0
1120+
};
1121+
1122+
// Prepare TLAS handle and buffer
1123+
vk::AccelerationStructureBuildGeometryInfoKHR tlasBuildInfo{
1124+
.type = vk::AccelerationStructureTypeKHR::eTopLevel,
1125+
.flags = vk::BuildAccelerationStructureFlagBitsKHR::eAllowUpdate,
1126+
.mode = vk::BuildAccelerationStructureModeKHR::eBuild,
1127+
.geometryCount = 1,
1128+
.pGeometries = &tlasGeometry
1129+
};
1130+
1131+
auto tlasSizeInfo = device.getAccelerationStructureBuildSizesKHR(
1132+
vk::AccelerationStructureBuildTypeKHR::eDevice,
1133+
tlasBuildInfo,
1134+
{ tlasRangeInfo.primitiveCount }
1135+
);
1136+
1137+
createBuffer(
1138+
tlasSizeInfo.accelerationStructureSize,
1139+
vk::BufferUsageFlagBits::eAccelerationStructureStorageKHR |
1140+
vk::BufferUsageFlagBits::eShaderDeviceAddress |
1141+
vk::BufferUsageFlagBits::eAccelerationStructureBuildInputReadOnlyKHR,
1142+
vk::MemoryPropertyFlagBits::eDeviceLocal,
1143+
tlasBuffer, tlasMemory
1144+
);
1145+
1146+
vk::AccelerationStructureCreateInfoKHR tlasCreateInfo{
1147+
.buffer = tlasBuffer,
1148+
.offset = 0,
1149+
.size = tlasSizeInfo.accelerationStructureSize,
1150+
.type = vk::AccelerationStructureTypeKHR::eTopLevel
1151+
};
1152+
1153+
tlas = device.createAccelerationStructureKHR(tlasCreateInfo);
1154+
1155+
// Allocate a scratch buffer for TLAS and build it
1156+
createBuffer(
1157+
tlasSizeInfo.buildScratchSize,
1158+
vk::BufferUsageFlagBits::eStorageBuffer |
1159+
vk::BufferUsageFlagBits::eShaderDeviceAddress,
1160+
vk::MemoryPropertyFlagBits::eDeviceLocal,
1161+
tlasScratchBuffer, tlasScratchMemory
1162+
);
1163+
1164+
vk::BufferDeviceAddressInfo scratchAddressInfo{ .buffer = *tlasScratchBuffer };
1165+
vk::DeviceAddress scratchAddr = device.getBufferAddressKHR(scratchAddressInfo);
1166+
1167+
auto cmd = beginSingleTimeCommands();
1168+
tlasBuildInfo.dstAccelerationStructure = tlas;
1169+
tlasBuildInfo.scratchData.deviceAddress = scratchAddr;
1170+
cmd->buildAccelerationStructuresKHR({ tlasBuildInfo }, { &tlasRangeInfo });
1171+
endSingleTimeCommands(*cmd);
1172+
}
1173+
9441174
void createDescriptorPool() {
9451175
std::array poolSize {
9461176
vk::DescriptorPoolSize( vk::DescriptorType::eUniformBuffer, MAX_FRAMES_IN_FLIGHT),
1177+
vk::DescriptorPoolSize( vk::DescriptorType::eAccelerationStructureKHR, MAX_FRAMES_IN_FLIGHT),
9471178
vk::DescriptorPoolSize( vk::DescriptorType::eSampler, MAX_FRAMES_IN_FLIGHT),
9481179
vk::DescriptorPoolSize( vk::DescriptorType::eSampledImage, (uint32_t)materials.size())
9491180
};
@@ -987,7 +1218,22 @@ class HelloTriangleApplication {
9871218
.pBufferInfo = &bufferInfo
9881219
};
9891220

990-
std::array<vk::WriteDescriptorSet, 1> descriptorWrites{bufferWrite};
1221+
// Acceleration structure
1222+
vk::WriteDescriptorSetAccelerationStructureKHR asInfo{
1223+
.accelerationStructureCount = 1,
1224+
.pAccelerationStructures = {&*tlas}
1225+
};
1226+
1227+
vk::WriteDescriptorSet asWrite{
1228+
.pNext = &asInfo,
1229+
.dstSet = globalDescriptorSets[i],
1230+
.dstBinding = 1,
1231+
.dstArrayElement = 0,
1232+
.descriptorCount = 1,
1233+
.descriptorType = vk::DescriptorType::eAccelerationStructureKHR
1234+
};
1235+
1236+
std::array<vk::WriteDescriptorSet, 2> descriptorWrites{bufferWrite, asWrite};
9911237

9921238
device.updateDescriptorSets(descriptorWrites, {});
9931239
}
@@ -1061,6 +1307,11 @@ class HelloTriangleApplication {
10611307
.allocationSize = memRequirements.size,
10621308
.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties)
10631309
};
1310+
vk::MemoryAllocateFlagsInfo allocFlagsInfo{};
1311+
if (usage & vk::BufferUsageFlagBits::eShaderDeviceAddress) {
1312+
allocFlagsInfo.flags = vk::MemoryAllocateFlagBits::eDeviceAddress;
1313+
allocInfo.pNext = &allocFlagsInfo;
1314+
}
10641315
bufferMemory = vk::raii::DeviceMemory(device, allocInfo);
10651316
buffer.bindMemory(bufferMemory, 0);
10661317
}

‎attachments/38_ray_tracing.slang‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@ struct UniformBuffer {
1313
[[vk::binding(0,0)]]
1414
ConstantBuffer<UniformBuffer> ubo;
1515

16+
[[vk::binding(1,0)]]
17+
RaytracingAccelerationStructure accelerationStructure;
18+
1619
struct VSOutput
1720
{
1821
float4 pos : SV_Position;

0 commit comments

Comments
 (0)