ao_fp64.comp

#version 420
#extension GL_ARB_compute_shader : enable

//
// aobench in OpenGL compute shader.
// Based on 
// http://glsl.heroku.com/e#3159.0
//

uniform float time;
uniform image2D destTex;
layout (local_size_x = 16, local_size_y = 16) in;

layout (binding=0, offset=0) uniform atomic_uint ac;


struct Ray
{
    dvec3 org;
    dvec3 dir;
};

struct Sphere
{
    dvec3 center;
    double radius;
};

struct Plane
{
    dvec3 p;
    dvec3 n;
};

struct Intersection
{
    double t;
    dvec3 p;     // hit point
    dvec3 n;     // normal
    int hit;
};

void shpere_intersect(Sphere s,  Ray ray, inout Intersection isect)
{
    // rs = ray.org - sphere.center
    dvec3 rs = ray.org - s.center;
    double B = dot(rs, ray.dir);
    double C = dot(rs, rs) - (s.radius * s.radius);
    double D = B * B - C;

    if (D > 0.0)
    {
        double t = -B - sqrt(D);
        if ( (t > 0.0) && (t < isect.t) )
        {
            isect.t = t;
            isect.hit = 1;

            // calculate normal.
            dvec3 p = dvec3(ray.org.x + ray.dir.x * t,
                          ray.org.y + ray.dir.y * t,
                          ray.org.z + ray.dir.z * t);
            dvec3 n = p - s.center;
            n = normalize(n);
            isect.n = n;
            isect.p = p;
        }
    }
}

void plane_intersect(Plane pl, Ray ray, inout Intersection isect)
{
    // d = -(p . n)
    // t = -(ray.org . n + d) / (ray.dir . n)
    double d = -dot(pl.p, pl.n);
    double v = dot(ray.dir, pl.n);

    if (abs(v) < 1.0e-6)
        return; // the plane is parallel to the ray.

    double t = -(dot(ray.org, pl.n) + d) / v;

    if ( (t > 0.0) && (t < isect.t) )
    {
        isect.hit = 1;
        isect.t   = t;
        isect.n   = pl.n;

        dvec3 p = dvec3(ray.org.x + t * ray.dir.x,
                      ray.org.y + t * ray.dir.y,
                      ray.org.z + t * ray.dir.z);
        isect.p = p;
    }
}

Sphere sphere[3];
Plane plane;
void Intersect(Ray r, inout Intersection i)
{
    for (int c = 0; c < 3; c++)
    {
        shpere_intersect(sphere[c], r, i);
    }
    plane_intersect(plane, r, i);
}

void orthoBasis(out dvec3 basis[3], dvec3 n)
{
    basis[2] = dvec3(n.x, n.y, n.z);
    basis[1] = dvec3(0.0, 0.0, 0.0);

    if ((n.x < 0.6) && (n.x > -0.6))
        basis[1].x = 1.0;
    else if ((n.y < 0.6) && (n.y > -0.6))
        basis[1].y = 1.0;
    else if ((n.z < 0.6) && (n.z > -0.6))
        basis[1].z = 1.0;
    else
        basis[1].x = 1.0;


    basis[0] = cross(basis[1], basis[2]);
    basis[0] = normalize(basis[0]);

    basis[1] = cross(basis[2], basis[0]);
    basis[1] = normalize(basis[1]);

}

//int seed = 0;
double random(double seed)
{
    double seed2 = (mod((seed)*1364.0+626.0, 509.0));
    return seed2/509.0;
}

dvec3 computeAO(inout Intersection isect)
{
    int i, j;
    int ntheta = 16;//irritating, guess I cant use a variable in a loop in webGL? Maybe I heard this already. -gtoledo
    int nphi   = 16;
    double eps  = 0.0001;

    // Slightly move ray org towards ray dir to avoid numerical probrem.
    dvec3 p = dvec3(isect.p.x + eps * isect.n.x,
                  isect.p.y + eps * isect.n.y,
                  isect.p.z + eps * isect.n.z);

    // Calculate orthogonal basis.
    dvec3 basis[3];
    orthoBasis(basis, isect.n);

    double occlusion = 0.0;

    for (int j = 0; j < 16; j++)//woulda used ntheta.-gtoledo
    {
        for (int i = 0; i < 16; i++)//woulda used nphi.-gtoledo
        {
            double s = isect.p.x += isect.p.y * 57. + isect.p.z * 21.;
            s = sin(cos(float(s)) * float(s))*.5+.5;
            // Pick a random ray direction with importance sampling.
            // p = cos(theta) / 3.141592
            double r = random(s*0.7);
            double phi = 2.0 * 3.141592 * random(s*3000.);

            dvec3 ref;
            ref.x = cos(float(phi)) * sqrt(1.0 - r);
            ref.y = sin(float(phi)) * sqrt(1.0 - r);
            ref.z = sqrt(r);

            // local -> global
            dvec3 rray;
            rray.x = ref.x * basis[0].x + ref.y * basis[1].x + ref.z * basis[2].x;
            rray.y = ref.x * basis[0].y + ref.y * basis[1].y + ref.z * basis[2].y;
            rray.z = ref.x * basis[0].z + ref.y * basis[1].z + ref.z * basis[2].z;

            dvec3 raydir = dvec3(rray.x, rray.y, rray.z);

            Ray ray;
            ray.org = p;
            ray.dir = raydir;

            Intersection occIsect;
            occIsect.hit = 0;
            occIsect.t = 1.0e+30;
            occIsect.n = occIsect.p = dvec3(0, 0, 0);
            Intersect(ray, occIsect);
            if (occIsect.hit != 0)
                occlusion += 1.0;
        }
    }

    // [0.0, 1.0]
    occlusion = (double(ntheta * nphi) - occlusion) / double(ntheta * nphi);

    //occlusion = (double(6 * 6) - occlusion) / double(6 * 6);
    return dvec3(occlusion, occlusion, occlusion);
}

void main() {
    ivec2 storePos = ivec2(gl_GlobalInvocationID.xy);

    dvec3    org=dvec3(0,0,0);
    vec2 resolution = vec2(gl_NumWorkGroups.xy * gl_WorkGroupSize.xy);
    vec2 pixel = -1.0 + 2.0 * vec2(gl_GlobalInvocationID.xy) / resolution;

    // compute ray origin and direction
    double asp = 1.0;
    dvec3 dir = normalize(dvec3(asp*pixel.x, pixel.y, -1.0));

    sphere[0].center = dvec3(-1.0+(sin(time*1.)), 1., -4.5+(cos(time*1.)));
    sphere[0].radius = 1.5;
    sphere[1].center = dvec3(1.5+(cos(time*1.)), .5, -3.5+(sin(time*1.)));
    sphere[1].radius = 1.;
    sphere[2].center = dvec3(.5+(cos(time*2.2)), 0.0, -3.2+(sin(time*2.3)));
    sphere[2].radius = 0.5;
    plane.p = dvec3(0,-0.5, 0);
    plane.n = dvec3(0, 1.0, 0);
    
    Intersection i;
    i.hit = 0;
    i.t = 1.0e+30;
    i.n = i.p = dvec3(0, 0, 0);
        
    Ray r;
    r.org = org;
    r.dir = normalize(dir);
    //double seed = (mod(dir.x * dir.y * 4525434.0, 65536.0));
    
    vec4 col = vec4(0,0,0,0);
    Intersect(r, i);
    if (i.hit != 0)
    {
        col.rgb = vec3(computeAO(i));
    }

    uint counter = atomicCounterIncrement(ac);
   
    imageStore(destTex, storePos, vec4(col.rgb, 0.0));
    //imageAtomicAdd(destTex, storePos, i.hit);
}