1#define_import_path bevy_pbr::meshlet_visibility_buffer_resolve
2
3#import bevy_pbr::{
4    meshlet_bindings::{
5        Meshlet,
6        meshlet_visibility_buffer,
7        meshlet_raster_clusters,
8        meshlets,
9        meshlet_instance_uniforms,
10        get_meshlet_vertex_id,
11        get_meshlet_vertex_position,
12        get_meshlet_vertex_normal,
13        get_meshlet_vertex_uv,
14    },
15    mesh_view_bindings::view,
16    mesh_functions::mesh_position_local_to_world,
17    mesh_types::Mesh,
18    view_transformations::{position_world_to_clip, frag_coord_to_ndc},
19}
20#import bevy_render::maths::{affine3_to_square, mat2x4_f32_to_mat3x3_unpack}
21
22#ifdef PREPASS_FRAGMENT
23#ifdef MOTION_VECTOR_PREPASS
24#import bevy_pbr::{
25    prepass_bindings::previous_view_uniforms,
26    pbr_prepass_functions::calculate_motion_vector,
27}
28#endif
29#endif
30
31/// Functions to be used by materials for reading from a meshlet visibility buffer texture.
32
33#ifdef MESHLET_MESH_MATERIAL_PASS
34struct PartialDerivatives {
35    barycentrics: vec3<f32>,
36    ddx: vec3<f32>,
37    ddy: vec3<f32>,
38}
39
40// https://github.com/ConfettiFX/The-Forge/blob/9d43e69141a9cd0ce2ce2d2db5122234d3a2d5b5/Common_3/Renderer/VisibilityBuffer2/Shaders/FSL/vb_shading_utilities.h.fsl#L90-L150
41fn compute_partial_derivatives(vertex_world_positions: array<vec4<f32>, 3>, ndc_uv: vec2<f32>, half_screen_size: vec2<f32>) -> PartialDerivatives {
42    var result: PartialDerivatives;
43
44    let vertex_clip_position_0 = position_world_to_clip(vertex_world_positions[0].xyz);
45    let vertex_clip_position_1 = position_world_to_clip(vertex_world_positions[1].xyz);
46    let vertex_clip_position_2 = position_world_to_clip(vertex_world_positions[2].xyz);
47
48    let inv_w = 1.0 / vec3(vertex_clip_position_0.w, vertex_clip_position_1.w, vertex_clip_position_2.w);
49    let ndc_0 = vertex_clip_position_0.xy * inv_w[0];
50    let ndc_1 = vertex_clip_position_1.xy * inv_w[1];
51    let ndc_2 = vertex_clip_position_2.xy * inv_w[2];
52
53    let inv_det = 1.0 / determinant(mat2x2(ndc_2 - ndc_1, ndc_0 - ndc_1));
54    result.ddx = vec3(ndc_1.y - ndc_2.y, ndc_2.y - ndc_0.y, ndc_0.y - ndc_1.y) * inv_det * inv_w;
55    result.ddy = vec3(ndc_2.x - ndc_1.x, ndc_0.x - ndc_2.x, ndc_1.x - ndc_0.x) * inv_det * inv_w;
56
57    var ddx_sum = dot(result.ddx, vec3(1.0));
58    var ddy_sum = dot(result.ddy, vec3(1.0));
59
60    let delta_v = ndc_uv - ndc_0;
61    let interp_inv_w = inv_w.x + delta_v.x * ddx_sum + delta_v.y * ddy_sum;
62    let interp_w = 1.0 / interp_inv_w;
63
64    result.barycentrics = vec3(
65        interp_w * (inv_w[0] + delta_v.x * result.ddx.x + delta_v.y * result.ddy.x),
66        interp_w * (delta_v.x * result.ddx.y + delta_v.y * result.ddy.y),
67        interp_w * (delta_v.x * result.ddx.z + delta_v.y * result.ddy.z),
68    );
69
70    result.ddx *= half_screen_size.x;
71    result.ddy *= half_screen_size.y;
72    ddx_sum *= half_screen_size.x;
73    ddy_sum *= half_screen_size.y;
74
75    result.ddy *= -1.0;
76    ddy_sum *= -1.0;
77
78    let interp_ddx_w = 1.0 / (interp_inv_w + ddx_sum);
79    let interp_ddy_w = 1.0 / (interp_inv_w + ddy_sum);
80
81    result.ddx = interp_ddx_w * (result.barycentrics * interp_inv_w + result.ddx) - result.barycentrics;
82    result.ddy = interp_ddy_w * (result.barycentrics * interp_inv_w + result.ddy) - result.barycentrics;
83    return result;
84}
85
86struct VertexOutput {
87    position: vec4<f32>,
88    world_position: vec4<f32>,
89    world_normal: vec3<f32>,
90    uv: vec2<f32>,
91    ddx_uv: vec2<f32>,
92    ddy_uv: vec2<f32>,
93    world_tangent: vec4<f32>,
94    mesh_flags: u32,
95    cluster_id: u32,
96    material_bind_group_slot: u32,
97#ifdef PREPASS_FRAGMENT
98#ifdef MOTION_VECTOR_PREPASS
99    motion_vector: vec2<f32>,
100#endif
101#endif
102}
103
104/// Load the visibility buffer texture and resolve it into a VertexOutput.
105fn resolve_vertex_output(frag_coord: vec4<f32>) -> VertexOutput {
106    let packed_ids = u32(textureLoad(meshlet_visibility_buffer, vec2<u32>(frag_coord.xy)).r);
107    let cluster_id = packed_ids >> 7u;
108    let instanced_offset = meshlet_raster_clusters[cluster_id];
109    let meshlet_id = instanced_offset.offset;
110    var meshlet = meshlets[meshlet_id];
111
112    let triangle_id = extractBits(packed_ids, 0u, 7u);
113    let index_ids = meshlet.start_index_id + (triangle_id * 3u) + vec3(0u, 1u, 2u);
114    let vertex_ids = vec3(get_meshlet_vertex_id(index_ids[0]), get_meshlet_vertex_id(index_ids[1]), get_meshlet_vertex_id(index_ids[2]));
115    let vertex_0 = load_vertex(&meshlet, vertex_ids[0]);
116    let vertex_1 = load_vertex(&meshlet, vertex_ids[1]);
117    let vertex_2 = load_vertex(&meshlet, vertex_ids[2]);
118
119    let instance_id = instanced_offset.instance_id;
120    var instance_uniform = meshlet_instance_uniforms[instance_id];
121
122    let world_from_local = affine3_to_square(instance_uniform.world_from_local);
123    let world_position_0 = mesh_position_local_to_world(world_from_local, vec4(vertex_0.position, 1.0));
124    let world_position_1 = mesh_position_local_to_world(world_from_local, vec4(vertex_1.position, 1.0));
125    let world_position_2 = mesh_position_local_to_world(world_from_local, vec4(vertex_2.position, 1.0));
126
127    let frag_coord_ndc = frag_coord_to_ndc(frag_coord).xy;
128    let partial_derivatives = compute_partial_derivatives(
129        array(world_position_0, world_position_1, world_position_2),
130        frag_coord_ndc,
131        view.viewport.zw / 2.0,
132    );
133
134    let world_position = mat3x4(world_position_0, world_position_1, world_position_2) * partial_derivatives.barycentrics;
135    let world_positions_camera_relative = mat3x3(
136        world_position_0.xyz - view.world_position,
137        world_position_1.xyz - view.world_position,
138        world_position_2.xyz - view.world_position,
139    );
140    let ddx_world_position = world_positions_camera_relative * partial_derivatives.ddx;
141    let ddy_world_position = world_positions_camera_relative * partial_derivatives.ddy;
142
143    let world_normal = mat3x3(
144        normal_local_to_world(vertex_0.normal, &instance_uniform),
145        normal_local_to_world(vertex_1.normal, &instance_uniform),
146        normal_local_to_world(vertex_2.normal, &instance_uniform),
147    ) * partial_derivatives.barycentrics;
148
149    let uv = mat3x2(vertex_0.uv, vertex_1.uv, vertex_2.uv) * partial_derivatives.barycentrics;
150    let ddx_uv = mat3x2(vertex_0.uv, vertex_1.uv, vertex_2.uv) * partial_derivatives.ddx;
151    let ddy_uv = mat3x2(vertex_0.uv, vertex_1.uv, vertex_2.uv) * partial_derivatives.ddy;
152
153    let world_tangent = calculate_world_tangent(world_normal, ddx_world_position, ddy_world_position, ddx_uv, ddy_uv);
154
155#ifdef PREPASS_FRAGMENT
156#ifdef MOTION_VECTOR_PREPASS
157    let previous_world_from_local = affine3_to_square(instance_uniform.previous_world_from_local);
158    let previous_world_position_0 = mesh_position_local_to_world(previous_world_from_local, vec4(vertex_0.position, 1.0));
159    let previous_world_position_1 = mesh_position_local_to_world(previous_world_from_local, vec4(vertex_1.position, 1.0));
160    let previous_world_position_2 = mesh_position_local_to_world(previous_world_from_local, vec4(vertex_2.position, 1.0));
161    let previous_world_position = mat3x4(previous_world_position_0, previous_world_position_1, previous_world_position_2) * partial_derivatives.barycentrics;
162    let motion_vector = calculate_motion_vector(world_position, previous_world_position);
163#endif
164#endif
165
166    return VertexOutput(
167        frag_coord,
168        world_position,
169        world_normal,
170        uv,
171        ddx_uv,
172        ddy_uv,
173        world_tangent,
174        instance_uniform.flags,
175        instance_id ^ meshlet_id,
176        instance_uniform.material_and_lightmap_bind_group_slot & 0xffffu,
177#ifdef PREPASS_FRAGMENT
178#ifdef MOTION_VECTOR_PREPASS
179        motion_vector,
180#endif
181#endif
182    );
183}
184
185struct MeshletVertex {
186    position: vec3<f32>,
187    normal: vec3<f32>,
188    uv: vec2<f32>,
189}
190
191fn load_vertex(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> MeshletVertex {
192    return MeshletVertex(
193        get_meshlet_vertex_position(meshlet, vertex_id),
194        get_meshlet_vertex_normal(meshlet, vertex_id),
195        get_meshlet_vertex_uv(meshlet, vertex_id),
196    );
197}
198
199fn normal_local_to_world(vertex_normal: vec3<f32>, instance_uniform: ptr<function, Mesh>) -> vec3<f32> {
200    if any(vertex_normal != vec3<f32>(0.0)) {
201        return normalize(
202            mat2x4_f32_to_mat3x3_unpack(
203                (*instance_uniform).local_from_world_transpose_a,
204                (*instance_uniform).local_from_world_transpose_b,
205            ) * vertex_normal
206        );
207    } else {
208        return vertex_normal;
209    }
210}
211
212// https://www.jeremyong.com/graphics/2023/12/16/surface-gradient-bump-mapping/#surface-gradient-from-a-tangent-space-normal-vector-without-an-explicit-tangent-basis
213fn calculate_world_tangent(
214    world_normal: vec3<f32>,
215    ddx_world_position: vec3<f32>,
216    ddy_world_position: vec3<f32>,
217    ddx_uv: vec2<f32>,
218    ddy_uv: vec2<f32>,
219) -> vec4<f32> {
220    // Project the position gradients onto the tangent plane
221    let ddx_world_position_s = ddx_world_position - dot(ddx_world_position, world_normal) * world_normal;
222    let ddy_world_position_s = ddy_world_position - dot(ddy_world_position, world_normal) * world_normal;
223
224    // Compute the jacobian matrix to leverage the chain rule
225    let jacobian_sign = sign(ddx_uv.x * ddy_uv.y - ddx_uv.y * ddy_uv.x);
226
227    var world_tangent = jacobian_sign * (ddy_uv.y * ddx_world_position_s - ddx_uv.y * ddy_world_position_s);
228
229    // The sign intrinsic returns 0 if the argument is 0
230    if jacobian_sign != 0.0 {
231        world_tangent = normalize(world_tangent);
232    }
233
234    // The second factor here ensures a consistent handedness between
235    // the tangent frame and surface basis w.r.t. screenspace.
236    let w = jacobian_sign * sign(dot(ddy_world_position, cross(world_normal, ddx_world_position)));
237
238    return vec4(world_tangent, -w); // TODO: Unclear why we need to negate this to match mikktspace generated tangents
239}
240#endif