forked from bevyengine/bevy
-
Notifications
You must be signed in to change notification settings - Fork 1
/
specialized_mesh_pipeline.rs
354 lines (330 loc) · 14.8 KB
/
specialized_mesh_pipeline.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
//! Demonstrates how to define and use specialized mesh pipeline
//!
//! This example shows how to use the built-in [`SpecializedMeshPipeline`]
//! functionality with a custom [`RenderCommand`] to allow custom mesh rendering with
//! more flexibility than the material api.
//!
//! [`SpecializedMeshPipeline`] let's you customize the entire pipeline used when rendering a mesh.
use bevy::{
core_pipeline::core_3d::{Opaque3d, Opaque3dBinKey, CORE_3D_DEPTH_FORMAT},
math::{vec3, vec4},
pbr::{
DrawMesh, MeshPipeline, MeshPipelineKey, MeshPipelineViewLayoutKey, RenderMeshInstances,
SetMeshBindGroup, SetMeshViewBindGroup,
},
prelude::*,
render::{
extract_component::{ExtractComponent, ExtractComponentPlugin},
mesh::{Indices, MeshVertexBufferLayoutRef, PrimitiveTopology, RenderMesh},
render_asset::{RenderAssetUsages, RenderAssets},
render_phase::{
AddRenderCommand, BinnedRenderPhaseType, DrawFunctions, SetItemPipeline,
ViewBinnedRenderPhases,
},
render_resource::{
ColorTargetState, ColorWrites, CompareFunction, DepthStencilState, Face, FragmentState,
FrontFace, MultisampleState, PipelineCache, PolygonMode, PrimitiveState,
RenderPipelineDescriptor, SpecializedMeshPipeline, SpecializedMeshPipelineError,
SpecializedMeshPipelines, TextureFormat, VertexState,
},
view::{self, ExtractedView, RenderVisibleEntities, ViewTarget, VisibilitySystems},
Render, RenderApp, RenderSet,
},
};
const SHADER_ASSET_PATH: &str = "shaders/specialized_mesh_pipeline.wgsl";
fn main() {
App::new()
.add_plugins(DefaultPlugins)
.add_plugins(CustomRenderedMeshPipelinePlugin)
.add_systems(Startup, setup)
.run();
}
/// Spawns the objects in the scene.
fn setup(mut commands: Commands, mut meshes: ResMut<Assets<Mesh>>) {
// Build a custom triangle mesh with colors
// We define a custom mesh because the examples only uses a limited
// set of vertex attributes for simplicity
let mesh = Mesh::new(
PrimitiveTopology::TriangleList,
RenderAssetUsages::default(),
)
.with_inserted_indices(Indices::U32(vec![0, 1, 2]))
.with_inserted_attribute(
Mesh::ATTRIBUTE_POSITION,
vec![
vec3(-0.5, -0.5, 0.0),
vec3(0.5, -0.5, 0.0),
vec3(0.0, 0.25, 0.0),
],
)
.with_inserted_attribute(
Mesh::ATTRIBUTE_COLOR,
vec![
vec4(1.0, 0.0, 0.0, 1.0),
vec4(0.0, 1.0, 0.0, 1.0),
vec4(0.0, 0.0, 1.0, 1.0),
],
);
// spawn 3 triangles to show that batching works
for (x, y) in [-0.5, 0.0, 0.5].into_iter().zip([-0.25, 0.5, -0.25]) {
// Spawn an entity with all the required components for it to be rendered with our custom pipeline
commands.spawn((
// We use a marker component to identify the mesh that will be rendered
// with our specialized pipeline
CustomRenderedEntity,
// We need to add the mesh handle to the entity
Mesh3d(meshes.add(mesh.clone())),
Transform::from_xyz(x, y, 0.0),
));
}
// Spawn the camera.
commands.spawn((
Camera3d::default(),
// Move the camera back a bit to see all the triangles
Transform::from_xyz(0.0, 0.0, 3.0).looking_at(Vec3::ZERO, Vec3::Y),
));
}
// When writing custom rendering code it's generally recommended to use a plugin.
// The main reason for this is that it gives you access to the finish() hook
// which is called after rendering resources are initialized.
struct CustomRenderedMeshPipelinePlugin;
impl Plugin for CustomRenderedMeshPipelinePlugin {
fn build(&self, app: &mut App) {
app.add_plugins(ExtractComponentPlugin::<CustomRenderedEntity>::default())
.add_systems(
PostUpdate,
// Make sure to tell Bevy to check our entity for visibility. Bevy won't
// do this by default, for efficiency reasons.
// This will do things like frustum culling and hierarchy visibility
view::check_visibility::<WithCustomRenderedEntity>
.in_set(VisibilitySystems::CheckVisibility),
);
// We make sure to add these to the render app, not the main app.
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app
// This is needed to tell bevy about your custom pipeline
.init_resource::<SpecializedMeshPipelines<CustomMeshPipeline>>()
// We need to use a custom draw command so we need to register it
.add_render_command::<Opaque3d, DrawSpecializedPipelineCommands>()
.add_systems(Render, queue_custom_mesh_pipeline.in_set(RenderSet::Queue));
}
fn finish(&self, app: &mut App) {
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
// Creating this pipeline needs the RenderDevice and RenderQueue
// which are only available once rendering plugins are initialized.
render_app.init_resource::<CustomMeshPipeline>();
}
}
/// A marker component that represents an entity that is to be rendered using
/// our specialized pipeline.
///
/// Note the [`ExtractComponent`] trait implementation. This is necessary to
/// tell Bevy that this object should be pulled into the render world.
#[derive(Clone, Component, ExtractComponent)]
struct CustomRenderedEntity;
/// The custom draw commands that Bevy executes for each entity we enqueue into
/// the render phase.
type DrawSpecializedPipelineCommands = (
// Set the pipeline
SetItemPipeline,
// Set the view uniform at bind group 0
SetMeshViewBindGroup<0>,
// Set the mesh uniform at bind group 1
SetMeshBindGroup<1>,
// Draw the mesh
DrawMesh,
);
/// A query filter that tells [`view::check_visibility`] about our custom
/// rendered entity.
type WithCustomRenderedEntity = With<CustomRenderedEntity>;
// This contains the state needed to speciazlize a mesh pipeline
#[derive(Resource)]
struct CustomMeshPipeline {
/// The base mesh pipeline defined by bevy
///
/// This isn't required, but if you want to use a bevy `Mesh` it's easier when you
/// have access to the base `MeshPipeline` that bevy already defines
mesh_pipeline: MeshPipeline,
/// Stores the shader used for this pipeline directly on the pipeline.
/// This isn't required, it's only done like this for simplicity.
shader_handle: Handle<Shader>,
}
impl FromWorld for CustomMeshPipeline {
fn from_world(world: &mut World) -> Self {
// Load the shader
let shader_handle: Handle<Shader> = world.resource::<AssetServer>().load(SHADER_ASSET_PATH);
Self {
mesh_pipeline: MeshPipeline::from_world(world),
shader_handle,
}
}
}
impl SpecializedMeshPipeline for CustomMeshPipeline {
/// Pipeline use keys to determine how to specialize it.
/// The key is also used by the pipeline cache to determine if
/// it needs to create a new pipeline or not
///
/// In this example we just use the base `MeshPipelineKey` defined by bevy, but this could be anything.
/// For example, if you want to make a pipeline with a procedural shader you could add the Handle<Shader> to the key.
type Key = MeshPipelineKey;
fn specialize(
&self,
mesh_key: Self::Key,
layout: &MeshVertexBufferLayoutRef,
) -> Result<RenderPipelineDescriptor, SpecializedMeshPipelineError> {
// Define the vertex attributes based on a standard bevy [`Mesh`]
let mut vertex_attributes = Vec::new();
if layout.0.contains(Mesh::ATTRIBUTE_POSITION) {
// Make sure this matches the shader location
vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0));
}
if layout.0.contains(Mesh::ATTRIBUTE_COLOR) {
// Make sure this matches the shader location
vertex_attributes.push(Mesh::ATTRIBUTE_COLOR.at_shader_location(1));
}
// This will automatically generate the correct `VertexBufferLayout` based on the vertex attributes
let vertex_buffer_layout = layout.0.get_layout(&vertex_attributes)?;
Ok(RenderPipelineDescriptor {
label: Some("Specialized Mesh Pipeline".into()),
layout: vec![
// Bind group 0 is the view uniform
self.mesh_pipeline
.get_view_layout(MeshPipelineViewLayoutKey::from(mesh_key))
.clone(),
// Bind group 1 is the mesh uniform
self.mesh_pipeline.mesh_layouts.model_only.clone(),
],
push_constant_ranges: vec![],
vertex: VertexState {
shader: self.shader_handle.clone(),
shader_defs: vec![],
entry_point: "vertex".into(),
// Customize how to store the meshes' vertex attributes in the vertex buffer
buffers: vec![vertex_buffer_layout],
},
fragment: Some(FragmentState {
shader: self.shader_handle.clone(),
shader_defs: vec![],
entry_point: "fragment".into(),
targets: vec![Some(ColorTargetState {
// This isn't required, but bevy supports HDR and non-HDR rendering
// so it's generally recommended to specialize the pipeline for that
format: if mesh_key.contains(MeshPipelineKey::HDR) {
ViewTarget::TEXTURE_FORMAT_HDR
} else {
TextureFormat::bevy_default()
},
// For this example we only use opaque meshes,
// but if you wanted to use alpha blending you would need to set it here
blend: None,
write_mask: ColorWrites::ALL,
})],
}),
primitive: PrimitiveState {
topology: mesh_key.primitive_topology(),
front_face: FrontFace::Ccw,
cull_mode: Some(Face::Back),
polygon_mode: PolygonMode::Fill,
..default()
},
// Note that if your view has no depth buffer this will need to be
// changed.
depth_stencil: Some(DepthStencilState {
format: CORE_3D_DEPTH_FORMAT,
depth_write_enabled: true,
depth_compare: CompareFunction::GreaterEqual,
stencil: default(),
bias: default(),
}),
// It's generally recommended to specialize your pipeline for MSAA,
// but it's not always possible
multisample: MultisampleState {
count: mesh_key.msaa_samples(),
..MultisampleState::default()
},
zero_initialize_workgroup_memory: false,
})
}
}
/// A render-world system that enqueues the entity with custom rendering into
/// the opaque render phases of each view.
#[allow(clippy::too_many_arguments)]
fn queue_custom_mesh_pipeline(
pipeline_cache: Res<PipelineCache>,
custom_mesh_pipeline: Res<CustomMeshPipeline>,
mut opaque_render_phases: ResMut<ViewBinnedRenderPhases<Opaque3d>>,
opaque_draw_functions: Res<DrawFunctions<Opaque3d>>,
mut specialized_mesh_pipelines: ResMut<SpecializedMeshPipelines<CustomMeshPipeline>>,
views: Query<(Entity, &RenderVisibleEntities, &ExtractedView, &Msaa), With<ExtractedView>>,
render_meshes: Res<RenderAssets<RenderMesh>>,
render_mesh_instances: Res<RenderMeshInstances>,
) {
// Get the id for our custom draw function
let draw_function_id = opaque_draw_functions
.read()
.id::<DrawSpecializedPipelineCommands>();
// Render phases are per-view, so we need to iterate over all views so that
// the entity appears in them. (In this example, we have only one view, but
// it's good practice to loop over all views anyway.)
for (view_entity, view_visible_entities, view, msaa) in views.iter() {
let Some(opaque_phase) = opaque_render_phases.get_mut(&view_entity) else {
continue;
};
// Create the key based on the view. In this case we only care about MSAA and HDR
let view_key = MeshPipelineKey::from_msaa_samples(msaa.samples())
| MeshPipelineKey::from_hdr(view.hdr);
// Find all the custom rendered entities that are visible from this
// view.
for &(render_entity, visible_entity) in view_visible_entities
.get::<WithCustomRenderedEntity>()
.iter()
{
// Get the mesh instance
let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(visible_entity)
else {
continue;
};
// Get the mesh data
let Some(mesh) = render_meshes.get(mesh_instance.mesh_asset_id) else {
continue;
};
// Specialize the key for the current mesh entity
// For this example we only specialize based on the mesh topology
// but you could have more complex keys and that's where you'd need to create those keys
let mut mesh_key = view_key;
mesh_key |= MeshPipelineKey::from_primitive_topology(mesh.primitive_topology());
// Finally, we can specialize the pipeline based on the key
let pipeline_id = specialized_mesh_pipelines
.specialize(
&pipeline_cache,
&custom_mesh_pipeline,
mesh_key,
&mesh.layout,
)
// This should never with this example, but if your pipeline specialization
// can fail you need to handle the error here
.expect("Failed to specialize mesh pipeline");
// Add the mesh with our specialized pipeline
opaque_phase.add(
Opaque3dBinKey {
draw_function: draw_function_id,
pipeline: pipeline_id,
// The asset ID is arbitrary; we simply use [`AssetId::invalid`],
// but you can use anything you like. Note that the asset ID need
// not be the ID of a [`Mesh`].
asset_id: AssetId::<Mesh>::invalid().untyped(),
material_bind_group_id: None,
lightmap_image: None,
},
(render_entity, visible_entity),
// This example supports batching, but if your pipeline doesn't
// support it you can use `BinnedRenderPhaseType::UnbatchableMesh`
BinnedRenderPhaseType::BatchableMesh,
);
}
}
}