From 9dc57617e45b29ad3e7b29091fe7b8659d83f7f7 Mon Sep 17 00:00:00 2001
From: wicast <wicastchen@hotmail.com>
Date: Tue, 24 Oct 2023 07:44:04 +0800
Subject: [PATCH] feat(spv): shader debug option (#4028)

Co-authored-by: Connor Fitzgerald <connorwadefitzgerald@gmail.com>
---
 wgpu-core/src/device/resource.rs  | 21 +++++++++++++++++++--
 wgpu-hal/examples/halmark/main.rs |  1 +
 wgpu-hal/src/lib.rs               |  8 ++++++++
 wgpu-hal/src/vulkan/device.rs     | 20 +++++++++++++++++++-
 4 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/wgpu-core/src/device/resource.rs b/wgpu-core/src/device/resource.rs
index 560f9890d4..fd85fd6a77 100644
--- a/wgpu-core/src/device/resource.rs
+++ b/wgpu-core/src/device/resource.rs
@@ -1317,6 +1317,19 @@ impl<A: HalApi> Device<A> {
                 .contains(wgt::DownlevelFlags::CUBE_ARRAY_TEXTURES),
         );
 
+        let debug_source = if self.instance_flags.contains(wgt::InstanceFlags::DEBUG) {
+            Some(hal::DebugSource {
+                file_name: Cow::Owned(
+                    desc.label
+                        .as_ref()
+                        .map_or("shader".to_string(), |l| l.to_string()),
+                ),
+                source_code: Cow::Owned(source.clone()),
+            })
+        } else {
+            None
+        };
+
         let info = naga::valid::Validator::new(naga::valid::ValidationFlags::all(), caps)
             .validate(&module)
             .map_err(|inner| {
@@ -1326,10 +1339,14 @@ impl<A: HalApi> Device<A> {
                     inner: Box::new(inner),
                 })
             })?;
+
         let interface =
             validation::Interface::new(&module, &info, self.limits.clone(), self.features);
-        let hal_shader = hal::ShaderInput::Naga(hal::NagaShader { module, info });
-
+        let hal_shader = hal::ShaderInput::Naga(hal::NagaShader {
+            module,
+            info,
+            debug_source,
+        });
         let hal_desc = hal::ShaderModuleDescriptor {
             label: desc.label.to_hal(self.instance_flags),
             runtime_checks: desc.shader_bound_checks.runtime_checks(),
diff --git a/wgpu-hal/examples/halmark/main.rs b/wgpu-hal/examples/halmark/main.rs
index b751a85e96..527c45e2aa 100644
--- a/wgpu-hal/examples/halmark/main.rs
+++ b/wgpu-hal/examples/halmark/main.rs
@@ -156,6 +156,7 @@ impl<A: hal::Api> Example<A> {
             hal::NagaShader {
                 module: Cow::Owned(module),
                 info,
+                debug_source: None,
             }
         };
         let shader_desc = hal::ShaderModuleDescriptor {
diff --git a/wgpu-hal/src/lib.rs b/wgpu-hal/src/lib.rs
index 5b4d0f45fd..2e989499e4 100644
--- a/wgpu-hal/src/lib.rs
+++ b/wgpu-hal/src/lib.rs
@@ -1097,6 +1097,8 @@ pub struct NagaShader {
     pub module: Cow<'static, naga::Module>,
     /// Analysis information of the module.
     pub info: naga::valid::ModuleInfo,
+    /// Source codes for debug
+    pub debug_source: Option<DebugSource>,
 }
 
 // Custom implementation avoids the need to generate Debug impl code
@@ -1119,6 +1121,12 @@ pub struct ShaderModuleDescriptor<'a> {
     pub runtime_checks: bool,
 }
 
+#[derive(Debug, Clone)]
+pub struct DebugSource {
+    pub file_name: Cow<'static, str>,
+    pub source_code: Cow<'static, str>,
+}
+
 /// Describes a programmable pipeline stage.
 #[derive(Debug)]
 pub struct ProgrammableStage<'a, A: Api> {
diff --git a/wgpu-hal/src/vulkan/device.rs b/wgpu-hal/src/vulkan/device.rs
index f541fd1032..d88b48ef73 100644
--- a/wgpu-hal/src/vulkan/device.rs
+++ b/wgpu-hal/src/vulkan/device.rs
@@ -723,7 +723,9 @@ impl super::Device {
                     entry_point: stage.entry_point.to_string(),
                     shader_stage: naga_stage,
                 };
-                let needs_temp_options = !runtime_checks || !binding_map.is_empty();
+                let needs_temp_options = !runtime_checks
+                    || !binding_map.is_empty()
+                    || naga_shader.debug_source.is_some();
                 let mut temp_options;
                 let options = if needs_temp_options {
                     temp_options = self.naga_options.clone();
@@ -739,6 +741,14 @@ impl super::Device {
                     if !binding_map.is_empty() {
                         temp_options.binding_map = binding_map.clone();
                     }
+
+                    if let Some(ref debug) = naga_shader.debug_source {
+                        temp_options.debug_info = Some(naga::back::spv::DebugInfo {
+                            source_code: &debug.source_code,
+                            file_name: debug.file_name.as_ref().as_ref(),
+                        })
+                    }
+
                     &temp_options
                 } else {
                     &self.naga_options
@@ -1513,6 +1523,14 @@ impl crate::Device<super::Api> for super::Device {
                     });
                 }
                 let mut naga_options = self.naga_options.clone();
+                naga_options.debug_info =
+                    naga_shader
+                        .debug_source
+                        .as_ref()
+                        .map(|d| naga::back::spv::DebugInfo {
+                            source_code: d.source_code.as_ref(),
+                            file_name: d.file_name.as_ref().as_ref(),
+                        });
                 if !desc.runtime_checks {
                     naga_options.bounds_check_policies = naga::proc::BoundsCheckPolicies {
                         index: naga::proc::BoundsCheckPolicy::Unchecked,