Initial DLSS implementation

Cargo.toml

@@ -463,6 +463,9 @@ tonemapping_luts = ["bevy_internal/tonemapping_luts", "ktx2", "bevy_image/zstd"]
 # Include SMAA Look Up Tables KTX2 Files
 smaa_luts = ["bevy_internal/smaa_luts"]
 
+# NVIDIA Deep Learning Super Sampling
+dlss = ["bevy_internal/dlss"]
+
 # Enable AccessKit on Unix backends (currently only works with experimental screen readers and forks.)
 accesskit_unix = ["bevy_internal/accesskit_unix"]
 

crates/bevy_anti_aliasing/Cargo.toml

@@ -13,6 +13,7 @@ trace = []
 webgl = []
 webgpu = []
 smaa_luts = ["bevy_render/ktx2", "bevy_image/ktx2", "bevy_image/zstd"]
+dlss = ["dep:dlss_wgpu"]
 
 [dependencies]
 # bevy
@@ -30,6 +31,7 @@ bevy_diagnostic = { path = "../bevy_diagnostic", version = "0.17.0-dev" }
 
 # other
 tracing = { version = "0.1", default-features = false, features = ["std"] }
+dlss_wgpu = { path = "../../../dlss_wgpu", optional = true }
 
 [lints]
 workspace = true

crates/bevy_anti_aliasing/src/dlss/extract.rs

@@ -0,0 +1,29 @@
+use super::{prepare::ViewDlssContext, Dlss};
+use bevy_ecs::{
+    query::With,
+    system::{Commands, ResMut},
+};
+use bevy_render::{
+    camera::{Camera, MainPassResolutionOverride, Projection},
+    sync_world::RenderEntity,
+    view::Hdr,
+    MainWorld,
+};
+
+pub fn extract_dlss(mut commands: Commands, mut main_world: ResMut<MainWorld>) {
+    let mut cameras_3d = main_world
+        .query_filtered::<(RenderEntity, &Camera, &Projection, Option<&mut Dlss>), With<Hdr>>();
+
+    for (entity, camera, camera_projection, mut dlss) in cameras_3d.iter_mut(&mut main_world) {
+        let has_perspective_projection = matches!(camera_projection, Projection::Perspective(_));
+        let mut entity_commands = commands
+            .get_entity(entity)
+            .expect("Camera entity wasn't synced.");
+        if dlss.is_some() && camera.is_active && has_perspective_projection {
+            entity_commands.insert(dlss.as_deref().unwrap().clone());
+            dlss.as_mut().unwrap().reset = false;
+        } else {
+            entity_commands.remove::<(Dlss, ViewDlssContext, MainPassResolutionOverride)>();
+        }
+    }
+}

crates/bevy_anti_aliasing/src/dlss/mod.rs

@@ -0,0 +1,105 @@
+mod extract;
+mod node;
+mod prepare;
+
+use bevy_app::{App, Plugin};
+use bevy_core_pipeline::{
+    core_3d::graph::{Core3d, Node3d},
+    prepass::{DepthPrepass, MotionVectorPrepass},
+};
+use bevy_ecs::{
+    component::Component, prelude::ReflectComponent, resource::Resource,
+    schedule::IntoScheduleConfigs,
+};
+use bevy_reflect::{prelude::ReflectDefault, reflect_remote, Reflect};
+use bevy_render::{
+    camera::{MipBias, TemporalJitter},
+    render_graph::{RenderGraphApp, ViewNodeRunner},
+    renderer::RenderDevice,
+    view::{prepare_view_targets, prepare_view_uniforms, Hdr},
+    ExtractSchedule, Render, RenderApp, RenderSystems,
+};
+use std::sync::{Arc, Mutex};
+use tracing::info;
+
+pub use bevy_render::{DlssProjectId, DlssSupported};
+pub use dlss_wgpu::DlssPerfQualityMode;
+
+pub struct DlssPlugin;
+
+impl Plugin for DlssPlugin {
+    fn build(&self, app: &mut App) {
+        app.register_type::<Dlss>();
+    }
+
+    fn finish(&self, app: &mut App) {
+        if app.world().get_resource::<DlssSupported>().is_none() {
+            info!("DLSS is not supported on this system");
+            return;
+        }
+
+        let dlss_project_id = app.world().resource::<DlssProjectId>().0;
+
+        let render_app = app.get_sub_app_mut(RenderApp).unwrap();
+        let render_device = render_app.world().resource::<RenderDevice>().clone();
+
+        let dlss_sdk =
+            dlss_wgpu::DlssSdk::new(dlss_project_id, render_device.wgpu_device().clone());
+        if dlss_sdk.is_err() {
+            app.world_mut().remove_resource::<DlssSupported>();
+            info!("DLSS is not supported on this system");
+            return;
+        }
+
+        render_app
+            .insert_resource(DlssSdk(dlss_sdk.unwrap()))
+            .add_systems(ExtractSchedule, extract::extract_dlss)
+            .add_systems(
+                Render,
+                prepare::configure_dlss_view_targets
+                    .in_set(RenderSystems::ManageViews)
+                    .before(prepare_view_targets),
+            )
+            .add_systems(
+                Render,
+                prepare::prepare_dlss
+                    .in_set(RenderSystems::PrepareResources)
+                    .before(prepare_view_uniforms),
+            )
+            .add_render_graph_node::<ViewNodeRunner<node::DlssNode>>(Core3d, Node3d::Dlss)
+            .add_render_graph_edges(
+                Core3d,
+                (
+                    Node3d::EndMainPass,
+                    Node3d::MotionBlur, // Running before DLSS reduces edge artifacts and noise
+                    Node3d::Dlss,
+                    Node3d::Bloom,
+                    Node3d::Tonemapping,
+                ),
+            );
+    }
+}
+
+#[derive(Component, Reflect, Clone, Default)]
+#[reflect(Component, Default)]
+#[require(TemporalJitter, MipBias, DepthPrepass, MotionVectorPrepass, Hdr)]
+pub struct Dlss {
+    #[reflect(remote = DlssPerfQualityModeRemoteReflect)]
+    pub perf_quality_mode: DlssPerfQualityMode,
+    pub reset: bool,
+}
+
+#[reflect_remote(DlssPerfQualityMode)]
+#[derive(Default)]
+enum DlssPerfQualityModeRemoteReflect {
+    #[default]
+    Auto,
+    Dlaa,
+    Quality,
+    Balanced,
+    Performance,
+    UltraPerformance,
+}
+
+#[derive(Resource)]
+struct DlssSdk(Arc<Mutex<dlss_wgpu::DlssSdk>>);

crates/bevy_anti_aliasing/src/dlss/node.rs

@@ -0,0 +1,85 @@
+use super::{prepare::ViewDlssContext, Dlss};
+use bevy_core_pipeline::prepass::ViewPrepassTextures;
+use bevy_ecs::{query::QueryItem, world::World};
+use bevy_render::{
+    camera::{MainPassResolutionOverride, TemporalJitter},
+    render_graph::{NodeRunError, RenderGraphContext, ViewNode},
+    renderer::{RenderAdapter, RenderContext},
+    view::ViewTarget,
+};
+use dlss_wgpu::{DlssExposure, DlssRenderParameters, DlssTexture};
+
+#[derive(Default)]
+pub struct DlssNode;
+
+impl ViewNode for DlssNode {
+    type ViewQuery = (
+        &'static Dlss,
+        &'static ViewDlssContext,
+        &'static MainPassResolutionOverride,
+        &'static TemporalJitter,
+        &'static ViewTarget,
+        &'static ViewPrepassTextures,
+    );
+
+    fn run(
+        &self,
+        _graph: &mut RenderGraphContext,
+        render_context: &mut RenderContext,
+        (
+            dlss,
+            dlss_context,
+            viewport_override,
+            temporal_jitter,
+            view_target,
+            prepass_textures,
+        ): QueryItem<Self::ViewQuery>,
+        world: &World,
+    ) -> Result<(), NodeRunError> {
+        let adapter = world.resource::<RenderAdapter>();
+        let (Some(prepass_motion_vectors_texture), Some(prepass_depth_texture)) =
+            (&prepass_textures.motion_vectors, &prepass_textures.depth)
+        else {
+            return Ok(());
+        };
+
+        let view_target = view_target.post_process_write();
+
+        let render_resolution = viewport_override.0;
+        let render_parameters = DlssRenderParameters {
+            color: DlssTexture {
+                texture: &view_target.source_texture,
+                view: &view_target.source,
+            },
+            depth: DlssTexture {
+                texture: &prepass_depth_texture.texture.texture,
+                view: &prepass_depth_texture.texture.default_view,
+            },
+            motion_vectors: DlssTexture {
+                texture: &prepass_motion_vectors_texture.texture.texture,
+                view: &prepass_motion_vectors_texture.texture.default_view,
+            },
+            exposure: DlssExposure::Automatic, // TODO
+            bias: None,                        // TODO
+            dlss_output: DlssTexture {
+                texture: &view_target.destination_texture,
+                view: &view_target.destination,
+            },
+            reset: dlss.reset,
+            jitter_offset: -temporal_jitter.offset,
+            partial_texture_size: Some(render_resolution),
+            motion_vector_scale: Some(-render_resolution.as_vec2()),
+        };
+
+        let command_encoder = render_context.command_encoder();
+        let mut dlss_context = dlss_context.context.lock().unwrap();
+
+        command_encoder.push_debug_group("dlss");
+        dlss_context
+            .render(render_parameters, command_encoder, &adapter)
+            .expect("Failed to render DLSS");
+        command_encoder.pop_debug_group();
+
+        Ok(())
+    }
+}

crates/bevy_anti_aliasing/src/dlss/prepare.rs

@@ -0,0 +1,111 @@
+use super::{Dlss, DlssSdk};
+use bevy_core_pipeline::{
+    core_3d::Camera3d,
+    prepass::{DepthPrepass, MotionVectorPrepass},
+};
+use bevy_diagnostic::FrameCount;
+use bevy_ecs::{
+    component::Component,
+    entity::Entity,
+    query::With,
+    system::{Commands, Query, Res},
+};
+use bevy_math::Vec4Swizzles;
+use bevy_render::{
+    camera::{CameraMainTextureUsages, MainPassResolutionOverride, MipBias, TemporalJitter},
+    render_resource::TextureUsages,
+    renderer::{RenderDevice, RenderQueue},
+    view::ExtractedView,
+};
+use dlss_wgpu::{DlssContext, DlssFeatureFlags, DlssPerfQualityMode};
+use std::sync::{Arc, Mutex};
+
+pub fn prepare_dlss(
+    mut query: Query<
+        (
+            Entity,
+            &ExtractedView,
+            &Dlss,
+            &mut TemporalJitter,
+            &mut MipBias,
+            Option<&mut ViewDlssContext>,
+        ),
+        (
+            With<Camera3d>,
+            With<TemporalJitter>,
+            With<DepthPrepass>,
+            With<MotionVectorPrepass>,
+        ),
+    >,
+    dlss_sdk: Res<DlssSdk>,
+    render_device: Res<RenderDevice>,
+    render_queue: Res<RenderQueue>,
+    frame_count: Res<FrameCount>,
+    mut commands: Commands,
+) {
+    for (entity, view, dlss, mut temporal_jitter, mut mip_bias, mut dlss_context) in &mut query {
+        let upscaled_resolution = view.viewport.zw();
+
+        let dlss_feature_flags = DlssFeatureFlags::LowResolutionMotionVectors
+            | DlssFeatureFlags::InvertedDepth
+            | DlssFeatureFlags::HighDynamicRange
+            | DlssFeatureFlags::AutoExposure; // TODO
+
+        match dlss_context.as_deref_mut() {
+            Some(dlss_context)
+                if upscaled_resolution
+                    == dlss_context.context.lock().unwrap().upscaled_resolution()
+                    && dlss.perf_quality_mode == dlss_context.perf_quality_mode
+                    && dlss_feature_flags == dlss_context.feature_flags =>
+            {
+                let dlss_context = dlss_context.context.lock().unwrap();
+                temporal_jitter.offset =
+                    dlss_context.suggested_jitter(frame_count.0, dlss_context.render_resolution());
+            }
+            _ => {
+                let dlss_context = DlssContext::new(
+                    upscaled_resolution,
+                    dlss.perf_quality_mode,
+                    dlss_feature_flags,
+                    Arc::clone(&dlss_sdk.0),
+                    render_device.wgpu_device(),
+                    &render_queue,
+                )
+                .expect("Failed to create DlssContext");
+
+                let render_resolution = dlss_context.render_resolution();
+                temporal_jitter.offset =
+                    dlss_context.suggested_jitter(frame_count.0, render_resolution);
+                mip_bias.0 = dlss_context.suggested_mip_bias(render_resolution);
+
+                commands.entity(entity).insert((
+                    ViewDlssContext {
+                        context: Mutex::new(dlss_context),
+                        perf_quality_mode: dlss.perf_quality_mode,
+                        feature_flags: dlss_feature_flags,
+                    },
+                    MainPassResolutionOverride(render_resolution),
+                ));
+            }
+        }
+    }
+}
+
+#[derive(Component)]
+pub struct ViewDlssContext {
+    pub context: Mutex<DlssContext>,
+    pub perf_quality_mode: DlssPerfQualityMode,
+    pub feature_flags: DlssFeatureFlags,
+}
+
+pub fn configure_dlss_view_targets(
+    mut view_targets: Query<(&mut Camera3d, &mut CameraMainTextureUsages), With<Dlss>>,
+) {
+    for (mut camera_3d, mut camera_main_texture_usages) in view_targets.iter_mut() {
+        camera_main_texture_usages.0 |= TextureUsages::STORAGE_BINDING;
+
+        let mut depth_texture_usages = TextureUsages::from(camera_3d.depth_texture_usages);
+        depth_texture_usages |= TextureUsages::TEXTURE_BINDING;
+        camera_3d.depth_texture_usages = depth_texture_usages.into();
+    }
+}

crates/bevy_anti_aliasing/src/lib.rs

@@ -13,6 +13,8 @@ use smaa::SmaaPlugin;
 use taa::TemporalAntiAliasPlugin;
 
 pub mod contrast_adaptive_sharpening;
+#[cfg(feature = "dlss")]
+pub mod dlss;
 pub mod fxaa;
 pub mod smaa;
 pub mod taa;
@@ -21,6 +23,13 @@ pub mod taa;
 pub struct AntiAliasingPlugin;
 impl Plugin for AntiAliasingPlugin {
     fn build(&self, app: &mut bevy_app::App) {
-        app.add_plugins((FxaaPlugin, SmaaPlugin, TemporalAntiAliasPlugin, CasPlugin));
+        app.add_plugins((
+            FxaaPlugin,
+            SmaaPlugin,
+            TemporalAntiAliasPlugin,
+            CasPlugin,
+            #[cfg(feature = "dlss")]
+            dlss::DlssPlugin,
+        ));
     }
 }

crates/bevy_core_pipeline/src/core_3d/mod.rs

@@ -31,6 +31,7 @@ pub mod graph {
         Wireframe,
         LateDownsampleDepth,
         Taa,
+        Dlss,
         MotionBlur,
         Bloom,
         AutoExposure,