wgpu_examples/storage_texture/
mod.rs

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            experimental_features: wgpu::ExperimentalFeatures::disabled(),
38            memory_hints: wgpu::MemoryHints::MemoryUsage,
39            trace: wgpu::Trace::Off,
40        })
41        .await
42        .unwrap();
43
44    let shader = device.create_shader_module(wgpu::include_wgsl!("shader.wgsl"));
45
46    let storage_texture = device.create_texture(&wgpu::TextureDescriptor {
47        label: None,
48        size: wgpu::Extent3d {
49            width: TEXTURE_DIMS.0 as u32,
50            height: TEXTURE_DIMS.1 as u32,
51            depth_or_array_layers: 1,
52        },
53        mip_level_count: 1,
54        sample_count: 1,
55        dimension: wgpu::TextureDimension::D2,
56        format: wgpu::TextureFormat::Rgba8Unorm,
57        usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC,
58        view_formats: &[],
59    });
60    let storage_texture_view = storage_texture.create_view(&wgpu::TextureViewDescriptor::default());
61    let output_staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
62        label: None,
63        size: size_of_val(&texture_data[..]) as u64,
64        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
65        mapped_at_creation: false,
66    });
67
68    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
69        label: None,
70        entries: &[wgpu::BindGroupLayoutEntry {
71            binding: 0,
72            visibility: wgpu::ShaderStages::COMPUTE,
73            ty: wgpu::BindingType::StorageTexture {
74                access: wgpu::StorageTextureAccess::WriteOnly,
75                format: wgpu::TextureFormat::Rgba8Unorm,
76                view_dimension: wgpu::TextureViewDimension::D2,
77            },
78            count: None,
79        }],
80    });
81    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
82        label: None,
83        layout: &bind_group_layout,
84        entries: &[wgpu::BindGroupEntry {
85            binding: 0,
86            resource: wgpu::BindingResource::TextureView(&storage_texture_view),
87        }],
88    });
89
90    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
91        label: None,
92        bind_group_layouts: &[&bind_group_layout],
93        push_constant_ranges: &[],
94    });
95    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
96        label: None,
97        layout: Some(&pipeline_layout),
98        module: &shader,
99        entry_point: Some("main"),
100        compilation_options: Default::default(),
101        cache: None,
102    });
103
104    log::info!("Wgpu context set up.");
105    //----------------------------------------
106
107    let mut command_encoder =
108        device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
109    {
110        let mut compute_pass = command_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
111            label: None,
112            timestamp_writes: None,
113        });
114        compute_pass.set_bind_group(0, &bind_group, &[]);
115        compute_pass.set_pipeline(&pipeline);
116        compute_pass.dispatch_workgroups(TEXTURE_DIMS.0 as u32, TEXTURE_DIMS.1 as u32, 1);
117    }
118    command_encoder.copy_texture_to_buffer(
119        wgpu::TexelCopyTextureInfo {
120            texture: &storage_texture,
121            mip_level: 0,
122            origin: wgpu::Origin3d::ZERO,
123            aspect: wgpu::TextureAspect::All,
124        },
125        wgpu::TexelCopyBufferInfo {
126            buffer: &output_staging_buffer,
127            layout: wgpu::TexelCopyBufferLayout {
128                offset: 0,
129                // This needs to be padded to 256.
130                bytes_per_row: Some((TEXTURE_DIMS.0 * 4) as u32),
131                rows_per_image: Some(TEXTURE_DIMS.1 as u32),
132            },
133        },
134        wgpu::Extent3d {
135            width: TEXTURE_DIMS.0 as u32,
136            height: TEXTURE_DIMS.1 as u32,
137            depth_or_array_layers: 1,
138        },
139    );
140    queue.submit(Some(command_encoder.finish()));
141
142    let buffer_slice = output_staging_buffer.slice(..);
143    let (sender, receiver) = flume::bounded(1);
144    buffer_slice.map_async(wgpu::MapMode::Read, move |r| sender.send(r).unwrap());
145    device.poll(wgpu::PollType::wait()).unwrap();
146    receiver.recv_async().await.unwrap().unwrap();
147    log::info!("Output buffer mapped");
148    {
149        let view = buffer_slice.get_mapped_range();
150        texture_data.copy_from_slice(&view[..]);
151    }
152    log::info!("GPU data copied to local.");
153    output_staging_buffer.unmap();
154
155    #[cfg(not(target_arch = "wasm32"))]
156    output_image_native(texture_data.to_vec(), TEXTURE_DIMS, _path.unwrap());
157    #[cfg(target_arch = "wasm32")]
158    output_image_wasm(texture_data.to_vec(), TEXTURE_DIMS);
159    log::info!("Done.")
160}
161
162pub fn main() {
163    #[cfg(not(target_arch = "wasm32"))]
164    {
165        env_logger::builder()
166            .filter_level(log::LevelFilter::Info)
167            .format_timestamp_nanos()
168            .init();
169
170        let path = std::env::args()
171            .nth(2)
172            .unwrap_or_else(|| "please_don't_git_push_me.png".to_string());
173        pollster::block_on(run(Some(path)));
174    }
175    #[cfg(target_arch = "wasm32")]
176    {
177        std::panic::set_hook(Box::new(console_error_panic_hook::hook));
178        console_log::init_with_level(log::Level::Info).expect("could not initialize logger");
179        wasm_bindgen_futures::spawn_local(run(None));
180    }
181}
wgpu_examples/storage_texture/mod.rs

wgpu_examples/storage_texture/
mod.rs
