vello/tests/src/transform.rs

// Copyright 2021 The piet-gpu authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Also licensed under MIT license, at your choice.

//! Tests for piet-gpu shaders.

use crate::{Config, Runner, TestResult};

use kurbo::Affine;
use piet_gpu::stages::{self, Transform, TransformCode, TransformStage};
use piet_gpu_hal::BufferUsage;
use rand::Rng;

struct AffineTestData {
    input_data: Vec<Transform>,
    expected: Vec<Affine>,
}

pub unsafe fn transform_test(runner: &mut Runner, config: &Config) -> TestResult {
    let mut result = TestResult::new("transform");
    let n_elements: u64 = config.size.choose(1 << 12, 1 << 18, 1 << 24);
    // TODO: would be nice to validate with real transform.
    let data = AffineTestData::new(n_elements as usize);
    let data_buf = runner
        .session
        .create_buffer_init(&data.input_data, BufferUsage::STORAGE)
        .unwrap();
    let memory = runner.buf_down(data_buf.size() + 24, BufferUsage::empty());
    let stage_config = stages::Config {
        n_trans: n_elements as u32,
        // This is a hack to get elements aligned.
        trans_alloc: 16,
        ..Default::default()
    };
    let config_buf = runner
        .session
        .create_buffer_init(std::slice::from_ref(&stage_config), BufferUsage::STORAGE)
        .unwrap();

    let code = TransformCode::new(&runner.session);
    let stage = TransformStage::new(&runner.session, &code);
    let binding = stage.bind(
        &runner.session,
        &code,
        &config_buf,
        &data_buf,
        &memory.dev_buf,
    );
    let mut total_elapsed = 0.0;
    let n_iter = config.n_iter;
    for i in 0..n_iter {
        let mut commands = runner.commands();
        commands.write_timestamp(0);
        stage.record(&mut commands.cmd_buf, &code, &binding, n_elements);
        commands.write_timestamp(1);
        if i == 0 || config.verify_all {
            commands.cmd_buf.memory_barrier();
            commands.download(&memory);
        }
        total_elapsed += runner.submit(commands);
        if i == 0 || config.verify_all {
            let mut dst: Vec<Transform> = Default::default();
            memory.read(&mut dst);
            if let Some(failure) = data.verify(&dst[1..]) {
                result.fail(failure);
            }
        }
    }
    result.timing(total_elapsed, n_elements * n_iter);
    result
}

impl AffineTestData {
    fn new(n: usize) -> AffineTestData {
        let mut rng = rand::thread_rng();
        let mut a = Affine::default();
        let mut b;
        let mut input_data = Vec::with_capacity(n);
        let mut expected = Vec::with_capacity(n);
        for _ in 0..n {
            loop {
                b = Affine::new([
                    rng.gen_range(-10.0, 10.0),
                    rng.gen_range(-10.0, 10.0),
                    rng.gen_range(-10.0, 10.0),
                    rng.gen_range(-10.0, 10.0),
                    rng.gen_range(-10.0, 10.0),
                    rng.gen_range(-10.0, 10.0),
                ]);
                if b.determinant() >= 1.0 {
                    break;
                }
            }
            expected.push(b);
            let c = a.inverse() * b;
            input_data.push(Transform::from_kurbo(c));
            a = b;
        }
        AffineTestData {
            input_data,
            expected,
        }
    }

    fn verify(&self, actual: &[Transform]) -> Option<String> {
        for (i, (actual, expected)) in actual.iter().zip(&self.expected).enumerate() {
            let error: f64 = actual
                .to_kurbo()
                .as_coeffs()
                .iter()
                .zip(expected.as_coeffs())
                .map(|(actual, expected)| (actual - expected).powi(2))
                .sum();
            let tolerance = 1e-6 * (i + 1) as f64;
            if error > tolerance {
                return Some(format!("{}: {} {}", i, error, tolerance));
            }
        }
        None
    }
}