librashader_test/render/
wgpu.rs

1use crate::render::{CommonFrameOptions, RenderTest};
2use anyhow::anyhow;
3use image::RgbaImage;
4use librashader::runtime::wgpu::*;
5use librashader::runtime::{Size, Viewport};
6use librashader_runtime::image::{Image, UVDirection};
7use std::io::{Cursor, Write};
8use std::ops::DerefMut;
9use std::path::Path;
10use std::sync::Arc;
11use wgpu::{Adapter, Device, Instance, Queue, TexelCopyBufferInfo, Texture};
12use wgpu_types::{
13    BufferAddress, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, TexelCopyBufferLayout,
14    TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
15};
16
17use librashader::presets::ShaderPreset;
18use librashader::runtime::{FilterChainParameters, RuntimeParameters};
19use parking_lot::Mutex;
20
21pub struct Wgpu {
22    _instance: Instance,
23    _adapter: Adapter,
24    device: Arc<Device>,
25    queue: Arc<Queue>,
26    image: Image,
27    texture: Arc<Texture>,
28}
29
30struct BufferDimensions {
31    height: usize,
32    unpadded_bytes_per_row: usize,
33    padded_bytes_per_row: usize,
34}
35
36impl BufferDimensions {
37    fn new(width: usize, height: usize) -> Self {
38        let bytes_per_pixel = std::mem::size_of::<u32>();
39        let unpadded_bytes_per_row = width * bytes_per_pixel;
40        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize;
41        let padded_bytes_per_row_padding = (align - unpadded_bytes_per_row % align) % align;
42        let padded_bytes_per_row = unpadded_bytes_per_row + padded_bytes_per_row_padding;
43        Self {
44            height,
45            unpadded_bytes_per_row,
46            padded_bytes_per_row,
47        }
48    }
49}
50
51impl RenderTest for Wgpu {
52    fn new(path: &Path) -> anyhow::Result<Self>
53    where
54        Self: Sized,
55    {
56        Wgpu::new(path)
57    }
58
59    fn image_size(&self) -> Size<u32> {
60        self.image.size
61    }
62
63    fn render_with_preset_and_params(
64        &mut self,
65        preset: ShaderPreset,
66        frame_count: usize,
67        output_size: Option<Size<u32>>,
68        param_setter: Option<&dyn Fn(&RuntimeParameters)>,
69        frame_options: Option<CommonFrameOptions>,
70    ) -> anyhow::Result<image::RgbaImage> {
71        let mut chain = FilterChain::load_from_preset(
72            preset,
73            &self.device,
74            &self.queue,
75            Some(&FilterChainOptions {
76                force_no_mipmaps: false,
77                enable_cache: true,
78                adapter_info: None,
79                ..Default::default()
80            }),
81        )?;
82        if let Some(setter) = param_setter {
83            setter(&chain.parameters());
84        }
85        let mut cmd = self
86            .device
87            .create_command_encoder(&CommandEncoderDescriptor { label: None });
88
89        let output_tex = self.device.create_texture(&TextureDescriptor {
90            label: None,
91            size: output_size.map_or(self.texture.size(), |size| size.into()),
92            mip_level_count: 1,
93            sample_count: 1,
94            dimension: TextureDimension::D2,
95            format: TextureFormat::Rgba8Unorm,
96            usage: TextureUsages::RENDER_ATTACHMENT | TextureUsages::COPY_SRC,
97            view_formats: &[wgpu::TextureFormat::Rgba8Unorm],
98        });
99
100        let buffer_dimensions =
101            BufferDimensions::new(output_tex.width() as usize, output_tex.height() as usize);
102        let output_buf = Arc::new(self.device.create_buffer(&BufferDescriptor {
103            label: None,
104            size: (buffer_dimensions.padded_bytes_per_row * buffer_dimensions.height)
105                as BufferAddress, // 4bpp
106            usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
107            mapped_at_creation: false,
108        }));
109
110        let view = output_tex.create_view(&wgpu::TextureViewDescriptor::default());
111        let output = WgpuOutputView::new_from_raw(
112            &view,
113            output_tex.size().into(),
114            TextureFormat::Rgba8Unorm,
115        );
116
117        let viewport = Viewport::new_render_target_sized_origin(output, None)?;
118        let options = frame_options.map(|options| FrameOptions {
119            clear_history: options.clear_history,
120            frame_direction: options.frame_direction,
121            rotation: options.rotation,
122            total_subframes: options.total_subframes,
123            current_subframe: options.current_subframe,
124            aspect_ratio: options.aspect_ratio,
125            frametime_delta: options.frametime_delta,
126            frames_per_second: options.frames_per_second,
127            ..Default::default()
128        });
129
130        for frame in 0..=frame_count {
131            chain.frame(&self.texture, &viewport, &mut cmd, frame, options.as_ref())?;
132        }
133
134        cmd.copy_texture_to_buffer(
135            output_tex.as_image_copy(),
136            TexelCopyBufferInfo {
137                buffer: &output_buf,
138                layout: TexelCopyBufferLayout {
139                    offset: 0,
140                    bytes_per_row: Some(buffer_dimensions.padded_bytes_per_row as u32),
141                    rows_per_image: None,
142                },
143            },
144            output_tex.size(),
145        );
146
147        let si = self.queue.submit([cmd.finish()]);
148        self.device.poll(wgpu::PollType::Wait {
149            submission_index: Some(si),
150            timeout: None,
151        })?;
152
153        let capturable = Arc::clone(&output_buf);
154
155        let pixels = Arc::new(Mutex::new(Vec::new()));
156
157        let pixels_async = Arc::clone(&pixels);
158        output_buf
159            .slice(..)
160            .map_async(wgpu::MapMode::Read, move |r| {
161                if r.is_ok() {
162                    if let Ok(buffer) = capturable.slice(..).get_mapped_range() {
163                        let mut pixels = pixels_async.lock();
164                        pixels.resize(buffer.len(), 0);
165
166                        let mut cursor = Cursor::new(pixels.deref_mut());
167                        for chunk in buffer.chunks(buffer_dimensions.padded_bytes_per_row) {
168                            cursor
169                                .write_all(&chunk[..buffer_dimensions.unpadded_bytes_per_row])
170                                .unwrap()
171                        }
172
173                        cursor.into_inner();
174                    }
175                }
176                capturable.unmap();
177            });
178
179        self.device.poll(wgpu::PollType::Wait {
180            submission_index: None,
181            timeout: None,
182        })?;
183
184        if pixels.lock().len() == 0 {
185            return Err(anyhow!("failed to copy pixels from buffer"));
186        }
187
188        let image = RgbaImage::from_raw(
189            output_tex.width(),
190            output_tex.height(),
191            pixels.lock().to_vec(),
192        )
193        .ok_or(anyhow!("Unable to create image from data"))?;
194
195        Ok(image)
196    }
197}
198
199impl Wgpu {
200    pub fn new(image: &Path) -> anyhow::Result<Self> {
201        pollster::block_on(async {
202            let instance = wgpu::Instance::default();
203            let adapter = instance
204                .request_adapter(&wgpu::RequestAdapterOptions {
205                    power_preference: wgpu::PowerPreference::default(),
206                    compatible_surface: None,
207                    force_fallback_adapter: false,
208                    apply_limit_buckets: false,
209                })
210                .await?;
211
212            let (device, queue) = adapter
213                .request_device(&wgpu::DeviceDescriptor {
214                    required_features: wgpu::Features::ADDRESS_MODE_CLAMP_TO_BORDER
215                        | wgpu::Features::PIPELINE_CACHE
216                        | wgpu::Features::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES
217                        | wgpu::Features::FLOAT32_FILTERABLE,
218                    required_limits: wgpu::Limits::default(),
219                    label: None,
220                    memory_hints: Default::default(),
221                    ..Default::default()
222                })
223                .await?;
224            let (image, texture) = Self::load_image(&device, &queue, image)?;
225
226            Ok(Self {
227                _instance: instance,
228                _adapter: adapter,
229                device: Arc::new(device),
230                queue: Arc::new(queue),
231                image,
232                texture: Arc::new(texture),
233            })
234        })
235    }
236
237    fn load_image(device: &Device, queue: &Queue, path: &Path) -> anyhow::Result<(Image, Texture)> {
238        let image = Image::load(path, UVDirection::TopLeft)?;
239        let texture = device.create_texture(&TextureDescriptor {
240            size: image.size.into(),
241            mip_level_count: 1,
242            sample_count: 1,
243            dimension: TextureDimension::D2,
244            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
245            format: wgpu::TextureFormat::Rgba8Unorm,
246            view_formats: &[wgpu::TextureFormat::Rgba8Unorm],
247            label: None,
248        });
249
250        queue.write_texture(
251            wgpu::TexelCopyTextureInfo {
252                texture: &texture,
253                mip_level: 0,
254                origin: wgpu::Origin3d::ZERO,
255                aspect: wgpu::TextureAspect::All,
256            },
257            &image.bytes,
258            wgpu::TexelCopyBufferLayout {
259                offset: 0,
260                bytes_per_row: Some(4 * image.size.width),
261                rows_per_image: None,
262            },
263            image.size.into(),
264        );
265
266        let si = queue.submit([]);
267
268        device.poll(wgpu::PollType::Wait {
269            submission_index: Some(si),
270            timeout: None,
271        })?;
272
273        Ok((image, texture))
274    }
275}
librashader_test/render/wgpu.rs

librashader_test/render/
wgpu.rs
