Support hot-reloading + shader overwrite (#1)

* feat: hot-reloading

* feat: support dynamic paths for hot-reloading + shader overwriting

* chore: add more wgcore examples

* fix hot-reloading boilerplate generation

* chore: clippy fixes

Author: sebcrozet

.gitignore

@@ -18,6 +18,7 @@ dist
 assets
 
 *_bg.wasm
+website
 
 # JetBrain IDEs
 #

.idea/vcs.xml

@@ -8,11 +8,11 @@ resolver = "2"
 
 [workspace.dependencies]
 nalgebra = { version = "0.33.1", features = ["convert-bytemuck"] }
-wgpu = { version = "22.1", features = ["naga-ir"] }
+wgpu = { version = "23", features = ["naga-ir"] }
 bytemuck = { version = "1", features = ["derive"] }
 anyhow = "1"
 async-channel = "2"
-naga_oil = "0.15"
+naga_oil = "0.16"
 thiserror = "1"
 encase = { version = "0.10.0", features = ["nalgebra"] }
 

.idea/wg-math.iml

@@ -73,28 +73,43 @@ pub fn derive_shader(item: TokenStream) -> TokenStream {
 
             let shader_defs = derive_shaders.shader_defs.map(|defs| quote! { #defs() })
                 .unwrap_or_else(|| quote! { Default::default() });
-            let src = derive_shaders.src_fn.map(|f| quote! { #f(include_str!(#src_path)) })
-                .unwrap_or_else(|| quote! { include_str!(#src_path).to_string() });
+            let raw_src = quote! {
+                // First try to find a path from the shader registry.
+                // If doesn’t exist in the registry, try the absolute path.
+                // If it doesn’t exist in the absolute path, load the embedded string.
+                if let Some(path) = Self::absolute_path() {
+                    // TODO: handle error
+                    std::fs::read_to_string(path).unwrap()
+                } else {
+                    include_str!(#src_path).to_string()
+                }
+            };
+
+            let src = derive_shaders.src_fn.map(|f| quote! { #f(&#raw_src) })
+                .unwrap_or_else(|| quote! { #raw_src });
+            let naga_module = quote! {
+                Self::composer().and_then(|mut c|
+                    c.make_naga_module(wgcore::re_exports::naga_oil::compose::NagaModuleDescriptor {
+                        source: &Self::src(),
+                        file_path: Self::FILE_PATH,
+                        shader_defs: #shader_defs,
+                        ..Default::default()
+                    })
+                )
+            };
 
             let from_device = if !kernels_to_build.is_empty() {
                 quote! {
-                    let module = Self::composer()
-                        .make_naga_module(wgcore::re_exports::naga_oil::compose::NagaModuleDescriptor {
-                            source: &Self::src(),
-                            file_path: Self::FILE_PATH,
-                            shader_defs: #shader_defs,
-                            ..Default::default()
-                        })
-                        .unwrap();
-                    Self {
+                    let module = #naga_module?;
+                    Ok(Self {
                         #(
                             #kernels_to_build
                         )*
-                    }
+                    })
                 }
             } else {
                 quote ! {
-                    Self
+                    Ok(Self)
                 }
             };
 
@@ -112,19 +127,37 @@ pub fn derive_shader(item: TokenStream) -> TokenStream {
                 impl wgcore::shader::Shader for #struct_identifier {
                     const FILE_PATH: &'static str = #src_path;
 
-                    fn from_device(device: &wgcore::re_exports::Device) -> Self {
+                    fn from_device(device: &wgcore::re_exports::Device) -> Result<Self, wgcore::re_exports::ComposerError> {
                         #from_device
                     }
 
-                    // TODO: could we avoid the String allocation here?
                     fn src() -> String {
                         #src
                     }
 
-                    fn compose(composer: &mut wgcore::re_exports::Composer) -> &mut wgcore::re_exports::Composer {
+                    fn naga_module() -> Result<wgcore::re_exports::wgpu::naga::Module, wgcore::re_exports::ComposerError> {
+                        #naga_module
+                    }
+
+                    fn absolute_path() -> Option<std::path::PathBuf> {
+                        if let Some(path) = wgcore::ShaderRegistry::get().get_path::<#struct_identifier>() {
+                            Some(path.clone())
+                        } else {
+                            // NOTE: this is a bit fragile, and won’t work if the current working directory
+                            //       isn’t the root of the workspace the binary crate is being run from.
+                            //       Ideally we need `proc_macro2::Span::source_file` but it is currently unstable.
+                            //       See: https://users.rust-lang.org/t/how-to-get-the-macro-called-file-path-in-a-rust-procedural-macro/109613/5
+                            std::path::Path::new(file!())
+                                .parent()?
+                                .join(Self::FILE_PATH)
+                                .canonicalize().ok()
+                        }
+                    }
+
+                    fn compose(composer: &mut wgcore::re_exports::Composer) -> Result<(), wgcore::re_exports::ComposerError> {
                         use wgcore::composer::ComposerExt;
                         #(
-                            #to_derive::compose(composer);
+                            #to_derive::compose(composer)?;
                         )*
 
                         if #composable {
@@ -134,10 +167,37 @@ pub fn derive_shader(item: TokenStream) -> TokenStream {
                                     file_path: Self::FILE_PATH,
                                     shader_defs: #shader_defs,
                                     ..Default::default()
-                                })
-                                .unwrap();
+                                })?;
+                        }
+
+                        Ok(())
+                    }
+
+                    /*
+                     * Hot reloading.
+                     */
+                    fn watch_sources(state: &mut wgcore::hot_reloading::HotReloadState) -> wgcore::re_exports::notify::Result<()> {
+                        #(
+                            #to_derive::watch_sources(state)?;
+                        )*
+
+                        if let Some(path) = Self::absolute_path() {
+                            state.watch_file(&path)?;
                         }
-                        composer
+
+                        Ok(())
+                    }
+
+                    fn needs_reload(state: &wgcore::hot_reloading::HotReloadState) -> bool {
+                        #(
+                            if #to_derive::needs_reload(state) {
+                                return true;
+                            }
+                        )*
+
+                        Self::absolute_path()
+                            .map(|path| state.file_changed(&path))
+                            .unwrap_or_default()
                     }
                 }
             }

Cargo.toml

@@ -11,10 +11,11 @@ license = "MIT OR Apache-2.0"
 
 [features]
 derive = ["wgcore-derive"]
+hot_reloading = [] # "notify"]
 
 [dependencies]
 nalgebra = { workspace = true }
-wgpu = { workspace = true }
+wgpu = { workspace = true, features = ["wgsl"] }
 bytemuck = { workspace = true }
 anyhow = { workspace = true }
 async-channel = { workspace = true }
@@ -24,6 +25,7 @@ encase = { workspace = true }
 wgcore-derive = { version = "0.1", path = "../wgcore-derive", optional = true }
 
 dashmap = "5"
+notify = { version = "7" } # , optional = true }
 
 
 [dev-dependencies]

crates/wgcore-derive/src/lib.rs

@@ -0,0 +1,36 @@
+use nalgebra::DVector;
+use wgcore::composer::ComposerExt;
+use wgcore::gpu::GpuInstance;
+use wgcore::kernel::{KernelInvocationBuilder, KernelInvocationQueue};
+use wgcore::tensor::GpuVector;
+use wgcore::Shader;
+use wgpu::{BufferUsages, ComputePipeline};
+
+#[async_std::main]
+async fn main() -> anyhow::Result<()> {
+    // Initialize the gpu device and its queue.
+    //
+    // Note that `GpuInstance` is just a simple helper struct for initializing the gpu resources.
+    // You are free to initialize them independently if more control is needed, or reuse the ones
+    // that were already created/owned by e.g., a game engine.
+    let gpu = GpuInstance::new().await?;
+
+    // Create the buffers.
+    const LEN: u32 = 1000;
+    let buffer_data = DVector::from_fn(LEN as usize, |i, _| i as u32);
+    let buffer = GpuVector::init(gpu.device(), &buffer_data, BufferUsages::STORAGE | BufferUsages::COPY_SRC);
+    let staging = GpuVector::uninit(gpu.device(), LEN, BufferUsages::COPY_DST | BufferUsages::MAP_READ);
+
+    // Queue the operation.
+    // Encode & submit the operation to the gpu.
+    let mut encoder = gpu.device().create_command_encoder(&Default::default());
+    // Copy the buffer to the staging buffer.
+    staging.copy_from(&mut encoder, &buffer);
+    gpu.queue().submit(Some(encoder.finish()));
+
+    let read = staging.read(gpu.device()).await?;
+    assert_eq!(buffer_data, DVector::from(read));
+    println!("Buffer copy & read succeeded!");
+
+    Ok(())
+}

crates/wgcore/Cargo.toml

@@ -10,17 +10,19 @@ use wgpu::{BufferUsages, ComputePipeline};
 // Note that we don’t build any compute pipeline from this wgsl file.
 #[derive(Shader)]
 #[shader(
-    src = "composable.wgsl" // Shader source code, will be embedded in the exe with `include_str!`
+    src = "compose_dependency.wgsl" // Shader source code, will be embedded in the exe with `include_str!`
 )]
 struct Composable;
 
 #[derive(Shader)]
 #[shader(
     derive(Composable), // This shader depends on the `Composable` shader.
-    src = "kernel.wgsl",  // Shader source code, will be embedded in the exe with `include_str!`.
+    src = "compose_kernel.wgsl",  // Shader source code, will be embedded in the exe with `include_str!`.
     composable = false    // This shader doesn’t export any symbols reusable from other wgsl shaders.
 )]
 struct WgKernel {
+    // This ComputePipeline field indicates that the Shader macro needs to generate the boilerplate
+    // for loading the compute pipeline in `WgKernel::from_device`.
     main: ComputePipeline,
 }
 
@@ -42,7 +44,11 @@ async fn main() -> anyhow::Result<()> {
     // Load and compile our kernel. The `from_device` function was generated by the `Shader` derive.
     // Note that its dependency to `Composable` is automatically resolved by the `Shader` derive
     // too.
-    let kernel = WgKernel::from_device(gpu.device());
+    let kernel = WgKernel::from_device(gpu.device())?;
+    println!("######################################");
+    println!("###### Composed shader sources: ######");
+    println!("######################################");
+    println!("{}", WgKernel::flat_wgsl()?);
 
     // Create the buffers.
     const LEN: u32 = 1000;

crates/wgcore/examples/buffer_readback.rs

@@ -0,0 +1,68 @@
+use nalgebra::{DVector, Vector4};
+use wgcore::composer::ComposerExt;
+use wgcore::gpu::GpuInstance;
+use wgcore::kernel::{KernelInvocationBuilder, KernelInvocationQueue};
+use wgcore::tensor::GpuVector;
+use wgcore::Shader;
+use wgpu::{BufferUsages, ComputePipeline};
+
+#[derive(Copy, Clone, PartialEq, Debug, Default, bytemuck::Pod, bytemuck::Zeroable)]
+#[repr(C)]
+pub struct BytemuckStruct {
+    value: f32,
+}
+
+#[derive(Copy, Clone, PartialEq, Debug, Default, encase::ShaderType)]
+#[repr(C)]
+pub struct EncaseStruct {
+    value: f32,
+    // This implies some internal padding, so we can’t rely on bytemuck.
+    // Encase will handle that properly.
+    value2: Vector4<f32>,
+}
+
+#[derive(Shader)]
+#[shader(
+    src = "encase.wgsl",
+    composable = false
+)]
+struct ShaderEncase {
+    main: ComputePipeline
+}
+
+#[async_std::main]
+async fn main() -> anyhow::Result<()> {
+    // Initialize the gpu device and its queue.
+    //
+    // Note that `GpuInstance` is just a simple helper struct for initializing the gpu resources.
+    // You are free to initialize them independently if more control is needed, or reuse the ones
+    // that were already created/owned by e.g., a game engine.
+    let gpu = GpuInstance::new().await?;
+
+    // Load and compile our kernel. The `from_device` function was generated by the `Shader` derive.
+    // Note that its dependency to `Composable` is automatically resolved by the `Shader` derive
+    // too.
+    let kernel = ShaderEncase::from_device(gpu.device())?;
+
+    // Create the buffers.
+    const LEN: u32 = 1000;
+    let a_data = (0..LEN).map(|x| EncaseStruct { value: x as f32, value2: Vector4::repeat(x as f32 * 10.0)}).collect::<Vec<_>>();
+    let b_data = (0..LEN).map(|x| BytemuckStruct { value: x as f32}).collect::<Vec<_>>();
+    // Call `encase` instead of `init` because `EncaseStruct` isn’t `Pod`.
+    // The `encase` function has a bit of overhead so bytemuck should be preferred whenever possible.
+    let a_buf = GpuVector::encase(gpu.device(), &a_data, BufferUsages::STORAGE);
+    let b_buf = GpuVector::init(gpu.device(), &b_data, BufferUsages::STORAGE);
+
+    // Queue the operation.
+    let mut queue = KernelInvocationQueue::new(gpu.device());
+    KernelInvocationBuilder::new(&mut queue, &kernel.main)
+        .bind0([a_buf.buffer(), b_buf.buffer()])
+        .queue(LEN.div_ceil(64));
+
+    // Encode & submit the operation to the gpu.
+    let mut encoder = gpu.device().create_command_encoder(&Default::default());
+    queue.encode(&mut encoder, None);
+    gpu.queue().submit(Some(encoder.finish()));
+
+    Ok(())
+}

crates/wgcore/examples/compose.rs

@@ -0,0 +1,22 @@
+@group(0) @binding(0)
+var<storage, read_write> a: array<EncaseStruct>;
+@group(0) @binding(1)
+var<storage, read> b: array<BytemuckStruct>;
+
+struct BytemuckStruct {
+    value: f32,
+}
+
+struct EncaseStruct {
+    value: f32,
+    value2: vec4<f32>
+}
+
+@compute @workgroup_size(64, 1, 1)
+fn main(@builtin(global_invocation_id) invocation_id: vec3<u32>) {
+    let i = invocation_id.x;
+    if i < arrayLength(&a) {
+        a[i].value += b[i].value;
+        a[i].value2 += vec4(b[i].value);
+    }
+}