3sixty2cubemap/src/gpu/mod.rs

//! WebGPU (`wgpu`) compute layer.
//!
//! `GpuContext` owns the wgpu instance/device/queue so every GPU stage
//! (extract sharpness, cubemap projection, later mask inference) shares one
//! device and one queue. `wgpu` maps to Vulkan/Metal/DX12 on native and WebGPU
//! in the browser — one WGSL codebase, cross-vendor.
//!
//! Added for stage 3 (cubemap): texture upload, storage-texture output,
//! and a `read_texture_rgba` readback that correctly handles the WebGPU
//! 256-byte `bytes_per_row` alignment requirement.

use wgpu::util::DeviceExt;

/// A ready-to-use wgpu device + queue plus the instance/adapter used to get it.
pub struct GpuContext {
    pub device: wgpu::Device,
    pub queue: wgpu::Queue,
}

/// WebGPU requires `bytes_per_row` to be a multiple of this for
/// buffer<->texture copies.
pub const COPY_BYTES_PER_ROW_ALIGNMENT: u64 = 256;

/// Round `n` up to the nearest multiple of `256` (WebGPU copy alignment).
pub fn align_bytes_per_row(n: u64) -> u64 {
    n.div_ceil(COPY_BYTES_PER_ROW_ALIGNMENT) * COPY_BYTES_PER_ROW_ALIGNMENT
}

impl GpuContext {
    /// Request a compute-capable adapter, preferring the *discrete* GPU.
    ///
    /// `RequestAdapterOptions::default()` uses `power_preference: LowPower`, which
    /// on a dual-GPU system (integrated + discrete) often selects the weak iGPU
    /// or a software (llvmpipe) adapter — making every GPU stage look like it
    /// "isn't using the GPU". We explicitly ask for `HighPerformance` and, as a
    /// fallback, enumerate all adapters and pick the first discrete one.
    pub async fn new(label: &str) -> anyhow::Result<GpuContext> {
        let instance = wgpu::Instance::default();
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                force_fallback_adapter: false,
                compatible_surface: None,
            })
            .await
            .or_else(|_| {
                // enumerate_adapter is blocking-ish; do it on a blocking context.
                // wgpu exposes enumerate_adapters synchronously.
                let mut best: Option<wgpu::Adapter> = None;
                for a in instance.enumerate_adapters(wgpu::Backends::all()) {
                    let info = a.get_info();
                    if info.device_type == wgpu::DeviceType::DiscreteGpu {
                        best = Some(a);
                        break;
                    }
                    if best.is_none() {
                        best = Some(a);
                    }
                }
                best.ok_or_else(|| {
                    anyhow::anyhow!("no wgpu adapter found (no Vulkan/Metal/DX12 device?)")
                })
            })?;
        let info = adapter.get_info();
        eprintln!(
            "[gpu] using {} ({:?}) via {:?}",
            info.name, info.device_type, info.backend
        );
        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                label: Some(label),
                required_features: wgpu::Features::empty(),
                required_limits: wgpu::Limits::default(),
                memory_hints: wgpu::MemoryHints::default(),
                trace: wgpu::Trace::Off,
            })
            .await?;
        Ok(GpuContext { device, queue })
    }

    /// Build a compute pipeline from an embedded WGSL source string.
    pub fn compute_pipeline(
        &self,
        label: &str,
        source: &'static str,
        layout: &wgpu::PipelineLayout,
    ) -> wgpu::ComputePipeline {
        let module = self
            .device
            .create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some(label),
                source: wgpu::ShaderSource::Wgsl(source.into()),
            });
        self.device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(label),
                layout: Some(layout),
                module: &module,
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            })
    }

    /// Create a storage buffer of `size` bytes (used for readback or GPU output).
    pub fn create_storage_buffer(
        &self,
        label: &str,
        size: u64,
        contents: Option<&[u8]>,
    ) -> wgpu::Buffer {
        let usage = wgpu::BufferUsages::STORAGE
            | wgpu::BufferUsages::COPY_SRC
            | wgpu::BufferUsages::COPY_DST;
        match contents {
            Some(data) => self
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some(label),
                    contents: data,
                    usage,
                }),
            None => self.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(label),
                size,
                usage,
                mapped_at_creation: false,
            }),
        }
    }

    /// Create a uniform buffer initialised from `data`.
    pub fn create_uniform_buffer(&self, label: &str, data: &[u8]) -> wgpu::Buffer {
        let usage = wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST;
        self.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some(label),
                contents: data,
                usage,
            })
    }

    /// Create a read-only RGBA8 input texture (sampled by the compute kernels).
    pub fn create_texture_2d(&self, label: &str, w: u32, h: u32) -> wgpu::Texture {
        self.device.create_texture(&wgpu::TextureDescriptor {
            label: Some(label),
            size: wgpu::Extent3d {
                width: w,
                height: h,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        })
    }

    /// Create a non-square output storage texture (written by kernels, read back).
    pub fn create_storage_texture_2d_wh(&self, label: &str, w: u32, h: u32) -> wgpu::Texture {
        self.device.create_texture(&wgpu::TextureDescriptor {
            label: Some(label),
            size: wgpu::Extent3d {
                width: w,
                height: h,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC,
            view_formats: &[],
        })
    }

    /// Create an output storage texture (written by the kernels, read back).
    pub fn create_storage_texture_2d(&self, label: &str, n: u32) -> wgpu::Texture {
        self.device.create_texture(&wgpu::TextureDescriptor {
            label: Some(label),
            size: wgpu::Extent3d {
                width: n,
                height: n,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC,
            view_formats: &[],
        })
    }

    /// Upload RGBA8 `data` (width*height*4 bytes) into `tex`.
    pub fn upload_rgba(&self, tex: &wgpu::Texture, w: u32, h: u32, data: &[u8]) {
        self.queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: tex,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            data,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(w * 4),
                rows_per_image: Some(h),
            },
            wgpu::Extent3d {
                width: w,
                height: h,
                depth_or_array_layers: 1,
            },
        );
    }

    /// Copy an RGBA storage texture to a tightly packed `Vec<u8>` (w*h*4).
    ///
    /// Handles the WebGPU 256-byte `bytes_per_row` padding: a padded staging
    /// buffer is mapped and then destrided into the tight layout expected by
    /// the `image` crate.
    pub fn read_texture_rgba(
        &self,
        tex: &wgpu::Texture,
        w: u32,
        h: u32,
    ) -> anyhow::Result<Vec<u8>> {
        let bytes_per_row = u64::from(w) * 4;
        let padded_row = align_bytes_per_row(bytes_per_row);
        let total = padded_row * u64::from(h);

        let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("texture-readback-staging"),
            size: total,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        encoder.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: tex,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: &staging,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(padded_row as u32),
                    rows_per_image: Some(h),
                },
            },
            wgpu::Extent3d {
                width: w,
                height: h,
                depth_or_array_layers: 1,
            },
        );
        self.queue.submit(Some(encoder.finish()));

        let slice = staging.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |res| {
            let _ = tx.send(res);
        });
        let _ = self.device.poll(wgpu::PollType::Wait);
        rx.recv()??;
        let mapped = slice.get_mapped_range();
        let mut tight = vec![0u8; (bytes_per_row * u64::from(h)) as usize];
        if padded_row == bytes_per_row {
            tight.copy_from_slice(&mapped);
        } else {
            let row_w = bytes_per_row as usize;
            let row_p = padded_row as usize;
            for r in 0..h as usize {
                let src_start = r * row_p;
                tight[r * row_w..(r + 1) * row_w]
                    .copy_from_slice(&mapped[src_start..src_start + row_w]);
            }
        }
        Ok(tight)
    }

    /// Map several already-submitted buffers in one `poll(Wait)` and return their
    /// contents in order. Amortises the per-readback poll latency across a batch.
    #[allow(clippy::needless_range_loop)]
    pub fn readback_many(
        &self,
        srcs: &[&wgpu::Buffer],
        sizes: &[u64],
    ) -> anyhow::Result<Vec<Vec<u8>>> {
        assert_eq!(srcs.len(), sizes.len(), "readback_many: len mismatch");
        let n = srcs.len();
        if n == 0 {
            return Ok(Vec::new());
        }
        let stagers: Vec<wgpu::Buffer> = (0..n)
            .map(|i| {
                self.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("rb-many-staging"),
                    size: sizes[i],
                    usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
                    mapped_at_creation: false,
                })
            })
            .collect();

        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        for i in 0..n {
            encoder.copy_buffer_to_buffer(srcs[i], 0, &stagers[i], 0, sizes[i]);
        }
        self.queue.submit(Some(encoder.finish()));

        let (tx, rx) = std::sync::mpsc::channel();
        for i in 0..n {
            let tx = tx.clone();
            stagers[i]
                .slice(..)
                .map_async(wgpu::MapMode::Read, move |res| {
                    let _ = tx.send((i, res));
                });
        }
        let _ = self.device.poll(wgpu::PollType::Wait);

        let mut results: Vec<Vec<u8>> = (0..n).map(|_| Vec::new()).collect();
        for _ in 0..n {
            let (i, res) = rx
                .recv()
                .map_err(|_| anyhow::anyhow!("readback_many: channel closed"))?;
            res?;
            results[i] = stagers[i].slice(..).get_mapped_range().to_vec();
        }
        Ok(results)
    }

    /// Copy several storage textures to tightly packed RGBA `Vec<u8>`s with one
    /// submit + one poll. Handles the 256-byte `bytes_per_row` padding per tex.
    #[allow(clippy::needless_range_loop)]
    pub fn read_textures_rgba(
        &self,
        texs: &[&wgpu::Texture],
        w: u32,
        h: u32,
    ) -> anyhow::Result<Vec<Vec<u8>>> {
        let n = texs.len();
        if n == 0 {
            return Ok(Vec::new());
        }
        let bytes_per_row = u64::from(w) * 4;
        let padded_row = align_bytes_per_row(bytes_per_row);
        let total = padded_row * u64::from(h);

        let stagers: Vec<wgpu::Buffer> = (0..n)
            .map(|_| {
                self.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("tex-rb-many-staging"),
                    size: total,
                    usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
                    mapped_at_creation: false,
                })
            })
            .collect();

        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        for i in 0..n {
            encoder.copy_texture_to_buffer(
                wgpu::TexelCopyTextureInfo {
                    texture: texs[i],
                    mip_level: 0,
                    origin: wgpu::Origin3d::ZERO,
                    aspect: wgpu::TextureAspect::All,
                },
                wgpu::TexelCopyBufferInfo {
                    buffer: &stagers[i],
                    layout: wgpu::TexelCopyBufferLayout {
                        offset: 0,
                        bytes_per_row: Some(padded_row as u32),
                        rows_per_image: Some(h),
                    },
                },
                wgpu::Extent3d {
                    width: w,
                    height: h,
                    depth_or_array_layers: 1,
                },
            );
        }
        self.queue.submit(Some(encoder.finish()));

        let (tx, rx) = std::sync::mpsc::channel();
        for i in 0..n {
            let tx = tx.clone();
            stagers[i]
                .slice(..)
                .map_async(wgpu::MapMode::Read, move |res| {
                    let _ = tx.send((i, res));
                });
        }
        let _ = self.device.poll(wgpu::PollType::Wait);

        let mut results: Vec<Vec<u8>> = (0..n).map(|_| Vec::new()).collect();
        for _ in 0..n {
            let (i, res) = rx
                .recv()
                .map_err(|_| anyhow::anyhow!("read_textures_rgba: channel closed"))?;
            res?;
            let mapped = stagers[i].slice(..).get_mapped_range();
            let mut tight = vec![0u8; (bytes_per_row * u64::from(h)) as usize];
            if padded_row == bytes_per_row {
                tight.copy_from_slice(&mapped);
            } else {
                let row_w = bytes_per_row as usize;
                let row_p = padded_row as usize;
                for r in 0..h as usize {
                    let src_start = r * row_p;
                    tight[r * row_w..(r + 1) * row_w]
                        .copy_from_slice(&mapped[src_start..src_start + row_w]);
                }
            }
            results[i] = tight;
        }
        Ok(results)
    }

    /// Snapshot the contents of a (storage or uniform) buffer into a `Vec<u8>`.
    pub fn readback(&self, src: &wgpu::Buffer, size: u64) -> anyhow::Result<Vec<u8>> {
        let snapshot = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("readback"),
            size,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        encoder.copy_buffer_to_buffer(src, 0, &snapshot, 0, size);
        self.queue.submit(Some(encoder.finish()));
        let slice = snapshot.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |res| {
            let _ = tx.send(res);
        });
        let _ = self.device.poll(wgpu::PollType::Wait);
        rx.recv()??;
        let data = slice.get_mapped_range().to_vec();
        Ok(data)
    }
}

// Shader sources (embedded at compile time).
pub mod shaders {
    /// Equirect -> 6 cube faces (images: bilinear; masks: nearest).
    pub const EQUIRECT_TO_CUBEMAP: &str = include_str!("shaders/equirect_to_cubemap.wgsl");
    /// Downscale + grayscale + Laplacian variance (sharpness scoring).
    pub const SHARPNESS: &str = include_str!("shaders/sharpness.wgsl");
    /// Cubemap -> equirect (nearest, for mask reprojection).
    pub const C2E: &str = include_str!("shaders/c2e.wgsl");
}