Meshlet screenspace-derived tangents (#15084)

* Save 16 bytes per vertex by calculating tangents in the shader at
runtime, rather than storing them in the vertex data.
* Based on https://jcgt.org/published/0009/03/04,
https://www.jeremyong.com/graphics/2023/12/16/surface-gradient-bump-mapping.
* Fixed visbuffer resolve to use the updated algorithm that flips ddy
correctly
* Added some more docs about meshlet material limitations, and some
TODOs about transforming UV coordinates for the future.


![image](https://github.com/user-attachments/assets/222d8192-8c82-4d77-945d-53670a503761)

For testing add a normal map to the bunnies with StandardMaterial like
below, and then test that on both main and this PR (make sure to
download the correct bunny for each). Results should be mostly
identical.

```rust
normal_map_texture: Some(asset_server.load_with_settings(
    "textures/BlueNoise-Normal.png",
    |settings: &mut ImageLoaderSettings| settings.is_srgb = false,
)),
```

Cargo.toml

@@ -1146,7 +1146,7 @@ setup = [
  "curl",
  "-o",
  "assets/models/bunny.meshlet_mesh",
- "https://raw.githubusercontent.com/JMS55/bevy_meshlet_asset/e3da1533b4c69fb967f233c817e9b0921134d317/bunny.meshlet_mesh",
+ "https://raw.githubusercontent.com/JMS55/bevy_meshlet_asset/854eb98353ad94aea1104f355fc24dbe4fda679d/bunny.meshlet_mesh",
  ],
 ]
 

crates/bevy_pbr/src/material.rs

@@ -183,6 +183,8 @@ pub trait Material: Asset + AsBindGroup + Clone + Sized {
  /// the default meshlet mesh fragment shader will be used.
  ///
  /// This is part of an experimental feature, and is unnecessary to implement unless you are using `MeshletMesh`'s.
+ ///
+ /// See [`crate::meshlet::MeshletMesh`] for limitations.
  #[allow(unused_variables)]
  #[cfg(feature = "meshlet")]
  fn meshlet_mesh_fragment_shader() -> ShaderRef {
@@ -193,6 +195,8 @@ pub trait Material: Asset + AsBindGroup + Clone + Sized {
  /// the default meshlet mesh prepass fragment shader will be used.
  ///
  /// This is part of an experimental feature, and is unnecessary to implement unless you are using `MeshletMesh`'s.
+ ///
+ /// See [`crate::meshlet::MeshletMesh`] for limitations.
  #[allow(unused_variables)]
  #[cfg(feature = "meshlet")]
  fn meshlet_mesh_prepass_fragment_shader() -> ShaderRef {
@@ -203,6 +207,8 @@ pub trait Material: Asset + AsBindGroup + Clone + Sized {
  /// the default meshlet mesh deferred fragment shader will be used.
  ///
  /// This is part of an experimental feature, and is unnecessary to implement unless you are using `MeshletMesh`'s.
+ ///
+ /// See [`crate::meshlet::MeshletMesh`] for limitations.
  #[allow(unused_variables)]
  #[cfg(feature = "meshlet")]
  fn meshlet_mesh_deferred_fragment_shader() -> ShaderRef {

crates/bevy_pbr/src/meshlet/asset.rs

@@ -26,9 +26,14 @@ pub const MESHLET_MESH_ASSET_VERSION: u64 = 1;
 /// There are restrictions on the [`crate::Material`] functionality that can be used with this type of mesh.
 /// * Materials have no control over the vertex shader or vertex attributes.
 /// * Materials must be opaque. Transparent, alpha masked, and transmissive materials are not supported.
+/// * Do not use normal maps baked from higher-poly geometry. Use the high-poly geometry directly and skip the normal map.
+/// * If additional detail is needed, a smaller tiling normal map not baked from a mesh is ok.
+/// * Material shaders must not use builtin functions that automatically calculate derivatives <https://gpuweb.github.io/gpuweb/wgsl/#derivatives>.
+/// * Use `pbr_functions::sample_texture` to sample textures instead.
+/// * Performing manual arithmetic on texture coordinates (UVs) is forbidden. Use the chain-rule version of arithmetic functions instead (TODO: not yet implemented).
+/// * Limited control over [`bevy_render::render_resource::RenderPipelineDescriptor`] attributes.
 /// * Materials must use the [`crate::Material::meshlet_mesh_fragment_shader`] method (and similar variants for prepass/deferred shaders)
 /// which requires certain shader patterns that differ from the regular material shaders.
-/// * Limited control over [`bevy_render::render_resource::RenderPipelineDescriptor`] attributes.
 ///
 /// See also [`super::MaterialMeshletMeshBundle`] and [`super::MeshletPlugin`].
 #[derive(Asset, TypePath, Clone)]

crates/bevy_pbr/src/meshlet/from_mesh.rs

@@ -24,7 +24,7 @@ impl MeshletMesh {
  /// The input mesh must:
  /// 1. Use [`PrimitiveTopology::TriangleList`]
  /// 2. Use indices
- /// 3. Have the exact following set of vertex attributes: `{POSITION, NORMAL, UV_0, TANGENT}`
+ /// 3. Have the exact following set of vertex attributes: `{POSITION, NORMAL, UV_0}` (tangents can be used in material shaders, but are calculated at runtime and are not stored in the mesh)
  pub fn from_mesh(mesh: &Mesh) -> Result<Self, MeshToMeshletMeshConversionError> {
  // Validate mesh format
  let indices = validate_input_mesh(mesh)?;
@@ -152,7 +152,6 @@ fn validate_input_mesh(mesh: &Mesh) -> Result<Cow<'_, [u32]>, MeshToMeshletMeshC
  Mesh::ATTRIBUTE_POSITION.id,
  Mesh::ATTRIBUTE_NORMAL.id,
  Mesh::ATTRIBUTE_UV_0.id,
- Mesh::ATTRIBUTE_TANGENT.id,
  ]) {
  return Err(MeshToMeshletMeshConversionError::WrongMeshVertexAttributes);
  }
@@ -336,7 +335,7 @@ fn convert_meshlet_bounds(bounds: meshopt_Bounds) -> MeshletBoundingSphere {
 pub enum MeshToMeshletMeshConversionError {
  #[error("Mesh primitive topology is not TriangleList")]
  WrongMeshPrimitiveTopology,
- #[error("Mesh attributes are not {{POSITION, NORMAL, UV_0, TANGENT}}")]
+ #[error("Mesh attributes are not {{POSITION, NORMAL, UV_0}}")]
  WrongMeshVertexAttributes,
  #[error("Mesh has no indices")]
  MeshMissingIndices,

crates/bevy_pbr/src/meshlet/meshlet_bindings.wgsl

@@ -7,23 +7,19 @@
 struct PackedMeshletVertex {
  a: vec4<f32>,
  b: vec4<f32>,
- tangent: vec4<f32>,
 }
 
-// TODO: Octahedral encode normal, remove tangent and derive from UV derivatives
 struct MeshletVertex {
  position: vec3<f32>,
  normal: vec3<f32>,
  uv: vec2<f32>,
- tangent: vec4<f32>,
 }
 
 fn unpack_meshlet_vertex(packed: PackedMeshletVertex) -> MeshletVertex {
  var vertex: MeshletVertex;
  vertex.position = packed.a.xyz;
  vertex.normal = vec3(packed.a.w, packed.b.xy);
  vertex.uv = packed.b.zw;
- vertex.tangent = packed.tangent;
  return vertex;
 }
 

crates/bevy_pbr/src/meshlet/persistent_buffer_impls.rs

@@ -4,7 +4,7 @@ use super::{
 };
 use alloc::sync::Arc;
 
-const MESHLET_VERTEX_SIZE_IN_BYTES: u32 = 48;
+const MESHLET_VERTEX_SIZE_IN_BYTES: u32 = 32;
 
 impl PersistentGpuBufferable for Arc<[u8]> {
  type Metadata = ();

crates/bevy_pbr/src/meshlet/visibility_buffer_resolve.wgsl

@@ -13,8 +13,8 @@
  unpack_meshlet_vertex,
  },
  mesh_view_bindings::view,
- mesh_functions::{mesh_position_local_to_world, sign_determinant_model_3x3m},
- mesh_types::{Mesh, MESH_FLAGS_SIGN_DETERMINANT_MODEL_3X3_BIT},
+ mesh_functions::mesh_position_local_to_world,
+ mesh_types::Mesh,
  view_transformations::{position_world_to_clip, frag_coord_to_ndc},
 }
 #import bevy_render::maths::{affine3_to_square, mat2x4_f32_to_mat3x3_unpack}
@@ -37,14 +37,18 @@ struct PartialDerivatives {
  ddy: vec3<f32>,
 }
 
-// https://github.com/ConfettiFX/The-Forge/blob/2d453f376ef278f66f97cbaf36c0d12e4361e275/Examples_3/Visibility_Buffer/src/Shaders/FSL/visibilityBuffer_shade.frag.fsl#L83-L139
-fn compute_partial_derivatives(vertex_clip_positions: array<vec4<f32>, 3>, ndc_uv: vec2<f32>, screen_size: vec2<f32>) -> PartialDerivatives {
+// https://github.com/ConfettiFX/The-Forge/blob/9d43e69141a9cd0ce2ce2d2db5122234d3a2d5b5/Common_3/Renderer/VisibilityBuffer2/Shaders/FSL/vb_shading_utilities.h.fsl#L90-L150
+fn compute_partial_derivatives(vertex_world_positions: array<vec4<f32>, 3>, ndc_uv: vec2<f32>, half_screen_size: vec2<f32>) -> PartialDerivatives {
  var result: PartialDerivatives;
 
- let inv_w = 1.0 / vec3(vertex_clip_positions[0].w, vertex_clip_positions[1].w, vertex_clip_positions[2].w);
- let ndc_0 = vertex_clip_positions[0].xy * inv_w[0];
- let ndc_1 = vertex_clip_positions[1].xy * inv_w[1];
- let ndc_2 = vertex_clip_positions[2].xy * inv_w[2];
+ let vertex_clip_position_0 = position_world_to_clip(vertex_world_positions[0].xyz);
+ let vertex_clip_position_1 = position_world_to_clip(vertex_world_positions[1].xyz);
+ let vertex_clip_position_2 = position_world_to_clip(vertex_world_positions[2].xyz);
+
+ let inv_w = 1.0 / vec3(vertex_clip_position_0.w, vertex_clip_position_1.w, vertex_clip_position_2.w);
+ let ndc_0 = vertex_clip_position_0.xy * inv_w[0];
+ let ndc_1 = vertex_clip_position_1.xy * inv_w[1];
+ let ndc_2 = vertex_clip_position_2.xy * inv_w[2];
 
  let inv_det = 1.0 / determinant(mat2x2(ndc_2 - ndc_1, ndc_0 - ndc_1));
  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;
@@ -58,15 +62,18 @@ fn compute_partial_derivatives(vertex_clip_positions: array<vec4<f32>, 3>, ndc_u
  let interp_w = 1.0 / interp_inv_w;
 
  result.barycentrics = vec3(
- interp_w * (delta_v.x * result.ddx.x + delta_v.y * result.ddy.x + inv_w.x),
+ interp_w * (inv_w[0] + delta_v.x * result.ddx.x + delta_v.y * result.ddy.x),
  interp_w * (delta_v.x * result.ddx.y + delta_v.y * result.ddy.y),
  interp_w * (delta_v.x * result.ddx.z + delta_v.y * result.ddy.z),
  );
 
- result.ddx *= 2.0 / screen_size.x;
- result.ddy *= 2.0 / screen_size.y;
- ddx_sum *= 2.0 / screen_size.x;
- ddy_sum *= 2.0 / screen_size.y;
+ result.ddx *= half_screen_size.x;
+ result.ddy *= half_screen_size.y;
+ ddx_sum *= half_screen_size.x;
+ ddy_sum *= half_screen_size.y;
+
+ result.ddy *= -1.0;
+ ddy_sum *= -1.0;
 
  let interp_ddx_w = 1.0 / (interp_inv_w + ddx_sum);
  let interp_ddy_w = 1.0 / (interp_inv_w + ddy_sum);
@@ -117,30 +124,33 @@ fn resolve_vertex_output(frag_coord: vec4<f32>) -> VertexOutput {
  let world_position_2 = mesh_position_local_to_world(world_from_local, vec4(vertex_2.position, 1.0));
  let world_position_3 = mesh_position_local_to_world(world_from_local, vec4(vertex_3.position, 1.0));
 
- let clip_position_1 = position_world_to_clip(world_position_1.xyz);
- let clip_position_2 = position_world_to_clip(world_position_2.xyz);
- let clip_position_3 = position_world_to_clip(world_position_3.xyz);
  let frag_coord_ndc = frag_coord_to_ndc(frag_coord).xy;
  let partial_derivatives = compute_partial_derivatives(
- array(clip_position_1, clip_position_2, clip_position_3),
+ array(world_position_1, world_position_2, world_position_3),
  frag_coord_ndc,
- view.viewport.zw,
+ view.viewport.zw / 2.0,
  );
 
  let world_position = mat3x4(world_position_1, world_position_2, world_position_3) * partial_derivatives.barycentrics;
+ let world_positions_camera_relative = mat3x3(
+ world_position_1.xyz - view.world_position,
+ world_position_2.xyz - view.world_position,
+ world_position_3.xyz - view.world_position,
+ );
+ let ddx_world_position = world_positions_camera_relative * partial_derivatives.ddx;
+ let ddy_world_position = world_positions_camera_relative * partial_derivatives.ddy;
+
  let world_normal = mat3x3(
  normal_local_to_world(vertex_1.normal, &instance_uniform),
  normal_local_to_world(vertex_2.normal, &instance_uniform),
  normal_local_to_world(vertex_3.normal, &instance_uniform),
  ) * partial_derivatives.barycentrics;
+
  let uv = mat3x2(vertex_1.uv, vertex_2.uv, vertex_3.uv) * partial_derivatives.barycentrics;
  let ddx_uv = mat3x2(vertex_1.uv, vertex_2.uv, vertex_3.uv) * partial_derivatives.ddx;
  let ddy_uv = mat3x2(vertex_1.uv, vertex_2.uv, vertex_3.uv) * partial_derivatives.ddy;
- let world_tangent = mat3x4(
- tangent_local_to_world(vertex_1.tangent, world_from_local, instance_uniform.flags),
- tangent_local_to_world(vertex_2.tangent, world_from_local, instance_uniform.flags),
- tangent_local_to_world(vertex_3.tangent, world_from_local, instance_uniform.flags),
- ) * partial_derivatives.barycentrics;
+
+ let world_tangent = calculate_world_tangent(world_normal, ddx_world_position, ddy_world_position, ddx_uv, ddy_uv);
 
 #ifdef PREPASS_FRAGMENT
 #ifdef MOTION_VECTOR_PREPASS
@@ -184,20 +194,32 @@ fn normal_local_to_world(vertex_normal: vec3<f32>, instance_uniform: ptr<functio
  }
 }
 
-fn tangent_local_to_world(vertex_tangent: vec4<f32>, world_from_local: mat4x4<f32>, mesh_flags: u32) -> vec4<f32> {
- if any(vertex_tangent != vec4<f32>(0.0)) {
- return vec4<f32>(
- normalize(
- mat3x3<f32>(
- world_from_local[0].xyz,
- world_from_local[1].xyz,
- world_from_local[2].xyz,
- ) * vertex_tangent.xyz
- ),
- vertex_tangent.w * sign_determinant_model_3x3m(mesh_flags)
- );
- } else {
- return vertex_tangent;
+// 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
+fn calculate_world_tangent(
+ world_normal: vec3<f32>,
+ ddx_world_position: vec3<f32>,
+ ddy_world_position: vec3<f32>,
+ ddx_uv: vec2<f32>,
+ ddy_uv: vec2<f32>,
+) -> vec4<f32> {
+ // Project the position gradients onto the tangent plane
+ let ddx_world_position_s = ddx_world_position - dot(ddx_world_position, world_normal) * world_normal;
+ let ddy_world_position_s = ddy_world_position - dot(ddy_world_position, world_normal) * world_normal;
+
+ // Compute the jacobian matrix to leverage the chain rule
+ let jacobian_sign = sign(ddx_uv.x * ddy_uv.y - ddx_uv.y * ddy_uv.x);
+
+ var world_tangent = jacobian_sign * (ddy_uv.y * ddx_world_position_s - ddx_uv.y * ddy_world_position_s);
+
+ // The sign intrinsic returns 0 if the argument is 0
+ if jacobian_sign != 0.0 {
+ world_tangent = normalize(world_tangent);
  }
+
+ // The second factor here ensures a consistent handedness between
+ // the tangent frame and surface basis w.r.t. screenspace.
+ let w = jacobian_sign * sign(dot(ddy_world_position, cross(world_normal, ddx_world_position)));
+
+ return vec4(world_tangent, -w); // TODO: Unclear why we need to negate this to match mikktspace generated tangents
 }
 #endif

crates/bevy_pbr/src/render/pbr_fragment.wgsl

@@ -89,12 +89,14 @@ fn pbr_input_from_standard_material(
  bias.mip_bias = view.mip_bias;
 #endif // MESHLET_MESH_MATERIAL_PASS
 
+// TODO: Transforming UVs mean we need to apply derivative chain rule for meshlet mesh material pass
 #ifdef VERTEX_UVS
  let uv_transform = pbr_bindings::material.uv_transform;
 #ifdef VERTEX_UVS_A
  var uv = (uv_transform * vec3(in.uv, 1.0)).xy;
 #endif
 
+// TODO: Transforming UVs mean we need to apply derivative chain rule for meshlet mesh material pass
 #ifdef VERTEX_UVS_B
  var uv_b = (uv_transform * vec3(in.uv_b, 1.0)).xy;
 #else
@@ -104,12 +106,14 @@ fn pbr_input_from_standard_material(
 #ifdef VERTEX_TANGENTS
  if ((pbr_bindings::material.flags & pbr_types::STANDARD_MATERIAL_FLAGS_DEPTH_MAP_BIT) != 0u) {
  let V = pbr_input.V;
- let N = in.world_normal;
- let T = in.world_tangent.xyz;
- let B = in.world_tangent.w * cross(N, T);
+ let TBN = pbr_functions::calculate_tbn_mikktspace(in.world_normal, in.world_tangent);
+ let T = TBN[0];
+ let B = TBN[1];
+ let N = TBN[2];
  // Transform V from fragment to camera in world space to tangent space.
  let Vt = vec3(dot(V, T), dot(V, B), dot(V, N));
 #ifdef VERTEX_UVS_A
+ // TODO: Transforming UVs mean we need to apply derivative chain rule for meshlet mesh material pass
  uv = parallaxed_uv(
  pbr_bindings::material.parallax_depth_scale,
  pbr_bindings::material.max_parallax_layer_count,
@@ -123,6 +127,7 @@ fn pbr_input_from_standard_material(
 #endif
 
 #ifdef VERTEX_UVS_B
+ // TODO: Transforming UVs mean we need to apply derivative chain rule for meshlet mesh material pass
  uv_b = parallaxed_uv(
  pbr_bindings::material.parallax_depth_scale,
  pbr_bindings::material.max_parallax_layer_count,

crates/bevy_pbr/src/render/pbr_functions.wgsl

@@ -172,6 +172,14 @@ fn calculate_tbn_mikktspace(world_normal: vec3<f32>, world_tangent: vec4<f32>) -
  var T: vec3<f32> = world_tangent.xyz;
  var B: vec3<f32> = world_tangent.w * cross(N, T);
 
+#ifdef MESHLET_MESH_MATERIAL_PASS
+ // https://www.jeremyong.com/graphics/2023/12/16/surface-gradient-bump-mapping/#a-note-on-mikktspace-usage
+ let inverse_length_n = 1.0 / length(N);
+ T *= inverse_length_n;
+ B *= inverse_length_n;
+ N *= inverse_length_n;
+#endif
+
  return mat3x3(T, B, N);
 }
 

crates/bevy_pbr/src/render/pbr_prepass.wgsl

@@ -53,6 +53,7 @@ fn fragment(
 #ifdef VERTEX_TANGENTS
 #ifdef STANDARD_MATERIAL_NORMAL_MAP
 
+// TODO: Transforming UVs mean we need to apply derivative chain rule for meshlet mesh material pass
 #ifdef STANDARD_MATERIAL_NORMAL_MAP_UV_B
  let uv = (material.uv_transform * vec3(in.uv_b, 1.0)).xy;
 #else

examples/3d/meshlet.rs

@@ -17,7 +17,7 @@ use camera_controller::{CameraController, CameraControllerPlugin};
 use std::{f32::consts::PI, path::Path, process::ExitCode};
 
 const ASSET_URL: &str =
- "https://raw.githubusercontent.com/JMS55/bevy_meshlet_asset/e3da1533b4c69fb967f233c817e9b0921134d317/bunny.meshlet_mesh";
+ "https://raw.githubusercontent.com/JMS55/bevy_meshlet_asset/854eb98353ad94aea1104f355fc24dbe4fda679d/bunny.meshlet_mesh";
 
 fn main() -> ExitCode {
  if !Path::new("./assets/models/bunny.meshlet_mesh").exists() {