|
| 1 | +extern crate clap; |
| 2 | +extern crate png; |
| 3 | +extern crate wyvern; |
| 4 | + |
| 5 | +use clap::{App, Arg}; |
| 6 | +use png::{BitDepth, ColorType, Encoder, HasParameters}; |
| 7 | +use std::fs::File; |
| 8 | +use std::io::BufWriter; |
| 9 | +use std::ops::{Add, Div, Mul, Sub}; |
| 10 | +use std::path::PathBuf; |
| 11 | +use std::time::{Duration, Instant}; |
| 12 | +use wyvern::core::builder::ProgramBuilder; |
| 13 | +use wyvern::core::executor::{Executable, Executor, Resource, IO}; |
| 14 | +use wyvern::core::program::{ConstantVector, TokenValue}; |
| 15 | +use wyvern::core::types::{Array, Constant, Variable}; |
| 16 | +use wyvern::vk::executor::VkExecutor; |
| 17 | + |
| 18 | +const WIDTH: u32 = 3840; |
| 19 | +const HEIGHT: u32 = 2160; |
| 20 | +const OUTFILE: &str = "out.png"; |
| 21 | +const CENTER_X: f32 = -0.75; |
| 22 | +const CENTER_Y: f32 = 0.0; |
| 23 | +const ZOOM: f32 = HEIGHT as f32 / 2.5; |
| 24 | +const ITERATIONS: usize = 2000; |
| 25 | + |
| 26 | +enum Mode { |
| 27 | + Native, |
| 28 | + Cpu, |
| 29 | + Vk, |
| 30 | +} |
| 31 | + |
| 32 | +trait Number: |
| 33 | + Copy |
| 34 | + + Add<Self, Output = Self> |
| 35 | + + Sub<Self, Output = Self> |
| 36 | + + Mul<Self, Output = Self> |
| 37 | + + Div<Self, Output = Self> |
| 38 | +{ |
| 39 | +} |
| 40 | + |
| 41 | +impl Number for f32 {} |
| 42 | +impl<'a> Number for Constant<'a, f32> {} |
| 43 | + |
| 44 | +fn get_opts() -> (Mode, u32, u32, PathBuf) { |
| 45 | + let args = App::new("mandelbrot") |
| 46 | + .author("Dario Ostuni <[email protected]>") |
| 47 | + .arg( |
| 48 | + Arg::with_name("mode") |
| 49 | + .short("m") |
| 50 | + .long("mode") |
| 51 | + .help("Compute engine") |
| 52 | + .possible_values(&["native", "cpu", "vulkan"]) |
| 53 | + .required(true) |
| 54 | + .takes_value(true), |
| 55 | + ) |
| 56 | + .arg( |
| 57 | + Arg::with_name("width") |
| 58 | + .short("w") |
| 59 | + .long("width") |
| 60 | + .help("Width of the output") |
| 61 | + .requires("height") |
| 62 | + .takes_value(true), |
| 63 | + ) |
| 64 | + .arg( |
| 65 | + Arg::with_name("height") |
| 66 | + .short("h") |
| 67 | + .long("height") |
| 68 | + .help("Height of the output") |
| 69 | + .requires("width") |
| 70 | + .takes_value(true), |
| 71 | + ) |
| 72 | + .arg( |
| 73 | + Arg::with_name("output") |
| 74 | + .short("o") |
| 75 | + .long("output") |
| 76 | + .help("Output file") |
| 77 | + .takes_value(true), |
| 78 | + ) |
| 79 | + .get_matches(); |
| 80 | + let mode = match args.value_of("mode").unwrap() { |
| 81 | + "native" => Mode::Native, |
| 82 | + "cpu" => Mode::Cpu, |
| 83 | + "vulkan" => Mode::Vk, |
| 84 | + _ => unreachable!(), |
| 85 | + }; |
| 86 | + let width = match args.value_of("width") { |
| 87 | + None => WIDTH, |
| 88 | + Some(s) => s.parse().unwrap_or(WIDTH), |
| 89 | + }; |
| 90 | + let height = match args.value_of("height") { |
| 91 | + None => HEIGHT, |
| 92 | + Some(s) => s.parse().unwrap_or(HEIGHT), |
| 93 | + }; |
| 94 | + let outfile = args.value_of("output").unwrap_or(OUTFILE); |
| 95 | + (mode, width, height, PathBuf::from(outfile)) |
| 96 | +} |
| 97 | + |
| 98 | +fn main() { |
| 99 | + let (mode, width, height, outfile_path) = get_opts(); |
| 100 | + let mut outfile = BufWriter::new(File::create(outfile_path).unwrap()); |
| 101 | + let mut encoder = Encoder::new(&mut outfile, width, height); |
| 102 | + encoder.set(ColorType::Grayscale).set(BitDepth::Eight); |
| 103 | + let mut writer = encoder.write_header().unwrap(); |
| 104 | + let (data, time) = &match mode { |
| 105 | + Mode::Native => native(width, height), |
| 106 | + Mode::Cpu => unimplemented!(), |
| 107 | + Mode::Vk => vk(width, height), |
| 108 | + }; |
| 109 | + let data = colorize(data); |
| 110 | + writer.write_image_data(&data).unwrap(); |
| 111 | + println!("{:?}", time); |
| 112 | +} |
| 113 | + |
| 114 | +fn colorize(data: &[f32]) -> Vec<u8> { |
| 115 | + data.iter() |
| 116 | + .map(|&x| if x <= 2.0 { 0 } else { 255 }) |
| 117 | + .collect() |
| 118 | +} |
| 119 | + |
| 120 | +fn native(width: u32, height: u32) -> (Vec<f32>, Duration) { |
| 121 | + let mut v = Vec::new(); |
| 122 | + let now = Instant::now(); |
| 123 | + for y in 0..height { |
| 124 | + for x in 0..width { |
| 125 | + let (x, y) = pixel2coordinates( |
| 126 | + x as f32, |
| 127 | + y as f32, |
| 128 | + CENTER_X, |
| 129 | + CENTER_Y, |
| 130 | + width as f32, |
| 131 | + height as f32, |
| 132 | + ZOOM, |
| 133 | + 2.0, |
| 134 | + ); |
| 135 | + v.push(mandelbrot(x, y, 0.0, 2.0, ITERATIONS)); |
| 136 | + } |
| 137 | + } |
| 138 | + (v, now.elapsed()) |
| 139 | +} |
| 140 | + |
| 141 | +fn vk(width: u32, height: u32) -> (Vec<f32>, Duration) { |
| 142 | + let builder = ProgramBuilder::new(); |
| 143 | + wyvern_program(&builder); |
| 144 | + let program = builder.finalize().unwrap(); |
| 145 | + let executor = VkExecutor::new(Default::default()).unwrap(); |
| 146 | + let mut executable = executor.compile(program).unwrap(); |
| 147 | + let input = executor.new_resource().unwrap(); |
| 148 | + let output = executor.new_resource().unwrap(); |
| 149 | + input.set_data(TokenValue::Vector(ConstantVector::U32(vec![width, height]))); |
| 150 | + output.set_data(TokenValue::Vector(ConstantVector::F32(vec![ |
| 151 | + 0.0; |
| 152 | + (width * height) |
| 153 | + as usize |
| 154 | + ]))); |
| 155 | + executable.bind("input", IO::Input, input.clone()); |
| 156 | + executable.bind("output", IO::Output, output.clone()); |
| 157 | + let now = Instant::now(); |
| 158 | + executable.run().unwrap(); |
| 159 | + let time = now.elapsed(); |
| 160 | + ( |
| 161 | + match output.get_data() { |
| 162 | + TokenValue::Vector(ConstantVector::F32(x)) => x, |
| 163 | + _ => unreachable!(), |
| 164 | + }, |
| 165 | + time, |
| 166 | + ) |
| 167 | +} |
| 168 | + |
| 169 | +fn wyvern_program(b: &ProgramBuilder) { |
| 170 | + let zero = Constant::new(0_u32, b); |
| 171 | + let fzero = Constant::new(0_f32, b); |
| 172 | + let one = Constant::new(1_u32, b); |
| 173 | + let ftwo = Constant::new(2_f32, b); |
| 174 | + let input = Array::new(zero, 0, true, b).mark_as_input("input"); |
| 175 | + let output = Array::new(zero, 0, true, b).mark_as_output("output"); |
| 176 | + let width = input.at(zero).load(); |
| 177 | + let height = input.at(one).load(); |
| 178 | + let fwidth = Constant::from(width); |
| 179 | + let fheight = Constant::from(height); |
| 180 | + let center_x = Constant::new(CENTER_X, b); |
| 181 | + let center_y = Constant::new(CENTER_Y, b); |
| 182 | + let zoom = Constant::new(ZOOM, b); |
| 183 | + let cells = width * height; |
| 184 | + let id = b.worker_id(); |
| 185 | + let size = b.num_workers(); |
| 186 | + let cell = Variable::new(b); |
| 187 | + cell.store(id); |
| 188 | + b.while_loop( |
| 189 | + |_| cell.load().lt(cells), |
| 190 | + |_| { |
| 191 | + let cell_id = cell.load(); |
| 192 | + let local_y = Constant::from(cell_id / width); |
| 193 | + let local_x = Constant::from(cell_id % width); |
| 194 | + let (local_x, local_y) = pixel2coordinates( |
| 195 | + local_x, local_y, center_x, center_y, fwidth, fheight, zoom, ftwo, |
| 196 | + ); |
| 197 | + let value = mandelbrot(local_x, local_y, fzero, ftwo, ITERATIONS); |
| 198 | + output.at(cell_id).store(value); |
| 199 | + cell.store(cell_id + size); |
| 200 | + }, |
| 201 | + ); |
| 202 | +} |
| 203 | + |
| 204 | +fn pixel2coordinates<T: Number>( |
| 205 | + mut x: T, |
| 206 | + mut y: T, |
| 207 | + x0: T, |
| 208 | + y0: T, |
| 209 | + width: T, |
| 210 | + height: T, |
| 211 | + zoom: T, |
| 212 | + two: T, |
| 213 | +) -> (T, T) { |
| 214 | + x = x - width / two; |
| 215 | + y = y - height / two; |
| 216 | + x = x / zoom; |
| 217 | + y = y / zoom; |
| 218 | + (x + x0, y + y0) |
| 219 | +} |
| 220 | + |
| 221 | +fn mandelbrot<T: Number>(a0: T, b0: T, zero: T, two: T, iterations: usize) -> T { |
| 222 | + let mut a = zero; |
| 223 | + let mut b = zero; |
| 224 | + for _ in 0..iterations { |
| 225 | + let tmp_a = a * a - b * b + a0; |
| 226 | + b = two * a * b + b0; |
| 227 | + a = tmp_a; |
| 228 | + } |
| 229 | + a * a + b * b |
| 230 | +} |
0 commit comments