demo_vk/graphics/streaming_renderer/texture/
loader.rs

1use {
2    super::{super::utility::round_to_power_of_two, Texture},
3    crate::graphics::vulkan::{CPUBuffer, SyncCommands, VulkanContext},
4    anyhow::{Context, Result},
5    ash::vk::{self},
6    image::{imageops::FilterType, DynamicImage, RgbaImage},
7    std::{path::PathBuf, sync::Arc},
8};
9
10/// A utility for loading textures from image files.
11pub struct TextureLoader {
12    sync_commands: SyncCommands,
13    transfer_buffer: CPUBuffer<u8>,
14    ctx: Arc<VulkanContext>,
15}
16
17impl TextureLoader {
18    /// Creates a new texture loader instance.
19    pub fn new(ctx: Arc<VulkanContext>) -> Result<Self> {
20        Ok(Self {
21            sync_commands: SyncCommands::new(ctx.clone()).context(
22                "Unable to create SyncCommands for the TextureLoader!",
23            )?,
24            transfer_buffer: CPUBuffer::allocate(
25                &ctx,
26                1024 * 1024,
27                vk::BufferUsageFlags::TRANSFER_SRC,
28            )
29            .context("Unable to allocate transfer buffer")?,
30            ctx,
31        })
32    }
33
34    /// Synchronously loads a texture from an image file. Supports .png, .bmp,
35    /// and .jpg files
36    pub fn load_from_file(
37        &mut self,
38        path: impl Into<PathBuf>,
39        generate_mipmaps: bool,
40    ) -> Result<Texture> {
41        let path: PathBuf = path.into();
42
43        let image = image::ImageReader::open(&path)
44            .context(format!("Unable to open image at {:?}", &path))?
45            .decode()
46            .context(format!("Unable to decode image from file {:?}", &path))?;
47
48        let mipmaps = if generate_mipmaps {
49            self.compute_generated_mipmaps(image)
50        } else {
51            vec![image.to_rgba8()]
52        };
53
54        let texture = Texture::builder()
55            .ctx(&self.ctx)
56            .dimensions(mipmaps[0].dimensions())
57            .format(vk::Format::R8G8B8A8_UNORM)
58            .image_usage_flags(
59                vk::ImageUsageFlags::TRANSFER_DST
60                    | vk::ImageUsageFlags::SAMPLED,
61            )
62            .memory_property_flags(vk::MemoryPropertyFlags::DEVICE_LOCAL)
63            .mip_levels(mipmaps.len() as u32)
64            .build()
65            .context(format!(
66                "Unable to create texture for image file {:?}",
67                &path
68            ))?;
69
70        self.copy_mipmaps_to_transfer_buffer(&mipmaps)
71            .context("Unable to upload texture data!")?;
72
73        self.copy_transfer_buffer_mipmaps_to_device_memory(&texture, &mipmaps)
74            .context("Unable to copy transfer buffer to texture memory!")?;
75
76        Ok(texture)
77    }
78
79    pub fn tex_sub_image(
80        &mut self,
81        ctx: &VulkanContext,
82        texture: &Texture,
83        rgba_data: &[u8],
84        offset: [u32; 2],
85        size: [u32; 2],
86    ) -> Result<()> {
87        debug_assert!(
88            (size[0] * size[1] * 4) as usize == rgba_data.len(),
89            "RGBA data and image size do not match!"
90        );
91
92        self.maybe_resize_transfer_buffer(rgba_data.len())?;
93        unsafe {
94            self.transfer_buffer.write_data(0, rgba_data).context(
95                "Unable to write tex_sub_image data to transfer buffer",
96            )?;
97        }
98
99        self.sync_commands.submit_and_wait(|cmd| {
100            texture
101                .pipeline_barrier()
102                .ctx(ctx)
103                .command_buffer(cmd)
104                .old_layout(vk::ImageLayout::UNDEFINED)
105                .new_layout(vk::ImageLayout::TRANSFER_DST_OPTIMAL)
106                .src_access_mask(vk::AccessFlags::SHADER_READ)
107                .src_stage_mask(vk::PipelineStageFlags::ALL_COMMANDS)
108                .dst_access_mask(vk::AccessFlags::TRANSFER_WRITE)
109                .dst_stage_mask(vk::PipelineStageFlags::ALL_COMMANDS)
110                .call();
111            unsafe {
112                ctx.cmd_copy_buffer_to_image(
113                    cmd,
114                    self.transfer_buffer.buffer(),
115                    texture.image().raw,
116                    vk::ImageLayout::TRANSFER_DST_OPTIMAL,
117                    &[vk::BufferImageCopy {
118                        buffer_offset: 0,
119                        buffer_row_length: 0,
120                        buffer_image_height: 0,
121                        image_subresource: vk::ImageSubresourceLayers {
122                            aspect_mask: vk::ImageAspectFlags::COLOR,
123                            mip_level: 0,
124                            base_array_layer: 0,
125                            layer_count: 1,
126                        },
127                        image_offset: vk::Offset3D {
128                            x: offset[0] as i32,
129                            y: offset[1] as i32,
130                            z: 0,
131                        },
132                        image_extent: vk::Extent3D {
133                            width: size[0],
134                            height: size[1],
135                            depth: 1,
136                        },
137                    }],
138                );
139            }
140            texture
141                .pipeline_barrier()
142                .ctx(ctx)
143                .command_buffer(cmd)
144                .old_layout(vk::ImageLayout::TRANSFER_DST_OPTIMAL)
145                .new_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL)
146                .src_access_mask(vk::AccessFlags::TRANSFER_WRITE)
147                .src_stage_mask(vk::PipelineStageFlags::TRANSFER)
148                .dst_access_mask(vk::AccessFlags::SHADER_READ)
149                .dst_stage_mask(vk::PipelineStageFlags::ALL_COMMANDS)
150                .call();
151            Ok(())
152        })?;
153        Ok(())
154    }
155}
156
157// Private Functions
158impl TextureLoader {
159    /// Rezizes the transfer buffer if needed to hold `total_size` bytes. No-op
160    /// if the transfer buffer is big enough to hold `total_size` bytes
161    /// already.
162    fn maybe_resize_transfer_buffer(
163        &mut self,
164        total_size_in_bytes: usize,
165    ) -> Result<()> {
166        if self.transfer_buffer.size_in_bytes() as usize >= total_size_in_bytes
167        {
168            return Ok(());
169        }
170
171        log::trace!(
172            "Existing transfer buffer {:?} needs reallocated",
173            self.transfer_buffer
174        );
175
176        self.transfer_buffer = CPUBuffer::allocate(
177            &self.ctx,
178            round_to_power_of_two(total_size_in_bytes),
179            vk::BufferUsageFlags::TRANSFER_SRC,
180        )
181        .context("Unable to reallocate transfer buffer")?;
182        log::trace!(
183            "New transfer buffer allocated: {:?}",
184            self.transfer_buffer
185        );
186        Ok(())
187    }
188
189    /// Generates mipmaps for the image.
190    fn compute_generated_mipmaps(
191        &self,
192        mut image: DynamicImage,
193    ) -> Vec<RgbaImage> {
194        let mut mipmaps = vec![image.to_rgba8()];
195        let mut width = image.width();
196        let mut height = image.height();
197        while width > 2 && height > 2 {
198            width /= 2;
199            height /= 2;
200            image = image.resize(width, height, FilterType::Lanczos3);
201            mipmaps.push(image.to_rgba8());
202        }
203        mipmaps
204    }
205
206    /// Copies data to the transfer buffer.
207    ///
208    /// This function may reallocate a new transfer buffer if the existing one
209    /// is too small for the data.
210    fn copy_mipmaps_to_transfer_buffer(
211        &mut self,
212        mipmaps: &[RgbaImage],
213    ) -> Result<()> {
214        self.maybe_resize_transfer_buffer(
215            mipmaps.iter().map(|image| image.as_raw().len()).sum(),
216        )?;
217
218        let mut offset = 0;
219        for mipmap in mipmaps {
220            // # Safety
221            //
222            // Safe because the transfer buffer is not shared and is only used
223            // for synchronous commands sent to the GPU and because
224            // the transfer buffer is always resized to have enough
225            // capacity for all data.
226            unsafe {
227                self.transfer_buffer
228                    .write_data(offset, mipmap.as_raw())
229                    .context("Unable to write data to transfer buffer")?;
230            }
231            offset += mipmap.as_raw().len();
232        }
233        Ok(())
234    }
235
236    /// Copies the contents of the transfer buffer into the texture's device
237    /// memory.
238    fn copy_transfer_buffer_mipmaps_to_device_memory(
239        &self,
240        texture: &Texture,
241        mipmaps: &[RgbaImage],
242    ) -> Result<()> {
243        let ctx = &self.ctx;
244        self.sync_commands.submit_and_wait(|command_buffer| unsafe {
245            ctx.cmd_pipeline_barrier(
246                command_buffer,
247                vk::PipelineStageFlags::TOP_OF_PIPE,
248                vk::PipelineStageFlags::TRANSFER,
249                vk::DependencyFlags::empty(),
250                &[],
251                &[],
252                &[vk::ImageMemoryBarrier {
253                    src_access_mask: vk::AccessFlags::empty(),
254                    dst_access_mask: vk::AccessFlags::TRANSFER_WRITE,
255                    old_layout: vk::ImageLayout::UNDEFINED,
256                    new_layout: vk::ImageLayout::TRANSFER_DST_OPTIMAL,
257                    src_queue_family_index: ctx.graphics_queue_family_index,
258                    dst_queue_family_index: ctx.graphics_queue_family_index,
259                    image: texture.image().raw,
260                    subresource_range: vk::ImageSubresourceRange {
261                        aspect_mask: vk::ImageAspectFlags::COLOR,
262                        base_mip_level: 0,
263                        level_count: mipmaps.len() as u32,
264                        base_array_layer: 0,
265                        layer_count: 1,
266                    },
267                    ..Default::default()
268                }],
269            );
270            let buffer_image_copies: Vec<vk::BufferImageCopy> = {
271                let mut offset: u64 = 0;
272                mipmaps
273                    .iter()
274                    .enumerate()
275                    .map(|(mip_level, mipmap)| {
276                        let buffer_image_copy = vk::BufferImageCopy {
277                            buffer_offset: offset,
278                            buffer_row_length: 0,
279                            buffer_image_height: 0,
280                            image_subresource: vk::ImageSubresourceLayers {
281                                aspect_mask: vk::ImageAspectFlags::COLOR,
282                                mip_level: mip_level as u32,
283                                base_array_layer: 0,
284                                layer_count: 1,
285                            },
286                            image_offset: vk::Offset3D { x: 0, y: 0, z: 0 },
287                            image_extent: vk::Extent3D {
288                                width: mipmap.width(),
289                                height: mipmap.height(),
290                                depth: 1,
291                            },
292                        };
293                        offset += mipmap.as_raw().len() as u64;
294                        buffer_image_copy
295                    })
296                    .collect()
297            };
298            ctx.cmd_copy_buffer_to_image(
299                command_buffer,
300                self.transfer_buffer.buffer(),
301                texture.image().raw,
302                vk::ImageLayout::TRANSFER_DST_OPTIMAL,
303                &buffer_image_copies,
304            );
305            ctx.cmd_pipeline_barrier(
306                command_buffer,
307                vk::PipelineStageFlags::TRANSFER,
308                vk::PipelineStageFlags::BOTTOM_OF_PIPE,
309                vk::DependencyFlags::empty(),
310                &[],
311                &[],
312                &[vk::ImageMemoryBarrier {
313                    src_access_mask: vk::AccessFlags::TRANSFER_WRITE,
314                    dst_access_mask: vk::AccessFlags::empty(),
315                    old_layout: vk::ImageLayout::TRANSFER_DST_OPTIMAL,
316                    new_layout: vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
317                    src_queue_family_index: ctx.graphics_queue_family_index,
318                    dst_queue_family_index: ctx.graphics_queue_family_index,
319                    image: texture.image().raw,
320                    subresource_range: vk::ImageSubresourceRange {
321                        aspect_mask: vk::ImageAspectFlags::COLOR,
322                        base_mip_level: 0,
323                        level_count: mipmaps.len() as u32,
324                        base_array_layer: 0,
325                        layer_count: 1,
326                    },
327                    ..Default::default()
328                }],
329            );
330            Ok(())
331        }) // sync_commands end
332    }
333}
334
335#[cfg(test)]
336mod test {
337    use super::*;
338
339    #[test]
340    pub fn power_of_two_should_round_up() {
341        assert_eq!(round_to_power_of_two(1), 1);
342        assert_eq!(round_to_power_of_two(2), 2);
343        assert_eq!(round_to_power_of_two(3), 4);
344        assert_eq!(round_to_power_of_two(6), 8);
345        assert_eq!(round_to_power_of_two(9), 16);
346        assert_eq!(round_to_power_of_two(20), 32);
347        assert_eq!(round_to_power_of_two(50), 64);
348        assert_eq!(round_to_power_of_two(93), 128);
349        assert_eq!(round_to_power_of_two(200), 256);
350        assert_eq!(round_to_power_of_two(500), 512);
351        assert_eq!(round_to_power_of_two(10_000), 16384);
352    }
353}