[backend-comparison] Burnbench CLI (#1260)

Cargo.lock

@@ -162,10 +162,15 @@ version = "0.13.0"
 dependencies = [
  "burn",
  "burn-common",
+ "clap",
+ "crossterm",
  "derive-new",
  "dirs 5.0.1",
  "rand",
+ "ratatui",
  "serde_json",
+ "strum",
+ "strum_macros",
 ]
 
 [[package]]

Cargo.toml

@@ -43,9 +43,8 @@ license = "MIT OR Apache-2.0"
 async-trait = "0.1.74"
 bytemuck = "1.14"
 candle-core = { version = "0.3.3" }
-clap = "4.4.11"
+clap = { version = "4.4.11", features = ["derive"] }
 console_error_panic_hook = "0.1.7"
-const-random = "0.1.17"
 csv = "1.3.0"
 dashmap = "5.5.3"
 dirs = "5.0.1"

backend-comparison/Cargo.toml

@@ -22,16 +22,22 @@ ndarray-blas-netlib = ["burn/ndarray", "burn/blas-netlib"]
 ndarray-blas-openblas = ["burn/ndarray", "burn/openblas"]
 tch-cpu = ["burn/tch"]
 tch-gpu = ["burn/tch"]
+tui = ["ratatui", "crossterm"]
 wgpu = ["burn/wgpu"]
 wgpu-fusion = ["burn/default", "burn/wgpu", "burn/fusion"]
 
 [dependencies]
 burn = { path = "../burn" }
-derive-new = { workspace = true }
-rand = { workspace = true }
 burn-common = { path = "../burn-common", version = "0.13.0" }
+clap = { workspace = true }
+crossterm = { workspace = true, optional = true }
+derive-new = { workspace = true }
+dirs = { workspace = true }
+rand = { workspace = true }
+ratatui = { workspace = true, optional = true }
 serde_json = { workspace = true }
-dirs = "5.0.1"
+strum = { workspace = true }
+strum_macros = { workspace = true }
 
 [dev-dependencies]
 
@@ -55,3 +61,7 @@ harness = false
 [[bench]]
 name = "custom_gelu"
 harness = false
+
+[[bin]]
+name = "burnbench"
+path = "src/bin/burnbench.rs"

backend-comparison/README.md

@@ -1,8 +1,134 @@
 # Burn Benchmark
 
-This crate is used with `cargo bench --features <backend>`
-to compare backend computation times, from tensor operations to complex models.
+This crate allows to compare backend computation times, from tensor operations
+to complex models.
 
 Note: in order to compare different backend-specific tensor operation
 implementations (for autotuning purposes, for instance), this should be done
 within the corresponding backend crate.
+
+## burnbench CLI
+
+This crate comes with a CLI binary called `burnbench` which can be executed via
+`cargo run --bin burnbench`.
+
+The end of options argument `--` is used to pass arguments to the `burnbench`
+application. For instance `cargo run --bin burnbench -- list` passes the `list`
+argument to `burnbench` effectively calling `burnbench list`.
+
+To list all the available benches and backends use the `list` command:
+
+```sh
+> cargo run --bin burnbench -- list
+    Finished dev [unoptimized] target(s) in 0.10s
+     Running `target/debug/burnbench list`
+Available Backends:
+- candle-cpu
+- candle-cuda
+- candle-metal
+- ndarray
+- ndarray-blas-accelerate
+- ndarray-blas-netlib
+- ndarray-blas-openblas
+- tch-cpu
+- tch-gpu
+- wgpu
+- wgpu-fusion
+
+Available Benchmarks:
+- binary
+- custom-gelu
+- data
+- matmul
+- unary
+```
+
+To execute a given benchmark against a specific backend we use the `run` command
+with the arguments `--benches` and `--backends` respectively. In the following
+example we execute the `unary` benchmark against the `wgpu-fusion` backend:
+
+```sh
+> cargo run --bin burnbench -- run --benches unary --backends wgpu-fusion
+```
+
+Shorthands can be used, the following command line is the same:
+
+```sh
+> cargo run --bin burnbench -- run -b unary -B wgpu-fusion
+```
+
+Multiple benchmarks and backends can be passed on the same command line. In this
+case, all the combinations of benchmarks with backends will be executed.
+
+```sh
+> cargo run --bin burnbench -- run --benches unary binary --backends wgpu-fusion tch-gpu
+    Finished dev [unoptimized] target(s) in 0.09s
+     Running `target/debug/burnbench run --benches unary binary --backends wgpu-fusion wgpu`
+Executing the following benchmark and backend combinations (Total: 4):
+- Benchmark: unary, Backend: wgpu-fusion
+- Benchmark: binary, Backend: wgpu-fusion
+- Benchmark: unary, Backend: tch-gpu
+- Benchmark: binary, Backend: tch-gpu
+Running benchmarks...
+```
+
+### Terminal UI
+
+This is a work in progress.
+
+## Execute benchmarks with cargo
+
+To execute a benchmark against a given backend using only cargo is done with the
+`bench` command. In this case the backend is a feature of this crate.
+
+```sh
+> cargo bench --features wgpu-fusion
+```
+
+## Add a new benchmark
+
+To add a new benchmark it must be first declared in the `Cargo.toml` file of this
+crate:
+
+```toml
+[[bench]]
+name = "mybench"
+harness = false
+```
+
+Then it must be registered in the `BenchmarkValues` enumeration:
+
+```rs
+
+#[derive(Debug, Clone, PartialEq, Eq, ValueEnum, Display, EnumIter)]
+pub(crate) enum BackendValues {
+    // ...
+    #[strum(to_string = "mybench")]
+    MyBench,
+    // ...
+}
+```
+
+Create a new file `mybench.rs` in the `benches` directory and implement the
+`Benchmark` trait over your benchmark structure. Then implement the `bench`
+function. At last call the macro `backend_comparison::bench_on_backend!()` in
+the `main` function.
+
+## Add a new backend
+
+You can easily register and new backend in the `BackendValues` enumeration:
+
+```rs
+
+#[derive(Debug, Clone, PartialEq, Eq, ValueEnum, Display, EnumIter)]
+pub(crate) enum BackendValues {
+    // ...
+    #[strum(to_string = "mybackend")]
+    MyBackend,
+    // ...
+}
+```
+
+Then update the macro `bench_on_backend` to support the newly registered
+backend.
+

backend-comparison/src/bin/burnbench.rs

@@ -0,0 +1,5 @@
+use backend_comparison::burnbenchapp;
+
+fn main() {
+    burnbenchapp::run()
+}

backend-comparison/src/burnbenchapp/base.rs

@@ -0,0 +1,140 @@
+use clap::{Parser, Subcommand, ValueEnum};
+use std::process::{Command, Stdio};
+use strum::IntoEnumIterator;
+use strum_macros::{Display, EnumIter};
+
+use super::App;
+
+/// Base trait to define an application
+pub(crate) trait Application {
+    fn init(&mut self) {}
+
+    #[allow(unused)]
+    fn run(&mut self, benches: &[BenchmarkValues], backends: &[BackendValues]) {}
+
+    fn cleanup(&mut self) {}
+}
+
+#[derive(Parser, Debug)]
+#[clap(author, version, about, long_about = None)]
+struct Args {
+    #[clap(subcommand)]
+    command: Commands,
+}
+
+#[derive(Subcommand, Debug)]
+enum Commands {
+    /// List all available benchmarks and backends
+    List,
+    /// Runs benchmarks
+    Run(RunArgs),
+}
+
+#[derive(Parser, Debug)]
+struct RunArgs {
+    /// Comma-separated list of backends to include
+    #[clap(short = 'B', long = "backends", value_name = "BACKEND,BACKEND,...", num_args(0..))]
+    backends: Vec<BackendValues>,
+
+    /// Comma-separated list of benches to run
+    #[clap(short = 'b', long = "benches", value_name = "BACKEND,BACKEND,...", num_args(0..))]
+    benches: Vec<BenchmarkValues>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, ValueEnum, Display, EnumIter)]
+pub(crate) enum BackendValues {
+    #[strum(to_string = "candle-cpu")]
+    CandleCpu,
+    #[strum(to_string = "candle-cuda")]
+    CandleCuda,
+    #[strum(to_string = "candle-metal")]
+    CandleMetal,
+    #[strum(to_string = "ndarray")]
+    Ndarray,
+    #[strum(to_string = "ndarray-blas-accelerate")]
+    NdarrayBlasAccelerate,
+    #[strum(to_string = "ndarray-blas-netlib")]
+    NdarrayBlasNetlib,
+    #[strum(to_string = "ndarray-blas-openblas")]
+    NdarrayBlasOpenblas,
+    #[strum(to_string = "tch-cpu")]
+    TchCpu,
+    #[strum(to_string = "tch-gpu")]
+    TchGpu,
+    #[strum(to_string = "wgpu")]
+    Wgpu,
+    #[strum(to_string = "wgpu-fusion")]
+    WgpuFusion,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, ValueEnum, Display, EnumIter)]
+pub(crate) enum BenchmarkValues {
+    #[strum(to_string = "binary")]
+    Binary,
+    #[strum(to_string = "custom-gelu")]
+    CustomGelu,
+    #[strum(to_string = "data")]
+    Data,
+    #[strum(to_string = "matmul")]
+    Matmul,
+    #[strum(to_string = "unary")]
+    Unary,
+}
+
+pub fn run() {
+    let args = Args::parse();
+
+    match args.command {
+        Commands::List => {
+            println!("Available Backends:");
+            for backend in BackendValues::iter() {
+                println!("- {}", backend);
+            }
+
+            println!("\nAvailable Benchmarks:");
+            for bench in BenchmarkValues::iter() {
+                println!("- {}", bench);
+            }
+        }
+        Commands::Run(run_args) => {
+            if run_args.backends.is_empty() || run_args.benches.is_empty() {
+                println!("No backends or benchmarks specified. Please select at least one backend and one benchmark.");
+                return;
+            }
+
+            let total_combinations = run_args.backends.len() * run_args.benches.len();
+            println!(
+                "Executing the following benchmark and backend combinations (Total: {}):",
+                total_combinations
+            );
+            for backend in &run_args.backends {
+                for bench in &run_args.benches {
+                    println!("- Benchmark: {}, Backend: {}", bench, backend);
+                }
+            }
+
+            let mut app = App::new();
+            app.init();
+            println!("Running benchmarks...");
+            app.run(&run_args.benches, &run_args.backends);
+            app.cleanup();
+            println!("Cleanup completed. Benchmark run(s) finished.");
+        }
+    }
+}
+
+#[allow(unused)] // for tui as this is WIP
+pub(crate) fn run_cargo(command: &str, params: &[&str]) {
+    let mut cargo = Command::new("cargo")
+        .arg(command)
+        .arg("--color=always")
+        .args(params)
+        .stdout(Stdio::inherit())
+        .stderr(Stdio::inherit())
+        .spawn()
+        .expect("cargo process should run");
+    let status = cargo.wait().expect("");
+    if !status.success() {
+        std::process::exit(status.code().unwrap_or(1));
+    }
+}

backend-comparison/src/burnbenchapp/mod.rs

@@ -0,0 +1,12 @@
+mod base;
+pub use base::*;
+
+#[cfg(feature = "tui")]
+mod tui;
+#[cfg(feature = "tui")]
+use tui::TuiApplication as App;
+
+#[cfg(not(feature = "tui"))]
+mod term;
+#[cfg(not(feature = "tui"))]
+use term::TermApplication as App;

backend-comparison/src/burnbenchapp/term/base.rs

@@ -0,0 +1,29 @@
+use crate::burnbenchapp::{run_cargo, Application, BackendValues, BenchmarkValues};
+
+use derive_new::new;
+
+#[derive(new)]
+pub struct TermApplication;
+
+impl Application for TermApplication {
+    fn init(&mut self) {}
+
+    fn run(&mut self, benches: &[BenchmarkValues], backends: &[BackendValues]) {
+        // Iterate over each combination of backend and bench
+        for backend in backends.iter() {
+            for bench in benches.iter() {
+                run_cargo(
+                    "bench",
+                    &[
+                        "--bench",
+                        &bench.to_string(),
+                        "--features",
+                        &backend.to_string(),
+                    ],
+                );
+            }
+        }
+    }
+
+    fn cleanup(&mut self) {}
+}

backend-comparison/src/burnbenchapp/term/mod.rs

@@ -0,0 +1,2 @@
+mod base;
+pub use base::*;

backend-comparison/src/burnbenchapp/tui/base.rs

@@ -0,0 +1,107 @@
+use crossterm::{
+    event::{self, Event, KeyCode},
+    execute,
+    terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
+};
+use ratatui::{
+    backend::CrosstermBackend,
+    layout::{Alignment, Margin},
+    prelude::Frame,
+    widgets::Paragraph,
+    Terminal,
+};
+use std::{io, time::Duration};
+
+use crate::burnbenchapp::{
+    tui::components::regions::*, Application, BackendValues, BenchmarkValues,
+};
+
+type BenchTerminal = Terminal<CrosstermBackend<io::Stdout>>;
+
+#[derive(PartialEq)]
+enum Message {
+    QuitApplication,
+}
+
+pub struct TuiApplication {
+    terminal: BenchTerminal,
+    regions: Regions<LeftRegion, RightRegion>,
+}
+
+impl Application for TuiApplication {
+    fn init(&mut self) {}
+
+    #[allow(unused)]
+    fn run(&mut self, benches: &[BenchmarkValues], backends: &[BackendValues]) {
+        // TODO initialize widgets given passed benches and backends on the command line
+        loop {
+            self.terminal
+                .draw(|f| TuiApplication::render_app(&mut self.regions, f))
+                .expect("frame should be drawn");
+            let mut current_msg = self.handle_event();
+            if let Some(Message::QuitApplication) = current_msg {
+                break;
+            } else {
+                while current_msg.is_some() {
+                    current_msg = self.update(current_msg.unwrap());
+                }
+            }
+        }
+    }
+
+    fn cleanup(&mut self) {
+        disable_raw_mode().expect("Terminal raw mode should be disabled");
+        execute!(self.terminal.backend_mut(), LeaveAlternateScreen)
+            .expect("Alternate screen should be disabled");
+        self.terminal
+            .show_cursor()
+            .expect("Terminal cursor should be made visible");
+    }
+}
+
+impl TuiApplication {
+    pub fn new() -> Self {
+        TuiApplication {
+            terminal: Self::setup_terminal(),
+            regions: Regions::new(),
+        }
+    }
+
+    fn setup_terminal() -> BenchTerminal {
+        let mut stdout = io::stdout();
+        enable_raw_mode().expect("Terminal raw mode should be enabled");
+        execute!(stdout, EnterAlternateScreen).expect("Alternate screen should be enabled");
+        BenchTerminal::new(CrosstermBackend::new(stdout)).unwrap()
+    }
+
+    fn handle_event(&mut self) -> Option<Message> {
+        if event::poll(Duration::from_millis(250)).unwrap() {
+            if let Event::Key(key) = event::read().unwrap() {
+                match key.code {
+                    KeyCode::Char('q') => return Some(Message::QuitApplication),
+                    _ => {
+                        self.regions.set_focus(key.code);
+                    }
+                }
+            }
+        }
+        None
+    }
+
+    fn update(&mut self, _msg: Message) -> Option<Message> {
+        None
+    }
+
+    fn render_app(regions: &mut Regions<LeftRegion, RightRegion>, frame: &mut Frame) {
+        regions.draw(frame);
+        let greeting =
+            Paragraph::new("Work in Progress\n\n(press 'q' to quit)").alignment(Alignment::Center);
+        frame.render_widget(
+            greeting,
+            regions.right.rect(&RightRegion::Top).inner(&Margin {
+                horizontal: 1,
+                vertical: 10,
+            }),
+        );
+    }
+}

backend-comparison/src/burnbenchapp/tui/components/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod regions;

backend-comparison/src/burnbenchapp/tui/components/regions.rs

@@ -0,0 +1,250 @@
+/// Define a left and a right region for the application.
+/// Each region is divided in vertically stacked rectangles.
+use std::marker::PhantomData;
+use std::rc::Rc;
+
+use crossterm::event::KeyCode;
+use ratatui::{
+    layout::{Alignment, Constraint, Direction, Layout, Rect},
+    style::{Color, Style},
+    widgets::{block::Position, Block, BorderType, Borders, Padding},
+    Frame,
+};
+
+// Region Base ---------------------------------------------------------------
+
+pub(crate) struct RegionInfo {
+    width_percentage: u16,
+}
+
+pub(crate) struct RegionSectionInfo {
+    index: usize,
+    title: &'static str,
+    height_percentage: u16,
+    pub hotkey: char,
+}
+
+pub(crate) trait GetRegionInfo {
+    fn get_region_info() -> RegionInfo;
+    fn get_section_info(&self) -> RegionSectionInfo;
+}
+
+pub(crate) struct Region<I: GetRegionInfo> {
+    rects: Rc<[Rect]>,
+    info: RegionInfo,
+    _i: PhantomData<I>,
+}
+
+impl<T: GetRegionInfo> Region<T> {
+    fn new() -> Self {
+        Self {
+            rects: [].into(),
+            info: T::get_region_info(),
+            _i: PhantomData,
+        }
+    }
+
+    pub fn rect(&self, region_section: &T) -> Rect {
+        self.rects[region_section.get_section_info().index]
+    }
+
+    /// Widget to draw the style of a region
+    fn block(&self, region_section: &T, is_focused: bool) -> Block {
+        let border_color = if is_focused {
+            Color::LightRed
+        } else {
+            Color::DarkGray
+        };
+        Block::default()
+            .title(format!(
+                "{} ({})",
+                region_section.get_section_info().title,
+                region_section.get_section_info().hotkey
+            ))
+            .title_position(Position::Top)
+            .title_alignment(Alignment::Center)
+            .borders(Borders::all())
+            .border_style(Style::default().fg(border_color))
+            .border_type(BorderType::Rounded)
+            .padding(Padding {
+                left: 10,
+                right: 10,
+                top: 2,
+                bottom: 2,
+            })
+            .style(Style::default().bg(Color::Black))
+    }
+}
+
+// Left Region --------------------------------------------------------------
+
+#[derive(PartialEq)]
+pub(crate) enum LeftRegion {
+    Top,
+    Middle,
+    Bottom,
+}
+
+impl LeftRegion {
+    fn variants() -> [LeftRegion; 3] {
+        [LeftRegion::Top, LeftRegion::Middle, LeftRegion::Bottom]
+    }
+}
+
+impl GetRegionInfo for LeftRegion {
+    fn get_region_info() -> RegionInfo {
+        RegionInfo {
+            width_percentage: 25,
+        }
+    }
+
+    fn get_section_info(&self) -> RegionSectionInfo {
+        match self {
+            LeftRegion::Top => RegionSectionInfo {
+                index: 0,
+                title: "Backend",
+                height_percentage: 30,
+                hotkey: 'b',
+            },
+            LeftRegion::Middle => RegionSectionInfo {
+                index: 1,
+                title: "Benches",
+                height_percentage: 60,
+                hotkey: 'n',
+            },
+            LeftRegion::Bottom => RegionSectionInfo {
+                index: 2,
+                title: "Action",
+                height_percentage: 10,
+                hotkey: 'a',
+            },
+        }
+    }
+}
+
+// Right Region --------------------------------------------------------------
+
+#[derive(PartialEq)]
+pub(crate) enum RightRegion {
+    Top,
+    Bottom,
+}
+
+impl RightRegion {
+    fn variants() -> [RightRegion; 2] {
+        [RightRegion::Top, RightRegion::Bottom]
+    }
+}
+
+impl GetRegionInfo for RightRegion {
+    fn get_region_info() -> RegionInfo {
+        RegionInfo {
+            width_percentage: 100 - LeftRegion::get_region_info().width_percentage,
+        }
+    }
+
+    fn get_section_info(&self) -> RegionSectionInfo {
+        match self {
+            RightRegion::Top => RegionSectionInfo {
+                index: 0,
+                title: "Results",
+                height_percentage: 90,
+                hotkey: 'r',
+            },
+            RightRegion::Bottom => RegionSectionInfo {
+                index: 1,
+                title: "Progress",
+                height_percentage: 10,
+                hotkey: 'p',
+            },
+        }
+    }
+}
+
+// Regions definition --------------------------------------------------------
+
+pub enum FocusedRegion<L: GetRegionInfo, R: GetRegionInfo> {
+    Left(L),
+    Right(R),
+}
+
+pub(crate) struct Regions<L: GetRegionInfo, R: GetRegionInfo> {
+    pub left: Region<L>,
+    pub right: Region<R>,
+    focused_region: FocusedRegion<L, R>,
+}
+
+impl Regions<LeftRegion, RightRegion> {
+    pub fn new() -> Self {
+        Self {
+            left: Region::<LeftRegion>::new(),
+            right: Region::<RightRegion>::new(),
+            focused_region: FocusedRegion::Left(LeftRegion::Top),
+        }
+    }
+
+    pub fn set_focus(&mut self, key: KeyCode) -> bool {
+        self.focused_region = if key == KeyCode::Char(LeftRegion::Top.get_section_info().hotkey) {
+            FocusedRegion::Left(LeftRegion::Top)
+        } else if key == KeyCode::Char(LeftRegion::Middle.get_section_info().hotkey) {
+            FocusedRegion::Left(LeftRegion::Middle)
+        } else if key == KeyCode::Char(LeftRegion::Bottom.get_section_info().hotkey) {
+            FocusedRegion::Left(LeftRegion::Bottom)
+        } else if key == KeyCode::Char(RightRegion::Top.get_section_info().hotkey) {
+            FocusedRegion::Right(RightRegion::Top)
+        } else if key == KeyCode::Char(RightRegion::Bottom.get_section_info().hotkey) {
+            FocusedRegion::Right(RightRegion::Bottom)
+        } else {
+            return false;
+        };
+        true
+    }
+
+    pub fn draw(&mut self, frame: &mut Frame) {
+        // compute rects boundaries and update the regions accordingly
+        let outer_layout = Layout::default()
+            .direction(Direction::Horizontal)
+            .constraints(vec![
+                Constraint::Percentage(self.left.info.width_percentage),
+                Constraint::Percentage(self.right.info.width_percentage),
+            ])
+            .split(frame.size());
+        let left_rects = Layout::default()
+            .direction(Direction::Vertical)
+            .constraints(vec![
+                Constraint::Percentage(LeftRegion::Top.get_section_info().height_percentage),
+                Constraint::Percentage(LeftRegion::Middle.get_section_info().height_percentage),
+                Constraint::Percentage(LeftRegion::Bottom.get_section_info().height_percentage),
+            ])
+            .split(outer_layout[0]);
+        let right_rects = Layout::default()
+            .direction(Direction::Vertical)
+            .constraints(vec![
+                Constraint::Percentage(RightRegion::Top.get_section_info().height_percentage),
+                Constraint::Percentage(RightRegion::Bottom.get_section_info().height_percentage),
+            ])
+            .split(outer_layout[1]);
+        self.set_rects(left_rects, right_rects);
+        // Draw left region
+        for region_variant in LeftRegion::variants() {
+            let is_focused = matches!(&self.focused_region, FocusedRegion::Left(ref lr) if *lr == region_variant);
+            frame.render_widget(
+                self.left.block(&region_variant, is_focused),
+                self.left.rect(&region_variant),
+            );
+        }
+        // Draw right region
+        for region_variant in RightRegion::variants() {
+            let is_focused = matches!(&self.focused_region, FocusedRegion::Right(ref rr) if *rr == region_variant);
+            frame.render_widget(
+                self.right.block(&region_variant, is_focused),
+                self.right.rect(&region_variant),
+            );
+        }
+    }
+
+    fn set_rects(&mut self, left_rects: Rc<[Rect]>, right_rects: Rc<[Rect]>) {
+        self.left.rects = left_rects;
+        self.right.rects = right_rects;
+    }
+}

backend-comparison/src/burnbenchapp/tui/mod.rs

@@ -0,0 +1,5 @@
+mod base;
+mod components;
+
+pub use base::*;
+// pub use components::*;

backend-comparison/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod burnbenchapp;
 pub mod persistence;
 
 #[macro_export]