using System; using Unity.Mathematics; using UnityEngine.PathTracing.Integration; using UnityEngine.Rendering; using UnityEngine.Rendering.UnifiedRayTracing; namespace UnityEngine.PathTracing.Lightmapping { internal static class BakeLightmapDriver { public class LightmapBakeState { public uint SampleIndex; public UInt64 TexelIndex; public void Init() { SampleIndex = 0; TexelIndex = 0; } public void Tick(uint passSampleCount, uint totalSampleCount, UInt64 chunkTexelCount, UInt64 totalTexelCount, out bool instanceIsDone, out bool chunkIsDone) { instanceIsDone = false; chunkIsDone = false; SampleIndex += passSampleCount; Debug.Assert(SampleIndex <= totalSampleCount); if (SampleIndex == totalSampleCount) { // a chunk is done since we have reached `totalSampleCount` chunkIsDone = true; TexelIndex += chunkTexelCount; SampleIndex = 0; } Debug.Assert(TexelIndex <= totalTexelCount); if (TexelIndex == totalTexelCount) { // an instance is done since we have reached `totalTexelCount` instanceIsDone = true; } } } public struct IntegrationSettings { public RayTracingBackend Backend; public uint MaxDispatchesPerFlush; // how many dispatches to do before flushing the GPU public bool DebugDispatches; public static readonly IntegrationSettings Default = new IntegrationSettings { Backend = RayTracingBackend.Compute, MaxDispatchesPerFlush = 1, DebugDispatches = false }; } public class LightmapBakeSettings { public uint AOSampleCount = 32; public uint DirectSampleCount = 32; public uint IndirectSampleCount = 512; public uint ValiditySampleCount = 512; public AntiAliasingType AOAntiAliasingType = AntiAliasingType.Stochastic; public AntiAliasingType DirectAntiAliasingType = AntiAliasingType.SuperSampling; public AntiAliasingType IndirectAntiAliasingType = AntiAliasingType.Stochastic; public AntiAliasingType ValidityAntiAliasingType = AntiAliasingType.Stochastic; public uint BounceCount = 4; public uint DirectLightingEvaluationCount = 4; public uint IndirectLightingEvaluationCount = 1; public float AOMaxDistance = 1.0f; public float PushOff = 0.00001f; public UInt64 ExpandedBufferSize = 262144; public uint GetSampleCount(IntegratedOutputType integratedOutputType) { switch (integratedOutputType) { case IntegratedOutputType.AO: return AOSampleCount; case IntegratedOutputType.Direct: return DirectSampleCount; case IntegratedOutputType.DirectionalityDirect: return DirectSampleCount; case IntegratedOutputType.Indirect: return IndirectSampleCount; case IntegratedOutputType.DirectionalityIndirect: return IndirectSampleCount; case IntegratedOutputType.Validity: return ValiditySampleCount; case IntegratedOutputType.ShadowMask: return DirectSampleCount; default: Debug.Assert(false, "Unexpected case."); return 0; } } public AntiAliasingType GetAntiAliasingType(IntegratedOutputType integratedOutputType) { switch (integratedOutputType) { case IntegratedOutputType.AO: return AOAntiAliasingType; case IntegratedOutputType.Direct: return DirectAntiAliasingType; case IntegratedOutputType.DirectionalityDirect: return DirectAntiAliasingType; case IntegratedOutputType.Indirect: return IndirectAntiAliasingType; case IntegratedOutputType.DirectionalityIndirect: return IndirectAntiAliasingType; case IntegratedOutputType.Validity: return ValidityAntiAliasingType; case IntegratedOutputType.ShadowMask: return DirectAntiAliasingType; default: Debug.Assert(false, "Unexpected case."); return 0; } } } static bool IsNewChunkStarted( uint maxChunkSize, uint instanceWidth, uint instanceHeight, uint currentChunkTexelOffset, // current starting texel index for the chunk, linear offset into instanceWidth*instanceHeight uint currentSampleIndex, // current sample index for the chunk uint maxSamplesPerTexel, // total sample count per texel out uint chunkSize, // number of texels to process in a single pass out uint expandedSampleWidth, // number of expanded samples per texel, power of two, this might exceed the required sample count out uint passSampleCount, // the actual number of samples to take, this might be smaller than the expanded sample width out uint2 chunkOffset // the chunk offset in 2D ) { // this function should only be called when there is work to do Debug.Assert(currentSampleIndex < maxSamplesPerTexel); Debug.Assert(maxChunkSize > 0); Debug.Assert(instanceWidth > 0); Debug.Assert(instanceHeight > 0); Debug.Assert(currentChunkTexelOffset < instanceWidth * instanceHeight); chunkOffset = new uint2((uint)(currentChunkTexelOffset % (UInt64)instanceWidth), (uint)(currentChunkTexelOffset / (UInt64)instanceWidth)); Debug.Assert(chunkOffset.x < instanceWidth); Debug.Assert(chunkOffset.y < instanceHeight); uint remainingTexels = (uint)instanceWidth - chunkOffset.x + ((uint)(instanceHeight - 1) - chunkOffset.y) * (uint)instanceWidth; Debug.Assert(remainingTexels > 0); uint remainingSampleCount = math.max(0, maxSamplesPerTexel - currentSampleIndex); Debug.Assert(remainingSampleCount > 0); // Choose the size of chunk to take chunkSize = math.min(remainingTexels, maxChunkSize); // Take as many *texels* as possible in a single pass - this is done to reduce the number of dispatches for compaction, reduction and `copy to lightmap` Debug.Assert(chunkSize <= maxChunkSize); // Calculate the maximum number of samples that can be taken in a single pass uint maxSamplesPerChunk = math.min(maxChunkSize / chunkSize, maxSamplesPerTexel); // The maximum number of samples we can take for the current chunk Debug.Assert(chunkSize * maxSamplesPerChunk <= maxChunkSize); // Sample count expansion needs to be a power of 2 - calculate the expansion width expandedSampleWidth = math.ceilpow2(maxSamplesPerChunk); if (expandedSampleWidth > maxSamplesPerChunk) expandedSampleWidth /= 2; Debug.Assert(expandedSampleWidth >= 1); Debug.Assert(expandedSampleWidth * chunkSize <= maxChunkSize); // Calculate how many samples we can take in this pass passSampleCount = math.min(maxSamplesPerChunk, math.min(remainingSampleCount, expandedSampleWidth)); Debug.Assert(passSampleCount > 0); Debug.Assert(passSampleCount <= expandedSampleWidth); Debug.Assert(passSampleCount + currentSampleIndex <= maxSamplesPerTexel); return currentSampleIndex == 0; // When the sample count has rolled back to zero a new chunk has started } internal static uint AccumulateLightmapInstance( LightmapBakeState bakeState, BakeInstance instance, LightmapBakeSettings lightmapBakeSettings, IntegratedOutputType integratedOutputType, LightmappingContext lightmappingContext, UVAccelerationStructure uvAS, UVFallbackBuffer uvFallbackBuffer, bool doDirectionality, out uint chunkSize, out bool instanceIsDone) { CommandBuffer cmd = lightmappingContext.GetCommandBuffer(); GraphicsBuffer traceScratchBuffer = lightmappingContext.TraceScratchBuffer; var ctx = lightmappingContext.IntegratorContext; var expansionShaders = ctx.ExpansionShaders; bool doDirectional = doDirectionality || integratedOutputType == IntegratedOutputType.ShadowMask; // Shadowmask uses the directional buffer to store the sample count - so also process that Vector2Int instanceTexelOffset = instance.TexelOffset; { var maxSampleCountPerTexel = lightmapBakeSettings.GetSampleCount(integratedOutputType); var instanceWidth = instance.TexelSize.x; var instanceHeight = instance.TexelSize.y; var instanceTexelCount = (UInt64)instanceWidth * (UInt64)instanceHeight; var sampleOffset = bakeState.SampleIndex; var maxChunkSize = (uint)lightmappingContext.ExpandedOutput.count; bool newChunkStarted = IsNewChunkStarted( maxChunkSize, (uint)instanceWidth, (uint)instanceHeight, (uint)bakeState.TexelIndex, sampleOffset, maxSampleCountPerTexel, out chunkSize, out uint expandedSampleWidth, out uint passSampleCount, out uint2 chunkOffset); if (newChunkStarted) { // compact the texels ExpansionHelpers.CompactGBuffer(cmd, expansionShaders, ctx.CompactGBufferKernel, (uint)instanceWidth, chunkSize, chunkOffset, uvFallbackBuffer, ctx.CompactedGBufferLength, lightmappingContext.CompactedTexelIndices); // clear the expanded output buffer // Populate the expanded clear indirect dispatch buffer - using the compacted size. expansionShaders.GetKernelThreadGroupSizes(ctx.ClearBufferKernel, out uint clearThreadGroupSizeX, out uint clearThreadGroupSizeY, out uint clearThreadGroupSizeZ); Debug.Assert(clearThreadGroupSizeY == 1 && clearThreadGroupSizeZ == 1); ExpansionHelpers.PopulateClearExpandedOutputIndirectDispatch(cmd, expansionShaders, ctx.PopulateClearDispatchKernel, clearThreadGroupSizeX, expandedSampleWidth, ctx.CompactedGBufferLength, ctx.ClearDispatchBuffer); // Clear the output buffers. ExpansionHelpers.ClearExpandedOutput(cmd, expansionShaders, ctx.ClearBufferKernel, lightmappingContext.ExpandedOutput, ctx.ClearDispatchBuffer); if (doDirectional) ExpansionHelpers.ClearExpandedOutput(cmd, expansionShaders, ctx.ClearBufferKernel, lightmappingContext.ExpandedDirectional, ctx.ClearDispatchBuffer); } // Work out the super sampling resolution. It's the width of the N x N supersampling kernel. Find the largest perfect square that is less than or equal to the max sample count per texel. uint superSampleWidth = Math.Max(1, (uint)Math.Sqrt(maxSampleCountPerTexel)); // generate the GBuffer ExpansionHelpers.GenerateGBuffer( cmd, lightmappingContext.IntegratorContext.GBufferShader, lightmappingContext.GBuffer, traceScratchBuffer, lightmappingContext.IntegratorContext.SamplingResources, uvAS, uvFallbackBuffer, ctx.CompactedGBufferLength, lightmappingContext.CompactedTexelIndices, instanceTexelOffset, chunkOffset, chunkSize, expandedSampleWidth, passSampleCount, sampleOffset, lightmapBakeSettings.GetAntiAliasingType(integratedOutputType), superSampleWidth ); GraphicsBuffer expandedDirectional = lightmappingContext.ExpandedDirectional; var instanceGeometryIndex = lightmappingContext.World.PathTracingWorld.GetAccelerationStructure().GeometryPool.GetInstanceGeometryIndex(instance.Mesh); bool debugGBuffer = false; if (debugGBuffer) { Debug.Log($"Lightmap resolution: {lightmappingContext.Width} x {lightmappingContext.Height}"); Debug.Log($"Instance resolution: {instanceWidth} x {instanceHeight}"); Debug.Log($"Instance offset: {instanceTexelOffset}"); Debug.Log($"Sample count: {maxSampleCountPerTexel}"); var occupancy = (double)(passSampleCount * chunkSize) / (double)maxChunkSize * 100.0; Debug.Log(string.Format(System.Globalization.CultureInfo.InvariantCulture, "Occupancy: {0:F2}%", occupancy)); // write out the lightmap UV samples var uvSampleData = ExpansionHelpers.DebugGBuffer(cmd, instance, lightmappingContext, expandedSampleWidth, passSampleCount); string sampleOutput = new(""); foreach (var sample in uvSampleData) sampleOutput += string.Format(System.Globalization.CultureInfo.InvariantCulture, "float2({0}, {1})\n", sample.x, sample.y); System.Console.WriteLine(sampleOutput); } // accumulate the lightmap texel cmd.BeginSample("AccumulateLightmapInstance"); switch (integratedOutputType) { case IntegratedOutputType.AO: { lightmappingContext.IntegratorContext.LightmapAOIntegrator.Accumulate( cmd, passSampleCount, bakeState.SampleIndex, instance.LocalToWorldMatrix, instance.LocalToWorldMatrixNormals, instanceGeometryIndex, instance.TexelSize, chunkOffset, lightmappingContext.World.PathTracingWorld, traceScratchBuffer, lightmappingContext.GBuffer, expandedSampleWidth, lightmappingContext.ExpandedOutput, lightmappingContext.CompactedTexelIndices, lightmappingContext.IntegratorContext.CompactedGBufferLength, lightmapBakeSettings.PushOff, lightmapBakeSettings.AOMaxDistance, newChunkStarted ); break; } case IntegratedOutputType.Validity: { lightmappingContext.IntegratorContext.LightmapValidityIntegrator.Accumulate( cmd, passSampleCount, bakeState.SampleIndex, instance.LocalToWorldMatrix, instance.LocalToWorldMatrixNormals, instanceGeometryIndex, instance.TexelSize, chunkOffset, lightmappingContext.World.PathTracingWorld, traceScratchBuffer, lightmappingContext.GBuffer, expandedSampleWidth, lightmappingContext.ExpandedOutput, lightmappingContext.CompactedTexelIndices, lightmappingContext.IntegratorContext.CompactedGBufferLength, lightmapBakeSettings.PushOff, newChunkStarted ); break; } case IntegratedOutputType.Direct: case IntegratedOutputType.DirectionalityDirect: { lightmappingContext.IntegratorContext.LightmapDirectIntegrator.Accumulate( cmd, passSampleCount, bakeState.SampleIndex, instance.LocalToWorldMatrix, instance.LocalToWorldMatrixNormals, instanceGeometryIndex, instance.TexelSize, chunkOffset, lightmappingContext.World.PathTracingWorld, traceScratchBuffer, lightmappingContext.GBuffer, expandedSampleWidth, lightmappingContext.ExpandedOutput, expandedDirectional, lightmappingContext.CompactedTexelIndices, lightmappingContext.IntegratorContext.CompactedGBufferLength, instance.ReceiveShadows, lightmapBakeSettings.PushOff, lightmapBakeSettings.DirectLightingEvaluationCount, newChunkStarted ); break; } case IntegratedOutputType.Indirect: case IntegratedOutputType.DirectionalityIndirect: { lightmappingContext.IntegratorContext.LightmapIndirectIntegrator.Accumulate( cmd, passSampleCount, bakeState.SampleIndex, lightmapBakeSettings.BounceCount, instance.LocalToWorldMatrix, instance.LocalToWorldMatrixNormals, instanceGeometryIndex, instance.TexelSize, chunkOffset, lightmappingContext.World.PathTracingWorld, traceScratchBuffer, lightmappingContext.GBuffer, expandedSampleWidth, lightmappingContext.ExpandedOutput, expandedDirectional, lightmappingContext.CompactedTexelIndices, lightmappingContext.IntegratorContext.CompactedGBufferLength, lightmapBakeSettings.PushOff, lightmapBakeSettings.IndirectLightingEvaluationCount, newChunkStarted ); break; } case IntegratedOutputType.ShadowMask: { lightmappingContext.IntegratorContext.LightmapShadowMaskIntegrator.Accumulate( cmd, passSampleCount, bakeState.SampleIndex, instance.LocalToWorldMatrix, instance.LocalToWorldMatrixNormals, instanceGeometryIndex, instance.TexelSize, chunkOffset, lightmappingContext.World.PathTracingWorld, traceScratchBuffer, lightmappingContext.GBuffer, expandedSampleWidth, lightmappingContext.ExpandedOutput, expandedDirectional, lightmappingContext.CompactedTexelIndices, lightmappingContext.IntegratorContext.CompactedGBufferLength, instance.ReceiveShadows, lightmapBakeSettings.PushOff, lightmapBakeSettings.DirectLightingEvaluationCount, newChunkStarted ); break; } } cmd.EndSample("AccumulateLightmapInstance"); //LightmapIntegrationHelpers.LogGraphicsBuffer(cmd, lightmappingContext.ExpandedOutput, "expandedOutput", LightmapIntegrationHelpers.LogBufferType.Float4); // Update the baking state bakeState.Tick( passSampleCount, maxSampleCountPerTexel, chunkSize, instanceTexelCount, out instanceIsDone, out bool chunkIsDone); if (chunkIsDone) { int maxExpandedDispatchSize = instanceWidth * instanceHeight * (int)expandedSampleWidth; // Gather to lightmap -> first reduce to output resolution // Populate the reduce indirect dispatch buffer - using the compacted size. expansionShaders.GetKernelThreadGroupSizes(ctx.ReductionKernel, out uint reduceThreadGroupSizeX, out uint reduceThreadGroupSizeY, out uint reduceThreadGroupSizeZ); Debug.Assert(reduceThreadGroupSizeY == 1 && reduceThreadGroupSizeZ == 1); ExpansionHelpers.PopulateReduceExpandedOutputIndirectDispatch(cmd, expansionShaders, ctx.PopulateReduceDispatchKernel, reduceThreadGroupSizeX, expandedSampleWidth, ctx.CompactedGBufferLength, ctx.ReduceDispatchBuffer); ExpansionHelpers.ReduceExpandedOutput(cmd, expansionShaders, ctx.ReductionKernel, lightmappingContext.ExpandedOutput, maxExpandedDispatchSize, expandedSampleWidth, ctx.ReduceDispatchBuffer); if (doDirectional) ExpansionHelpers.ReduceExpandedOutput(cmd, expansionShaders, ctx.ReductionKernel, lightmappingContext.ExpandedDirectional, maxExpandedDispatchSize, expandedSampleWidth, ctx.ReduceDispatchBuffer); // Populate the copy indirect dispatch buffer - using the compacted size. expansionShaders.GetKernelThreadGroupSizes(ctx.CopyToLightmapKernel, out uint copyThreadGroupSizeX, out uint copyThreadGroupSizeY, out uint copyThreadGroupSizeZ); Debug.Assert(copyThreadGroupSizeY == 1 && copyThreadGroupSizeZ == 1); ExpansionHelpers.PopulateCopyToLightmapIndirectDispatch(cmd, expansionShaders, ctx.PopulateCopyDispatchKernel, copyThreadGroupSizeX, ctx.CompactedGBufferLength, ctx.CopyDispatchBuffer); ExpansionHelpers.CopyToLightmap(cmd, expansionShaders, ctx.CopyToLightmapKernel, expandedSampleWidth, instanceWidth, instanceTexelOffset, chunkOffset, ctx.CompactedGBufferLength, lightmappingContext.CompactedTexelIndices, lightmappingContext.ExpandedOutput, ctx.CopyDispatchBuffer, lightmappingContext.AccumulatedOutput); if (doDirectional) ExpansionHelpers.CopyToLightmap(cmd, expansionShaders, ctx.CopyToLightmapKernel, expandedSampleWidth, instanceWidth, instanceTexelOffset, chunkOffset, ctx.CompactedGBufferLength, lightmappingContext.CompactedTexelIndices, lightmappingContext.ExpandedDirectional, ctx.CopyDispatchBuffer, lightmappingContext.AccumulatedDirectionalOutput); } return passSampleCount; } } } }