Skip to content

Commit

Permalink
Merge pull request #31 from corporateshark/rtx_blas
Browse files Browse the repository at this point in the history 
Implemented BLAS via `VK_KHR_acceleration_structure`
  • Loading branch information
@corporateshark
corporateshark authored Sep 8, 2024
2 parents b811a18 + 0af37bb commit 4c91bac
Show file tree
Hide file tree
Showing 6 changed files with 257 additions and 1 deletion.
6 changes: 6 additions & 0 deletions lvk/LVK.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,12 @@ void lvk::destroy(lvk::IContext* ctx, lvk::QueryPoolHandle handle) {
}
}

void lvk::destroy(lvk::IContext* ctx, lvk::AccelStructHandle handle) {
if (ctx) {
ctx->destroy(handle);
}
}

// Logs GLSL shaders with line numbers annotation
void lvk::logShaderSource(const char* text) {
uint32_t line = 0;
Expand Down
54 changes: 54 additions & 0 deletions lvk/LVK.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@ using SamplerHandle = lvk::Handle<struct Sampler>;
using BufferHandle = lvk::Handle<struct Buffer>;
using TextureHandle = lvk::Handle<struct Texture>;
using QueryPoolHandle = lvk::Handle<struct QueryPool>;
using AccelStructHandle = lvk::Handle<struct AccelerationStructure>;

// forward declarations to access incomplete type IContext
void destroy(lvk::IContext* ctx, lvk::ComputePipelineHandle handle);
Expand All @@ -119,6 +120,7 @@ void destroy(lvk::IContext* ctx, lvk::SamplerHandle handle);
void destroy(lvk::IContext* ctx, lvk::BufferHandle handle);
void destroy(lvk::IContext* ctx, lvk::TextureHandle handle);
void destroy(lvk::IContext* ctx, lvk::QueryPoolHandle handle);
void destroy(lvk::IContext* ctx, lvk::AccelStructHandle handle);

template<typename HandleType>
class Holder final {
Expand Down Expand Up @@ -781,6 +783,55 @@ struct TextureDesc {
const char* debugName = "";
};

enum AccelStructType : uint8_t {
AccelStructType_Invalid = 0,
AccelStructType_TLAS = 1,
AccelStructType_BLAS = 2,
};

enum AccelStructGeomType : uint8_t {
AccelStructGeomType_Triangles = 0,
AccelStructGeomType_AABBs = 1,
AccelStructGeomType_Instances = 2,
};

enum AccelStructBuildFlagBits : uint8_t {
AccelStructBuildFlagBits_AllowUpdate = 1 << 0,
AccelStructBuildFlagBits_AllowCompaction = 1 << 1,
AccelStructBuildFlagBits_PreferFastTrace = 1 << 2,
AccelStructBuildFlagBits_PreferFastBuild = 1 << 3,
AccelStructBuildFlagBits_LowMemory = 1 << 4,
};

enum AccelStructGeometryFlagBits : uint8_t {
AccelStructGeometryFlagBits_Opaque = 1 << 0,
AccelStructGeometryFlagBits_NoDuplicateAnyHit = 1 << 1,
};

struct AccelStructBuildRange {
uint32_t primitiveCount = 0;
uint32_t primitiveOffset = 0;
uint32_t firstVertex = 0;
uint32_t transformOffset = 0;
};

struct AccelStructDesc {
AccelStructType type = AccelStructType_Invalid;
AccelStructGeomType geometryType = AccelStructGeomType_Triangles;
uint8_t geometryFlags = AccelStructGeometryFlagBits_Opaque;

VertexFormat vertexFormat = VertexFormat::Invalid;
BufferHandle vertexBuffer;
uint32_t vertexStride = 0; // zero means the size of `vertexFormat`
uint32_t numVertices = 0;
IndexFormat indexFormat = IndexFormat_UI32;
BufferHandle indexBuffer;
BufferHandle transformBuffer;
AccelStructBuildRange buildRange = {};
uint8_t buildFlags = AccelStructBuildFlagBits_PreferFastTrace;
const char* debugName = "";
};

struct Dependencies {
enum { LVK_MAX_SUBMIT_DEPENDENCIES = 4 };
TextureHandle textures[LVK_MAX_SUBMIT_DEPENDENCIES] = {};
Expand Down Expand Up @@ -905,13 +956,16 @@ class IContext {
const char* debugName,
Result* outResult = nullptr) = 0;

[[nodiscard]] virtual Holder<AccelStructHandle> createAccelerationStructure(const AccelStructDesc& desc, Result* outResult = nullptr) = 0;

virtual void destroy(ComputePipelineHandle handle) = 0;
virtual void destroy(RenderPipelineHandle handle) = 0;
virtual void destroy(ShaderModuleHandle handle) = 0;
virtual void destroy(SamplerHandle handle) = 0;
virtual void destroy(BufferHandle handle) = 0;
virtual void destroy(TextureHandle handle) = 0;
virtual void destroy(QueryPoolHandle handle) = 0;
virtual void destroy(AccelStructHandle handle) = 0;
virtual void destroy(Framebuffer& fb) = 0;

#pragma region Buffer functions
Expand Down
177 changes: 177 additions & 0 deletions lvk/vulkan/VulkanClasses.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3271,6 +3271,36 @@ lvk::Holder<lvk::QueryPoolHandle> lvk::VulkanContext::createQueryPool(uint32_t n
return {this, handle};
}

lvk::Holder<lvk::AccelStructHandle> lvk::VulkanContext::createAccelerationStructure(const AccelStructDesc& desc, Result* outResult) {
LVK_PROFILER_FUNCTION();

Result result;

AccelStructHandle handle;

switch (desc.type) {
case AccelStructType_BLAS:
handle = createBLAS(desc, &result);
break;
case AccelStructType_TLAS:
LVK_ASSERT_MSG(false, "Not implemented (yet)");
break;
default:
LVK_ASSERT_MSG(false, "Invalid acceleration structure type");
Result::setResult(outResult, Result(Result::Code::ArgumentOutOfRange, "Invalid acceleration structure type"));
return {};
}

if (!LVK_VERIFY(result.isOk())) {
Result::setResult(outResult, Result(Result::Code::RuntimeError, "Cannot create AccelerationStructure"));
return {};
}

Result::setResult(outResult, result);

return {this, handle};
}

lvk::Holder<lvk::SamplerHandle> lvk::VulkanContext::createSampler(const SamplerStateDesc& desc, Result* outResult) {
LVK_PROFILER_FUNCTION();

Expand Down Expand Up @@ -3582,6 +3612,142 @@ lvk::Holder<lvk::TextureHandle> lvk::VulkanContext::createTexture(const TextureD
return {this, handle};
}

lvk::AccelStructHandle lvk::VulkanContext::createBLAS(const AccelStructDesc& desc, Result* outResult) {
LVK_ASSERT(desc.type == AccelStructType_BLAS);
LVK_ASSERT(desc.geometryType == AccelStructGeomType_Triangles);
LVK_ASSERT(desc.numVertices);
LVK_ASSERT(desc.indexBuffer.valid());
LVK_ASSERT(desc.vertexBuffer.valid());
LVK_ASSERT(desc.transformBuffer.valid());
LVK_ASSERT(desc.buildRange.primitiveCount);

VkGeometryFlagsKHR geometryFlags = 0;

if (desc.geometryFlags & AccelStructGeometryFlagBits_Opaque) {
geometryFlags |= VK_GEOMETRY_OPAQUE_BIT_KHR;
}
if (desc.geometryFlags & AccelStructGeometryFlagBits_NoDuplicateAnyHit) {
geometryFlags |= VK_GEOMETRY_NO_DUPLICATE_ANY_HIT_INVOCATION_BIT_KHR;
}

const VkAccelerationStructureGeometryKHR accelerationStructureGeometry{
.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
.geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR,
.geometry =
{
.triangles =
{
.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR,
.vertexFormat = vertexFormatToVkFormat(desc.vertexFormat),
.vertexData = {.deviceAddress = gpuAddress(desc.vertexBuffer)},
.vertexStride = desc.vertexStride ? desc.vertexStride : lvk::getVertexFormatSize(desc.vertexFormat),
.maxVertex = desc.numVertices - 1,
.indexType = VK_INDEX_TYPE_UINT32,
.indexData = {.deviceAddress = gpuAddress(desc.indexBuffer)},
.transformData = {.deviceAddress = gpuAddress(desc.transformBuffer)},
},
},
.flags = VK_GEOMETRY_OPAQUE_BIT_KHR,
};

const VkAccelerationStructureBuildGeometryInfoKHR accelerationStructureBuildGeometryInfo{
.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR,
.geometryCount = 1,
.pGeometries = &accelerationStructureGeometry,
};

VkAccelerationStructureBuildSizesInfoKHR accelerationStructureBuildSizesInfo = {
.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR,
};
vkGetAccelerationStructureBuildSizesKHR(vkDevice_,
VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
&accelerationStructureBuildGeometryInfo,
&desc.buildRange.primitiveCount,
&accelerationStructureBuildSizesInfo);
char debugNameBuffer[256] = {0};
if (desc.debugName) {
snprintf(debugNameBuffer, sizeof(debugNameBuffer) - 1, "Buffer: %s", desc.debugName);
}
lvk::AccelerationStructure accelStruct = {
.buffer = createBuffer(
{
.usage = lvk::BufferUsageBits_AccelStructStorage,
.storage = lvk::StorageType_Device,
.size = accelerationStructureBuildSizesInfo.accelerationStructureSize,
.debugName = debugNameBuffer,
},
outResult),
};
const VkAccelerationStructureCreateInfoKHR ciAccelerationStructure = {
.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR,
.buffer = getVkBuffer(this, accelStruct.buffer),
.size = accelerationStructureBuildSizesInfo.accelerationStructureSize,
.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
};
VK_ASSERT(vkCreateAccelerationStructureKHR(vkDevice_, &ciAccelerationStructure, nullptr, &accelStruct.vkHandle));

lvk::Holder<lvk::BufferHandle> scratchBuffer = createBuffer(
{
.usage = lvk::BufferUsageBits_Storage,
.storage = lvk::StorageType_Device,
.size = accelerationStructureBuildSizesInfo.buildScratchSize,
.debugName = "Buffer: BLAS scratch",
},
outResult);

VkBuildAccelerationStructureFlagsKHR buildFlags = 0;

if (desc.buildFlags & AccelStructBuildFlagBits_AllowUpdate) {
buildFlags |= VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHR;
}
if (desc.buildFlags & AccelStructBuildFlagBits_AllowCompaction) {
buildFlags |= VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_COMPACTION_BIT_KHR;
}
if (desc.buildFlags & AccelStructBuildFlagBits_PreferFastTrace) {
buildFlags |= VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
}
if (desc.buildFlags & AccelStructBuildFlagBits_PreferFastBuild) {
buildFlags |= VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR;
}
if (desc.buildFlags & AccelStructBuildFlagBits_LowMemory) {
buildFlags |= VK_BUILD_ACCELERATION_STRUCTURE_LOW_MEMORY_BIT_KHR;
}

const VkAccelerationStructureBuildGeometryInfoKHR accelerationBuildGeometryInfo{
.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
.flags = buildFlags,
.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR,
.dstAccelerationStructure = accelStruct.vkHandle,
.geometryCount = 1,
.pGeometries = &accelerationStructureGeometry,
.scratchData = {.deviceAddress = gpuAddress(scratchBuffer)},
};

const VkAccelerationStructureBuildRangeInfoKHR accelerationStructureBuildRangeInfo = {
.primitiveCount = desc.buildRange.primitiveCount,
.primitiveOffset = desc.buildRange.primitiveOffset,
.firstVertex = desc.buildRange.firstVertex,
.transformOffset = desc.buildRange.transformOffset,
};
const VkAccelerationStructureBuildRangeInfoKHR* accelerationBuildStructureRangeInfos[] = {&accelerationStructureBuildRangeInfo};

lvk::ICommandBuffer& buffer = acquireCommandBuffer();
vkCmdBuildAccelerationStructuresKHR(
lvk::getVkCommandBuffer(buffer), 1, &accelerationBuildGeometryInfo, accelerationBuildStructureRangeInfos);
wait(submit(buffer, {}));

const VkAccelerationStructureDeviceAddressInfoKHR accelerationDeviceAddressInfo{
.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR,
.accelerationStructure = accelStruct.vkHandle,
};
accelStruct.deviceAddress = vkGetAccelerationStructureDeviceAddressKHR(vkDevice_, &accelerationDeviceAddressInfo);

return accelStructuresPool_.create(std::move(accelStruct));
}

static_assert(1 << (sizeof(lvk::Format) * 8) <= LVK_ARRAY_NUM_ELEMENTS(lvk::VulkanContextImpl::ycbcrConversionData_),
"There aren't enough elements in `ycbcrConversionData_` to be accessed by lvk::Format");

Expand Down Expand Up @@ -4180,6 +4346,17 @@ void lvk::VulkanContext::destroy(lvk::QueryPoolHandle handle) {
deferredTask(std::packaged_task<void()>([device = vkDevice_, pool = pool]() { vkDestroyQueryPool(device, pool, nullptr); }));
}

void lvk::VulkanContext::destroy(lvk::AccelStructHandle handle) {
AccelerationStructure* accelStruct = accelStructuresPool_.get(handle);

SCOPE_EXIT {
accelStructuresPool_.destroy(handle);
};

deferredTask(std::packaged_task<void()>(
[device = vkDevice_, as = accelStruct->vkHandle]() { vkDestroyAccelerationStructureKHR(device, as, nullptr); }));
}

void lvk::VulkanContext::destroy(Framebuffer& fb) {
auto destroyFbTexture = [this](TextureHandle& handle) {
{
Expand Down
13 changes: 12 additions & 1 deletion lvk/vulkan/VulkanClasses.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,12 @@ struct ShaderModuleState final {
uint32_t pushConstantsSize = 0;
};

struct AccelerationStructure {
VkAccelerationStructureKHR vkHandle = VK_NULL_HANDLE;
uint64_t deviceAddress = 0;
lvk::Holder<lvk::BufferHandle> buffer;
};

class CommandBuffer final : public ICommandBuffer {
public:
CommandBuffer() = default;
Expand Down Expand Up @@ -458,18 +464,21 @@ class VulkanContext final : public IContext {

Holder<QueryPoolHandle> createQueryPool(uint32_t numQueries, const char* debugName, Result* outResult) override;

Holder<AccelStructHandle> createAccelerationStructure(const AccelStructDesc& desc, Result* outResult) override;

void destroy(ComputePipelineHandle handle) override;
void destroy(RenderPipelineHandle handle) override;
void destroy(ShaderModuleHandle handle) override;
void destroy(SamplerHandle handle) override;
void destroy(BufferHandle handle) override;
void destroy(TextureHandle handle) override;
void destroy(QueryPoolHandle handle) override;
void destroy(AccelStructHandle handle) override;
void destroy(Framebuffer& fb) override;

Result upload(BufferHandle handle, const void* data, size_t size, size_t offset) override;
uint8_t* getMappedPtr(BufferHandle handle) const override;
uint64_t gpuAddress(BufferHandle handle, size_t offset) const override;
uint64_t gpuAddress(BufferHandle handle, size_t offset = 0) const override;
void flushMappedMemory(BufferHandle handle, size_t offset, size_t size) const override;

Result upload(TextureHandle handle, const TextureRangeDesc& range, const void* data) override;
Expand Down Expand Up @@ -507,6 +516,7 @@ class VulkanContext final : public IContext {
lvk::Result* outResult,
lvk::Format yuvFormat = Format_Invalid,
const char* debugName = nullptr);
AccelStructHandle createBLAS(const AccelStructDesc& desc, Result* outResult);

bool hasSwapchain() const noexcept {
return swapchain_ != nullptr;
Expand Down Expand Up @@ -629,6 +639,7 @@ class VulkanContext final : public IContext {
lvk::Pool<lvk::Buffer, lvk::VulkanBuffer> buffersPool_;
lvk::Pool<lvk::Texture, lvk::VulkanImage> texturesPool_;
lvk::Pool<lvk::QueryPool, VkQueryPool> queriesPool_;
lvk::Pool<lvk::AccelerationStructure, lvk::AccelerationStructure> accelStructuresPool_;
};

} // namespace lvk
7 changes: 7 additions & 0 deletions lvk/vulkan/VulkanUtils.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1005,3 +1005,10 @@ VkShaderModule lvk::getVkShaderModule(const IContext* ctx, ShaderModuleHandle sh

return static_cast<const VulkanContext*>(ctx)->shaderModulesPool_.get(shader)->sm;
}

VkDeviceAddress lvk::getAccelerationStructureDeviceAddress(const IContext* ctx, AccelStructHandle accelStruct) {
if (!ctx || accelStruct.empty())
return 0;

return static_cast<const VulkanContext*>(ctx)->accelStructuresPool_.get(accelStruct)->deviceAddress;
}
1 change: 1 addition & 0 deletions lvk/vulkan/VulkanUtils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,5 +124,6 @@ VkCommandBuffer getVkCommandBuffer(const ICommandBuffer& buffer);
VkBuffer getVkBuffer(const IContext* ctx, BufferHandle buffer);
VkImageView getVkImageView(const IContext* ctx, TextureHandle texture);
VkShaderModule getVkShaderModule(const IContext* ctx, ShaderModuleHandle shader);
VkDeviceAddress getAccelerationStructureDeviceAddress(const IContext* ctx, AccelStructHandle accelStruct);

} // namespace lvk

0 comments on commit 4c91bac

Please sign in to comment.