LumixEngine/data/pipelines/instancing.shd

include "pipelines/common.glsl"

compute_shader [[
	struct Indirect {
		uint vertex_count;
		uint instance_count;
		uint first_index;
		uint base_vertex;
		uint base_instance;
	};
	
	struct OutputData {
		vec4 rot_lod;
		vec4 pos_scale;
	};

	struct InputData {
		vec4 rot_lod;
		vec4 pos_scale;
	};

	layout(local_size_x = 256, local_size_y = 1) in;

	#ifdef UPDATE_LODS
		layout(binding = 0, std430) buffer InData {
			InputData b_input[];
		};
	#else
		layout(binding = 0, std430) readonly buffer InData {
			InputData b_input[];
		};
	#endif

	layout(binding = 1, std430) buffer OutData {
		uvec4 b_batch_offset;
		uvec4 b_lod_count;
		uvec4 b_lod_offset;
		OutputData b_output[];
	};

	layout(binding = 2, std430) writeonly buffer OutData2 {
		Indirect b_indirect[];
	};

	layout(std140, binding = 4) uniform UniformData {
		vec4 u_camera_offset;
		vec4 u_lod_distances;
		ivec4 u_lod_indices;
		uint u_indirect_offset;
		float u_radius;
		vec2 padding;
		vec4 u_camera_planes[6];
		uvec4 u_indices_count[32];
	};

	layout(std140, binding = 5) uniform UniformData2 {
		uint u_from_instance;
		uint u_instance_count;
	};

	shared uint group_count[5];

	uint getLOD(uint id) {
		vec3 p = b_input[id].pos_scale.xyz + u_camera_offset.xyz;
		float d = dot(p, p);
		if (d > u_lod_distances.w) return 4;
		else if (d > u_lod_distances.z) return 3;
		else if (d > u_lod_distances.y) return 2;
		else if (d > u_lod_distances.x) return 1;
		return 0;
	}

	bool cull(uint id) {
		float scale = b_input[id].pos_scale.w;
		vec4 cullp = vec4(b_input[id].pos_scale.xyz + u_camera_offset.xyz, 1);
		for (int i = 0; i < 6; ++i) {
			if (dot(u_camera_planes[i], cullp) < -u_radius * scale) {
				return false;
			}
		}
		return true;
	}

	void main() {
		uint id = gl_GlobalInvocationID.y * gl_WorkGroupSize.x + gl_GlobalInvocationID.x;
		#ifdef PASS0
			if (id == 0) {
				b_batch_offset.x += uint(dot(b_lod_count, vec4(1)));
				b_lod_offset = uvec4(b_batch_offset.x);
				b_lod_count = uvec4(0);
			}
		#elif defined PASS1
			bool master = gl_LocalInvocationID.x == 0;
			if (master) {
				group_count[0] = 0;
				group_count[1] = 0;
				group_count[2] = 0;
				group_count[3] = 0;
			}
			groupMemoryBarrier();
			barrier();

			if (id < u_instance_count) {
				id += u_from_instance;
				#ifdef UPDATE_LODS
					float dst_lod = getLOD(id);
					float src_lod = b_input[id].rot_lod.w;
					float d = dst_lod - src_lod;
					float td = Global.frame_time_delta * 2;
					float lod = abs(d) < td ? dst_lod : src_lod + td * sign(d);
					b_input[id].rot_lod.w = lod;
				#else
					float lod = b_input[id].rot_lod.w;
				#endif
				if (lod <= 3 && cull(id)) {
					float t = fract(lod);
					uint ilod = uint(lod);
					atomicAdd(group_count[ilod], 1);
					if (t > 0.01) {
						atomicAdd(group_count[ilod + 1], 1);
					}
				}
			}

			groupMemoryBarrier();
			barrier();

			if (master) {
				atomicAdd(b_lod_count.x, group_count[0]);
				atomicAdd(b_lod_count.y, group_count[1]);
				atomicAdd(b_lod_count.z, group_count[2]);
				atomicAdd(b_lod_count.w, group_count[3]);
				atomicAdd(b_lod_offset.y, group_count[0]);
				atomicAdd(b_lod_offset.z, group_count[1] + group_count[0]);
				atomicAdd(b_lod_offset.w, group_count[2] + group_count[1] + group_count[0]);
			}
		#elif defined PASS2
			int iid = int(id);
			if (iid > u_lod_indices.w) return;

			if (iid <= u_lod_indices.x) {
				b_indirect[id + u_indirect_offset].instance_count = b_lod_count.x;
				b_indirect[id + u_indirect_offset].base_instance = b_lod_offset.x;
			}
			else if (iid <= u_lod_indices.y) {
				b_indirect[id + u_indirect_offset].instance_count = b_lod_count.y;
				b_indirect[id + u_indirect_offset].base_instance = b_lod_offset.y;
			}
			else if (iid <= u_lod_indices.z) {
				b_indirect[id + u_indirect_offset].instance_count = b_lod_count.z;
				b_indirect[id + u_indirect_offset].base_instance = b_lod_offset.z;
			}
			else {
				b_indirect[id + u_indirect_offset].instance_count = b_lod_count.w;
				b_indirect[id + u_indirect_offset].base_instance = b_lod_offset.w;
			}

			b_indirect[id + u_indirect_offset].base_vertex = 0;
			b_indirect[id + u_indirect_offset].first_index = 0;
			b_indirect[id + u_indirect_offset].vertex_count = u_indices_count[id].x;
		#elif defined PASS3
			if (id >= u_instance_count) return;

			id += u_from_instance;
			InputData inp = b_input[id];

			float lod = inp.rot_lod.w;

			uint idx;
			if (lod > 3 || !cull(id)) return;

			float t = fract(lod);
			uint ilod = uint(lod);

			if (ilod == 0) {
				idx = atomicAdd(b_lod_offset.x, 1);
			}
			else if (ilod == 1) {
				idx = atomicAdd(b_lod_offset.y, 1);
			}
			else if (ilod == 2) {
				idx = atomicAdd(b_lod_offset.z, 1);
			}
			else if (ilod == 3) {
				idx = atomicAdd(b_lod_offset.w, 1);
			}
			else return;

			b_output[idx].rot_lod.xyz = inp.rot_lod.xyz;
			b_output[idx].rot_lod.w = t;
			b_output[idx].pos_scale = inp.pos_scale + vec4(u_camera_offset.xyz, 0);

			if (t > 0.01) {
				if (ilod == 0) {
					idx = atomicAdd(b_lod_offset.y, 1);
				}
				else if (ilod == 1) {
					idx = atomicAdd(b_lod_offset.z, 1);
				}
				else if (ilod == 2) {
					idx = atomicAdd(b_lod_offset.w, 1);
				}
				b_output[idx].rot_lod.xyz = inp.rot_lod.xyz;
				b_output[idx].rot_lod.w = t - 1;
				b_output[idx].pos_scale = inp.pos_scale + vec4(u_camera_offset.xyz, 0);
			}
		#elif defined UPDATE_LODS
			if (id < u_instance_count) {
				id += u_from_instance;
				b_input[id].rot_lod.w = getLOD(id);
			}
		#endif
	}
]]