BRDF LUT

#include <DirectXTex.h>

#include <algorithm>

#include <DirectXMath.h>


using namespace DirectX;
using namespace std;




#define IMAGE_SIZE 64
#define BRDF_SAMPLES 1024
#define FP_BITS 16


float RadicalInverse_VdC(unsigned int bits)
{
    bits = (bits << 16u) | (bits >> 16u);
    bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
    bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
    bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
    bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
    return float(bits) * 2.3283064365386963e-10; // / 0x100000000
}


XMFLOAT2 Hammersley(unsigned int i, unsigned int N)
{
    return XMFLOAT2(float(i) / float(N), RadicalInverse_VdC(i));
}


XMFLOAT3 ImportanceSampleGGX(XMFLOAT2 Xi, XMFLOAT3 N, float roughness)
{
    XMVECTOR N_simd = XMLoadFloat3(&N);

    float a = roughness * roughness;

    float phi = 2.0 * XM_PI * Xi.x;
    float cosTheta = sqrt((1.0 - Xi.y) / (1.0 + (a * a - 1.0) * Xi.y));
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);

    // from spherical coordinates to cartesian coordinates - halfway vector
    XMFLOAT3 H;
    H.x = cos(phi) * sinTheta;
    H.y = sin(phi) * sinTheta;
    H.z = cosTheta;

    // from tangent-space H vector to world-space sample vector
    XMFLOAT3 up = fabs(N.z) < 0.999 ? XMFLOAT3(0.0, 0.0, 1.0) : XMFLOAT3(1.0, 0.0, 0.0);
    XMVECTOR up_simd = XMLoadFloat3(&up);
    XMVECTOR tangent_simd = XMVector3Normalize(XMVector3Cross(up_simd, N_simd));
    XMVECTOR bitangent_simd = XMVector3Cross(N_simd, tangent_simd);

    XMVECTOR sampleVec_simd = tangent_simd * H.x + bitangent_simd * H.y + N_simd * H.z;
    sampleVec_simd = XMVector3Normalize(sampleVec_simd);

    XMFLOAT3 sampleVec;
    XMStoreFloat3(&sampleVec, sampleVec_simd);
    return sampleVec;
}


float GeometrySchlickGGX(float NdotV, float roughness)
{
    // note that we use a different k for IBL
    float a = roughness;
    float k = (a * a) / 2.0;

    float nom = NdotV;
    float denom = NdotV * (1.0 - k) + k;

    return nom / denom;
}


float GeometrySmith(XMFLOAT3 N, XMFLOAT3 V, XMFLOAT3 L, float roughness)
{
    XMVECTOR N_simd = XMLoadFloat3(&N);
    XMVECTOR V_simd = XMLoadFloat3(&V);
    XMVECTOR L_simd = XMLoadFloat3(&L);

    XMVECTOR NdotV_simd = XMVector3Dot(N_simd, V_simd);
    float NdotV;
    XMStoreFloat(&NdotV, NdotV_simd);
    NdotV = max(NdotV, 0.0f);
    XMVECTOR NdotL_simd = XMVector3Dot(N_simd, L_simd);
    float NdotL;
    XMStoreFloat(&NdotL, NdotL_simd);
    NdotL = max(NdotL, 0.0f);
    float ggx2 = GeometrySchlickGGX(NdotV, roughness);
    float ggx1 = GeometrySchlickGGX(NdotL, roughness);

    return ggx1 * ggx2;
}


int main(int argc, char* argv[])
{
    CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE);

    ScratchImage image;
    image.Initialize2D(DXGI_FORMAT_R32G32B32A32_FLOAT, IMAGE_SIZE, IMAGE_SIZE, 1, 1);

    for (size_t y = 0; y < IMAGE_SIZE; ++y)
    {
        for (size_t x = 0; x < IMAGE_SIZE; ++x)
        {
            float NdotV = (x + 0.5f) / static_cast<float>(IMAGE_SIZE);
            float roughness = (y + 0.5f) / static_cast<float>(IMAGE_SIZE);

            XMFLOAT3 V;
            V.x = sqrt(1.0 - NdotV * NdotV);
            V.y = 0.0;
            V.z = NdotV;
            XMVECTOR V_simd = XMLoadFloat3(&V);

            float A = 0.0;
            float B = 0.0;

            XMFLOAT3 N = XMFLOAT3(0.0, 0.0, 1.0);

            for (unsigned int i = 0; i < BRDF_SAMPLES; ++i)
            {
                XMFLOAT2 Xi = Hammersley(i, BRDF_SAMPLES);
                XMFLOAT3 H = ImportanceSampleGGX(Xi, N, roughness);
                XMVECTOR H_simd = XMLoadFloat3(&H);
                XMVECTOR L_simd = XMVector3Normalize(2.0 * XMVector3Dot(V_simd, H_simd) * H_simd - V_simd);
                XMFLOAT3 L;
                XMStoreFloat3(&L, L_simd);

                float NdotL = max(L.z, 0.0f);
                float NdotH = max(H.z, 0.0f);
                XMVECTOR VdotH_simd = XMVector3Dot(V_simd, H_simd);
                float VdotH;
                XMStoreFloat(&VdotH, VdotH_simd);
                VdotH = max(VdotH, 0.0f);

                if (NdotL > 0.0)
                {
                    float G = GeometrySmith(N, V, L, roughness);
                    float G_Vis = (G * VdotH) / (NdotH * NdotV);
                    float Fc = pow(1.0 - VdotH, 5.0);

                    A += (1.0 - Fc) * G_Vis;
                    B += Fc * G_Vis;
                }
            }
            A /= float(BRDF_SAMPLES);
            B /= float(BRDF_SAMPLES);

            float* pixelsPtr = reinterpret_cast<float*>(image.GetImage(0, 0, 0)->pixels) + (x + (IMAGE_SIZE - 1 - y) * IMAGE_SIZE) * 4;

            *pixelsPtr++ = A;
            *pixelsPtr++ = B;
            *pixelsPtr++ = 0.0f;
            *pixelsPtr++ = 1.0f;
        }
    }

    ScratchImage imageFinal;

    if (FP_BITS == 16)
    {
        Convert(image.GetImages(),
                image.GetImageCount(),
                image.GetMetadata(),
                DXGI_FORMAT_R16G16B16A16_FLOAT,
                TEX_FILTER_DEFAULT,
                TEX_THRESHOLD_DEFAULT,
                imageFinal);
    }
    else
    {
        imageFinal = move(image);
    }

    SaveToDDSFile(imageFinal.GetImages(), imageFinal.GetImageCount(), imageFinal.GetMetadata(), DDS_FLAGS_NONE, L"BrdfLut.dds");

    return EXIT_SUCCESS;
}