Add headless rendering feature for GPU testing without a display

Author: LegNeato

Cargo.lock

@@ -9,6 +9,7 @@ edition = "2021"
 default = ["use-compiled-tools"]
 use-installed-tools = ["spirv-builder/use-installed-tools"]
 use-compiled-tools = ["spirv-builder/use-compiled-tools"]
+headless = ["png"]
 
 [dependencies]
 shared = { path = "shared" }
@@ -24,6 +25,7 @@ winit = "0.30.12"
 bytemuck = "1.20.0"
 env_logger = "0.11.6"
 ouroboros = "0.18.5"
+png = { version = "0.17", optional = true }
 
 [build-dependencies]
 spirv-builder.workspace = true

Cargo.toml

@@ -0,0 +1,205 @@
+use futures::executor::block_on;
+use shared::ShaderConstants;
+use std::error::Error;
+use wgpu::{self, InstanceDescriptor};
+use wgpu::{include_spirv, include_spirv_raw};
+
+pub fn run() -> Result<(), Box<dyn Error>> {
+    let width = 1280u32;
+    let height = 720u32;
+
+    let instance = wgpu::Instance::new(&InstanceDescriptor::default());
+    let adapter = block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
+        power_preference: wgpu::PowerPreference::HighPerformance,
+        compatible_surface: None,
+        force_fallback_adapter: false,
+    }))?;
+
+    let mut required_features = wgpu::Features::PUSH_CONSTANTS;
+    if adapter
+        .features()
+        .contains(wgpu::Features::SPIRV_SHADER_PASSTHROUGH)
+    {
+        required_features |= wgpu::Features::SPIRV_SHADER_PASSTHROUGH;
+    }
+
+    let required_limits = wgpu::Limits {
+        max_push_constant_size: std::mem::size_of::<ShaderConstants>() as u32,
+        ..Default::default()
+    };
+    let (device, queue) = block_on(adapter.request_device(&wgpu::DeviceDescriptor {
+        label: None,
+        required_features,
+        required_limits,
+        ..Default::default()
+    }))?;
+
+    let shader_module = if device
+        .features()
+        .contains(wgpu::Features::SPIRV_SHADER_PASSTHROUGH)
+    {
+        let x = include_spirv_raw!(env!("shadertoys_shaders.spv"));
+        unsafe { device.create_shader_module_passthrough(x) }
+    } else {
+        device.create_shader_module(include_spirv!(env!("shadertoys_shaders.spv")))
+    };
+
+    let texture_format = wgpu::TextureFormat::Rgba8UnormSrgb;
+    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+        label: None,
+        bind_group_layouts: &[],
+        push_constant_ranges: &[wgpu::PushConstantRange {
+            stages: wgpu::ShaderStages::VERTEX_FRAGMENT,
+            range: 0..std::mem::size_of::<ShaderConstants>() as u32,
+        }],
+    });
+    let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+        label: None,
+        layout: Some(&pipeline_layout),
+        vertex: wgpu::VertexState {
+            module: &shader_module,
+            entry_point: Some("main_vs"),
+            buffers: &[],
+            compilation_options: Default::default(),
+        },
+        fragment: Some(wgpu::FragmentState {
+            module: &shader_module,
+            entry_point: Some("main_fs"),
+            targets: &[Some(wgpu::ColorTargetState {
+                format: texture_format,
+                blend: Some(wgpu::BlendState::REPLACE),
+                write_mask: wgpu::ColorWrites::ALL,
+            })],
+            compilation_options: Default::default(),
+        }),
+        primitive: wgpu::PrimitiveState {
+            topology: wgpu::PrimitiveTopology::TriangleList,
+            ..Default::default()
+        },
+        depth_stencil: None,
+        multisample: wgpu::MultisampleState::default(),
+        multiview: None,
+        cache: None,
+    });
+
+    let texture = device.create_texture(&wgpu::TextureDescriptor {
+        label: Some("headless render target"),
+        size: wgpu::Extent3d {
+            width,
+            height,
+            depth_or_array_layers: 1,
+        },
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: wgpu::TextureDimension::D2,
+        format: texture_format,
+        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
+        view_formats: &[],
+    });
+    let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+    let bytes_per_pixel = 4u32;
+    let unpadded_bytes_per_row = width * bytes_per_pixel;
+    let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
+    let padded_bytes_per_row = unpadded_bytes_per_row.div_ceil(align) * align;
+    let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
+        label: Some("headless output buffer"),
+        size: (padded_bytes_per_row * height) as u64,
+        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+        mapped_at_creation: false,
+    });
+
+    let push_constants = ShaderConstants {
+        width,
+        height,
+        time: 0.0,
+        cursor_x: 0.0,
+        cursor_y: 0.0,
+        drag_start_x: 0.0,
+        drag_start_y: 0.0,
+        drag_end_x: 0.0,
+        drag_end_y: 0.0,
+        mouse_left_pressed: 0,
+        mouse_left_clicked: 0,
+    };
+
+    let mut encoder =
+        device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
+    {
+        let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+            label: None,
+            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                view: &view,
+                resolve_target: None,
+                ops: wgpu::Operations {
+                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
+                    store: wgpu::StoreOp::Store,
+                },
+            })],
+            depth_stencil_attachment: None,
+            timestamp_writes: None,
+            occlusion_query_set: None,
+        });
+        rpass.set_viewport(0.0, 0.0, width as f32, height as f32, 0.0, 1.0);
+        rpass.set_pipeline(&render_pipeline);
+        rpass.set_push_constants(
+            wgpu::ShaderStages::VERTEX_FRAGMENT,
+            0,
+            bytemuck::bytes_of(&push_constants),
+        );
+        rpass.draw(0..3, 0..1);
+    }
+    encoder.copy_texture_to_buffer(
+        wgpu::TexelCopyTextureInfo {
+            texture: &texture,
+            mip_level: 0,
+            origin: wgpu::Origin3d::ZERO,
+            aspect: wgpu::TextureAspect::All,
+        },
+        wgpu::TexelCopyBufferInfo {
+            buffer: &output_buffer,
+            layout: wgpu::TexelCopyBufferLayout {
+                offset: 0,
+                bytes_per_row: Some(padded_bytes_per_row),
+                rows_per_image: Some(height),
+            },
+        },
+        wgpu::Extent3d {
+            width,
+            height,
+            depth_or_array_layers: 1,
+        },
+    );
+    queue.submit(Some(encoder.finish()));
+
+    let buffer_slice = output_buffer.slice(..);
+    let (tx, rx) = std::sync::mpsc::channel();
+    buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
+        tx.send(result).unwrap();
+    });
+    device.poll(wgpu::PollType::Wait)?;
+    rx.recv()??;
+
+    let data = buffer_slice.get_mapped_range();
+    let output_path = "output.png";
+    let file = std::fs::File::create(output_path)?;
+    let w = &mut std::io::BufWriter::new(file);
+    let mut encoder = png::Encoder::new(w, width, height);
+    encoder.set_color(png::ColorType::Rgba);
+    encoder.set_depth(png::BitDepth::Eight);
+    encoder.set_source_srgb(png::SrgbRenderingIntent::Perceptual);
+    let mut writer = encoder.write_header()?;
+    // Remove row padding
+    let mut unpadded = Vec::with_capacity((unpadded_bytes_per_row * height) as usize);
+    for row in 0..height as usize {
+        let start = row * padded_bytes_per_row as usize;
+        let end = start + unpadded_bytes_per_row as usize;
+        unpadded.extend_from_slice(&data[start..end]);
+    }
+    writer.write_image_data(&unpadded)?;
+    drop(data);
+    output_buffer.unmap();
+
+    eprintln!("Rendered frame to {output_path}");
+    Ok(())
+}