implement debugPrintf in msl & spv

CHANGELOG.md

@@ -58,6 +58,7 @@ Bottom level categories:
 - Unconditionally enable `Features::CLIP_DISTANCES`. By @ErichDonGubler in [#9270](https://github.com/gfx-rs/wgpu/pull/9270).
 - Added full support for mesh shaders, including in WGSL shaders. By @inner-daemons in [#8739](https://github.com/gfx-rs/wgpu/pull/8739).
 - Fixed structure field names incorrectly ignoring reserved keywords in the Metal (MSL) backend. By @39ali [#9379](https://github.com/gfx-rs/wgpu/pull/9379).
+- Implement `debugPrintf()` in Metal backend. By @39ali [#9389](https://github.com/gfx-rs/wgpu/pull/9389).
 
 #### GLES
 
@@ -73,6 +74,7 @@ Bottom level categories:
 - Add support for RawWindowHandle::Drm on unix, conditional on the `"drm"` feature.
   - DRM support by @rectalogic in [#9182](https://github.com/gfx-rs/wgpu/pull/9182).
   - Conditional compilation by @jimblandy in [#9390](https://github.com/gfx-rs/wgpu/pull/9390)
+  - Implement `debugPrintf()` in Vulkan backend. By @39ali [#9389](https://github.com/gfx-rs/wgpu/pull/9389).
 
 ### Changes
 

Cargo.toml

@@ -255,6 +255,7 @@ objc2-metal = { version = "0.3.2", default-features = false, features = [
     "MTLAccelerationStructureTypes",
     "MTLAccelerationStructureCommandEncoder",
     "MTLResidencySet",
+    "MTLLogState",
 ] }
 objc2-quartz-core = { version = "0.3.2", default-features = false, features = [
     "std",

examples/features/src/debug_printf/README.md

@@ -0,0 +1,9 @@
+# debugPrintf
+
+This example shows how to printf in shaders
+
+## To Run
+
+```
+cargo run --bin wgpu-examples debug_printf
+```

examples/features/src/debug_printf/mod.rs

@@ -0,0 +1,85 @@
+struct Example {}
+
+impl crate::framework::Example for Example {
+    fn required_features() -> wgpu::Features {
+        wgpu::Features::DEBUG_PRINTF
+    }
+
+    fn required_downlevel_capabilities() -> wgpu::DownlevelCapabilities {
+        wgpu::DownlevelCapabilities {
+            flags: wgpu::DownlevelFlags::COMPUTE_SHADERS,
+            ..Default::default()
+        }
+    }
+
+    fn required_limits() -> wgpu::Limits {
+        wgpu::Limits::default()
+    }
+
+    fn init(
+        _config: &wgpu::SurfaceConfiguration,
+        _adapter: &wgpu::Adapter,
+        device: &wgpu::Device,
+        queue: &wgpu::Queue,
+    ) -> Self {
+        let shader_source = r#"
+        enable wgpu_debug_printf;
+
+        @compute @workgroup_size(8, 1, 1)
+        fn main(@builtin(local_invocation_index) idx: u32) {
+             // We use a specific ID to identify our print in the logs
+             debugPrintf("WGSL_METAL_LOG: Thread index is %u", idx);
+        }
+    "#;
+
+        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+            label: Some("DebugPrintfShader"),
+            source: wgpu::ShaderSource::Wgsl(shader_source.into()),
+        });
+
+        // Create a simple pipeline
+        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+            label: Some("Debug Pipeline"),
+            layout: None,
+            module: &shader,
+            entry_point: Some("main"),
+            compilation_options: Default::default(),
+            cache: None,
+        });
+
+        let mut encoder =
+            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
+        {
+            let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
+                label: None,
+                timestamp_writes: None,
+            });
+            cpass.set_pipeline(&pipeline);
+            cpass.dispatch_workgroups(1, 1, 1);
+        }
+
+        queue.submit(Some(encoder.finish()));
+
+        device.poll(wgpu::PollType::wait_indefinitely()).unwrap();
+
+        Example {}
+    }
+
+    fn update(&mut self, _event: winit::event::WindowEvent) {
+        //empty
+    }
+
+    fn resize(
+        &mut self,
+        _config: &wgpu::SurfaceConfiguration,
+        _device: &wgpu::Device,
+        _queue: &wgpu::Queue,
+    ) {
+    }
+
+    fn render(&mut self, _view: &wgpu::TextureView, _device: &wgpu::Device, _queue: &wgpu::Queue) {}
+}
+
+pub fn main() {
+    crate::framework::run::<Example>("debug-printf");
+}

examples/features/src/lib.rs

@@ -10,6 +10,7 @@ pub mod bunnymark;
 pub mod conservative_raster;
 pub mod cooperative_matrix;
 pub mod cube;
+pub mod debug_printf;
 pub mod hello_synchronization;
 pub mod hello_triangle;
 pub mod hello_windows;

examples/features/src/main.rs

@@ -212,6 +212,12 @@ const EXAMPLES: &[ExampleDesc] = &[
         webgl: false,
         webgpu: false,
     },
+    ExampleDesc {
+        name: "debug_printf",
+        function: wgpu_examples::debug_printf::main,
+        webgl: false,
+        webgpu: false,
+    },
 ];
 
 fn get_example_name() -> Option<String> {

naga/src/back/dot/mod.rs

@@ -429,6 +429,15 @@ impl StatementGraph {
                         "TraceRay"
                     }
                 },
+                S::DebugPrintf {
+                    format: _,
+                    ref arguments,
+                } => {
+                    for &expr in arguments {
+                        self.dependencies.push((id, expr, "arg"));
+                    }
+                    "DebugPrintf"
+                }
             };
             // Set the last node to the merge node
             last_node = merge_id;

naga/src/back/glsl/writer.rs

@@ -2264,6 +2264,7 @@ impl<'a, W: Write> Writer<'a, W> {
             }
             Statement::CooperativeStore { .. } => unimplemented!(),
             Statement::RayPipelineFunction(_) => unimplemented!(),
+            Statement::DebugPrintf { .. } => unimplemented!(),
         }
 
         Ok(())

naga/src/back/hlsl/writer.rs

@@ -3047,6 +3047,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             }
             Statement::CooperativeStore { .. } => unimplemented!(),
             Statement::RayPipelineFunction(_) => unreachable!(),
+            Statement::DebugPrintf { .. } => unimplemented!(),
         }
 
         Ok(())

naga/src/back/msl/mod.rs

@@ -929,6 +929,7 @@ pub fn supported_capabilities() -> crate::valid::Capabilities {
         // No DRAW_INDEX
         // No MEMORY_DECORATION_VOLATILE
         | Caps::MEMORY_DECORATION_COHERENT
+        | Caps::DEBUG_PRINTF
 }
 
 #[test]

naga/src/back/msl/writer.rs

@@ -4373,6 +4373,21 @@ impl<W: Write> Writer<W> {
                     writeln!(self.out, ");")?;
                 }
                 crate::Statement::RayPipelineFunction(_) => unreachable!(),
+                crate::Statement::DebugPrintf {
+                    ref format,
+                    ref arguments,
+                } => {
+                    write!(
+                        self.out,
+                        "{level}metal::os_log_default.log_info(\"{}\"",
+                        format
+                    )?;
+                    for &arg in arguments {
+                        write!(self.out, ", ")?;
+                        self.put_expression(arg, &context.expression, true)?;
+                    }
+                    writeln!(self.out, ");")?;
+                }
             }
         }
 

naga/src/back/pipeline_constants.rs

@@ -921,6 +921,14 @@ fn adjust_stmt(new_pos: &HandleVec<Expression, Handle<Expression>>, stmt: &mut S
                 adjust(payload);
             }
         },
+        Statement::DebugPrintf {
+            format: _,
+            ref mut arguments,
+        } => {
+            for argument in arguments.iter_mut() {
+                adjust(argument);
+            }
+        }
         Statement::Break
         | Statement::Continue
         | Statement::Kill

naga/src/back/spv/block.rs

@@ -4178,6 +4178,19 @@ impl BlockContext<'_> {
                     };
                 }
                 Statement::RayPipelineFunction(_) => unreachable!(),
+                Statement::DebugPrintf {
+                    ref format,
+                    ref arguments,
+                } => {
+                    let mut format_params = Vec::with_capacity(arguments.len());
+                    for &arg in arguments {
+                        let word = self.cached[arg];
+                        format_params.push(word);
+                    }
+
+                    self.writer
+                        .write_debug_printf(&mut block, format, &format_params);
+                }
             }
         }
 

naga/src/back/spv/mod.rs

@@ -1195,4 +1195,5 @@ pub fn supported_capabilities() -> crate::valid::Capabilities {
         | Caps::DRAW_INDEX
         | Caps::MEMORY_DECORATION_COHERENT
         | Caps::MEMORY_DECORATION_VOLATILE
+        | Caps::DEBUG_PRINTF
 }

naga/src/back/wgsl/writer.rs

@@ -283,6 +283,38 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    /// Returns `true` if a `debugPrintf` is found anywhere within this statement, `false` otherwise.
+    ///
+    /// Does not traverse into function calls.
+    fn find_debug_printf(stmt: &crate::Statement) -> bool {
+        match *stmt {
+            crate::Statement::DebugPrintf { .. } => true,
+            crate::Statement::Block(ref b) => b.iter().any(Self::find_debug_printf),
+            crate::Statement::If {
+                ref accept,
+                ref reject,
+                ..
+            } => {
+                accept.iter().any(Self::find_debug_printf)
+                    || reject.iter().any(Self::find_debug_printf)
+            }
+            crate::Statement::Loop {
+                ref body,
+                ref continuing,
+                ..
+            } => {
+                body.iter().any(Self::find_debug_printf)
+                    || continuing.iter().any(Self::find_debug_printf)
+            }
+            crate::Statement::Switch { ref cases, .. } => cases
+                .iter()
+                .any(|c| c.body.iter().any(Self::find_debug_printf)),
+            // Note: does not match on function `Call`s to look inside function bodies recursively.
+            // This is fine because this is called on all functions and entrypoints in a module to check for `debugPrintf` usage.
+            _ => false,
+        }
+    }
+
     /// Helper method which writes all the `enable` declarations
     /// needed for a module.
     fn write_enable_declarations(&mut self, module: &Module) -> BackendResult {
@@ -298,6 +330,7 @@ impl<W: Write> Writer<W> {
             ray_tracing_pipeline: bool,
             per_vertex: bool,
             binding_array: bool,
+            debug_printf: bool,
         }
         let mut needed = RequiredEnabled {
             mesh_shaders: module.uses_mesh_shaders(),
@@ -427,6 +460,19 @@ impl<W: Write> Writer<W> {
             needed.ray_tracing_pipeline = true;
         }
 
+        // search for debugPrintf statements in all functions and entry points
+        let mut functions_to_check = module.functions.iter().map(|(_, f)| f).collect::<Vec<_>>();
+        functions_to_check.extend(module.entry_points.iter().map(|ep| &ep.function));
+
+        'outer: for func in functions_to_check {
+            for stmt in func.body.iter() {
+                if Self::find_debug_printf(stmt) {
+                    needed.debug_printf = true;
+                    break 'outer;
+                }
+            }
+        }
+
         // Write required declarations
         let mut any_written = false;
         if needed.f16 {
@@ -469,6 +515,10 @@ impl<W: Write> Writer<W> {
             writeln!(self.out, "enable wgpu_per_vertex;")?;
             any_written = true;
         }
+        if needed.debug_printf {
+            writeln!(self.out, "enable wgpu_debug_printf;")?;
+            any_written = true;
+        }
         if any_written {
             // Empty line for readability
             writeln!(self.out)?;
@@ -1210,6 +1260,17 @@ impl<W: Write> Writer<W> {
                     writeln!(self.out, ");")?
                 }
             },
+            Statement::DebugPrintf {
+                ref format,
+                ref arguments,
+            } => {
+                write!(self.out, "{level}debugPrintf(\"{}\"", format)?;
+                for &arg in arguments {
+                    write!(self.out, ", ")?;
+                    self.write_expr(module, arg, func_ctx)?;
+                }
+                writeln!(self.out, ");")?;
+            }
         }
 
         Ok(())

naga/src/compact/statements.rs

@@ -168,6 +168,14 @@ impl FunctionTracer<'_> {
                             self.expressions_used.insert(payload);
                         }
                     },
+                    St::DebugPrintf {
+                        format: _,
+                        ref arguments,
+                    } => {
+                        for &expr in arguments {
+                            self.expressions_used.insert(expr);
+                        }
+                    }
 
                     // Trivial statements.
                     St::Break
@@ -406,6 +414,14 @@ impl FunctionMap {
                             adjust(payload);
                         }
                     },
+                    St::DebugPrintf {
+                        format: _,
+                        ref mut arguments,
+                    } => {
+                        for expr in arguments {
+                            adjust(expr);
+                        }
+                    }
 
                     // Trivial statements.
                     St::Break

naga/src/front/spv/mod.rs

@@ -1664,6 +1664,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                 }
                 S::WorkGroupUniformLoad { .. } => unreachable!(),
                 S::CooperativeStore { .. } => unreachable!(),
+                S::DebugPrintf { .. } => unreachable!(),
             }
             i += 1;
         }

naga/src/front/wgsl/error.rs

@@ -501,6 +501,7 @@ impl<'a> Error<'a> {
                         Token::Attribute => "@".to_string(),
                         Token::Number(_) => "number".to_string(),
                         Token::Word(s) => s.to_string(),
+                        Token::String(_) => "a string literal".to_string(),
                         Token::Operation(c) => format!("operation (`{c}`)"),
                         Token::LogicalOperation(c) => format!("logical operation (`{c}`)"),
                         Token::ShiftOperation(c) => format!("bitshift (`{c}{c}`)"),