o3de-atom-sampleviewer/Shaders/RayTracing/RTAOGeneration.azsl

/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#include <Atom/Features/SrgSemantics.azsli>
#include <Atom/RPI/Math.azsli> // PI

#include "RTAODefines.azsli"

// GlobalSrg
ShaderResourceGroup RayTracingGlobalSrg : SRG_RayTracingGlobal
{
    RaytracingAccelerationStructure m_scene;

    Texture2D<float> m_depth;
    Texture2D<float4> m_worldNormalMap;

    RWTexture2D<float4> m_outputAO;

    float m_aoRadius;   // Ambient occlusion radius. Default: 0.4f
    float m_rayMinT;    // The minimum ray extent
    int m_frameCount;   // Used for unique random seeds each frame. Default: 0
    uint  m_numRays;    // Number of ray casted for each pixel

    // Copy of ViewSrg::m_viewProjectionInverseMatrix since we can't access ViewSrg in ray tracing shaders ATM.
    row_major float4x4 m_viewProjectionInverseMatrix;
};

// Generates a seed for RNG from 2 input values
uint InitRandomSeed(uint value1, uint value2)
{
    uint v0 = value1, v1 = value2, s0 = 0;

    [unroll]
    for (uint n = 0; n < 16; n++)
    {
        s0 += 0x9e3779b9;
        v0 += ((v1 << 4) + 0xa341316c) ^ (v1 + s0) ^ ((v1 >> 5) + 0xc8013ea4);
        v1 += ((v0 << 4) + 0xad90777d) ^ (v0 + s0) ^ ((v0 >> 5) + 0x7e95761e);
    }
    return v0;
}

// Get a vector perpendicular to an input vector 
// "Efficient Construction of Perpendicular Vectors Without Branching"
float3 GetPerpendicularVector(float3 input)
{
    float3 absInput = abs(input);
    uint x = ((absInput.x < absInput.y) && (absInput.x < absInput.z)) ? 1 : 0;
    uint y = (absInput.y < absInput.z)? (1 ^ x) : 0;
    uint z = 1 ^ (x | y);
    return normalize(cross(input, float3(x, y, z)));
}

// Get a cosine-weighted random vector centered around a specified normal direction.
float3 GetCosHemisphereSample(inout uint randSeed, float3 hitNorm)
{
    // Get 2 random numbers to select our sample with
    float2 randVal = float2(NextRandomFloatUniform(randSeed), NextRandomFloatUniform(randSeed));

    float3 bitangent = GetPerpendicularVector(hitNorm);
    float3 tangent = cross(bitangent, hitNorm);

    float radius = sqrt(randVal.x);
    float angle = 2.0f * PI * randVal.y;

    // Get our cosine-weighted hemisphere lobe sample direction
    return tangent * (radius * cos(angle)) + bitangent * (radius * sin(angle)) + hitNorm.xyz * sqrt(1 - randVal.x);
}

float ShootRay(float3 orig, float3 dir, float minT, float maxT)
{
    // Setup AO payload. Set the default to 0.0f (assume it's a hit)
    AORayPayload  rayPayload = { 0.0f };

    RayDesc rayAO = { orig, minT, dir, maxT };

    // Skip closest-hit shader and stop as soon as any intersection found
    uint rayFlags = RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH | RAY_FLAG_SKIP_CLOSEST_HIT_SHADER;

    TraceRay(RayTracingGlobalSrg::m_scene, rayFlags, 0xFF, 0, 1, 0, rayAO, rayPayload );

    return rayPayload.aoValue;
}

[shader("raygeneration")]
void AoRayGen()
{
    // Where this thread's ray is on the screen
    uint2 launchIndex = DispatchRaysIndex().xy;
    uint2 launchDim   = DispatchRaysDimensions().xy;

    // Initialize a random seed, per-pixel, based on a screen position and temporally varying count
    uint randSeed = InitRandomSeed(launchIndex.x + launchIndex.y * launchDim.x, RayTracingGlobalSrg::m_frameCount);

    // Get world position from screen position and depth
    float depth = RayTracingGlobalSrg::m_depth.Load(uint3(launchIndex, 0));
    // Position in native device coordinate
    float2 ndcPos = float2((float)launchIndex.x/(float)launchDim.x, 1.0f - (float)launchIndex.y/(float)launchDim.y) * 2.0f - 1.0f;
    float4 projectedPos = float4(ndcPos, depth, 1.0f);
    float4 worldPos = mul(RayTracingGlobalSrg::m_viewProjectionInverseMatrix, projectedPos);
    worldPos /= worldPos.w;

    float4 encodedNormal = RayTracingGlobalSrg::m_worldNormalMap.Load(uint3(launchIndex, 0));
    float3 worldNorm = DecodeNormalSignedOctahedron(encodedNormal.rgb);

    // Default ambient occlusion value if it hits the background
    float ambientOcclusion = 1.0f;

    // depth == 0 for background pixels; only shoot AO rays elsewhere
    if (depth != 0.0f)  
    {
        // Start accumulating from zero if we don't hit the background
        ambientOcclusion = 0.0f;

        for (int i = 0; i < RayTracingGlobalSrg::m_numRays; i++)
        {
            // Sample cosine-weighted hemisphere around surface normal to pick a random ray direction
            float3 worldDir = GetCosHemisphereSample(randSeed, worldNorm.xyz);

            // Shoot our ambient occlusion ray and update the value we'll output with the result
            float minT = RayTracingGlobalSrg::m_rayMinT;
            float maxT = RayTracingGlobalSrg::m_aoRadius;
            ambientOcclusion += ShootRay(worldPos.xyz, worldDir, minT, maxT);
        }
        ambientOcclusion = ambientOcclusion / float(RayTracingGlobalSrg::m_numRays);
    }
    
    // Save out AO color
    RayTracingGlobalSrg::m_outputAO[launchIndex].rgb = ambientOcclusion;
    RayTracingGlobalSrg::m_outputAO[launchIndex].a = 1.0f;
}