1use crate::{
2 core_3d::graph::{Core3d, Node3d},
3 fullscreen_vertex_shader::fullscreen_shader_vertex_state,
4 prelude::Camera3d,
5 prepass::{DepthPrepass, MotionVectorPrepass, ViewPrepassTextures},
6};
7use bevy_app::{App, Plugin};
8use bevy_asset::{load_internal_asset, weak_handle, Handle};
9use bevy_diagnostic::FrameCount;
10use bevy_ecs::{
11 prelude::{Component, Entity, ReflectComponent},
12 query::{QueryItem, With},
13 resource::Resource,
14 schedule::IntoScheduleConfigs,
15 system::{Commands, Query, Res, ResMut},
16 world::{FromWorld, World},
17};
18use bevy_image::BevyDefault as _;
19use bevy_math::vec2;
20use bevy_reflect::{std_traits::ReflectDefault, Reflect};
21use bevy_render::{
22 camera::{ExtractedCamera, MipBias, TemporalJitter},
23 prelude::{Camera, Projection},
24 render_graph::{NodeRunError, RenderGraphApp, RenderGraphContext, ViewNode, ViewNodeRunner},
25 render_resource::{
26 binding_types::{sampler, texture_2d, texture_depth_2d},
27 BindGroupEntries, BindGroupLayout, BindGroupLayoutEntries, CachedRenderPipelineId,
28 ColorTargetState, ColorWrites, Extent3d, FilterMode, FragmentState, MultisampleState,
29 Operations, PipelineCache, PrimitiveState, RenderPassColorAttachment, RenderPassDescriptor,
30 RenderPipelineDescriptor, Sampler, SamplerBindingType, SamplerDescriptor, Shader,
31 ShaderStages, SpecializedRenderPipeline, SpecializedRenderPipelines, TextureDescriptor,
32 TextureDimension, TextureFormat, TextureSampleType, TextureUsages,
33 },
34 renderer::{RenderContext, RenderDevice},
35 sync_component::SyncComponentPlugin,
36 sync_world::RenderEntity,
37 texture::{CachedTexture, TextureCache},
38 view::{ExtractedView, Msaa, ViewTarget},
39 ExtractSchedule, MainWorld, Render, RenderApp, RenderSet,
40};
41use tracing::warn;
42
43const TAA_SHADER_HANDLE: Handle<Shader> = weak_handle!("fea20d50-86b6-4069-aa32-374346aec00c");
44
45pub struct TemporalAntiAliasPlugin;
49
50impl Plugin for TemporalAntiAliasPlugin {
51 fn build(&self, app: &mut App) {
52 load_internal_asset!(app, TAA_SHADER_HANDLE, "taa.wgsl", Shader::from_wgsl);
53
54 app.register_type::<TemporalAntiAliasing>();
55
56 app.add_plugins(SyncComponentPlugin::<TemporalAntiAliasing>::default());
57
58 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
59 return;
60 };
61 render_app
62 .init_resource::<SpecializedRenderPipelines<TaaPipeline>>()
63 .add_systems(ExtractSchedule, extract_taa_settings)
64 .add_systems(
65 Render,
66 (
67 prepare_taa_jitter_and_mip_bias.in_set(RenderSet::ManageViews),
68 prepare_taa_pipelines.in_set(RenderSet::Prepare),
69 prepare_taa_history_textures.in_set(RenderSet::PrepareResources),
70 ),
71 )
72 .add_render_graph_node::<ViewNodeRunner<TemporalAntiAliasNode>>(Core3d, Node3d::Taa)
73 .add_render_graph_edges(
74 Core3d,
75 (
76 Node3d::EndMainPass,
77 Node3d::MotionBlur, Node3d::Taa,
79 Node3d::Bloom,
80 Node3d::Tonemapping,
81 ),
82 );
83 }
84
85 fn finish(&self, app: &mut App) {
86 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
87 return;
88 };
89
90 render_app.init_resource::<TaaPipeline>();
91 }
92}
93
94#[derive(Component, Reflect, Clone)]
134#[reflect(Component, Default, Clone)]
135#[require(TemporalJitter, DepthPrepass, MotionVectorPrepass)]
136#[doc(alias = "Taa")]
137pub struct TemporalAntiAliasing {
138 pub reset: bool,
146}
147
148impl Default for TemporalAntiAliasing {
149 fn default() -> Self {
150 Self { reset: true }
151 }
152}
153
154#[derive(Default)]
156pub struct TemporalAntiAliasNode;
157
158impl ViewNode for TemporalAntiAliasNode {
159 type ViewQuery = (
160 &'static ExtractedCamera,
161 &'static ViewTarget,
162 &'static TemporalAntiAliasHistoryTextures,
163 &'static ViewPrepassTextures,
164 &'static TemporalAntiAliasPipelineId,
165 &'static Msaa,
166 );
167
168 fn run(
169 &self,
170 _graph: &mut RenderGraphContext,
171 render_context: &mut RenderContext,
172 (camera, view_target, taa_history_textures, prepass_textures, taa_pipeline_id, msaa): QueryItem<
173 Self::ViewQuery,
174 >,
175 world: &World,
176 ) -> Result<(), NodeRunError> {
177 if *msaa != Msaa::Off {
178 warn!("Temporal anti-aliasing requires MSAA to be disabled");
179 return Ok(());
180 }
181
182 let (Some(pipelines), Some(pipeline_cache)) = (
183 world.get_resource::<TaaPipeline>(),
184 world.get_resource::<PipelineCache>(),
185 ) else {
186 return Ok(());
187 };
188 let (Some(taa_pipeline), Some(prepass_motion_vectors_texture), Some(prepass_depth_texture)) = (
189 pipeline_cache.get_render_pipeline(taa_pipeline_id.0),
190 &prepass_textures.motion_vectors,
191 &prepass_textures.depth,
192 ) else {
193 return Ok(());
194 };
195 let view_target = view_target.post_process_write();
196
197 let taa_bind_group = render_context.render_device().create_bind_group(
198 "taa_bind_group",
199 &pipelines.taa_bind_group_layout,
200 &BindGroupEntries::sequential((
201 view_target.source,
202 &taa_history_textures.read.default_view,
203 &prepass_motion_vectors_texture.texture.default_view,
204 &prepass_depth_texture.texture.default_view,
205 &pipelines.nearest_sampler,
206 &pipelines.linear_sampler,
207 )),
208 );
209
210 {
211 let mut taa_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
212 label: Some("taa_pass"),
213 color_attachments: &[
214 Some(RenderPassColorAttachment {
215 view: view_target.destination,
216 resolve_target: None,
217 ops: Operations::default(),
218 }),
219 Some(RenderPassColorAttachment {
220 view: &taa_history_textures.write.default_view,
221 resolve_target: None,
222 ops: Operations::default(),
223 }),
224 ],
225 depth_stencil_attachment: None,
226 timestamp_writes: None,
227 occlusion_query_set: None,
228 });
229 taa_pass.set_render_pipeline(taa_pipeline);
230 taa_pass.set_bind_group(0, &taa_bind_group, &[]);
231 if let Some(viewport) = camera.viewport.as_ref() {
232 taa_pass.set_camera_viewport(viewport);
233 }
234 taa_pass.draw(0..3, 0..1);
235 }
236
237 Ok(())
238 }
239}
240
241#[derive(Resource)]
242struct TaaPipeline {
243 taa_bind_group_layout: BindGroupLayout,
244 nearest_sampler: Sampler,
245 linear_sampler: Sampler,
246}
247
248impl FromWorld for TaaPipeline {
249 fn from_world(world: &mut World) -> Self {
250 let render_device = world.resource::<RenderDevice>();
251
252 let nearest_sampler = render_device.create_sampler(&SamplerDescriptor {
253 label: Some("taa_nearest_sampler"),
254 mag_filter: FilterMode::Nearest,
255 min_filter: FilterMode::Nearest,
256 ..SamplerDescriptor::default()
257 });
258 let linear_sampler = render_device.create_sampler(&SamplerDescriptor {
259 label: Some("taa_linear_sampler"),
260 mag_filter: FilterMode::Linear,
261 min_filter: FilterMode::Linear,
262 ..SamplerDescriptor::default()
263 });
264
265 let taa_bind_group_layout = render_device.create_bind_group_layout(
266 "taa_bind_group_layout",
267 &BindGroupLayoutEntries::sequential(
268 ShaderStages::FRAGMENT,
269 (
270 texture_2d(TextureSampleType::Float { filterable: true }),
272 texture_2d(TextureSampleType::Float { filterable: true }),
274 texture_2d(TextureSampleType::Float { filterable: true }),
276 texture_depth_2d(),
278 sampler(SamplerBindingType::NonFiltering),
280 sampler(SamplerBindingType::Filtering),
282 ),
283 ),
284 );
285
286 TaaPipeline {
287 taa_bind_group_layout,
288 nearest_sampler,
289 linear_sampler,
290 }
291 }
292}
293
294#[derive(PartialEq, Eq, Hash, Clone)]
295struct TaaPipelineKey {
296 hdr: bool,
297 reset: bool,
298}
299
300impl SpecializedRenderPipeline for TaaPipeline {
301 type Key = TaaPipelineKey;
302
303 fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
304 let mut shader_defs = vec![];
305
306 let format = if key.hdr {
307 shader_defs.push("TONEMAP".into());
308 ViewTarget::TEXTURE_FORMAT_HDR
309 } else {
310 TextureFormat::bevy_default()
311 };
312
313 if key.reset {
314 shader_defs.push("RESET".into());
315 }
316
317 RenderPipelineDescriptor {
318 label: Some("taa_pipeline".into()),
319 layout: vec![self.taa_bind_group_layout.clone()],
320 vertex: fullscreen_shader_vertex_state(),
321 fragment: Some(FragmentState {
322 shader: TAA_SHADER_HANDLE,
323 shader_defs,
324 entry_point: "taa".into(),
325 targets: vec![
326 Some(ColorTargetState {
327 format,
328 blend: None,
329 write_mask: ColorWrites::ALL,
330 }),
331 Some(ColorTargetState {
332 format,
333 blend: None,
334 write_mask: ColorWrites::ALL,
335 }),
336 ],
337 }),
338 primitive: PrimitiveState::default(),
339 depth_stencil: None,
340 multisample: MultisampleState::default(),
341 push_constant_ranges: Vec::new(),
342 zero_initialize_workgroup_memory: false,
343 }
344 }
345}
346
347fn extract_taa_settings(mut commands: Commands, mut main_world: ResMut<MainWorld>) {
348 let mut cameras_3d = main_world.query_filtered::<(
349 RenderEntity,
350 &Camera,
351 &Projection,
352 &mut TemporalAntiAliasing,
353 ), (
354 With<Camera3d>,
355 With<TemporalJitter>,
356 With<DepthPrepass>,
357 With<MotionVectorPrepass>,
358 )>();
359
360 for (entity, camera, camera_projection, mut taa_settings) in
361 cameras_3d.iter_mut(&mut main_world)
362 {
363 let has_perspective_projection = matches!(camera_projection, Projection::Perspective(_));
364 let mut entity_commands = commands
365 .get_entity(entity)
366 .expect("Camera entity wasn't synced.");
367 if camera.is_active && has_perspective_projection {
368 entity_commands.insert(taa_settings.clone());
369 taa_settings.reset = false;
370 } else {
371 entity_commands.remove::<(
373 TemporalAntiAliasing,
374 TemporalAntiAliasHistoryTextures,
376 TemporalAntiAliasPipelineId,
377 )>();
378 }
379 }
380}
381
382fn prepare_taa_jitter_and_mip_bias(
383 frame_count: Res<FrameCount>,
384 mut query: Query<(Entity, &mut TemporalJitter, Option<&MipBias>), With<TemporalAntiAliasing>>,
385 mut commands: Commands,
386) {
387 let halton_sequence = [
389 vec2(0.0, -0.16666666),
390 vec2(-0.25, 0.16666669),
391 vec2(0.25, -0.3888889),
392 vec2(-0.375, -0.055555552),
393 vec2(0.125, 0.2777778),
394 vec2(-0.125, -0.2777778),
395 vec2(0.375, 0.055555582),
396 vec2(-0.4375, 0.3888889),
397 ];
398
399 let offset = halton_sequence[frame_count.0 as usize % halton_sequence.len()];
400
401 for (entity, mut jitter, mip_bias) in &mut query {
402 jitter.offset = offset;
403
404 if mip_bias.is_none() {
405 commands.entity(entity).insert(MipBias(-1.0));
406 }
407 }
408}
409
410#[derive(Component)]
411pub struct TemporalAntiAliasHistoryTextures {
412 write: CachedTexture,
413 read: CachedTexture,
414}
415
416fn prepare_taa_history_textures(
417 mut commands: Commands,
418 mut texture_cache: ResMut<TextureCache>,
419 render_device: Res<RenderDevice>,
420 frame_count: Res<FrameCount>,
421 views: Query<(Entity, &ExtractedCamera, &ExtractedView), With<TemporalAntiAliasing>>,
422) {
423 for (entity, camera, view) in &views {
424 if let Some(physical_target_size) = camera.physical_target_size {
425 let mut texture_descriptor = TextureDescriptor {
426 label: None,
427 size: Extent3d {
428 depth_or_array_layers: 1,
429 width: physical_target_size.x,
430 height: physical_target_size.y,
431 },
432 mip_level_count: 1,
433 sample_count: 1,
434 dimension: TextureDimension::D2,
435 format: if view.hdr {
436 ViewTarget::TEXTURE_FORMAT_HDR
437 } else {
438 TextureFormat::bevy_default()
439 },
440 usage: TextureUsages::TEXTURE_BINDING | TextureUsages::RENDER_ATTACHMENT,
441 view_formats: &[],
442 };
443
444 texture_descriptor.label = Some("taa_history_1_texture");
445 let history_1_texture = texture_cache.get(&render_device, texture_descriptor.clone());
446
447 texture_descriptor.label = Some("taa_history_2_texture");
448 let history_2_texture = texture_cache.get(&render_device, texture_descriptor);
449
450 let textures = if frame_count.0 % 2 == 0 {
451 TemporalAntiAliasHistoryTextures {
452 write: history_1_texture,
453 read: history_2_texture,
454 }
455 } else {
456 TemporalAntiAliasHistoryTextures {
457 write: history_2_texture,
458 read: history_1_texture,
459 }
460 };
461
462 commands.entity(entity).insert(textures);
463 }
464 }
465}
466
467#[derive(Component)]
468pub struct TemporalAntiAliasPipelineId(CachedRenderPipelineId);
469
470fn prepare_taa_pipelines(
471 mut commands: Commands,
472 pipeline_cache: Res<PipelineCache>,
473 mut pipelines: ResMut<SpecializedRenderPipelines<TaaPipeline>>,
474 pipeline: Res<TaaPipeline>,
475 views: Query<(Entity, &ExtractedView, &TemporalAntiAliasing)>,
476) {
477 for (entity, view, taa_settings) in &views {
478 let mut pipeline_key = TaaPipelineKey {
479 hdr: view.hdr,
480 reset: taa_settings.reset,
481 };
482 let pipeline_id = pipelines.specialize(&pipeline_cache, &pipeline, pipeline_key.clone());
483
484 if pipeline_key.reset {
486 pipeline_key.reset = false;
487 pipelines.specialize(&pipeline_cache, &pipeline, pipeline_key);
488 }
489
490 commands
491 .entity(entity)
492 .insert(TemporalAntiAliasPipelineId(pipeline_id));
493 }
494}