Live reloading of shaders

[yrns]

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]

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.

[joshuajbouw]

And run solve clippy issues 🍺

[yrns]

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

[alec-deason]

This is an awesome bit of functionality.

[alec-deason]

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.

[alec-deason]

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.

[cart]

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.

[cart]

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

Just a few minor comments about code structure.

[cart]

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.

[cart]

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

[cart]

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"

[cart]

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]

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

bevy/crates/bevy_render/src/shader/shader.rs

Line 32 in e6867e9

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]

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

[yrns]

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]

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

[cart]

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]

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]

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.