gpu/renderer/src/lib.rs

// Copyright 2022 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use anyhow::Error;
use conveyor_sort::{SortContext, BLOCK_SIZE};
use painter::PaintContext;
pub use painter::{Color, Style, TILE_HEIGHT, TILE_WIDTH};
use rasterizer::RasterizerContext;
use std::{borrow::Cow, mem, num::NonZeroU32, slice, sync::Arc, time::Duration};
use surpass::{painter::Image, Lines};
use wgpu::{util::DeviceExt, Extent3d, ImageDataLayout, TextureAspect};

fn div_round_up(n: usize, d: usize) -> usize {
    (n + d - 1) / d
}

#[derive(Debug)]
pub struct Timings {
    pub rasterize: Duration,
    pub sort: Duration,
    pub paint: Duration,
    pub render: Duration,
}

impl Timings {
    pub(crate) const fn size() -> usize {
        mem::size_of::<Timings>() / mem::size_of::<u64>()
    }
}

#[derive(Debug)]
struct RenderContext {
    pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    sampler: wgpu::Sampler,
}

#[derive(Debug, Default)]
struct Resources {
    atlas: Option<wgpu::Texture>,
    texture: Option<(wgpu::Texture, u32, u32)>,
}

#[derive(Debug)]
pub struct Renderer {
    rasterizer: RasterizerContext,
    sort: SortContext,
    paint: PaintContext,
    render: RenderContext,
    common: Resources,
    has_timestamp_query: bool,
}

impl Renderer {
    pub fn minimum_device(adapter: &wgpu::Adapter) -> (wgpu::DeviceDescriptor, bool) {
        let adapter_features = adapter.features();
        let has_timestamp_query = adapter_features.contains(wgpu::Features::TIMESTAMP_QUERY);

        let desc = wgpu::DeviceDescriptor {
            label: None,
            features: wgpu::Features::TIMESTAMP_QUERY & adapter_features,
            limits: wgpu::Limits {
                max_texture_dimension_2d: 4096,
                max_storage_buffer_binding_size: 1 << 30,
                ..wgpu::Limits::downlevel_defaults()
            },
        };

        (desc, has_timestamp_query)
    }

    pub fn new(
        device: &wgpu::Device,
        swap_chain_format: wgpu::TextureFormat,
        has_timestamp_query: bool,
    ) -> Self {
        let rasterizer = rasterizer::init(device);

        let sort = conveyor_sort::init(device);

        let paint = painter::init(device);

        let shader = device.create_shader_module(&wgpu::ShaderModuleDescriptor {
            label: None,
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("draw_texture.wgsl"))),
        });

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: None,
            layout: None,
            vertex: wgpu::VertexState { module: &shader, entry_point: "vs_main", buffers: &[] },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: "fs_main",
                targets: &[swap_chain_format.into()],
            }),
            primitive: wgpu::PrimitiveState::default(),
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
        });

        let bind_group_layout = pipeline.get_bind_group_layout(0);

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor::default());

        let render = RenderContext { pipeline, bind_group_layout, sampler };

        Self { rasterizer, sort, paint, render, common: Resources::default(), has_timestamp_query }
    }

    #[allow(clippy::too_many_arguments)]
    pub fn render(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        texture: &wgpu::Texture,
        width: u32,
        height: u32,
        lines: &Lines,
        style_offsets: &[u32],
        styles: &[u32],
        new_allocations: &[(Arc<Image>, [u32; 4])],
        background_color: Color,
    ) -> Result<Option<Timings>, Error> {
        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());

        let timestamp_context = self.has_timestamp_query.then(|| {
            let timestamp = device.create_query_set(&wgpu::QuerySetDescriptor {
                label: None,
                count: Timings::size() as u32,
                ty: wgpu::QueryType::Timestamp,
            });

            let timestamp_period = queue.get_timestamp_period();

            let data_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: mem::size_of::<Timings>() as wgpu::BufferAddress,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
                mapped_at_creation: false,
            });

            (timestamp, timestamp_period, data_buffer)
        });

        if let Some((texture, current_width, current_height)) = self.common.texture.as_ref() {
            if *current_width != width || *current_height != height {
                texture.destroy();

                self.common.texture = None;
            }
        }

        let texture_view = self
            .common
            .texture
            .get_or_insert_with(|| {
                let texture = device.create_texture(&wgpu::TextureDescriptor {
                    label: None,
                    size: wgpu::Extent3d { width, height, depth_or_array_layers: 1 },
                    mip_level_count: 1,
                    sample_count: 1,
                    dimension: wgpu::TextureDimension::D2,
                    format: wgpu::TextureFormat::Rgba16Float,
                    usage: wgpu::TextureUsages::STORAGE_BINDING
                        | wgpu::TextureUsages::TEXTURE_BINDING,
                });

                (texture, width, height)
            })
            .0
            .create_view(&wgpu::TextureViewDescriptor::default());

        let mut encoder =
            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });

        let atlas_texture = self.common.atlas.get_or_insert_with(|| {
            device.create_texture(&wgpu::TextureDescriptor {
                label: Some("atlas"),
                size: wgpu::Extent3d { width: 4096, height: 4096, depth_or_array_layers: 1 },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: wgpu::TextureFormat::Rgba16Float,
                usage: wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::TEXTURE_BINDING,
            })
        });

        for (image, [xmin, ymin, _, _]) in new_allocations {
            queue.write_texture(
                wgpu::ImageCopyTexture {
                    texture: atlas_texture,
                    mip_level: 0,
                    origin: wgpu::Origin3d { x: *xmin, y: *ymin, z: 0 },
                    aspect: TextureAspect::All,
                },
                unsafe {
                    slice::from_raw_parts(
                        (image.data() as *const _) as *const u8,
                        mem::size_of_val(image.data()),
                    )
                },
                ImageDataLayout {
                    offset: 0,
                    bytes_per_row: Some(NonZeroU32::new(4 * 2 * image.width()).unwrap()),
                    rows_per_image: None,
                },
                Extent3d { width: image.width(), height: image.height(), depth_or_array_layers: 1 },
            );
        }

        // Number of segments to be generated, out of the prefix sum.
        let segments_count = lines.segments_count();
        if segments_count > 0 && width != 0 && height != 0 {
            let segments_blocks =
                div_round_up(segments_count as usize, BLOCK_SIZE.block_len as usize);

            let pixel_segment_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (segments_count * std::mem::size_of::<u64>()) as wgpu::BufferAddress,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });

            rasterizer::encode(
                device,
                &mut encoder,
                &self.rasterizer,
                lines,
                segments_count,
                &pixel_segment_buffer,
                timestamp_context.as_ref().map(|(timestamp, _, _)| (timestamp, 0, 1)),
            );

            let slice_size = segments_count * std::mem::size_of::<u64>();
            let size = slice_size as wgpu::BufferAddress;

            let storage_buffer1 = device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let offsets_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: ((segments_blocks as usize + 1) * std::mem::size_of::<u32>())
                    as wgpu::BufferAddress,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });

            let segments_buffer = conveyor_sort::encode(
                device,
                &mut encoder,
                &self.sort,
                segments_count,
                &pixel_segment_buffer,
                &storage_buffer1,
                &offsets_buffer,
                timestamp_context.as_ref().map(|(timestamp, _, _)| (timestamp, 2, 3)),
            );
            let style_offsets_buffer =
                device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(style_offsets),
                    usage: wgpu::BufferUsages::STORAGE,
                });
            let styles_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(styles),
                usage: wgpu::BufferUsages::STORAGE,
            });

            painter::encode(
                device,
                &mut encoder,
                &self.paint,
                &texture_view,
                segments_buffer,
                segments_count as u32,
                &style_offsets_buffer,
                &styles_buffer,
                &atlas_texture.create_view(&wgpu::TextureViewDescriptor::default()),
                width,
                height,
                background_color,
                timestamp_context.as_ref().map(|(timestamp, _, _)| (timestamp, 4, 5)),
            );
        }

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &self.render.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&texture_view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&self.render.sampler),
                },
            ],
            label: None,
        });

        {
            let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: None,
                color_attachments: &[wgpu::RenderPassColorAttachment {
                    view: &view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: true,
                    },
                }],
                depth_stencil_attachment: None,
            });

            if let Some((timestamp, _, _)) = timestamp_context.as_ref() {
                rpass.write_timestamp(timestamp, 4);
            }

            rpass.set_pipeline(&self.render.pipeline);
            rpass.set_bind_group(0, &bind_group, &[]);
            rpass.draw(0..3, 0..1);

            if let Some((timestamp, _, _)) = timestamp_context.as_ref() {
                rpass.write_timestamp(timestamp, 5);
            }
        }

        if let Some((timestamp, _, data_buffer)) = &timestamp_context {
            encoder.resolve_query_set(timestamp, 0..4, data_buffer, 0);
        }

        queue.submit(Some(encoder.finish()));

        let timings = timestamp_context.as_ref().map(|(_, timestamp_period, data_buffer)| {
            use bytemuck::{Pod, Zeroable};

            #[repr(C)]
            #[derive(Clone, Copy, Debug, Pod, Zeroable)]
            struct TimestampData {
                start: u64,
                end: u64,
            }

            let _ = data_buffer.slice(..).map_async(wgpu::MapMode::Read);

            device.poll(wgpu::Maintain::Wait);

            let view = data_buffer.slice(..).get_mapped_range();
            let timestamps: [TimestampData; Timings::size() / 2] = *bytemuck::from_bytes(&*view);
            let durations = timestamps.map(|timestamp| {
                let nanos = (timestamp.end - timestamp.start) as f32 * timestamp_period;
                Duration::from_nanos(nanos as u64)
            });

            Timings {
                rasterize: durations[0],
                sort: durations[1],
                paint: durations[2],
                render: durations[3],
            }
        });

        Ok(timings)
    }
}