Make ash backend use ComputeShader trait to reduce code duplication

Author: LegNeato

tests/difftests/lib/src/scaffold/compute/backend.rs

@@ -17,19 +17,32 @@ pub enum BufferUsage {
     Uniform,
 }
 
+use super::SpirvShader;
+
 /// A generic trait for compute backends
 pub trait ComputeBackend: Sized {
     /// Initialize the backend
     fn init() -> Result<Self>;
 
-    /// Create and run a compute shader with multiple buffers
+    /// Create and run a compute shader with multiple buffers from raw SPIRV bytes
     fn run_compute(
         &self,
         spirv_bytes: &[u8],
         entry_point: &str,
         dispatch: [u32; 3],
         buffers: Vec<BufferConfig>,
     ) -> Result<Vec<Vec<u8>>>;
+
+    /// Create and run a compute shader with multiple buffers from a shader object
+    fn run_compute_shader<S: SpirvShader>(
+        &self,
+        shader: &S,
+        dispatch: [u32; 3],
+        buffers: Vec<BufferConfig>,
+    ) -> Result<Vec<Vec<u8>>> {
+        let (spirv_bytes, entry_point) = shader.spirv_bytes()?;
+        self.run_compute(&spirv_bytes, &entry_point, dispatch, buffers)
+    }
 }
 
 /// A compute test that can run on any backend
@@ -82,3 +95,43 @@ impl<B: ComputeBackend> ComputeTest<B> {
         anyhow::bail!("No storage buffer output found")
     }
 }
+
+/// A compute test that can run on any backend using a shader object
+pub struct ComputeShaderTest<B: ComputeBackend, S: SpirvShader> {
+    backend: B,
+    shader: S,
+    dispatch: [u32; 3],
+    buffers: Vec<BufferConfig>,
+}
+
+impl<B: ComputeBackend, S: SpirvShader> ComputeShaderTest<B, S> {
+    pub fn new(shader: S, dispatch: [u32; 3], buffers: Vec<BufferConfig>) -> Result<Self> {
+        Ok(Self {
+            backend: B::init()?,
+            shader,
+            dispatch,
+            buffers,
+        })
+    }
+
+    pub fn run(self) -> Result<Vec<Vec<u8>>> {
+        self.backend
+            .run_compute_shader(&self.shader, self.dispatch, self.buffers)
+    }
+
+    pub fn run_test(self, config: &Config) -> Result<()> {
+        let buffers = self.buffers.clone();
+        let outputs = self.run()?;
+        // Write the first storage buffer output to the file
+        for (output, buffer_config) in outputs.iter().zip(&buffers) {
+            if matches!(buffer_config.usage, BufferUsage::Storage) && !output.is_empty() {
+                use std::fs::File;
+                use std::io::Write;
+                let mut f = File::create(&config.output_path)?;
+                f.write_all(output)?;
+                return Ok(());
+            }
+        }
+        anyhow::bail!("No storage buffer output found")
+    }
+}

tests/difftests/lib/src/scaffold/compute/mod.rs

@@ -3,8 +3,8 @@ mod backend;
 mod wgpu;
 
 pub use ash::AshBackend;
-pub use backend::{BufferConfig, BufferUsage, ComputeBackend, ComputeTest};
+pub use backend::{BufferConfig, BufferUsage, ComputeBackend, ComputeShaderTest, ComputeTest};
 pub use wgpu::{
-    RustComputeShader, WgpuBackend, WgpuComputeTest, WgpuComputeTestMultiBuffer,
-    WgpuComputeTestPushConstants, WgslComputeShader,
+    RustComputeShader, SpirvShader, WgpuBackend, WgpuComputeTest, WgpuComputeTestMultiBuffer,
+    WgpuComputeTestPushConstants, WgpuShader, WgslComputeShader,
 };

tests/difftests/lib/src/scaffold/compute/wgpu.rs

@@ -18,8 +18,15 @@ use super::backend::{self, ComputeBackend};
 pub type BufferConfig = backend::BufferConfig;
 pub type BufferUsage = backend::BufferUsage;
 
-/// Trait that creates a shader module and provides its entry point.
-pub trait ComputeShader {
+/// Trait for shaders that can provide SPIRV bytes.
+pub trait SpirvShader {
+    /// Returns the SPIRV bytes and entry point name.
+    fn spirv_bytes(&self) -> anyhow::Result<(Vec<u8>, String)>;
+}
+
+/// Trait for shaders that can create wgpu modules.
+pub trait WgpuShader {
+    /// Creates a wgpu shader module.
     fn create_module(
         &self,
         device: &wgpu::Device,
@@ -29,25 +36,46 @@ pub trait ComputeShader {
 /// A compute shader written in Rust compiled with spirv-builder.
 pub struct RustComputeShader {
     pub path: PathBuf,
+    pub target: String,
+    pub capabilities: Vec<spirv_builder::Capability>,
 }
 
 impl RustComputeShader {
     pub fn new<P: Into<PathBuf>>(path: P) -> Self {
-        Self { path: path.into() }
+        Self {
+            path: path.into(),
+            target: "spirv-unknown-vulkan1.1".to_string(),
+            capabilities: Vec::new(),
+        }
+    }
+
+    pub fn with_target<P: Into<PathBuf>>(path: P, target: impl Into<String>) -> Self {
+        Self {
+            path: path.into(),
+            target: target.into(),
+            capabilities: Vec::new(),
+        }
+    }
+
+    pub fn with_capability(mut self, capability: spirv_builder::Capability) -> Self {
+        self.capabilities.push(capability);
+        self
     }
 }
 
-impl ComputeShader for RustComputeShader {
-    fn create_module(
-        &self,
-        device: &wgpu::Device,
-    ) -> anyhow::Result<(wgpu::ShaderModule, Option<String>)> {
-        let builder = SpirvBuilder::new(&self.path, "spirv-unknown-vulkan1.1")
+impl SpirvShader for RustComputeShader {
+    fn spirv_bytes(&self) -> anyhow::Result<(Vec<u8>, String)> {
+        let mut builder = SpirvBuilder::new(&self.path, &self.target)
             .print_metadata(spirv_builder::MetadataPrintout::None)
             .release(true)
             .multimodule(false)
             .shader_panic_strategy(spirv_builder::ShaderPanicStrategy::SilentExit)
             .preserve_bindings(true);
+
+        for capability in &self.capabilities {
+            builder = builder.capability(*capability);
+        }
+
         let artifact = builder.build().context("SpirvBuilder::build() failed")?;
 
         if artifact.entry_points.len() != 1 {
@@ -66,6 +94,17 @@ impl ComputeShader for RustComputeShader {
             }
         };
 
+        Ok((shader_bytes, entry_point))
+    }
+}
+
+impl WgpuShader for RustComputeShader {
+    fn create_module(
+        &self,
+        device: &wgpu::Device,
+    ) -> anyhow::Result<(wgpu::ShaderModule, Option<String>)> {
+        let (shader_bytes, entry_point) = self.spirv_bytes()?;
+
         if shader_bytes.len() % 4 != 0 {
             anyhow::bail!("SPIR-V binary length is not a multiple of 4");
         }
@@ -93,7 +132,7 @@ impl WgslComputeShader {
     }
 }
 
-impl ComputeShader for WgslComputeShader {
+impl WgpuShader for WgslComputeShader {
     fn create_module(
         &self,
         device: &wgpu::Device,
@@ -133,7 +172,7 @@ pub struct WgpuComputeTestPushConstants<S> {
 
 impl<S> WgpuComputeTest<S>
 where
-    S: ComputeShader,
+    S: WgpuShader,
 {
     pub fn new(shader: S, dispatch: [u32; 3], output_bytes: u64) -> Self {
         Self {
@@ -544,7 +583,7 @@ impl Default for RustComputeShader {
 
 impl<S> WgpuComputeTestMultiBuffer<S>
 where
-    S: ComputeShader,
+    S: WgpuShader,
 {
     pub fn new(shader: S, dispatch: [u32; 3], buffers: Vec<BufferConfig>) -> Self {
         Self {
@@ -714,7 +753,7 @@ where
 
 impl<S> WgpuComputeTestPushConstants<S>
 where
-    S: ComputeShader,
+    S: WgpuShader,
 {
     pub fn new(
         shader: S,

tests/difftests/tests/arch/workgroup_memory/workgroup_memory-ash/src/main.rs

@@ -17,37 +17,10 @@ fn main() {
     // Run the actual test on other platforms
     #[cfg(not(target_os = "macos"))]
     {
-        use difftest::scaffold::compute::{AshBackend, BufferConfig, BufferUsage, ComputeTest};
-        use difftest::spirv_builder::{
-            Capability, MetadataPrintout, ModuleResult, ShaderPanicStrategy, SpirvBuilder,
-        };
-        use std::fs;
-
-        // Build the Rust shader to SPIR-V
-        let builder = SpirvBuilder::new(".", "spirv-unknown-vulkan1.2")
-            .print_metadata(MetadataPrintout::None)
-            .release(true)
-            .multimodule(false)
-            .shader_panic_strategy(ShaderPanicStrategy::SilentExit)
-            .preserve_bindings(true)
-            .capability(Capability::VulkanMemoryModel);
-
-        let artifact = builder.build().expect("Failed to build SPIR-V");
-
-        if artifact.entry_points.len() != 1 {
-            panic!(
-                "Expected exactly one entry point, found {}",
-                artifact.entry_points.len()
-            );
-        }
-        let entry_point = artifact.entry_points.into_iter().next().unwrap();
-
-        let spirv_bytes = match artifact.module {
-            ModuleResult::SingleModule(path) => {
-                fs::read(&path).expect("Failed to read SPIR-V file")
-            }
-            ModuleResult::MultiModule(_) => panic!("Unexpected multi-module result"),
+        use difftest::scaffold::compute::{
+            AshBackend, BufferConfig, BufferUsage, ComputeShaderTest, RustComputeShader,
         };
+        use difftest::spirv_builder::Capability;
 
         // Initialize input buffer with values to sum
         let input_data: Vec<u32> = (1..=64).collect();
@@ -66,9 +39,10 @@ fn main() {
             },
         ];
 
-        let test = ComputeTest::<AshBackend>::new(
-            spirv_bytes,
-            entry_point,
+        let shader = RustComputeShader::with_target(".", "spirv-unknown-vulkan1.2")
+            .with_capability(Capability::VulkanMemoryModel);
+        let test = ComputeShaderTest::<AshBackend, _>::new(
+            shader,
             [1, 1, 1], // Single workgroup with 64 threads
             buffers,
         )