xash3d-fwgs/ref/vk/shaders/bounce.comp

#version 460 core
#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_nonuniform_qualifier : enable
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_ray_query: require

#include "debug.glsl"

#define GLSL
#include "ray_interop.h"
#undef GLSL

#define RAY_BOUNCE
#define RAY_QUERY
layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

layout(set = 0, binding = 1) uniform accelerationStructureEXT tlas;

layout(set = 0, binding = 2) uniform UBO { UniformBuffer ubo; } ubo;
layout(set = 0, binding = 6) uniform sampler2D textures[MAX_TEXTURES];
layout(set = 0, binding = 7) uniform samplerCube skybox;

layout (set = 0, binding = 8) readonly buffer SBOLights { LightsMetadata m; } lights;
layout (set = 0, binding = 9, align = 1) readonly buffer UBOLightClusters {
	LightCluster clusters_[MAX_LIGHT_CLUSTERS];
} light_grid;

layout(set = 0, binding = 10, rgba32f) uniform readonly image2D position_t;
layout(set = 0, binding = 11, rgba16f) uniform readonly image2D normals_gs;
layout(set = 0, binding = 12, rgba8) uniform readonly image2D material_rmxx;
layout(set = 0, binding = 13, rgba8) uniform readonly image2D base_color_a;

layout(set = 0, binding = 20, rgba16f) uniform writeonly image2D out_indirect_diffuse;
layout(set = 0, binding = 21, rgba16f) uniform writeonly image2D out_indirect_specular;

layout(set = 0, binding = 30, std430) readonly buffer ModelHeaders { ModelHeader a[]; } model_headers;
layout(set = 0, binding = 31, std430) readonly buffer Kusochki { Kusok a[]; } kusochki;
layout(set = 0, binding = 32, std430) readonly buffer Indices { uint16_t a[]; } indices;
layout(set = 0, binding = 33, std430) readonly buffer Vertices { Vertex a[]; } vertices;

#include "ray_primary_common.glsl"
#include "ray_primary_hit.glsl"
#include "noise.glsl"
#include "brdf.glsl"

#define LIGHT_POLYGON 1
#define LIGHT_POINT 1

#include "light.glsl"

#include "trace_simple_blending.glsl"

void readNormals(ivec2 uv, out vec3 geometry_normal, out vec3 shading_normal) {
	const vec4 n = imageLoad(normals_gs, uv);
	geometry_normal = normalDecode(n.xy);
	shading_normal = normalDecode(n.zw);
}

bool getHit(vec3 origin, vec3 direction, inout RayPayloadPrimary payload) {
	rayQueryEXT rq;
	const uint flags = 0
		| gl_RayFlagsCullFrontFacingTrianglesEXT
		//| gl_RayFlagsOpaqueEXT
		//| gl_RayFlagsTerminateOnFirstHitEXT
		//| gl_RayFlagsSkipClosestHitShaderEXT
		;
	const float L = 10000.; // TODO Why 10k?
	rayQueryInitializeEXT(rq, tlas, flags, GEOMETRY_BIT_OPAQUE | GEOMETRY_BIT_ALPHA_TEST, origin, 0., direction, L);
	while (rayQueryProceedEXT(rq)) {
		if (0 != (rayQueryGetRayFlagsEXT(rq) & gl_RayFlagsOpaqueEXT))
			continue;

		// alpha test
		// TODO check other possible ways of doing alpha test. They might be more efficient
		// (although in this particular primary ray case it's not taht important):
		// 1. Do a separate ray query for alpha masked geometry. Reason: here we might accidentally do the expensive
		//    texture sampling for geometry that's ultimately invisible (i.e. behind walls). Also, shader threads congruence.
		//    Separate pass could be more efficient as it'd be doing the same thing for every invocation.
		// 2. Same as the above, but also with a completely independent TLAS. Why: no need to mask-check geometry for opaque-vs-alpha
		const MiniGeometry geom = readCandidateMiniGeometry(rq);
		const uint tex_base_color = getKusok(geom.kusok_index).material.tex_base_color;
		// tex_base_color cannot be TEX_BASE_SKYBOX, as skybox is opaque
		const vec4 texture_color = texture(textures[nonuniformEXT(tex_base_color)], geom.uv);

		const float alpha_mask_threshold = .1f;
		if (texture_color.a >= alpha_mask_threshold) {
			rayQueryConfirmIntersectionEXT(rq);
		}
	}

	if (rayQueryGetIntersectionTypeEXT(rq, true) != gl_RayQueryCommittedIntersectionTriangleEXT)
		return false;

	primaryRayHit(rq, payload);
	//L = rayQueryGetIntersectionTEXT(rq, true);
	return true;
}
const int INDIRECT_SCALE = 2;

struct MaterialEx {
	MaterialProperties prop;
	vec3 geometry_normal, shading_normal;
	//TODO float alpha;
};

int max_bounces = 1;

ivec2 pix;
void computeBounces(MaterialEx mat, vec3 pos, vec3 direction, inout vec3 diffuse, inout vec3 specular) {
	const float ray_normal_fudge = .01;
	vec3 throughput = vec3(1.);
	vec3 bounce_direction;

	//const int max_bounces = 2;
	for (int i = 0; i < max_bounces; ++i) {
		// TODO blue noise
		const vec2 rnd = vec2(rand01(), rand01());
		const int brdf_type = brdfGetSample(rnd, mat.prop, -direction, mat.geometry_normal, mat.shading_normal/* TODO, mat.base_color_a.a*/, bounce_direction, throughput);

		if (brdf_type == BRDF_TYPE_NONE)
			return;

		RayPayloadPrimary payload;
		payload.base_color_a = vec4(0.);
		payload.emissive = vec4(0.);
		payload.material_rmxx = vec4(0.);
		payload.normals_gs = vec4(0.);

		vec3 lighting = vec3(0.);
		pos += mat.geometry_normal * ray_normal_fudge;
		const bool did_hit = getHit(pos, bounce_direction, payload);
		bool hit_skybox = false;
		if (did_hit) {
			const vec3 hit_pos = payload.hit_t.xyz;
			const vec3 hit_shading_normal = normalDecode(payload.normals_gs.zw);

			vec3 ldiffuse = vec3(0.);
			vec3 lspecular = vec3(0.);
			MaterialProperties hit_material;
			hit_material.base_color = payload.base_color_a.rgb;
			hit_material.metalness = payload.material_rmxx.g;
			hit_material.roughness = payload.material_rmxx.r;
			computeLighting(hit_pos, hit_shading_normal, -bounce_direction, hit_material, ldiffuse, lspecular);

			if (IS_INVALID3(ldiffuse)) {
#ifdef SHADER_DEBUG_ENABLE
				debugPrintfEXT("pix=(%d,%d) pos=(%f,%f,%f) dir=(%f,%f,%f) ldiffuse=invalid",
					pix.x, pix.y, pos.x, pos.y, pos.z, direction.x, direction.y, direction.z);
#endif
				ldiffuse = vec3(0.);
			}

			if (IS_INVALID3(lspecular)) {
#ifdef SHADER_DEBUG_ENABLE
				debugPrintfEXT("pix=(%d,%d) pos=(%f,%f,%f) dir=(%f,%f,%f) lspecular=invalid",
					pix.x, pix.y, pos.x, pos.y, pos.z, direction.x, direction.y, direction.z);
#endif
				lspecular = vec3(0.);
			}

			const vec3 background = mixFinalColor(payload.base_color_a.rgb, ldiffuse, lspecular, hit_material.metalness);
			// OLD
			//vec3 background = payload.base_color_a.rgb * ldiffuse;
			//background += lspecular * mix(vec3(1.), payload.base_color_a.rgb, hit_material.metalness);

			const vec4 blend = traceLegacyBlending(pos, bounce_direction, payload.hit_t.w);
			lighting = SRGBtoLINEAR(blend.rgb) + background * blend.a;

			// Prepare next bounce state
			pos = hit_pos;
			direction = bounce_direction;
			mat.prop = hit_material;
			mat.geometry_normal = normalDecode(payload.normals_gs.xy);
			mat.shading_normal = hit_shading_normal;

			hit_skybox = payload.hit_t.w < 0.;
		} else {
			lighting = texture(skybox, bounce_direction).rgb * ubo.ubo.skybox_exposure;
			hit_skybox = true;
		} // not hit

		/* TODO: also if last bounce (i.e. w/o direct light sampling) or */
		if (hit_skybox)
			lighting += payload.emissive.rgb;

		if (brdf_type == BRDF_TYPE_DIFFUSE)
			diffuse += throughput * lighting;
		else
			specular += throughput * lighting;

		if (IS_INVALID3(throughput)) {
#ifdef SHADER_DEBUG_ENABLE
			debugPrintfEXT("pix=(%d,%d) pos=(%f,%f,%f) dir=(%f,%f,%f) throughput=invalid",
				pix.x, pix.y, pos.x, pos.y, pos.z, direction.x, direction.y, direction.z);
#endif
			throughput = vec3(0.);
			//diffuse = 10.*vec3(1., 0., 0.);
			//specular = 10.*vec3(0., 0., 1.);
		} else {
			//diffuse = throughput;
		}

/*
		if (any(lessThan(throughput, vec3(0.)))) {
#ifdef SHADER_DEBUG_ENABLE
			debugPrintfEXT("pix=(%d,%d) pos=(%f,%f,%f) dir=(%f,%f,%f) throughput=(%f, %f, %f)",
				pix.x, pix.y, pos.x, pos.y, pos.z, direction.x, direction.y, direction.z,
				throughput.r, throughput.g, throughput.b
				);
#endif
		}
*/

		if (hit_skybox) {
			//diffuse = vec3(1.,1.,0.);
			break;
		} else {
			//diffuse = vec3(0.,1.,1.);
		}

		//if (should_include_emissive)
			//diffuse = lighting;// payload.emissive.rgb;
		//diffuse = bounce_direction * .5 + .5;
		//diffuse = throughput;

	} // for bounces
}


void main() {
	pix = ivec2(gl_GlobalInvocationID);
	const ivec2 res = ubo.ubo.res / INDIRECT_SCALE;
	if (any(greaterThanEqual(pix, res))) {
		return;
	}
	const vec2 uv = (gl_GlobalInvocationID.xy + .5) / res * 2. - 1.;

	const vec4 pos_t = imageLoad(position_t, pix * INDIRECT_SCALE);
	vec3 diffuse = vec3(0.), specular = vec3(0.);
	if (pos_t.w > 0.) {
		const vec3 origin    = (ubo.ubo.inv_view * vec4(0, 0, 0, 1)).xyz;
		const vec4 target    = ubo.ubo.inv_proj * vec4(uv.x, uv.y, 1, 1);
		const vec3 direction = normalize((ubo.ubo.inv_view * vec4(target.xyz, 0)).xyz);

		rand01_state = ubo.ubo.random_seed + pix.x * 1833 + pix.y * 31337 + 12;

		const vec4 material_data = imageLoad(material_rmxx, pix * INDIRECT_SCALE);
		const vec4 base_a = SRGBtoLINEAR(imageLoad(base_color_a, pix * INDIRECT_SCALE));

		vec3 geometry_normal, shading_normal;
		readNormals(pix * INDIRECT_SCALE, geometry_normal, shading_normal);

		MaterialEx mat;
		mat.prop.base_color = base_a.rgb;
		mat.prop.metalness = material_data.g;
		mat.prop.roughness = material_data.r;
		mat.geometry_normal = geometry_normal;
		mat.shading_normal = shading_normal;

		//max_bounces = (pix.x > res.x / 2.) ? 2 : 1;

		computeBounces(mat, pos_t.xyz, direction, diffuse, specular);
	}

	if (IS_INVALID3(specular)) {
#ifdef SHADER_DEBUG_ENABLE
		debugPrintfEXT("pix=(%d,%d) specular=inv", pix.x, pix.y);
#endif
		specular = vec3(0.);
	}

	if (IS_INVALID3(diffuse)) {
#ifdef SHADER_DEBUG_ENABLE
		debugPrintfEXT("pix=(%d,%d) diffuse=inv", pix.x, pix.y);
#endif
		diffuse = vec3(0.);
	}

	imageStore(out_indirect_diffuse, pix, vec4(diffuse, 0.f));
	imageStore(out_indirect_specular, pix, vec4(specular, 0.f));
}