wgpu_examples/storage_texture/
mod.rs

Help
1//! This example demonstrates the basic usage of storage textures for the purpose of
2//! creating a digital image of the Mandelbrot set
3//! (<https://en.wikipedia.org/wiki/Mandelbrot_set>).
4//!
5//! Storage textures work like normal textures but they operate similar to storage buffers
6//! in that they can be written to. The issue is that as it stands, write-only is the
7//! only valid access mode for storage textures in WGSL and although there is a WGPU feature
8//! to allow for read-write access, this is unfortunately a native-only feature and thus
9//! we won't be using it here. If we needed a reference texture, we would need to add a
10//! second texture to act as a reference and attach that as well. Luckily, we don't need
11//! to read anything in our shader except the dimensions of our texture, which we can
12//! easily get via `textureDimensions`.
13//!
14//! A lot of things aren't explained here via comments. See hello-compute and
15//! repeated-compute for code that is more thoroughly commented.
16
17#[cfg(not(target_arch = "wasm32"))]
18use crate::utils::output_image_native;
19#[cfg(target_arch = "wasm32")]
20use crate::utils::output_image_wasm;
21
22const TEXTURE_DIMS: (usize, usize) = (512, 512);
23
24async fn run(_path: Option<String>) {
25    let mut texture_data = vec![0u8; TEXTURE_DIMS.0 * TEXTURE_DIMS.1 * 4];
26
27    let instance = wgpu::Instance::default();
28    let adapter = instance
29        .request_adapter(&wgpu::RequestAdapterOptions::default())
30        .await
31        .unwrap();
32    let (device, queue) = adapter
33        .request_device(&wgpu::DeviceDescriptor {
34            label: None,
35            required_features: wgpu::Features::empty(),
36            required_limits: wgpu::Limits::downlevel_defaults(),
37            memory_hints: wgpu::MemoryHints::MemoryUsage,
38            trace: wgpu::Trace::Off,
39        })
40        .await
41        .unwrap();
42
43    let shader = device.create_shader_module(wgpu::include_wgsl!("shader.wgsl"));
44
45    let storage_texture = device.create_texture(&wgpu::TextureDescriptor {
46        label: None,
47        size: wgpu::Extent3d {
48            width: TEXTURE_DIMS.0 as u32,
49            height: TEXTURE_DIMS.1 as u32,
50            depth_or_array_layers: 1,
51        },
52        mip_level_count: 1,
53        sample_count: 1,
54        dimension: wgpu::TextureDimension::D2,
55        format: wgpu::TextureFormat::Rgba8Unorm,
56        usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC,
57        view_formats: &[],
58    });
59    let storage_texture_view = storage_texture.create_view(&wgpu::TextureViewDescriptor::default());
60    let output_staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
61        label: None,
62        size: size_of_val(&texture_data[..]) as u64,
63        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
64        mapped_at_creation: false,
65    });
66
67    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
68        label: None,
69        entries: &[wgpu::BindGroupLayoutEntry {
70            binding: 0,
71            visibility: wgpu::ShaderStages::COMPUTE,
72            ty: wgpu::BindingType::StorageTexture {
73                access: wgpu::StorageTextureAccess::WriteOnly,
74                format: wgpu::TextureFormat::Rgba8Unorm,
75                view_dimension: wgpu::TextureViewDimension::D2,
76            },
77            count: None,
78        }],
79    });
80    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
81        label: None,
82        layout: &bind_group_layout,
83        entries: &[wgpu::BindGroupEntry {
84            binding: 0,
85            resource: wgpu::BindingResource::TextureView(&storage_texture_view),
86        }],
87    });
88
89    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
90        label: None,
91        bind_group_layouts: &[&bind_group_layout],
92        push_constant_ranges: &[],
93    });
94    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
95        label: None,
96        layout: Some(&pipeline_layout),
97        module: &shader,
98        entry_point: Some("main"),
99        compilation_options: Default::default(),
100        cache: None,
101    });
102
103    log::info!("Wgpu context set up.");
104    //----------------------------------------
105
106    let mut command_encoder =
107        device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
108    {
109        let mut compute_pass = command_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
110            label: None,
111            timestamp_writes: None,
112        });
113        compute_pass.set_bind_group(0, &bind_group, &[]);
114        compute_pass.set_pipeline(&pipeline);
115        compute_pass.dispatch_workgroups(TEXTURE_DIMS.0 as u32, TEXTURE_DIMS.1 as u32, 1);
116    }
117    command_encoder.copy_texture_to_buffer(
118        wgpu::TexelCopyTextureInfo {
119            texture: &storage_texture,
120            mip_level: 0,
121            origin: wgpu::Origin3d::ZERO,
122            aspect: wgpu::TextureAspect::All,
123        },
124        wgpu::TexelCopyBufferInfo {
125            buffer: &output_staging_buffer,
126            layout: wgpu::TexelCopyBufferLayout {
127                offset: 0,
128                // This needs to be padded to 256.
129                bytes_per_row: Some((TEXTURE_DIMS.0 * 4) as u32),
130                rows_per_image: Some(TEXTURE_DIMS.1 as u32),
131            },
132        },
133        wgpu::Extent3d {
134            width: TEXTURE_DIMS.0 as u32,
135            height: TEXTURE_DIMS.1 as u32,
136            depth_or_array_layers: 1,
137        },
138    );
139    queue.submit(Some(command_encoder.finish()));
140
141    let buffer_slice = output_staging_buffer.slice(..);
142    let (sender, receiver) = flume::bounded(1);
143    buffer_slice.map_async(wgpu::MapMode::Read, move |r| sender.send(r).unwrap());
144    device.poll(wgpu::PollType::wait()).unwrap();
145    receiver.recv_async().await.unwrap().unwrap();
146    log::info!("Output buffer mapped");
147    {
148        let view = buffer_slice.get_mapped_range();
149        texture_data.copy_from_slice(&view[..]);
150    }
151    log::info!("GPU data copied to local.");
152    output_staging_buffer.unmap();
153
154    #[cfg(not(target_arch = "wasm32"))]
155    output_image_native(texture_data.to_vec(), TEXTURE_DIMS, _path.unwrap());
156    #[cfg(target_arch = "wasm32")]
157    output_image_wasm(texture_data.to_vec(), TEXTURE_DIMS);
158    log::info!("Done.")
159}
160
161pub fn main() {
162    #[cfg(not(target_arch = "wasm32"))]
163    {
164        env_logger::builder()
165            .filter_level(log::LevelFilter::Info)
166            .format_timestamp_nanos()
167            .init();
168
169        let path = std::env::args()
170            .nth(2)
171            .unwrap_or_else(|| "please_don't_git_push_me.png".to_string());
172        pollster::block_on(run(Some(path)));
173    }
174    #[cfg(target_arch = "wasm32")]
175    {
176        std::panic::set_hook(Box::new(console_error_panic_hook::hook));
177        console_log::init_with_level(log::Level::Info).expect("could not initialize logger");
178        wasm_bindgen_futures::spawn_local(run(None));
179    }
180}
wgpu_examples/storage_texture/mod.rs

wgpu_examples/storage_texture/
mod.rs
