spatialrust_gpu/image/
gpu_image.rs

1//! Owned GPU image textures and named upload/readback.
2
3use spatialrust_core::{SpatialError, SpatialResult};
4use spatialrust_image::{Image, ImageMetadata, ImageView};
5
6use crate::WgpuRuntime;
7
8const BYTES_PER_TEXEL: u64 = 4;
9
10/// Transfer accounting for GPU image uploads, device copies, and readbacks.
11#[derive(Clone, Debug, Default, PartialEq, Eq)]
12pub struct GpuImageReceipt {
13    host_to_device_bytes: u64,
14    gpu_to_gpu_bytes: u64,
15    device_to_host_bytes: u64,
16    stages: Vec<&'static str>,
17}
18
19impl GpuImageReceipt {
20    /// Returns physical bytes explicitly uploaded from host memory.
21    #[must_use]
22    pub const fn host_to_device_bytes(&self) -> u64 {
23        self.host_to_device_bytes
24    }
25
26    /// Returns physical bytes copied or written by device-side stages.
27    #[must_use]
28    pub const fn gpu_to_gpu_bytes(&self) -> u64 {
29        self.gpu_to_gpu_bytes
30    }
31
32    /// Returns physical bytes explicitly copied back to host memory.
33    #[must_use]
34    pub const fn device_to_host_bytes(&self) -> u64 {
35        self.device_to_host_bytes
36    }
37
38    /// Returns the ordered logical transfer and kernel stage names.
39    #[must_use]
40    pub fn stages(&self) -> &[&'static str] {
41        &self.stages
42    }
43
44    /// Verifies the transfer contract for a caller-uploaded resident chain.
45    ///
46    /// The expected byte count is the physical RGBA8 texture upload size.
47    pub fn validate_resident_chain(&self, expected_upload_bytes: u64) -> SpatialResult<()> {
48        if self.host_to_device_bytes != expected_upload_bytes {
49            return Err(SpatialError::InvalidArgument(format!(
50                "resident chain expected {expected_upload_bytes} host-to-device bytes, recorded {}",
51                self.host_to_device_bytes
52            )));
53        }
54        if self.device_to_host_bytes != 0 {
55            return Err(SpatialError::InvalidArgument(format!(
56                "resident chain forbids device-to-host transfers, recorded {} bytes",
57                self.device_to_host_bytes
58            )));
59        }
60        Ok(())
61    }
62
63    pub(crate) fn record_host_to_device(&mut self, bytes: u64, stage: &'static str) {
64        self.host_to_device_bytes = self.host_to_device_bytes.saturating_add(bytes);
65        self.stages.push(stage);
66    }
67
68    pub(crate) fn record_gpu_to_gpu(&mut self, bytes: u64, stage: &'static str) {
69        self.gpu_to_gpu_bytes = self.gpu_to_gpu_bytes.saturating_add(bytes);
70        self.stages.push(stage);
71    }
72
73    pub(crate) fn record_device_to_host(&mut self, bytes: u64, stage: &'static str) {
74        self.device_to_host_bytes = self.device_to_host_bytes.saturating_add(bytes);
75        self.stages.push(stage);
76    }
77
78    pub(crate) fn merge_from(&mut self, other: &Self) {
79        self.host_to_device_bytes =
80            self.host_to_device_bytes.saturating_add(other.host_to_device_bytes);
81        self.gpu_to_gpu_bytes = self.gpu_to_gpu_bytes.saturating_add(other.gpu_to_gpu_bytes);
82        self.device_to_host_bytes =
83            self.device_to_host_bytes.saturating_add(other.device_to_host_bytes);
84        self.stages.extend_from_slice(&other.stages);
85    }
86}
87
88/// GPU-resident packed image backed by an `rgba8uint` 2D texture.
89///
90/// The logical channel count remains 1..=4. Unused texture components are zero,
91/// so device storage is a predictable four bytes per pixel on every backend.
92pub struct GpuImage {
93    width: u32,
94    height: u32,
95    channels: u32,
96    device_key: usize,
97    texture: wgpu::Texture,
98    storage_bytes: u64,
99    metadata: ImageMetadata,
100    receipt: GpuImageReceipt,
101}
102
103impl GpuImage {
104    /// Explicitly uploads a packed or strided `u8` image into a texture.
105    pub fn upload_u8<const CHANNELS: usize>(
106        runtime: &WgpuRuntime,
107        view: ImageView<'_, u8, CHANNELS>,
108    ) -> SpatialResult<Self> {
109        if CHANNELS == 0 || CHANNELS > 4 {
110            return Err(SpatialError::InvalidArgument(
111                "GpuImage upload supports 1..=4 channels".to_owned(),
112            ));
113        }
114        let width = u32::try_from(view.width())
115            .map_err(|_| SpatialError::InvalidArgument("GpuImage width exceeds u32".to_owned()))?;
116        let height = u32::try_from(view.height())
117            .map_err(|_| SpatialError::InvalidArgument("GpuImage height exceeds u32".to_owned()))?;
118        if width == 0 || height == 0 {
119            return Err(SpatialError::InvalidArgument(
120                "GpuImage upload requires positive width and height".to_owned(),
121            ));
122        }
123        let texture = create_texture(runtime, width, height, "gpu-image-upload");
124        let texels = pack_rgba8(view);
125        runtime.queue().write_texture(
126            wgpu::TexelCopyTextureInfo {
127                texture: &texture,
128                mip_level: 0,
129                origin: wgpu::Origin3d::ZERO,
130                aspect: wgpu::TextureAspect::All,
131            },
132            &texels,
133            wgpu::TexelCopyBufferLayout {
134                offset: 0,
135                bytes_per_row: Some(width * BYTES_PER_TEXEL as u32),
136                rows_per_image: Some(height),
137            },
138            extent(width, height),
139        );
140        let storage_bytes = texture_bytes(width, height);
141        let mut receipt = GpuImageReceipt::default();
142        receipt.record_host_to_device(storage_bytes, "upload_u8_texture");
143        Ok(Self {
144            width,
145            height,
146            channels: CHANNELS as u32,
147            device_key: runtime_device_key(runtime),
148            texture,
149            storage_bytes,
150            metadata: view.metadata(),
151            receipt,
152        })
153    }
154
155    /// Explicitly reads the texture back into an owned packed host image.
156    pub fn readback_u8<const CHANNELS: usize>(
157        &mut self,
158        runtime: &WgpuRuntime,
159    ) -> SpatialResult<Image<u8, CHANNELS>> {
160        self.validate_runtime(runtime)?;
161        if self.channels as usize != CHANNELS {
162            return Err(SpatialError::InvalidArgument(format!(
163                "GpuImage has {} channels but readback requested {CHANNELS}",
164                self.channels
165            )));
166        }
167        let unpadded_row = self.width * BYTES_PER_TEXEL as u32;
168        let padded_row = unpadded_row.div_ceil(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT)
169            * wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
170        let staging_size = u64::from(padded_row) * u64::from(self.height);
171        let staging = runtime.device().create_buffer(&wgpu::BufferDescriptor {
172            label: Some("gpu-image-texture-readback"),
173            size: staging_size,
174            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
175            mapped_at_creation: false,
176        });
177        let mut encoder =
178            runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
179                label: Some("gpu-image-texture-readback-encoder"),
180            });
181        encoder.copy_texture_to_buffer(
182            wgpu::TexelCopyTextureInfo {
183                texture: &self.texture,
184                mip_level: 0,
185                origin: wgpu::Origin3d::ZERO,
186                aspect: wgpu::TextureAspect::All,
187            },
188            wgpu::TexelCopyBufferInfo {
189                buffer: &staging,
190                layout: wgpu::TexelCopyBufferLayout {
191                    offset: 0,
192                    bytes_per_row: Some(padded_row),
193                    rows_per_image: Some(self.height),
194                },
195            },
196            extent(self.width, self.height),
197        );
198        runtime.queue().submit(Some(encoder.finish()));
199        let rgba = read_staging_bytes(runtime.device(), &staging, staging_size as usize)?;
200        let mut data = Vec::with_capacity(self.width as usize * self.height as usize * CHANNELS);
201        for row in rgba.chunks_exact(padded_row as usize).take(self.height as usize) {
202            for texel in row[..unpadded_row as usize].chunks_exact(BYTES_PER_TEXEL as usize) {
203                data.extend_from_slice(&texel[..CHANNELS]);
204            }
205        }
206        self.receipt.record_device_to_host(self.storage_bytes, "readback_u8_texture");
207        Image::try_new_with_metadata(self.width as usize, self.height as usize, data, self.metadata)
208            .map_err(|error| SpatialError::InvalidArgument(error.to_string()))
209    }
210
211    /// Returns image width in pixels.
212    #[must_use]
213    pub const fn width(&self) -> u32 {
214        self.width
215    }
216
217    /// Returns image height in pixels.
218    #[must_use]
219    pub const fn height(&self) -> u32 {
220        self.height
221    }
222
223    /// Returns the logical component count.
224    #[must_use]
225    pub const fn channels(&self) -> u32 {
226        self.channels
227    }
228
229    /// Returns retained semantic metadata.
230    #[must_use]
231    pub const fn metadata(&self) -> ImageMetadata {
232        self.metadata
233    }
234
235    /// Returns cumulative transfer and device-stage accounting.
236    #[must_use]
237    pub const fn receipt(&self) -> &GpuImageReceipt {
238        &self.receipt
239    }
240
241    /// Drops texture storage after verifying the runtime ownership contract.
242    pub fn recycle(self, runtime: &WgpuRuntime) {
243        if self.validate_runtime(runtime).is_ok() {
244            runtime.recycle_image_texture(self.width, self.height, self.texture);
245        }
246    }
247
248    pub(crate) fn validate_runtime(&self, runtime: &WgpuRuntime) -> SpatialResult<()> {
249        if self.device_key != runtime_device_key(runtime) {
250            return Err(SpatialError::InvalidArgument(
251                "GpuImage belongs to a different runtime device".to_owned(),
252            ));
253        }
254        Ok(())
255    }
256
257    pub(crate) const fn storage_bytes(&self) -> u64 {
258        self.storage_bytes
259    }
260
261    pub(crate) const fn texture(&self) -> &wgpu::Texture {
262        &self.texture
263    }
264
265    pub(crate) fn view(&self) -> wgpu::TextureView {
266        self.texture.create_view(&wgpu::TextureViewDescriptor::default())
267    }
268
269    pub(crate) fn from_parts(
270        runtime: &WgpuRuntime,
271        width: u32,
272        height: u32,
273        channels: u32,
274        texture: wgpu::Texture,
275        metadata: ImageMetadata,
276        receipt: GpuImageReceipt,
277    ) -> SpatialResult<Self> {
278        if width == 0 || height == 0 || !(1..=4).contains(&channels) {
279            return Err(SpatialError::InvalidArgument(
280                "GpuImage dimensions/channels must be positive with 1..=4 channels".to_owned(),
281            ));
282        }
283        Ok(Self {
284            width,
285            height,
286            channels,
287            device_key: runtime_device_key(runtime),
288            texture,
289            storage_bytes: texture_bytes(width, height),
290            metadata,
291            receipt,
292        })
293    }
294}
295
296pub(crate) fn create_texture(
297    runtime: &WgpuRuntime,
298    width: u32,
299    height: u32,
300    label: &'static str,
301) -> wgpu::Texture {
302    runtime.acquire_image_texture(width, height, label)
303}
304
305fn pack_rgba8<const CHANNELS: usize>(view: ImageView<'_, u8, CHANNELS>) -> Vec<u8> {
306    let mut packed = vec![0_u8; view.width() * view.height() * BYTES_PER_TEXEL as usize];
307    for y in 0..view.height() {
308        let source = view.row(y).expect("input row in bounds");
309        let target = &mut packed[y * view.width() * 4..(y + 1) * view.width() * 4];
310        for (pixel, texel) in source.chunks_exact(CHANNELS).zip(target.chunks_exact_mut(4)) {
311            texel[..CHANNELS].copy_from_slice(pixel);
312        }
313    }
314    packed
315}
316
317const fn extent(width: u32, height: u32) -> wgpu::Extent3d {
318    wgpu::Extent3d { width, height, depth_or_array_layers: 1 }
319}
320
321const fn texture_bytes(width: u32, height: u32) -> u64 {
322    width as u64 * height as u64 * BYTES_PER_TEXEL
323}
324
325pub(crate) fn runtime_device_key(runtime: &WgpuRuntime) -> usize {
326    runtime.device() as *const wgpu::Device as usize
327}
328
329pub(crate) fn read_staging_bytes(
330    device: &wgpu::Device,
331    staging_buffer: &wgpu::Buffer,
332    len: usize,
333) -> SpatialResult<Vec<u8>> {
334    let slice = staging_buffer.slice(..);
335    let (sender, receiver) = std::sync::mpsc::channel();
336    slice.map_async(wgpu::MapMode::Read, move |result| {
337        let _ = sender.send(result);
338    });
339    device.poll(wgpu::Maintain::Wait);
340    receiver
341        .recv()
342        .map_err(|_| SpatialError::InvalidArgument("failed to receive wgpu map result".to_owned()))?
343        .map_err(|error| {
344            SpatialError::InvalidArgument(format!("failed to map wgpu texture: {error}"))
345        })?;
346    let data = slice.get_mapped_range();
347    let values = data[..len].to_vec();
348    drop(data);
349    staging_buffer.unmap();
350    Ok(values)
351}
spatialrust_gpu/image/gpu_image.rs

spatialrust_gpu/image/
gpu_image.rs
