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 std::{borrow::Cow, mem, time::Duration};

use anyhow::Error;
use conveyor_sort::{SortContext, BLOCK_SIZE};
use painter::PaintContext;
pub use painter::{Color, Style, TILE_HEIGHT, TILE_WIDTH};
use wgpu::util::DeviceExt;

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

#[derive(Debug)]
pub struct Timings {
    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 {
    texture: Option<(wgpu::Texture, u32, u32)>,
}

#[derive(Debug)]
pub struct Renderer {
    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 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 { sort, paint, render, common: Resources::default(), has_timestamp_query }
    }

    pub fn render(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        texture: &wgpu::Texture,
        width: u32,
        height: u32,
        segments: &[u64],
        style_offsets: &[u32],
        styles: &[u32],
        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: Timings::size() as u64 * mem::size_of::<u64>() 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 });

        if !segments.is_empty() && width != 0 && height != 0 {
            let mut segments = segments.to_vec();

            let old_len = segments.len();
            let len_in_blocks =
                u32::try_from(div_round_up(old_len, BLOCK_SIZE.block_len as usize)).unwrap();

            // Adds padding to the data to please the sort. `BLOCK_SIZE.block_len` should
            // be enough, but twice this amount is required for an unknown reason.
            // TODO remove this when the sort is replaced by a better version.
            segments.resize(
                std::cmp::max(
                    (len_in_blocks * BLOCK_SIZE.block_len) as usize,
                    (BLOCK_SIZE.block_len * 2) as usize,
                ),
                bytemuck::cast(u64::MAX),
            );

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

            let storage_buffer0 = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(segments.as_slice()),
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
            });
            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: ((len_in_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.len(),
                &storage_buffer0,
                &storage_buffer1,
                &offsets_buffer,
                timestamp_context.as_ref().map(|(timestamp, _, _)| (timestamp, 0, 1)),
            );
            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,
                old_len as u32,
                &style_offsets_buffer,
                &styles_buffer,
                width,
                height,
                background_color,
                timestamp_context.as_ref().map(|(timestamp, _, _)| (timestamp, 2, 3)),
            );
        }

        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 { sort: durations[0], paint: durations[1], render: durations[2] }
        });

        Ok(timings)
    }
}