Allow tests to set max_sets

Author: LegNeato

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

@@ -3,6 +3,21 @@ use anyhow::{Context, Result, bail};
 use ash::vk;
 use std::ffi::{CStr, CString};
 
+/// Configuration for the Ash backend
+#[derive(Clone)]
+pub struct AshConfig {
+    /// Maximum number of descriptor sets
+    pub max_descriptor_sets: u32,
+}
+
+impl Default for AshConfig {
+    fn default() -> Self {
+        Self {
+            max_descriptor_sets: 16,
+        }
+    }
+}
+
 pub struct AshBackend {
     instance: ash::Instance,
     device: ash::Device,
@@ -90,8 +105,8 @@ impl AshBackend {
     }
 }
 
-impl ComputeBackend for AshBackend {
-    fn init() -> Result<Self> {
+impl AshBackend {
+    pub fn new(config: AshConfig) -> Result<Self> {
         unsafe {
             let entry = ash::Entry::load().context("Failed to load Vulkan entry")?;
 
@@ -174,7 +189,7 @@ impl ComputeBackend for AshBackend {
 
             let descriptor_pool_create_info = vk::DescriptorPoolCreateInfo::default()
                 .pool_sizes(&descriptor_pool_sizes)
-                .max_sets(16);
+                .max_sets(config.max_descriptor_sets);
 
             let descriptor_pool = device
                 .create_descriptor_pool(&descriptor_pool_create_info, None)
@@ -191,7 +206,15 @@ impl ComputeBackend for AshBackend {
             })
         }
     }
+}
+
+impl Default for AshBackend {
+    fn default() -> Self {
+        Self::new(AshConfig::default()).expect("Failed to create AshBackend with default config")
+    }
+}
 
+impl ComputeBackend for AshBackend {
     fn run_compute(
         &self,
         spirv_bytes: &[u8],

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

@@ -19,9 +19,6 @@ pub enum BufferUsage {
 
 /// 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
     fn run_compute(
         &self,
@@ -42,14 +39,32 @@ pub struct ComputeTest<B: ComputeBackend> {
 }
 
 impl<B: ComputeBackend> ComputeTest<B> {
+    pub fn new_with_backend(
+        backend: B,
+        spirv_bytes: Vec<u8>,
+        entry_point: String,
+        dispatch: [u32; 3],
+        buffers: Vec<BufferConfig>,
+    ) -> Result<Self> {
+        Ok(Self {
+            backend,
+            spirv_bytes,
+            entry_point,
+            dispatch,
+            buffers,
+        })
+    }
+}
+
+impl<B: ComputeBackend + Default> ComputeTest<B> {
     pub fn new(
         spirv_bytes: Vec<u8>,
         entry_point: String,
         dispatch: [u32; 3],
         buffers: Vec<BufferConfig>,
     ) -> Result<Self> {
         Ok(Self {
-            backend: B::init()?,
+            backend: B::default(),
             spirv_bytes,
             entry_point,
             dispatch,

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

@@ -2,7 +2,7 @@ mod ash;
 mod backend;
 mod wgpu;
 
-pub use ash::AshBackend;
+pub use ash::{AshBackend, AshConfig};
 pub use backend::{BufferConfig, BufferUsage, ComputeBackend, ComputeTest};
 pub use wgpu::{
     RustComputeShader, WgpuBackend, WgpuComputeTest, WgpuComputeTestMultiBuffer,

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

@@ -336,15 +336,18 @@ pub struct WgpuBackend {
     queue: Arc<wgpu::Queue>,
 }
 
-impl ComputeBackend for WgpuBackend {
-    fn init() -> anyhow::Result<Self> {
-        let (device, queue) = WgpuComputeTest::<RustComputeShader>::init()?;
-        Ok(Self {
+impl Default for WgpuBackend {
+    fn default() -> Self {
+        let (device, queue) = WgpuComputeTest::<RustComputeShader>::init()
+            .expect("Failed to create WgpuBackend with default config");
+        Self {
             device: Arc::new(device),
             queue: Arc::new(queue),
-        })
+        }
     }
+}
 
+impl ComputeBackend for WgpuBackend {
     fn run_compute(
         &self,
         spirv_bytes: &[u8],