Skip to content

AccelerationStructure.kt

src/jvmMain/kotlin/afield/graphics/render/AccelerationStructure.kt · 459 lines

package afield.graphics.render

import afield.document.context.fromContext
import afield.graphics.vulkan.Renderer
import afield.graphics.vulkan.buffers.HasBuffer
import afield.graphics.vulkan.buffers.VulkanBuffer
import afield.graphics.vulkan.device.CommandBuffer
import afield.graphics.vulkan.device.Fence
import afield.graphics.vulkan.device.nameObject
import afield.graphics.vulkan.util.vkCheck
import afield.graphics.vulkan.util.AutoDisposer
import org.lwjgl.PointerBuffer
import org.lwjgl.system.MemoryStack
import org.lwjgl.vulkan.KHRAccelerationStructure.*
import org.lwjgl.vulkan.KHRSynchronization2.*
import org.lwjgl.vulkan.VK10.*
import org.lwjgl.vulkan.VK12.VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
import org.lwjgl.vulkan.VK13.*
import org.lwjgl.vulkan.*
import kotlin.math.max

Vulkan ray-tracing acceleration structures: AccelerationStructure builds a bottom-level structure from triangles or AABBs, or a top-level one from RayTracingInstances, on the Renderer found from context.

Geometry is described by AccelerationStructureTriangles and AccelerationStructureAabbs; identityTransform3x4 is the default instance transform. After a build, accelerationStructureBuildBarrier makes it visible to later commands that read it, and rayTracingMemoryBarrier orders ray-tracing shader writes before later shader reads.

interface AccelerationStructureAddress { fun deviceAddress(): Long }

A BLAS device address usable in a TLAS instance, including NV cluster BLAS without a Vk handle.

interface HasAccelerationStructure : AccelerationStructureAddress

val vkAccelerationStructure: Long

data class AccelerationStructureTriangles

data class AccelerationStructureTriangles(
    val vertexBuffer: HasBuffer,
    val vertexFormat: Int = VK_FORMAT_R32G32B32_SFLOAT,
    val vertexStride: Long = 12L,
    val maxVertex: Int,
    val primitiveCount: Int,
    val indexBuffer: HasBuffer? = null,
    val indexType: Int = VK_INDEX_TYPE_UINT32,
    val transformBuffer: HasBuffer? = null,
    val flags: Int = VK_GEOMETRY_OPAQUE_BIT_KHR,
    val primitiveOffset: Int = 0,
    val firstVertex: Int = 0,
    val transformOffset: Int = 0
)

data class AccelerationStructureAabbs

data class AccelerationStructureAabbs(
    val aabbBuffer: HasBuffer,
    val stride: Long = 24L,
    val primitiveCount: Int,
    val flags: Int = VK_GEOMETRY_OPAQUE_BIT_KHR,
    val primitiveOffset: Int = 0,
    val firstVertex: Int = 0,
    val transformOffset: Int = 0
)

data class RayTracingInstance

data class RayTracingInstance(
    val accelerationStructure: AccelerationStructureAddress,
    val transform: FloatArray = identityTransform3x4(),
    val customIndex: Int = 0,
    val mask: Int = 0xff,
    val shaderBindingTableRecordOffset: Int = 0,
    val flags: Int = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR
) {

init

    init {
        require(transform.size == 12) { "RayTracingInstance transform must be a row-major 3x4 float array." }
    }
}

fun identityTransform3x4(): FloatArray

fun identityTransform3x4(): FloatArray = floatArrayOf(
    1f, 0f, 0f, 0f,
    0f, 1f, 0f, 0f,
    0f, 0f, 1f, 0f
)

class AccelerationStructure

class AccelerationStructure(
    val renderer: Renderer = fromContext(),
    private val debugName: String = "acceleration-structure"
) : HasAccelerationStructure {

override var vkAccelerationStructure: Long

    override var vkAccelerationStructure: Long = 0L
        private set

    private var storage: VulkanBuffer? = null
    private var scratch: VulkanBuffer? = null
    private var instancesBuffer: VulkanBuffer? = null
    private var type: Int = VK_ACCELERATION_STRUCTURE_TYPE_GENERIC_KHR
    private var cachedDeviceAddress: Long = 0L
    private var handleDisposer: AutoDisposer.Token? = null
    private var gpuInstanceContract: Pair<Int,Int>? = null

override fun deviceAddress(): Long

    override fun deviceAddress(): Long {
        if (cachedDeviceAddress != 0L) return cachedDeviceAddress
        require(vkAccelerationStructure != 0L) { "Acceleration structure '$debugName' has not been built yet." }
        MemoryStack.stackPush().use { stack ->
            val info = VkAccelerationStructureDeviceAddressInfoKHR.calloc(stack)
                .sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR)
                .accelerationStructure(vkAccelerationStructure)
            cachedDeviceAddress = vkGetAccelerationStructureDeviceAddressKHR(renderer.device.vkDevice, info)
            return cachedDeviceAddress
        }
    }

fun buildBottomLevel

body — 67 lines, from line 110 
    fun buildBottomLevel(
        command: CommandBuffer,
        triangles: List<AccelerationStructureTriangles> = emptyList(),
        aabbs: List<AccelerationStructureAabbs> = emptyList(),
        flags: Int = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR
    ): AccelerationStructure {
        require(gpuInstanceContract==null) { "Use a separate acceleration structure for a GPU-instance TLAS" }
        require(renderer.device.supportsAccelerationStructures) { "Device has no enabled VK_KHR_acceleration_structure." }
        require(triangles.isNotEmpty() || aabbs.isNotEmpty()) { "BLAS build needs at least one triangles or AABB geometry." }

        MemoryStack.stackPush().use { stack ->
            val count = triangles.size + aabbs.size
            val geometries = VkAccelerationStructureGeometryKHR.calloc(count, stack)
            val ranges = VkAccelerationStructureBuildRangeInfoKHR.calloc(count, stack)
            val primitiveCounts = stack.mallocInt(count)

            var i = 0
            triangles.forEach { spec ->
                geometries[i]
                    .sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR)
                    .geometryType(VK_GEOMETRY_TYPE_TRIANGLES_KHR)
                    .flags(spec.flags)
                    .geometry { data ->
                        data.triangles { tri ->
                            tri.sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR)
                                .vertexFormat(spec.vertexFormat)
                                .vertexData { it.deviceAddress(spec.vertexBuffer.getGPUAddress(command)) }
                                .vertexStride(spec.vertexStride)
                                .maxVertex(spec.maxVertex)
                                .indexType(if (spec.indexBuffer == null) VK_INDEX_TYPE_NONE_KHR else spec.indexType)
                                .indexData { it.deviceAddress(spec.indexBuffer?.getGPUAddress(command) ?: 0L) }
                                .transformData { it.deviceAddress(spec.transformBuffer?.getGPUAddress(command) ?: 0L) }
                        }
                    }
                ranges[i]
                    .primitiveCount(spec.primitiveCount)
                    .primitiveOffset(spec.primitiveOffset)
                    .firstVertex(spec.firstVertex)
                    .transformOffset(spec.transformOffset)
                primitiveCounts.put(i, spec.primitiveCount)
                i++
            }
            aabbs.forEach { spec ->
                geometries[i]
                    .sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR)
                    .geometryType(VK_GEOMETRY_TYPE_AABBS_KHR)
                    .flags(spec.flags)
                    .geometry { data ->
                        data.aabbs { aabb ->
                            aabb.sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR)
                                .data { it.deviceAddress(spec.aabbBuffer.getGPUAddress(command)) }
                                .stride(spec.stride)
                        }
                    }
                ranges[i]
                    .primitiveCount(spec.primitiveCount)
                    .primitiveOffset(spec.primitiveOffset)
                    .firstVertex(spec.firstVertex)
                    .transformOffset(spec.transformOffset)
                primitiveCounts.put(i, spec.primitiveCount)
                i++
            }

            build(command, VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR, flags, geometries, primitiveCounts, ranges)
        }
        return this
    }

fun buildTopLevel

body — 60 lines, from line 178 
    fun buildTopLevel(
        command: CommandBuffer,
        instances: List<RayTracingInstance>,
        flags: Int = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR
    ): AccelerationStructure {
        require(gpuInstanceContract==null) { "Use a separate acceleration structure for a GPU-instance TLAS" }
        require(renderer.device.supportsAccelerationStructures) { "Device has no enabled VK_KHR_acceleration_structure." }
        require(instances.isNotEmpty()) { "TLAS build needs at least one instance." }

        val byteSize = instances.size * VkAccelerationStructureInstanceKHR.SIZEOF.toLong()
        val buffer = instancesBuffer
            ?.takeIf { it.requestedSize >= byteSize }
            ?: VulkanBuffer(
                renderer,
                byteSize,
                VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR or VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT or VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
                needsDeviceAddress = true
            ).also {
                instancesBuffer?.cleanup()
                instancesBuffer = it
            }

        buffer.mapAndWrite_bytes {
            val records = VkAccelerationStructureInstanceKHR.Buffer(this)
            records.limit(instances.size)
            instances.forEachIndexed { index, instance ->
                val record = records[index]
                for (n in 0 until 12) record.transform().matrix(n, instance.transform[n])
                record.instanceCustomIndex(instance.customIndex)
                    .mask(instance.mask)
                    .instanceShaderBindingTableRecordOffset(instance.shaderBindingTableRecordOffset)
                    .flags(instance.flags)
                    .accelerationStructureReference(instance.accelerationStructure.deviceAddress())
            }
        }

        MemoryStack.stackPush().use { stack ->
            val geometries = VkAccelerationStructureGeometryKHR.calloc(1, stack)
            geometries[0]
                .sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR)
                .geometryType(VK_GEOMETRY_TYPE_INSTANCES_KHR)
                .flags(VK_GEOMETRY_OPAQUE_BIT_KHR)
                .geometry { data ->
                    data.instances { inst ->
                        inst.sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_INSTANCES_DATA_KHR)
                            .arrayOfPointers(false)
                            .data { it.deviceAddress(buffer.getGPUAddress(command)) }
                    }
                }

            val ranges = VkAccelerationStructureBuildRangeInfoKHR.calloc(1, stack)
            ranges[0].primitiveCount(instances.size).primitiveOffset(0).firstVertex(0).transformOffset(0)
            val primitiveCounts = stack.mallocInt(1)
            primitiveCounts.put(0, instances.size)

            build(command, VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR, flags, geometries, primitiveCounts, ranges)
        }
        return this
    }

fun buildTopLevel(command: CommandBuffer, instances: VulkanBuffer, instanceCount: Int,

GPU-produced instance records, on a fresh owner. The first call fixes count and flags: recreate the TLAS to change them, retiring the old owner after GPU completion. This avoids resizing its storage/scratch while previous frames still reference them. Update transforms and addresses within those fixed records freely. Caller owns and synchronizes the producer buffer on Field's graphics queue. No host mapping or per-frame upload occurs.

    fun buildTopLevel(command: CommandBuffer, instances: VulkanBuffer, instanceCount: Int,
        flags: Int = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR): AccelerationStructure {
        val contract=instanceCount to flags
        require(gpuInstanceContract?.let { it==contract } ?: (vkAccelerationStructure==0L)) {
            "GPU-instance TLAS count and flags are fixed; recreate the owner and defer retiring the previous one"
        }
        require(renderer.device.supportsAccelerationStructures && instances.renderer === renderer)
        require(instanceCount > 0 && instances.requestedSize >= instanceCount.toLong()*VkAccelerationStructureInstanceKHR.SIZEOF)
        require(instances.usage and VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR != 0 && instances.gpuAddress > 0)
        MemoryStack.stackPush().use { stack ->
            val geometries=VkAccelerationStructureGeometryKHR.calloc(1,stack).`sType$Default`()
                .geometryType(VK_GEOMETRY_TYPE_INSTANCES_KHR).flags(VK_GEOMETRY_OPAQUE_BIT_KHR)
            geometries.geometry().instances().`sType$Default`().arrayOfPointers(false).data().deviceAddress(instances.getGPUAddress(command))
            val ranges=VkAccelerationStructureBuildRangeInfoKHR.calloc(1,stack).primitiveCount(instanceCount)
            build(command,VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR,flags,geometries,stack.ints(instanceCount),ranges)
        }
        gpuInstanceContract=contract
        return this
    }

fun buildBottomLevelNow

    fun buildBottomLevelNow(
        triangles: List<AccelerationStructureTriangles> = emptyList(),
        aabbs: List<AccelerationStructureAabbs> = emptyList(),
        flags: Int = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR
    ): AccelerationStructure = submitBuildNow { buildBottomLevel(it, triangles, aabbs, flags) }

fun buildTopLevelNow

body — 136 lines, from line 270 
    fun buildTopLevelNow(
        instances: List<RayTracingInstance>,
        flags: Int = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR
    ): AccelerationStructure = submitBuildNow { buildTopLevel(it, instances, flags) }

    private fun submitBuildNow(build: (CommandBuffer) -> Unit): AccelerationStructure {
        val cmd = CommandBuffer(renderer.commandPool, primary = true, oneTimeSubmit = true)
        val fence = Fence(renderer.device, signaled = true)
        try {
            cmd.beginRecording()
            build(cmd)
            cmd.endRecording()
            fence.reset()
            MemoryStack.stackPush().use { stack ->
                renderer.graphQueue.submit(stack.pointers(cmd.vkCommandBuffer), null, null, null, fence)
            }
            fence.fenceWait()
        } finally {
            cmd.cleanup()
            fence.cleanup()
        }
        return this
    }

    private fun build(
        command: CommandBuffer,
        buildType: Int,
        flags: Int,
        geometries: VkAccelerationStructureGeometryKHR.Buffer,
        primitiveCounts: java.nio.IntBuffer,
        ranges: VkAccelerationStructureBuildRangeInfoKHR.Buffer
    ) {
        type = buildType
        val info = VkAccelerationStructureBuildGeometryInfoKHR.calloc()
            .sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR)
            .type(buildType)
            .flags(flags)
            .mode(VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR)
            .geometryCount(geometries.remaining())
            .pGeometries(geometries)

        try {
            val sizes = MemoryStack.stackPush().use { stack ->
                val out = VkAccelerationStructureBuildSizesInfoKHR.calloc(stack)
                    .sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR)
                vkGetAccelerationStructureBuildSizesKHR(
                    renderer.device.vkDevice,
                    VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
                    info,
                    primitiveCounts,
                    out
                )
                BuildSizes(out.accelerationStructureSize(), out.buildScratchSize())
            }

            ensureStorage(sizes.accelerationStructureSize)
            val scratchAlignment = max(1, renderer.device.rayTracingLimits?.minAccelerationStructureScratchOffsetAlignment ?: 1)
            ensureScratch(sizes.buildScratchSize + scratchAlignment)

            val scratchAddress = alignUp(scratch!!.getGPUAddress(command), scratchAlignment.toLong())
            info.dstAccelerationStructure(vkAccelerationStructure)
                .scratchData { it.deviceAddress(scratchAddress) }

            MemoryStack.stackPush().use { stack ->
                val infos = VkAccelerationStructureBuildGeometryInfoKHR.calloc(1, stack)
                infos[0].set(info)
                val rangePointers: PointerBuffer = stack.mallocPointer(1)
                rangePointers.put(0, ranges.address())
                vkCmdBuildAccelerationStructuresKHR(command.vkCommandBuffer, infos, rangePointers)
            }
        } finally {
            info.free()
        }
        accelerationStructureBuildBarrier(command)
    }

    private fun ensureStorage(size: Long) {
        if (storage != null && storage!!.requestedSize >= size && vkAccelerationStructure != 0L) return
        destroyHandle()
        storage?.cleanup()
        storage = VulkanBuffer(
            renderer,
            size,
            VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR or VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
            needsDeviceAddress = true
        )
        MemoryStack.stackPush().use { stack ->
            val createInfo = VkAccelerationStructureCreateInfoKHR.calloc(stack)
                .sType(VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR)
                .type(type)
                .size(size)
                .buffer(storage!!.buffer)
            val out = stack.mallocLong(1)
            vkCheck(
                vkCreateAccelerationStructureKHR(renderer.device.vkDevice, createInfo, null, out),
                "Failed to create acceleration structure"
            )
            vkAccelerationStructure = out[0]
            cachedDeviceAddress = 0L
            // Register after the backing buffer so device shutdown destroys the handle first.
            // Capture only native values: the cleanup must not retain this object's graph.
            val device = renderer.device.vkDevice
            val handle = vkAccelerationStructure
            handleDisposer = AutoDisposer.register(this, debugName, owner = renderer.device) {
                vkDestroyAccelerationStructureKHR(device, handle, null)
            }
            renderer.device.nameObject(
                vkAccelerationStructure,
                VK_OBJECT_TYPE_ACCELERATION_STRUCTURE_KHR,
                debugName
            )
        }
    }

    private fun ensureScratch(size: Long) {
        if (scratch != null && scratch!!.requestedSize >= size) return
        scratch?.cleanup()
        scratch = VulkanBuffer(
            renderer,
            size,
            VK_BUFFER_USAGE_STORAGE_BUFFER_BIT or VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
            needsDeviceAddress = true
        )
    }

    private fun destroyHandle() {
        if (vkAccelerationStructure != 0L) {
            if (handleDisposer?.disable() == true)
                vkDestroyAccelerationStructureKHR(renderer.device.vkDevice, vkAccelerationStructure, null)
            handleDisposer = null
            vkAccelerationStructure = 0L
            cachedDeviceAddress = 0L
        }
    }

fun cleanup()

    fun cleanup() {
        destroyHandle()
        storage?.cleanup()
        scratch?.cleanup()
        instancesBuffer?.cleanup()
        storage = null
        scratch = null
        instancesBuffer = null
        gpuInstanceContract = null
    }

    private data class BuildSizes(val accelerationStructureSize: Long, val buildScratchSize: Long)
}

fun accelerationStructureBuildBarrier(command: CommandBuffer)

fun accelerationStructureBuildBarrier(command: CommandBuffer) {
    MemoryStack.stackPush().use { stack ->
        val barrier = VkMemoryBarrier2.calloc(1, stack)
        barrier[0]
            .sType(VK_STRUCTURE_TYPE_MEMORY_BARRIER_2)
            .srcStageMask(VK_PIPELINE_STAGE_2_ACCELERATION_STRUCTURE_BUILD_BIT_KHR)
            .srcAccessMask(VK_ACCESS_2_ACCELERATION_STRUCTURE_WRITE_BIT_KHR)
            .dstStageMask(VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT)
            .dstAccessMask(VK_ACCESS_2_ACCELERATION_STRUCTURE_READ_BIT_KHR or VK_ACCESS_2_SHADER_READ_BIT)

        val dep = VkDependencyInfo.calloc(stack)
            .sType(VK_STRUCTURE_TYPE_DEPENDENCY_INFO)
            .pMemoryBarriers(barrier)
        vkCmdPipelineBarrier2(command.vkCommandBuffer, dep)
    }
}

fun rayTracingMemoryBarrier(command: CommandBuffer)

fun rayTracingMemoryBarrier(command: CommandBuffer) {
    MemoryStack.stackPush().use { stack ->
        val barrier = VkMemoryBarrier2.calloc(1, stack)
        barrier[0]
            .sType(VK_STRUCTURE_TYPE_MEMORY_BARRIER_2)
            .srcStageMask(VK_PIPELINE_STAGE_2_RAY_TRACING_SHADER_BIT_KHR)
            .srcAccessMask(VK_ACCESS_2_SHADER_WRITE_BIT)
            .dstStageMask(VK_PIPELINE_STAGE_2_RAY_TRACING_SHADER_BIT_KHR or VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT or VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT)
            .dstAccessMask(VK_ACCESS_2_SHADER_READ_BIT or VK_ACCESS_2_SHADER_WRITE_BIT)

        val dep = VkDependencyInfo.calloc(stack)
            .sType(VK_STRUCTURE_TYPE_DEPENDENCY_INFO)
            .pMemoryBarriers(barrier)
        vkCmdPipelineBarrier2(command.vkCommandBuffer, dep)
    }
}

private fun alignUp(value: Long, alignment: Long): Long {
    if (alignment <= 1L) return value
    return ((value + alignment - 1L) / alignment) * alignment
}