using System;
using System.Runtime.InteropServices;
using Unity.Mathematics;
namespace UnityEngine.Rendering.RadeonRays
{
internal struct TopLevelAccelStruct : IDisposable
{
public const GraphicsBuffer.Target topLevelBvhTarget = GraphicsBuffer.Target.Structured;
public const GraphicsBuffer.Target instanceInfoTarget = GraphicsBuffer.Target.Structured;
public GraphicsBuffer topLevelBvh;
public GraphicsBuffer bottomLevelBvhs;
public GraphicsBuffer instanceInfos;
public uint instanceCount;
public void Dispose()
{
topLevelBvh?.Dispose();
instanceInfos?.Dispose();
}
}
internal class HlbvhTopLevelBuilder
{
readonly ComputeShader shaderBuildHlbvh;
readonly int kernelInit;
readonly int kernelCalculateAabb;
readonly int kernelCalculateMortonCodes;
readonly int kernelBuildTreeBottomUp;
readonly RadixSort radixSort;
const uint kTrianglesPerThread = 8u;
const uint kGroupSize = 256u;
const uint kTrianglesPerGroup = kTrianglesPerThread * kGroupSize;
public HlbvhTopLevelBuilder(RadeonRaysShaders shaders)
{
shaderBuildHlbvh = shaders.buildHlbvh;
kernelInit = shaderBuildHlbvh.FindKernel("Init");
kernelCalculateAabb = shaderBuildHlbvh.FindKernel("CalculateAabb");
kernelCalculateMortonCodes = shaderBuildHlbvh.FindKernel("CalculateMortonCodes");
kernelBuildTreeBottomUp = shaderBuildHlbvh.FindKernel("BuildTreeBottomUp");
radixSort = new RadixSort(shaders);
}
public ulong GetScratchDataSizeInDwords(uint instanceCount)
{
var scratchLayout = ScratchBufferLayout.Create(instanceCount);
return scratchLayout.TotalSize;
}
public static uint GetBvhNodeCount(uint leafCount)
{
return leafCount - 1;
}
public void AllocateResultBuffers(uint instanceCount, ref TopLevelAccelStruct accelStruct)
{
var bvhNodeCount = GetBvhNodeCount(instanceCount);
accelStruct.Dispose();
accelStruct.instanceInfos = new GraphicsBuffer(TopLevelAccelStruct.instanceInfoTarget, (int)instanceCount, Marshal.SizeOf<InstanceInfo>());
accelStruct.topLevelBvh = new GraphicsBuffer(TopLevelAccelStruct.topLevelBvhTarget, (int)bvhNodeCount + 1, Marshal.SizeOf<BvhNode>()); // plus one for header
}
public void CreateEmpty(ref TopLevelAccelStruct accelStruct)
{
accelStruct.Dispose();
accelStruct.topLevelBvh = new GraphicsBuffer(TopLevelAccelStruct.topLevelBvhTarget, 2, Marshal.SizeOf<BvhNode>());
accelStruct.instanceInfos = accelStruct.topLevelBvh;
accelStruct.bottomLevelBvhs = accelStruct.topLevelBvh;
accelStruct.instanceCount = 0;
var top = new BvhNode[2];
top[0].child0 = 0;
top[0].child1 = 0;
top[0].parent = 0xFFFFFFFF;
top[1].child0 = 0;
top[1].child1 = 0;
top[1].parent = 0xFFFFFFFF;
top[1].update = 0;
top[1].aabb0_min = new float3(float.NegativeInfinity, float.NegativeInfinity, float.NegativeInfinity);
top[1].aabb0_max = new float3(float.NegativeInfinity, float.NegativeInfinity, float.NegativeInfinity);
top[1].aabb1_min = new float3(float.NegativeInfinity, float.NegativeInfinity, float.NegativeInfinity);
top[1].aabb1_max = new float3(float.NegativeInfinity, float.NegativeInfinity, float.NegativeInfinity);
accelStruct.topLevelBvh.SetData(top);
}
public void Execute(CommandBuffer cmd, GraphicsBuffer scratch, ref TopLevelAccelStruct accelStruct)
{
Common.EnableKeyword(cmd, shaderBuildHlbvh, "TOP_LEVEL", true);
Common.EnableKeyword(cmd, shaderBuildHlbvh, "UINT16_INDICES", false);
uint instanceCount = accelStruct.instanceCount;
var scratchLayout = ScratchBufferLayout.Create(instanceCount);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_constants_vertex_stride, 0);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_constants_triangle_count, (int)instanceCount);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_bvh_offset, 0);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_internal_node_range_offset, (int)scratchLayout.InternalNodeRange);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_aabb_offset, (int)scratchLayout.Aabb);
BindKernelArguments(cmd, kernelInit, scratch, scratchLayout, accelStruct, false);
cmd.DispatchCompute(shaderBuildHlbvh, kernelInit, 1, 1, 1);
BindKernelArguments(cmd, kernelCalculateAabb, scratch, scratchLayout, accelStruct, false);
cmd.DispatchCompute(shaderBuildHlbvh, kernelCalculateAabb, (int)Common.CeilDivide(instanceCount, kTrianglesPerGroup), 1, 1);
BindKernelArguments(cmd, kernelCalculateMortonCodes, scratch, scratchLayout, accelStruct, false);
cmd.DispatchCompute(shaderBuildHlbvh, kernelCalculateMortonCodes, (int)Common.CeilDivide(instanceCount, kTrianglesPerGroup), 1, 1);
radixSort.Execute(cmd, scratch,
scratchLayout.MortonCodes, scratchLayout.SortedMortonCodes,
scratchLayout.PrimitiveRefs, scratchLayout.SortedPrimitiveRefs,
scratchLayout.SortMemory, instanceCount);
BindKernelArguments(cmd, kernelBuildTreeBottomUp, scratch, scratchLayout, accelStruct, true);
cmd.DispatchCompute(shaderBuildHlbvh, kernelBuildTreeBottomUp, (int)Common.CeilDivide(instanceCount, kTrianglesPerGroup), 1, 1);
}
struct ScratchBufferLayout
{
public uint Aabb;
public uint MortonCodes;
public uint PrimitiveRefs;
public uint SortedMortonCodes;
public uint SortedPrimitiveRefs;
public uint SortMemory;
public uint InternalNodeRange;
public uint TotalSize;
public static ScratchBufferLayout Create(uint instanceCount)
{
var result = new ScratchBufferLayout();
result.Aabb = result.Reserve(6);
result.MortonCodes = result.Reserve(instanceCount);
result.PrimitiveRefs = result.Reserve(instanceCount);
result.SortedMortonCodes = result.Reserve(instanceCount);
result.SortedPrimitiveRefs = result.Reserve(instanceCount);
result.SortMemory = result.Reserve((uint)RadixSort.GetScratchDataSizeInDwords(instanceCount));
// overlaps with MortonCodes and PrimitiveRefs
result.InternalNodeRange = result.MortonCodes;
return result;
}
uint Reserve(uint size)
{
var temp = TotalSize;
TotalSize += size;
return temp;
}
}
private void BindKernelArguments(
CommandBuffer cmd,
int kernel,
GraphicsBuffer scratch,
ScratchBufferLayout scratchLayout,
TopLevelAccelStruct accelStruct,
bool setSortedCodes)
{
cmd.SetComputeBufferParam(shaderBuildHlbvh, kernel, SID.g_scratch_buffer, scratch);
cmd.SetComputeBufferParam(shaderBuildHlbvh, kernel, SID.g_bvh, accelStruct.topLevelBvh);
cmd.SetComputeBufferParam(shaderBuildHlbvh, kernel, SID.g_bottom_bvhs, accelStruct.bottomLevelBvhs);
cmd.SetComputeBufferParam(shaderBuildHlbvh, kernel, SID.g_instance_infos, accelStruct.instanceInfos);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_aabb_offset, (int)scratchLayout.Aabb);
if (setSortedCodes)
{
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_morton_codes_offset, (int)scratchLayout.SortedMortonCodes);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_primitive_refs_offset, (int)scratchLayout.SortedPrimitiveRefs);
}
else
{
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_morton_codes_offset, (int)scratchLayout.MortonCodes);
cmd.SetComputeIntParam(shaderBuildHlbvh, SID.g_primitive_refs_offset, (int)scratchLayout.PrimitiveRefs);
}
}
}
}