Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Live reloading of shaders #937

Merged
merged 9 commits into from
Dec 7, 2020
Merged

Conversation

yrns
Copy link
Contributor

@yrns yrns commented Nov 28, 2020

The shaders have to be loaded via asset loader out of the assets directory, so the existing crates and examples don't work without changes. I'm not sure if that's desired or if I should add an new example specific to this. I modified shader_custom_material as a working example.

The scope of changes you can make to a shader before having to rebuild the engine is limited, of course, but I think it's useful!

@joshuajbouw
Copy link
Member

Please solve the conflicts.

Also, I do wonder if we will have compiled shaders at some point? Then, you would want this to be optional.

@yrns yrns force-pushed the shader-asset-loader branch from 57b0866 to 006e6e0 Compare November 28, 2020 08:39
@joshuajbouw
Copy link
Member

And run solve clippy issues 🍺

@yrns
Copy link
Contributor Author

yrns commented Nov 28, 2020

I thought rust-analyzer runs clippy by default, but apparently it doesn't. I updated the PR.

@Moxinilian Moxinilian added A-Assets Load files from disk to use for things like images, models, and sounds C-Feature A new feature, making something new possible A-Rendering Drawing game state to the screen labels Nov 28, 2020
Copy link
Member

@alec-deason alec-deason left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is an awesome bit of functionality.

@@ -282,4 +282,22 @@ impl PipelineCompiler {
})
.flatten()
}

/// Remove any specialized shaders or pipelines that match the
/// changed shader set.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would remove "changed" here. In this local scope we don't know why the shaders in the set should be removed, just that they should.

// Create a new shader pipeline
let pipeline_handle = pipelines.add(PipelineDescriptor::default_config(ShaderStages {
vertex: shaders.add(Shader::from_glsl(ShaderStage::Vertex, VERTEX_SHADER)),
fragment: Some(shaders.add(Shader::from_glsl(ShaderStage::Fragment, FRAGMENT_SHADER))),
fragment: Some(asset_server.load::<Shader, _>("shaders/custom.frag")),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's probably worth having a comment here to explain why the vertex shader and fragment shader are being loaded in different ways, ie. to demonstrate the shader-as-asset functionality.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lets just create a new "hot_shader_reloading" example. I'd rather have multiple simple examples that illustrate specific features than pile too many features into the same example.

Copy link
Member

@cart cart left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is an awesome change that I'm sure people will like 😄

Just a few minor comments about code structure.

// Create a new shader pipeline
let pipeline_handle = pipelines.add(PipelineDescriptor::default_config(ShaderStages {
vertex: shaders.add(Shader::from_glsl(ShaderStage::Vertex, VERTEX_SHADER)),
fragment: Some(shaders.add(Shader::from_glsl(ShaderStage::Fragment, FRAGMENT_SHADER))),
fragment: Some(asset_server.load::<Shader, _>("shaders/custom.frag")),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lets just create a new "hot_shader_reloading" example. I'd rather have multiple simple examples that illustrate specific features than pile too many features into the same example.

@@ -82,6 +83,27 @@ pub fn draw_render_pipelines_system(
meshes: Res<Assets<Mesh>>,
mut query: Query<(&mut Draw, &mut RenderPipelines, &Handle<Mesh>)>,
) {
// Pipelines will be rebuilt for shaders that have been modified.
let mut changed_shaders = HashSet::default();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should move shader updating to its own system (like the current mesh_resource_provider_system)

@@ -123,6 +124,8 @@ pub enum DrawError {
pub struct DrawContext<'a> {
pub pipelines: ResMut<'a, Assets<PipelineDescriptor>>,
pub shaders: ResMut<'a, Assets<Shader>>,
pub shader_events: ResMut<'a, Events<AssetEvent<Shader>>>,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

once we move shader updating to its own system, lets move these out of DrawContext (probably into a new Local resource used by the shader update system). I don't think they should live on DrawContext because "shader updating" is orthogonal to "drawing"


/// Remove any specialized shaders or pipelines that match the
/// changed shader set.
pub fn remove_shaders(
Copy link
Member

@cart cart Dec 3, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A few things:

  1. I think we should change this to a remove_shader operation instead of removing an arbitrary set of shaders
  2. Rather than doing scans, I think when we compile pipelines we should build maps that maintain specialized_pipeline <-> specialized_shader mappings. That way we can quickly clean up the relevant resources.
  3. We currently aren't cleaning up the specialized Assets and Assets. We should do that.

It doesn't quite work, but this is the general idea:

use super::{state_descriptors::PrimitiveTopology, IndexFormat, PipelineDescriptor};
use crate::{
    pipeline::{BindType, InputStepMode, VertexBufferDescriptor},
    renderer::RenderResourceContext,
    shader::{Shader, ShaderSource},
};
use bevy_asset::{Assets, Handle};
use bevy_reflect::Reflect;
use bevy_utils::{HashMap, HashSet};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};

#[derive(Clone, Eq, PartialEq, Debug, Reflect)]
pub struct PipelineSpecialization {
    pub shader_specialization: ShaderSpecialization,
    pub primitive_topology: PrimitiveTopology,
    pub dynamic_bindings: Vec<String>,
    pub index_format: IndexFormat,
    pub vertex_buffer_descriptor: VertexBufferDescriptor,
    pub sample_count: u32,
}

impl Default for PipelineSpecialization {
    fn default() -> Self {
        Self {
            sample_count: 1,
            index_format: IndexFormat::Uint32,
            shader_specialization: Default::default(),
            primitive_topology: Default::default(),
            dynamic_bindings: Default::default(),
            vertex_buffer_descriptor: Default::default(),
        }
    }
}

impl PipelineSpecialization {
    pub fn empty() -> &'static PipelineSpecialization {
        pub static EMPTY: Lazy<PipelineSpecialization> = Lazy::new(PipelineSpecialization::default);
        &EMPTY
    }
}

#[derive(Clone, Eq, PartialEq, Debug, Default, Reflect, Serialize, Deserialize)]
pub struct ShaderSpecialization {
    pub shader_defs: HashSet<String>,
}

#[derive(Debug)]
struct SpecializedShader {
    shader: Handle<Shader>,
    pipelines: Vec<Handle<PipelineDescriptor>>,
    specialization: ShaderSpecialization,
}

#[derive(Debug)]
struct SpecializedPipeline {
    pipeline: Handle<PipelineDescriptor>,
    specialization: PipelineSpecialization,
}

#[derive(Debug, Default)]
pub struct PipelineCompiler {
    specialized_shaders: HashMap<Handle<Shader>, Vec<SpecializedShader>>,
    specialized_shader_pipelines: HashMap<Handle<Shader>, Vec<Handle<PipelineDescriptor>>>,
    specialized_pipelines: HashMap<Handle<PipelineDescriptor>, Vec<SpecializedPipeline>>,
}

impl PipelineCompiler {
    fn compile_shader(
        &mut self,
        render_resource_context: &dyn RenderResourceContext,
        shaders: &mut Assets<Shader>,
        shader_handle: &Handle<Shader>,
        shader_specialization: &ShaderSpecialization,
    ) -> Handle<Shader> {
        let specialized_shaders = self
            .specialized_shaders
            .entry(shader_handle.clone_weak())
            .or_insert_with(Vec::new);

        let shader = shaders.get(shader_handle).unwrap();

        // don't produce new shader if the input source is already spirv
        if let ShaderSource::Spirv(_) = shader.source {
            return shader_handle.clone_weak();
        }

        if let Some(specialized_shader) =
            specialized_shaders
                .iter()
                .find(|current_specialized_shader| {
                    current_specialized_shader.specialization == *shader_specialization
                })
        {
            // if shader has already been compiled with current configuration, use existing shader
            specialized_shader.shader.clone_weak()
        } else {
            // if no shader exists with the current configuration, create new shader and compile
            let shader_def_vec = shader_specialization
                .shader_defs
                .iter()
                .cloned()
                .collect::<Vec<String>>();
            let compiled_shader =
                render_resource_context.get_specialized_shader(shader, Some(&shader_def_vec));
            let specialized_handle = shaders.add(compiled_shader);
            let weak_specialized_handle = specialized_handle.clone_weak();
            specialized_shaders.push(SpecializedShader {
                shader: specialized_handle,
                pipelines: Vec::new(),
                specialization: shader_specialization.clone(),
            });
            weak_specialized_handle
        }
    }

    pub fn get_specialized_pipeline(
        &self,
        pipeline: &Handle<PipelineDescriptor>,
        specialization: &PipelineSpecialization,
    ) -> Option<Handle<PipelineDescriptor>> {
        self.specialized_pipelines
            .get(pipeline)
            .and_then(|specialized_pipelines| {
                specialized_pipelines
                    .iter()
                    .find(|current_specialized_pipeline| {
                        &current_specialized_pipeline.specialization == specialization
                    })
            })
            .map(|specialized_pipeline| specialized_pipeline.pipeline.clone_weak())
    }

    pub fn compile_pipeline(
        &mut self,
        render_resource_context: &dyn RenderResourceContext,
        pipelines: &mut Assets<PipelineDescriptor>,
        shaders: &mut Assets<Shader>,
        source_pipeline: &Handle<PipelineDescriptor>,
        pipeline_specialization: &PipelineSpecialization,
    ) -> Handle<PipelineDescriptor> {
        let source_descriptor = pipelines.get(source_pipeline).unwrap();
        let mut specialized_descriptor = source_descriptor.clone();
        let specialized_vertex_shader = self.compile_shader(
            render_resource_context,
            shaders,
            &specialized_descriptor.shader_stages.vertex,
            &pipeline_specialization.shader_specialization,
        );
        specialized_descriptor.shader_stages.vertex = specialized_vertex_shader.clone_weak();
        let mut specialized_fragment_shader = None;
        specialized_descriptor.shader_stages.fragment = specialized_descriptor
            .shader_stages
            .fragment
            .as_ref()
            .map(|fragment| {
                let shader = self.compile_shader(
                    render_resource_context,
                    shaders,
                    fragment,
                    &pipeline_specialization.shader_specialization,
                );
                specialized_fragment_shader = Some(shader.clone_weak());
                shader
            });

        let mut layout = render_resource_context.reflect_pipeline_layout(
            &shaders,
            &specialized_descriptor.shader_stages,
            true,
        );

        if !pipeline_specialization.dynamic_bindings.is_empty() {
            // set binding uniforms to dynamic if render resource bindings use dynamic
            for bind_group in layout.bind_groups.iter_mut() {
                let mut binding_changed = false;
                for binding in bind_group.bindings.iter_mut() {
                    if pipeline_specialization
                        .dynamic_bindings
                        .iter()
                        .any(|b| b == &binding.name)
                    {
                        if let BindType::Uniform {
                            ref mut dynamic, ..
                        } = binding.bind_type
                        {
                            *dynamic = true;
                            binding_changed = true;
                        }
                    }
                }

                if binding_changed {
                    bind_group.update_id();
                }
            }
        }
        specialized_descriptor.layout = Some(layout);

        // create a vertex layout that provides all attributes from either the specialized vertex buffers or a zero buffer
        let mut pipeline_layout = specialized_descriptor.layout.as_mut().unwrap();
        // the vertex buffer descriptor of the mesh
        let mesh_vertex_buffer_descriptor = &pipeline_specialization.vertex_buffer_descriptor;

        // the vertex buffer descriptor that will be used for this pipeline
        let mut compiled_vertex_buffer_descriptor = VertexBufferDescriptor {
            step_mode: InputStepMode::Vertex,
            stride: mesh_vertex_buffer_descriptor.stride,
            ..Default::default()
        };

        for shader_vertex_attribute in pipeline_layout.vertex_buffer_descriptors.iter() {
            let shader_vertex_attribute = shader_vertex_attribute
                .attributes
                .get(0)
                .expect("Reflected layout has no attributes.");

            if let Some(target_vertex_attribute) = mesh_vertex_buffer_descriptor
                .attributes
                .iter()
                .find(|x| x.name == shader_vertex_attribute.name)
            {
                // copy shader location from reflected layout
                let mut compiled_vertex_attribute = target_vertex_attribute.clone();
                compiled_vertex_attribute.shader_location = shader_vertex_attribute.shader_location;
                compiled_vertex_buffer_descriptor
                    .attributes
                    .push(compiled_vertex_attribute);
            } else {
                panic!(
                    "Attribute {} is required by shader, but not supplied by mesh. Either remove the attribute from the shader or supply the attribute ({}) to the mesh. ",
                    shader_vertex_attribute.name,
                    shader_vertex_attribute.name,
                );
            }
        }

        //TODO: add other buffers (like instancing) here
        let mut vertex_buffer_descriptors = Vec::<VertexBufferDescriptor>::default();
        vertex_buffer_descriptors.push(compiled_vertex_buffer_descriptor);

        pipeline_layout.vertex_buffer_descriptors = vertex_buffer_descriptors;
        specialized_descriptor.sample_count = pipeline_specialization.sample_count;
        specialized_descriptor.primitive_topology = pipeline_specialization.primitive_topology;
        specialized_descriptor.index_format = pipeline_specialization.index_format;

        let specialized_pipeline_handle = pipelines.add(specialized_descriptor);
        render_resource_context.create_render_pipeline(
            specialized_pipeline_handle.clone_weak(),
            pipelines.get(&specialized_pipeline_handle).unwrap(),
            &shaders,
        );

        // track specialized shader pipelines
        self.specialized_shader_pipelines
            .entry(specialized_vertex_shader)
            .or_insert_with(Default::default)
            .push(specialized_pipeline_handle.clone_weak());
        if let Some(specialized_fragment_shader) = specialized_fragment_shader {
            self.specialized_shader_pipelines
                .entry(specialized_fragment_shader)
                .or_insert_with(Default::default)
                .push(specialized_pipeline_handle.clone_weak());
        }

        let specialized_pipelines = self
            .specialized_pipelines
            .entry(source_pipeline.clone_weak())
            .or_insert_with(Vec::new);
        let weak_specialized_pipeline_handle = specialized_pipeline_handle.clone_weak();
        specialized_pipelines.push(SpecializedPipeline {
            pipeline: specialized_pipeline_handle,
            specialization: pipeline_specialization.clone(),
        });

        weak_specialized_pipeline_handle
    }

    pub fn iter_compiled_pipelines(
        &self,
        pipeline_handle: Handle<PipelineDescriptor>,
    ) -> Option<impl Iterator<Item = &Handle<PipelineDescriptor>>> {
        if let Some(compiled_pipelines) = self.specialized_pipelines.get(&pipeline_handle) {
            Some(
                compiled_pipelines
                    .iter()
                    .map(|specialized_pipeline| &specialized_pipeline.pipeline),
            )
        } else {
            None
        }
    }

    pub fn iter_all_compiled_pipelines(&self) -> impl Iterator<Item = &Handle<PipelineDescriptor>> {
        self.specialized_pipelines
            .values()
            .map(|compiled_pipelines| {
                compiled_pipelines
                    .iter()
                    .map(|specialized_pipeline| &specialized_pipeline.pipeline)
            })
            .flatten()
    }

    /// Remove any specialized shaders or pipelines that match the
    /// changed shader set.
    pub fn remove_shader(
        &mut self,
        shader: &Handle<Shader>,
        pipelines: &mut Assets<PipelineDescriptor>,
        shaders: &mut Assets<Shader>,
    ) {
        if let Some(specialized_shaders) = self.specialized_shaders.remove(shader) {
            for specialized_shader in specialized_shaders {
                shaders.remove(&specialized_shader.shader);
                if let Some(specialized_pipelines) = self
                    .specialized_shader_pipelines
                    .get(&specialized_shader.shader)
                {
                    for specialized_pipeline in specialized_pipelines {
                        println!("remove pipeline");
                        self.specialized_pipelines.remove(specialized_pipeline);
                        pipelines.remove(specialized_pipeline);
                    }
                }
            }
        }
    }
}

@alec-deason
Copy link
Member

I've run into one frustration using this. When I write a shader that doesn't compile, this panics:

bevy_glsl_to_spirv::compile(glsl_source, stage.into(), shader_defs).unwrap()

When I'm not relying on reloading that doesn't bother me but when the shaders are hot reloading it's a really bad experience. I'm not sure what the right thing to do is but it seems like logging an error and continuing to use the old version of the shader is reasonable.

@cart
Copy link
Member

cart commented Dec 3, 2020

Yeah I think we might want to make asset load failures "soft failures" (ex: just log an error message instead of panic-ing).

@yrns yrns marked this pull request as draft December 4, 2020 23:08
@yrns
Copy link
Contributor Author

yrns commented Dec 5, 2020

I think I addressed all the feedback:

  • created a separate system for updating
  • new example
  • less panic
  • asset cleanup

I followed your example code and then things got a little weird trying to avoid the panic on compilation failure. I wish I had a snazzier example, but maybe that can be a separate PR.

@alec-deason
Copy link
Member

Works great for me, even when I mangle the shader.

@yrns yrns marked this pull request as ready for review December 5, 2020 07:35
@cart
Copy link
Member

cart commented Dec 7, 2020

Fantastic work! This is already usable as is, but I think its worth doing one last bit of cleanup now (if you're up for it). We currently leak the "gpu allocated" shader and pipeline resources when they are removed.

when we first load the hot_shader_reloading example, the wgpu resource counts look like this:

Diagnostics:
---------------------------------------------------------------------------------------------
render_pipelines                                                 : 1.000000    (avg 1.000000)

bind_group_layouts                                               : 3.000000    (avg 3.000000)

swap_chains                                                      : 1.000000    (avg 1.000000)

window_surfaces                                                  : 1.000000    (avg 1.000000)

textures                                                         : 1.000000    (avg 1.000000)

swap_chain_outputs                                               : 0.000000    (avg 0.000000)

buffers                                                          : 17.000000   (avg 17.000000)

bind_group_ids                                                   : 3.000000    (avg 3.000000)

texture_views                                                    : 1.000000    (avg 1.000000)

bind_groups                                                      : 3.000000    (avg 3.000000)

shader_modules                                                   : 2.000000    (avg 2.000000)

samplers                                                         : 0.000000    (avg 0.000000)

Each time we save a change to a shader, we create new versions of the shader module and pipeline groups. After saving 5 changes we get this:

Diagnostics:
---------------------------------------------------------------------------------------------
render_pipelines                                                 : 6.000000    (avg 6.000000)

bind_group_layouts                                               : 3.000000    (avg 3.000000)

swap_chains                                                      : 1.000000    (avg 1.000000)

window_surfaces                                                  : 1.000000    (avg 1.000000)

textures                                                         : 1.000000    (avg 1.000000)

swap_chain_outputs                                               : 0.000000    (avg 0.000000)

buffers                                                          : 17.000000   (avg 17.000000)

bind_group_ids                                                   : 3.000000    (avg 3.000000)

texture_views                                                    : 1.000000    (avg 1.000000)

bind_groups                                                      : 3.000000    (avg 3.000000)

shader_modules                                                   : 7.000000    (avg 7.000000)

samplers                                                         : 0.000000    (avg 0.000000)

@cart
Copy link
Member

cart commented Dec 7, 2020

diagnostics printed to console using these two plugins:

.add_plugin(bevy::wgpu::diagnostic::WgpuResourceDiagnosticsPlugin::default())
.add_plugin(bevy::diagnostic::PrintDiagnosticsPlugin::default())

However on second thought, adding support for cleanup involves extending the RenderResourceContext abstraction, which feels a bit out of scope for this pr. Ill just open an issue for it and we'll get to it eventually.

@yrns
Copy link
Contributor Author

yrns commented Dec 7, 2020

I noticed there was a second set of shaders stored on the backend when I tried to reuse the same handle with an updated shader asset. (It just kept using the old shader.) I don't know what they are saved for, but I assumed they were being managed already. Thanks everyone for reviewing.

@yrns yrns deleted the shader-asset-loader branch December 7, 2020 22:25
@fopsdev fopsdev mentioned this pull request Jan 24, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
A-Assets Load files from disk to use for things like images, models, and sounds A-Rendering Drawing game state to the screen C-Feature A new feature, making something new possible
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants