EnvironmentImportanceSamplingBuild.compute

// This implementation is adapted from BuildProbabilityTables.compute
#pragma only_renderers d3d11 vulkan metal glcore

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "EnvironmentImportanceSampling.hlsl"

#define KERNEL_WIDTH 64

uint _ConditionalResolution;
uint _MarginalResolution;
RWStructuredBuffer<float> _ConditionalBuffer;
RWStructuredBuffer<float> _MarginalBuffer;
TextureCube<float4> _EnvCubemap;
SamplerState sampler_EnvCubemap;

float3 SampleSkyTexture(float3 dir, float lod, int sliceIndex)
{
    return _EnvCubemap.SampleLevel(sampler_EnvCubemap, dir, lod).xyz;
}

// Performs a block-level parallel scan.
// Ref: GPU Gems 3, Chapter 39: "Parallel Prefix Sum (Scan) with CUDA".
void ParallelScan(uint i, uint w, uint iterCount, RWStructuredBuffer<float> buf, uint bufferOffset, out float sum)
{
    uint offset;

    // Execute the up-sweep phase.
    for (offset = 1; offset <= w / 2; offset *= 2)
    {
        AllMemoryBarrierWithGroupSync();

        for (uint iter = 0; iter < iterCount; iter++)
        {
            uint idx = i + iter * KERNEL_WIDTH;

            // a1 = (2 * i + 1) * offset - 1
            uint a1 = Mad24(Mad24(2u, idx, 1u), offset, -1);
            uint a2 = a1 + offset;

            if (a2 < w)
            {
                buf[bufferOffset + a2] += buf[bufferOffset + a1];
            }
        }
    }

    AllMemoryBarrierWithGroupSync();

    // Prevent NaNs arising from the division of 0 by 0.
    sum = max(buf[bufferOffset + w - 1], FLT_MIN);

    AllMemoryBarrierWithGroupSync();

    // The exclusive scan requires the last element to be 0.
    if (i == 0)
    {
        buf[bufferOffset + w - 1] = 0.0;
    }

    // Execute the down-sweep phase.
    for (offset = w / 2; offset > 0; offset /= 2)
    {
        AllMemoryBarrierWithGroupSync();

        for (uint iter = 0; iter < iterCount; iter++)
        {
            uint idx = i + iter * KERNEL_WIDTH;

            // a1 = (2 * i + 1) * offset - 1
            uint a1 = Mad24(Mad24(2u, idx, 1u), offset, -1);
            uint a2 = a1 + offset;

            if (a2 < w)
            {
                float t1 = buf[bufferOffset + a1];
                buf[bufferOffset + a1] = buf[bufferOffset + a2];
                buf[bufferOffset + a2] += t1;
            }
        }
    }

    AllMemoryBarrierWithGroupSync();
}

#pragma kernel ComputeConditional
[numthreads(KERNEL_WIDTH, 1, 1)]
void ComputeConditional(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    const uint i = dispatchThreadId.x;
    const uint j = dispatchThreadId.y;
    const uint iterCount = _ConditionalResolution / KERNEL_WIDTH;
    const uint conditionalBufferOffset = j * _ConditionalResolution;

    uint iter;

    float v = (j + 0.5) / _MarginalResolution;
    //float sinTheta = sin(v * PI);

    for (iter = 0; iter < iterCount; iter++)
    {
        uint idx = i + iter * KERNEL_WIDTH;
        float u = (idx + 0.5) / _ConditionalResolution;

        // No need for a sinTheta term in the PDF when using equiareal mapping
        float3 dir = MapUVToSkyDirection(float2(u, v));
        _ConditionalBuffer[conditionalBufferOffset + idx] = Luminance(SampleSkyTexture(dir, 0.0, 0).rgb); // * sinTheta;
    }

    float rowValSum = 0.0f;

    ParallelScan(i, _ConditionalResolution, iterCount, _ConditionalBuffer, conditionalBufferOffset, rowValSum);

    for (iter = 0; iter < iterCount; iter++)
    {
        uint idx = i + iter * KERNEL_WIDTH;
        _ConditionalBuffer[conditionalBufferOffset + idx] /= rowValSum;
    }

    if (i == 0)
    {
        float rowIntegralValue = rowValSum / _ConditionalResolution;
        _MarginalBuffer[j] = rowIntegralValue;
    }
}

#pragma kernel ComputeMarginal

[numthreads(KERNEL_WIDTH, 1, 1)]
void ComputeMarginal(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    const uint i = dispatchThreadId.x;
    const uint iterCount = _MarginalResolution / KERNEL_WIDTH;

    float rowValSum = 0.0f;

    ParallelScan(i, _MarginalResolution, iterCount, _MarginalBuffer, 0, rowValSum);

    for (uint iter = 0; iter < iterCount; iter++)
    {
        uint idx = i + iter * KERNEL_WIDTH;
        _MarginalBuffer[idx] /= rowValSum;
    }

    if (i == 0)
    {
        float imgIntegralValue = rowValSum / _MarginalResolution;
        float pdfNormalization = imgIntegralValue > 0.0 ? rcp(imgIntegralValue) * 0.25 * INV_PI : 0.0;
        _MarginalBuffer[0] = pdfNormalization;
    }
}