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

#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"

#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;
		// Should never happen: skybox is opaque if (tex_base_color == TEX_BASE_SKYBOX)
		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;

void computeBounce(ivec2 pix, vec3 direction, out vec3 diffuse, out vec3 specular) {
	diffuse = vec3(0.);
	specular = vec3(0.);

	const vec4 pos_t = imageLoad(position_t, pix);
	if (pos_t.w <= 0.)
		return;

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

	MaterialProperties material;
	material.baseColor = vec3(1.f);
	material.emissive = vec3(0.f);
	material.metalness = material_data.g;
	material.roughness = material_data.r;

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

	const float ray_normal_fudge = .01;
	vec3 throughput = vec3(1.);

	// 1. Make a "random" material-based ray for diffuse lighting
	vec3 bounce_direction = vec3(0.);
	int brdf_type = DIFFUSE_TYPE;

	const float alpha = (base_a.a);
	if (1. > alpha && rand01() > alpha) {
		brdf_type = SPECULAR_TYPE;
		// TODO: when not sampling randomly: throughput *= (1. - base_a.a);
		bounce_direction = normalize(refract(direction, geometry_normal, .95));
		geometry_normal = -geometry_normal;
		//throughput /= base_a.rgb;
	} else {
		if (material.metalness == 1.0f && material.roughness == 0.0f) {
			// Fast path for mirrors
			brdf_type = SPECULAR_TYPE;
		} else {
			// Decide whether to sample diffuse or specular BRDF (based on Fresnel term)
			const float brdf_probability = getBrdfProbability(material, -direction, shading_normal);
			if (rand01() < brdf_probability) {
				brdf_type = SPECULAR_TYPE;
				throughput /= brdf_probability;
			}
		}

		const vec2 u = vec2(rand01(), rand01());
		vec3 brdf_weight = vec3(0.);
		if (!evalIndirectCombinedBRDF(u, shading_normal, geometry_normal, -direction, material, brdf_type, bounce_direction, brdf_weight))
			return;
		throughput *= brdf_weight;
	}

	const float throughput_threshold = 1e-3;
	if (dot(throughput, throughput) < throughput_threshold)
		return;

	// 2. Rake yuri it, get hit
	// 3. Get relevant Geometry data
	RayPayloadPrimary payload;
	payload.base_color_a = vec4(0.);
	payload.emissive = vec4(0.);
	payload.material_rmxx = vec4(0.);

	vec3 lighting = vec3(0.);
	const vec3 pos = pos_t.xyz + geometry_normal * ray_normal_fudge;
	if (getHit(pos, bounce_direction, payload)) {
		// 4. Sample light sources
		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.baseColor = vec3(1.);
		hit_material.emissive = vec3(0.f);
		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);

		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;
	} else {
		lighting = texture(skybox, bounce_direction).rgb * ubo.ubo.skybox_exposure;
		//payload.emissive.rgb = texture(skybox, bounce_direction).rgb * ubo.ubo.skybox_exposure;
	}

	if (brdf_type == DIFFUSE_TYPE)
		diffuse += lighting;
	else
		specular += lighting + payload.emissive.rgb;
}


void main() {
	const ivec2 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 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;

	vec3 diffuse = vec3(0.), specular = vec3(0.);
	computeBounce(pix * INDIRECT_SCALE, direction, diffuse, specular);
	imageStore(out_indirect_diffuse, pix, vec4(diffuse, 0.f));
	imageStore(out_indirect_specular, pix, vec4(specular, 0.f));
}