gistfile1.txt

struct PickingParams
{
	OptixTraversableHandle sceneAccel;
	float4* accum_buffer;
	uchar4* frame_buffer;
	CUdeviceptr pickBuffer; // CUDA device pointer with room for at least two unsigned int.

	// 4-byte aligned
	uint32_t subframe_index;
	int picking;
	float3 rayOrigin;
	float3 rayDir;
	uint32_t screenX;
	uint32_t screenY;
	int32_t max_depth;

	OptixBufferView<OptixLight::Point> lights;
};

enum WhitedRayType
{
	WHITTED_RAY_TYPE_RADIANCE,
	WHITTED_RAY_TYPE_OCCLUSION,
	WHITTED_RAY_TYPE_COUNT,
};

struct WhittedPayload
{
	float3 result;
	float  importance;
	int    depth;
};

enum PickRayType
{
	PICK_RAY_TYPE_PICK,
	PICK_RAY_TYPE_PAD,
	PICK_RAY_TYPE_COUNT,
};

struct PickingPayload
{
	uint32_t meshID;
	uint32_t primitiveID;
};

const uint32_t NUM_PAYLOAD_VALUES = 6u;

extern "C" {
__constant__ PickingParams params;
}


static __forceinline__ __device__ void pick(
        OptixTraversableHandle      handle,
        float3                      ray_origin,
        float3                      ray_direction,
        float                       tmin,
        float                       tmax,
        PickingPayload*   pickPayload
        )
{
    uint32_t u0=0, u1=0;
    optixTrace(
            handle,
            ray_origin, ray_direction,
            tmin,
            tmax,
            0.0f, // rayTime
            OptixVisibilityMask( 1 ),
            OPTIX_RAY_FLAG_NONE,
            PICK_RAY_TYPE_PICK,     // SBT offset
            PICK_RAY_TYPE_COUNT,    // SBT stride
            PICK_RAY_TYPE_PICK,     // missSBTIndex
            u0, u1);

	pickPayload->meshID = u0;
	pickPayload->primitiveID = u1;
}

static __forceinline__ __device__ bool traceOcclusion(
        OptixTraversableHandle handle,
        float3                 ray_origin,
        float3                 ray_direction,
        float                  tmin,
        float                  tmax
        )
{
    uint32_t occluded = 0u;
    optixTrace(
            handle,
            ray_origin,
            ray_direction,
            tmin,
            tmax,
            0.0f,                    // rayTime
            OptixVisibilityMask( 1 ),
            OPTIX_RAY_FLAG_TERMINATE_ON_FIRST_HIT,
            WHITTED_RAY_TYPE_OCCLUSION,      // SBT offset
            WHITTED_RAY_TYPE_COUNT,          // SBT stride
            WHITTED_RAY_TYPE_OCCLUSION,      // missSBTIndex
            occluded );
    return occluded;
}

static __forceinline__ __device__ void traceRadiance(
        OptixTraversableHandle      handle,
        float3                      ray_origin,
        float3                      ray_direction,
        float                       tmin,
        float                       tmax,
        WhittedPayload*   whittedPayload
        )
{
    uint32_t u0=0, u1=0, u2=0, u3=0;
    optixTrace(
            handle,
            ray_origin, ray_direction,
            tmin,
            tmax,
            0.0f,                     // rayTime
            OptixVisibilityMask( 1 ),
            OPTIX_RAY_FLAG_NONE,
            WHITTED_RAY_TYPE_RADIANCE,        // SBT offset
            WHITTED_RAY_TYPE_COUNT,           // SBT stride
            WHITTED_RAY_TYPE_RADIANCE,        // missSBTIndex
            u0, u1, u2, u3 );

     whittedPayload->result.x = __int_as_float( u0 );
     whittedPayload->result.y = __int_as_float( u1 );
     whittedPayload->result.z = __int_as_float( u2 );
     whittedPayload->depth = u3;
}


__forceinline__ __device__ void setPayloadResult( float3 p )
{
	// payload slots 0 and 1 are used for picking
    optixSetPayload_2( float_as_int( p.x ) );
    optixSetPayload_3( float_as_int( p.y ) );
    optixSetPayload_4( float_as_int( p.z ) );
}

__forceinline__ __device__ uchar4 make_color( const float3&  c )
{
    return make_uchar4(
            static_cast<uint8_t>( clamp( c.x, 0.0f, 1.0f ) *255.0f ),
            static_cast<uint8_t>( clamp( c.y, 0.0f, 1.0f ) *255.0f ),
            static_cast<uint8_t>( clamp( c.z, 0.0f, 1.0f ) *255.0f ),
            255u
            );
}


extern "C" __global__ void __raygen__rg()
{
	const uint3 launch_idx = optixGetLaunchIndex();
    const uint3 launch_dims = optixGetLaunchDimensions();

   // pick when picking is on and the application mouse screen coords match the launch index
  if(params.picking == 1 && params.screenX == launch_idx.x && params.screenY == launch_idx.y)
  {
		// trace pick ray
		PickingPayload pickPayLoad;
		pickPayLoad.meshID = ~0;
		pickPayLoad.primitiveID = ~0;
	   
		pick(
				params.sceneAccel,
				params.rayOrigin,
				params.rayDir,
				0.01f,  
				1e16f,  
				&pickPayLoad );
				
		unsigned int* pickData = reinterpret_cast<unsigned int*>(params.pickBuffer);

		pickData[0] = pickPayLoad.meshID;
		pickData[1] = pickPayLoad.primitiveID;
	  
		printf("Picked mesh: %i\n", pickData[0] ); 
		printf("Picked triangle:  %i\n", pickData[1]); 
		//return;
	}
	
    const RayGenData* rtData = (RayGenData*)optixGetSbtDataPointer();
    const float3      U      = rtData->camera_u;
    const float3      V      = rtData->camera_v;
    const float3      W      = rtData->camera_w;
    
	const int subframe_index  = params.subframe_index;
	uint32_t seed = tea<4>( launch_idx.y*launch_dims.x + launch_idx.x, subframe_index );

    const float2 subpixel_jitter = subframe_index == 0 ?
									make_float2( 0.0f, 0.0f ) :
									make_float2( rnd( seed )-0.5f, rnd( seed )-0.5f );

    const float2 d = 2.0f * make_float2(
            ( static_cast<float>( launch_idx.x ) + subpixel_jitter.x ) / static_cast<float>( launch_dims.x ),
            ( static_cast<float>( launch_idx.y ) + subpixel_jitter.y ) / static_cast<float>( launch_dims.y )
            ) - 1.0f;
    const float3 ray_direction = normalize(d.x*U + d.y*V + W);
    const float3 ray_origin    = rtData->cam_eye;
	
    // Trace camera ray
    WhittedPayload whittedPayload;
    whittedPayload.result = make_float3(0.0f);
    whittedPayload.importance = 1.0f;
    whittedPayload.depth = 0.0f;
	
    traceRadiance(
            params.sceneAccel,
            ray_origin,
            ray_direction,
            0.01f,  // tmin       // TODO: smarter offset
            1e16f,  // tmax
            &whittedPayload );

    // Update results
    const uint32_t image_index  = launch_idx.y * launch_dims.x + launch_idx.x;
    float3 accum_color  = whittedPayload.result;

	if( subframe_index > 0 )
	{
		const float a = 1.0f / static_cast<float>( subframe_index+1 );
		const float3 accum_color_prev = make_float3( params.accum_buffer[ image_index ]);
		accum_color = lerp( accum_color_prev, accum_color, a );
	}
	params.accum_buffer[ image_index ] = make_float4( accum_color, 1.0f);
	params.frame_buffer[ image_index ] = make_color ( accum_color );
	
	
}

extern "C" __global__ void __miss__ms()
{
	if(params.picking == 1)
	{
		optixSetPayload_0( static_cast<unsigned int>( ~0 ) );
		optixSetPayload_1( static_cast<unsigned int>( ~0 ) );
		//return;
	}
	
    MissData* rt_data  = reinterpret_cast<MissData*>( optixGetSbtDataPointer() );
    setPayloadResult( make_float3( rt_data->r, rt_data->g, rt_data->b ) );
}

__forceinline__ __device__ void setPayloadOcclusion( bool occluded )
{
    optixSetPayload_2( static_cast<uint32_t>( occluded ) );
}


extern "C" __global__ void __closesthit__occlusion()
{
    setPayloadOcclusion( true );
}


extern "C" __global__ void __closesthit__ch()
{
	if(params.picking == 1)
	{
		unsigned int primID = optixGetPrimitiveIndex();
		 unsigned int instID = optixGetInstanceId();
		 
		 optixSetPayload_0( static_cast<unsigned int>( instID ) );
		 optixSetPayload_1( static_cast<unsigned int>( primID ) );
			//return;
	}

	const HitGroupSBT* hit_group_data = reinterpret_cast<HitGroupSBT*>( optixGetSbtDataPointer() );
	if( hit_group_data == nullptr) return;
	
    const LocalGeometry geom = getLocalGeometry( hit_group_data->geometry_data );

 //
    // Retrieve material data
    //
    float3 base_color = make_float3( hit_group_data->material_data.pbr.base_color );
    if( hit_group_data->material_data.pbr.base_color_tex )
        base_color *= linearize( make_float3(
                    tex2D<float4>( hit_group_data->material_data.pbr.base_color_tex, geom.UV.x, geom.UV.y )
                    ) );

    float metallic  = hit_group_data->material_data.pbr.metallic;
    float roughness = hit_group_data->material_data.pbr.roughness;
    float4 mr_tex = make_float4( 1.0f );
    if( hit_group_data->material_data.pbr.metallic_roughness_tex )
        // MR tex is (occlusion, roughness, metallic )
        mr_tex = tex2D<float4>( hit_group_data->material_data.pbr.metallic_roughness_tex, geom.UV.x, geom.UV.y );
    roughness *= mr_tex.y;
    metallic  *= mr_tex.z;
	
    // Convert to material params
    const float  F0         = 0.04f;
    const float3 diff_color = base_color*( 1.0f - F0 )*( 1.0f - metallic );
    const float3 spec_color = lerp( make_float3( F0 ), base_color, metallic );
    const float  alpha      = roughness*roughness;
    
    // compute direct lighting
    float3 N = geom.N;
    if( hit_group_data->material_data.pbr.normal_tex )
    {
        const float4 NN = 2.0f*tex2D<float4>( hit_group_data->material_data.pbr.normal_tex, geom.UV.x, geom.UV.y ) - make_float4(1.0f);
        N = normalize( NN.x*normalize( geom.dpdu ) + NN.y*normalize( geom.dpdv ) + NN.z*geom.N );
    }

    float3 result = make_float3( 0.0f );

    for( int i = 0; i < params.lights.count; ++i )
    {	
        OptixLight::Point light = params.lights[i];

        const float  L_dist  = length( light.position - geom.P );
        const float3 L       = ( light.position - geom.P ) / L_dist;
        const float3 V       = -normalize( optixGetWorldRayDirection() );
        const float3 H       = normalize( L + V );
        const float  N_dot_L = dot( N, L );
        const float  N_dot_V = dot( N, V );
        const float  N_dot_H = dot( N, H );
        const float  V_dot_H = dot( V, H );

        if( N_dot_L > 0.0f && N_dot_V > 0.0f )
        {
            const float tmin     = 0.001f;          // TODO
            const float tmax     = L_dist - 0.001f; // TODO
            const bool  occluded = traceOcclusion( params.sceneAccel, geom.P, L, tmin, tmax );
            if( !occluded )
            {
                const float3 F     = schlick( spec_color, V_dot_H );
                const float  G_vis = vis( N_dot_L, N_dot_V, alpha );
                const float  D     = ggxNormal( N_dot_H, alpha );

                const float3 diff = ( 1.0f - F )*diff_color / M_PIf;
                const float3 spec = F*G_vis*D;

                result += light.color*light.intensity*N_dot_L*( diff + spec );
            }
        }
    }
	
    setPayloadResult( result );
}