Initial DLSS implementation

Cargo.toml

@@ -467,6 +467,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"]
 
@@ -1012,7 +1015,7 @@ doc-scrape-examples = true
 
 [package.metadata.example.anti_aliasing]
 name = "Anti-aliasing"
-description = "Compares different anti-aliasing methods"
+description = "Compares different anti-aliasing techniques supported by Bevy"
 category = "3D Rendering"
 # TAA not supported by WebGL
 wasm = false

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 = { git = "https://github.com/bevyengine/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,116 @@
+//! NVIDIA Deep Learning Super Sampling (DLSS).
+//!
+//! DLSS uses machine learning models to upscale and anti-alias images.
+//!
+//! Requires a NVIDIA RTX GPU, and the Windows/Linux Vulkan rendering backend. Does not work on other platforms.
+//!
+//! See https://github.com/bevyengine/dlss_wgpu for licensing requirements and setup instructions.
+//!
+//! # Usage
+//! 1. Enable Bevy's `dlss` feature
+//! 2. During app setup, insert the `DlssProjectId` resource before `DefaultPlugins`
+//! 3. Check for the presence of `Option<Res<DlssSupported>>` at runtime to see if DLSS is supported on the current machine
+//! 4. Add the `Dlss` component to your camera entity, optionally setting a specific `DlssPerfQualityMode` (defaults to `Auto`)
+//! 5. Optionally add sharpening via `ContrastAdaptiveSharpening`
+//! 6. Custom rendering code, including third party crates, should account for the optional `MainPassResolutionOverride` to work with DLSS (see the `custom_render_phase` example)
+
+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::{RenderGraphExt, 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::prepare_dlss
+                    .in_set(RenderSystems::ManageViews)
+                    .before(prepare_view_targets),
+            )
+            .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,
+                ),
+            );
+    }
+}
+
+/// Camera component to enable DLSS.
+#[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,
+            resolution_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 = resolution_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,117 @@
+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};
+
+#[derive(Component)]
+pub struct ViewDlssContext {
+    pub context: Mutex<DlssContext>,
+    pub perf_quality_mode: DlssPerfQualityMode,
+    pub feature_flags: DlssFeatureFlags,
+}
+
+pub fn prepare_dlss(
+    mut query: Query<
+        (
+            Entity,
+            &ExtractedView,
+            &Dlss,
+            &mut Camera3d,
+            &mut CameraMainTextureUsages,
+            &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 camera_3d,
+        mut camera_main_texture_usages,
+        mut temporal_jitter,
+        mut mip_bias,
+        mut dlss_context,
+    ) in &mut query
+    {
+        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();
+
+        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),
+                ));
+            }
+        }
+    }
+}