bevy_pbr/meshlet/
asset.rs

1use alloc::sync::Arc;
2use bevy_asset::{
3    io::{Reader, Writer},
4    saver::{AssetSaver, SavedAsset},
5    Asset, AssetLoader, AsyncReadExt, AsyncWriteExt, LoadContext,
6};
7use bevy_math::{Vec2, Vec3};
8use bevy_reflect::TypePath;
9use bevy_render::render_resource::ShaderType;
10use bevy_tasks::block_on;
11use bytemuck::{Pod, Zeroable};
12use lz4_flex::frame::{FrameDecoder, FrameEncoder};
13use std::io::{Read, Write};
14use thiserror::Error;
15
16/// Unique identifier for the [`MeshletMesh`] asset format.
17const MESHLET_MESH_ASSET_MAGIC: u64 = 1717551717668;
18
19/// The current version of the [`MeshletMesh`] asset format.
20pub const MESHLET_MESH_ASSET_VERSION: u64 = 2;
21
22/// A mesh that has been pre-processed into multiple small clusters of triangles called meshlets.
23///
24/// A [`bevy_mesh::Mesh`] can be converted to a [`MeshletMesh`] using `MeshletMesh::from_mesh` when the `meshlet_processor` cargo feature is enabled.
25/// The conversion step is very slow, and is meant to be ran once ahead of time, and not during runtime. This type of mesh is not suitable for
26/// dynamically generated geometry.
27///
28/// There are restrictions on the [`crate::Material`] functionality that can be used with this type of mesh.
29/// * Materials have no control over the vertex shader or vertex attributes.
30/// * Materials must be opaque. Transparent, alpha masked, and transmissive materials are not supported.
31/// * Do not use normal maps baked from higher-poly geometry. Use the high-poly geometry directly and skip the normal map.
32///   * If additional detail is needed, a smaller tiling normal map not baked from a mesh is ok.
33/// * Material shaders must not use builtin functions that automatically calculate derivatives <https://gpuweb.github.io/gpuweb/wgsl/#derivatives>.
34///   * Performing manual arithmetic on texture coordinates (UVs) is forbidden. Use the chain-rule version of arithmetic functions instead (TODO: not yet implemented).
35/// * Limited control over [`bevy_render::render_resource::RenderPipelineDescriptor`] attributes.
36/// * Materials must use the [`crate::Material::meshlet_mesh_fragment_shader`] method (and similar variants for prepass/deferred shaders)
37///   which requires certain shader patterns that differ from the regular material shaders.
38///
39/// See also [`super::MeshletMesh3d`] and [`super::MeshletPlugin`].
40#[derive(Asset, TypePath, Clone)]
41pub struct MeshletMesh {
42    /// Quantized and bitstream-packed vertex positions for meshlet vertices.
43    pub(crate) vertex_positions: Arc<[u32]>,
44    /// Octahedral-encoded and 2x16snorm packed normals for meshlet vertices.
45    pub(crate) vertex_normals: Arc<[u32]>,
46    /// Uncompressed vertex texture coordinates for meshlet vertices.
47    pub(crate) vertex_uvs: Arc<[Vec2]>,
48    /// Triangle indices for meshlets.
49    pub(crate) indices: Arc<[u8]>,
50    /// The BVH8 used for culling and LOD selection of the meshlets. The root is at index 0.
51    pub(crate) bvh: Arc<[BvhNode]>,
52    /// The list of meshlets making up this mesh.
53    pub(crate) meshlets: Arc<[Meshlet]>,
54    /// Spherical bounding volumes.
55    pub(crate) meshlet_cull_data: Arc<[MeshletCullData]>,
56    /// The tight AABB of the meshlet mesh, used for frustum and occlusion culling at the instance
57    /// level.
58    pub(crate) aabb: MeshletAabb,
59    /// The depth of the culling BVH, used to determine the number of dispatches at runtime.
60    pub(crate) bvh_depth: u32,
61}
62
63/// A single BVH8 node in the BVH used for culling and LOD selection of a [`MeshletMesh`].
64#[derive(Copy, Clone, Default, Pod, Zeroable)]
65#[repr(C)]
66pub struct BvhNode {
67    /// The tight AABBs of this node's children, used for frustum and occlusion during BVH
68    /// traversal.
69    pub aabbs: [MeshletAabbErrorOffset; 8],
70    /// The LOD bounding spheres of this node's children, used for LOD selection during BVH
71    /// traversal.
72    pub lod_bounds: [MeshletBoundingSphere; 8],
73    /// If `u8::MAX`, it indicates that the child of each children is a BVH node, otherwise it is the number of meshlets in the group.
74    pub child_counts: [u8; 8],
75    pub _padding: [u32; 2],
76}
77
78/// A single meshlet within a [`MeshletMesh`].
79#[derive(Copy, Clone, Pod, Zeroable)]
80#[repr(C)]
81pub struct Meshlet {
82    /// The bit offset within the parent mesh's [`MeshletMesh::vertex_positions`] buffer where the vertex positions for this meshlet begin.
83    pub start_vertex_position_bit: u32,
84    /// The offset within the parent mesh's [`MeshletMesh::vertex_normals`] and [`MeshletMesh::vertex_uvs`] buffers
85    /// where non-position vertex attributes for this meshlet begin.
86    pub start_vertex_attribute_id: u32,
87    /// The offset within the parent mesh's [`MeshletMesh::indices`] buffer where the indices for this meshlet begin.
88    pub start_index_id: u32,
89    /// The amount of vertices in this meshlet.
90    pub vertex_count: u8,
91    /// The amount of triangles in this meshlet.
92    pub triangle_count: u8,
93    /// Unused.
94    pub padding: u16,
95    /// Number of bits used to store the X channel of vertex positions within this meshlet.
96    pub bits_per_vertex_position_channel_x: u8,
97    /// Number of bits used to store the Y channel of vertex positions within this meshlet.
98    pub bits_per_vertex_position_channel_y: u8,
99    /// Number of bits used to store the Z channel of vertex positions within this meshlet.
100    pub bits_per_vertex_position_channel_z: u8,
101    /// Power of 2 factor used to quantize vertex positions within this meshlet.
102    pub vertex_position_quantization_factor: u8,
103    /// Minimum quantized X channel value of vertex positions within this meshlet.
104    pub min_vertex_position_channel_x: f32,
105    /// Minimum quantized Y channel value of vertex positions within this meshlet.
106    pub min_vertex_position_channel_y: f32,
107    /// Minimum quantized Z channel value of vertex positions within this meshlet.
108    pub min_vertex_position_channel_z: f32,
109}
110
111/// Bounding spheres used for culling and choosing level of detail for a [`Meshlet`].
112#[derive(Copy, Clone, Pod, Zeroable)]
113#[repr(C)]
114pub struct MeshletCullData {
115    /// Tight bounding box, used for frustum and occlusion culling for this meshlet.
116    pub aabb: MeshletAabbErrorOffset,
117    /// Bounding sphere used for determining if this meshlet's group is at the correct level of detail for a given view.
118    pub lod_group_sphere: MeshletBoundingSphere,
119}
120
121/// An axis-aligned bounding box used for a [`Meshlet`].
122#[derive(Copy, Clone, Default, Pod, Zeroable, ShaderType)]
123#[repr(C)]
124pub struct MeshletAabb {
125    pub center: Vec3,
126    pub half_extent: Vec3,
127}
128
129// An axis-aligned bounding box used for a [`Meshlet`].
130#[derive(Copy, Clone, Default, Pod, Zeroable, ShaderType)]
131#[repr(C)]
132pub struct MeshletAabbErrorOffset {
133    pub center: Vec3,
134    pub error: f32,
135    pub half_extent: Vec3,
136    pub child_offset: u32,
137}
138
139/// A spherical bounding volume used for a [`Meshlet`].
140#[derive(Copy, Clone, Default, Pod, Zeroable)]
141#[repr(C)]
142pub struct MeshletBoundingSphere {
143    pub center: Vec3,
144    pub radius: f32,
145}
146
147/// An [`AssetSaver`] for `.meshlet_mesh` [`MeshletMesh`] assets.
148pub struct MeshletMeshSaver;
149
150impl AssetSaver for MeshletMeshSaver {
151    type Asset = MeshletMesh;
152    type Settings = ();
153    type OutputLoader = MeshletMeshLoader;
154    type Error = MeshletMeshSaveOrLoadError;
155
156    async fn save(
157        &self,
158        writer: &mut Writer,
159        asset: SavedAsset<'_, MeshletMesh>,
160        _settings: &(),
161    ) -> Result<(), MeshletMeshSaveOrLoadError> {
162        // Write asset magic number
163        writer
164            .write_all(&MESHLET_MESH_ASSET_MAGIC.to_le_bytes())
165            .await?;
166
167        // Write asset version
168        writer
169            .write_all(&MESHLET_MESH_ASSET_VERSION.to_le_bytes())
170            .await?;
171
172        writer.write_all(bytemuck::bytes_of(&asset.aabb)).await?;
173        writer
174            .write_all(bytemuck::bytes_of(&asset.bvh_depth))
175            .await?;
176
177        // Compress and write asset data
178        let mut writer = FrameEncoder::new(AsyncWriteSyncAdapter(writer));
179        write_slice(&asset.vertex_positions, &mut writer)?;
180        write_slice(&asset.vertex_normals, &mut writer)?;
181        write_slice(&asset.vertex_uvs, &mut writer)?;
182        write_slice(&asset.indices, &mut writer)?;
183        write_slice(&asset.bvh, &mut writer)?;
184        write_slice(&asset.meshlets, &mut writer)?;
185        write_slice(&asset.meshlet_cull_data, &mut writer)?;
186        // BUG: Flushing helps with an async_fs bug, but it still fails sometimes. https://github.com/smol-rs/async-fs/issues/45
187        // ERROR bevy_asset::server: Failed to load asset with asset loader MeshletMeshLoader: failed to fill whole buffer
188        writer.flush()?;
189        writer.finish()?;
190
191        Ok(())
192    }
193}
194
195/// An [`AssetLoader`] for `.meshlet_mesh` [`MeshletMesh`] assets.
196pub struct MeshletMeshLoader;
197
198impl AssetLoader for MeshletMeshLoader {
199    type Asset = MeshletMesh;
200    type Settings = ();
201    type Error = MeshletMeshSaveOrLoadError;
202
203    async fn load(
204        &self,
205        reader: &mut dyn Reader,
206        _settings: &(),
207        _load_context: &mut LoadContext<'_>,
208    ) -> Result<MeshletMesh, MeshletMeshSaveOrLoadError> {
209        // Load and check magic number
210        let magic = async_read_u64(reader).await?;
211        if magic != MESHLET_MESH_ASSET_MAGIC {
212            return Err(MeshletMeshSaveOrLoadError::WrongFileType);
213        }
214
215        // Load and check asset version
216        let version = async_read_u64(reader).await?;
217        if version != MESHLET_MESH_ASSET_VERSION {
218            return Err(MeshletMeshSaveOrLoadError::WrongVersion { found: version });
219        }
220
221        let mut bytes = [0u8; size_of::<MeshletAabb>()];
222        reader.read_exact(&mut bytes).await?;
223        let aabb = bytemuck::cast(bytes);
224        let mut bytes = [0u8; size_of::<u32>()];
225        reader.read_exact(&mut bytes).await?;
226        let bvh_depth = u32::from_le_bytes(bytes);
227
228        // Load and decompress asset data
229        let reader = &mut FrameDecoder::new(AsyncReadSyncAdapter(reader));
230        let vertex_positions = read_slice(reader)?;
231        let vertex_normals = read_slice(reader)?;
232        let vertex_uvs = read_slice(reader)?;
233        let indices = read_slice(reader)?;
234        let bvh = read_slice(reader)?;
235        let meshlets = read_slice(reader)?;
236        let meshlet_cull_data = read_slice(reader)?;
237
238        Ok(MeshletMesh {
239            vertex_positions,
240            vertex_normals,
241            vertex_uvs,
242            indices,
243            bvh,
244            meshlets,
245            meshlet_cull_data,
246            aabb,
247            bvh_depth,
248        })
249    }
250
251    fn extensions(&self) -> &[&str] {
252        &["meshlet_mesh"]
253    }
254}
255
256#[derive(Error, Debug)]
257pub enum MeshletMeshSaveOrLoadError {
258    #[error("file was not a MeshletMesh asset")]
259    WrongFileType,
260    #[error("expected asset version {MESHLET_MESH_ASSET_VERSION} but found version {found}")]
261    WrongVersion { found: u64 },
262    #[error("failed to compress or decompress asset data")]
263    CompressionOrDecompression(#[from] lz4_flex::frame::Error),
264    #[error(transparent)]
265    Io(#[from] std::io::Error),
266}
267
268async fn async_read_u64(reader: &mut dyn Reader) -> Result<u64, std::io::Error> {
269    let mut bytes = [0u8; 8];
270    reader.read_exact(&mut bytes).await?;
271    Ok(u64::from_le_bytes(bytes))
272}
273
274fn read_u64(reader: &mut dyn Read) -> Result<u64, std::io::Error> {
275    let mut bytes = [0u8; 8];
276    reader.read_exact(&mut bytes)?;
277    Ok(u64::from_le_bytes(bytes))
278}
279
280fn write_slice<T: Pod>(
281    field: &[T],
282    writer: &mut dyn Write,
283) -> Result<(), MeshletMeshSaveOrLoadError> {
284    writer.write_all(&(field.len() as u64).to_le_bytes())?;
285    writer.write_all(bytemuck::cast_slice(field))?;
286    Ok(())
287}
288
289fn read_slice<T: Pod>(reader: &mut dyn Read) -> Result<Arc<[T]>, std::io::Error> {
290    let len = read_u64(reader)? as usize;
291
292    let mut data: Arc<[T]> = core::iter::repeat_with(T::zeroed).take(len).collect();
293    let slice = Arc::get_mut(&mut data).unwrap();
294    reader.read_exact(bytemuck::cast_slice_mut(slice))?;
295
296    Ok(data)
297}
298
299// TODO: Use async for everything and get rid of this adapter
300struct AsyncWriteSyncAdapter<'a>(&'a mut Writer);
301
302impl Write for AsyncWriteSyncAdapter<'_> {
303    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
304        block_on(self.0.write(buf))
305    }
306
307    fn flush(&mut self) -> std::io::Result<()> {
308        block_on(self.0.flush())
309    }
310}
311
312// TODO: Use async for everything and get rid of this adapter
313struct AsyncReadSyncAdapter<'a>(&'a mut dyn Reader);
314
315impl Read for AsyncReadSyncAdapter<'_> {
316    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
317        block_on(self.0.read(buf))
318    }
319}