//! Benchmarks for matrix multiplication.
//!
//! Run with:
//! ```bash
//! cargo bench --bench matmul --features simd,gemm
//! ```
//!
//! Memory allocation tracking is enabled via divan's AllocProfiler.

#[path = "common/mod.rs"]
mod common;
use common::BencherExt;

use burn_tensor::{Tensor, TensorData};
use divan::{AllocProfiler, Bencher};

#[global_allocator]
static ALLOC: AllocProfiler = AllocProfiler::system();

fn main() {
    println!("Matrix Multiplication Benchmarks");
    println!("Memory allocation tracking enabled");
    println!();
    divan::main();
    common::report_failures();
}

fn make_matrix(rows: usize, cols: usize) -> Tensor<2> {
    let data: Vec<f32> = (0..rows * cols)
        .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5)
        .collect();
    Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default())
}

fn make_batch_matrix(batch: usize, rows: usize, cols: usize) -> Tensor<3> {
    let data: Vec<f32> = (0..batch * rows * cols)
        .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5)
        .collect();
    Tensor::from_data(
        TensorData::new(data, [batch, rows, cols]),
        &Default::default(),
    )
}

macro_rules! bench_backend {
    ($mod_name:ident, $backend_name:literal) => {
        #[divan::bench_group(name = $backend_name)]
        mod $mod_name {
            use super::*;

            // Square matrices
            #[divan::bench_group(name = "square")]
            mod square {
                use super::*;

                #[divan::bench]
                fn matmul_64x64(bencher: Bencher) {
                    let a = make_matrix(64, 64);
                    let b = make_matrix(64, 64);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn matmul_128x128(bencher: Bencher) {
                    let a = make_matrix(128, 128);
                    let b = make_matrix(128, 128);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn matmul_256x256(bencher: Bencher) {
                    let a = make_matrix(256, 256);
                    let b = make_matrix(256, 256);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn matmul_512x512(bencher: Bencher) {
                    let a = make_matrix(512, 512);
                    let b = make_matrix(512, 512);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn matmul_1024x1024(bencher: Bencher) {
                    let a = make_matrix(1024, 1024);
                    let b = make_matrix(1024, 1024);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }
            }

            // Rectangular matrices (common in neural networks)
            #[divan::bench_group(name = "rectangular")]
            mod rectangular {
                use super::*;

                // Linear layer: [batch, in] x [in, out]
                #[divan::bench]
                fn linear_256x512_512x256(bencher: Bencher) {
                    let a = make_matrix(256, 512);
                    let b = make_matrix(512, 256);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                // Attention: [seq, hidden] x [hidden, seq]
                #[divan::bench]
                fn attention_512x64_64x512(bencher: Bencher) {
                    let a = make_matrix(512, 64);
                    let b = make_matrix(64, 512);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                // Wide matrix (embedding lookup style)
                #[divan::bench]
                fn wide_128x1024_1024x128(bencher: Bencher) {
                    let a = make_matrix(128, 1024);
                    let b = make_matrix(1024, 128);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }
            }

            // Batched matmul (common in transformers)
            #[divan::bench_group(name = "batched")]
            mod batched {
                use super::*;

                #[divan::bench]
                fn batch8_64x64(bencher: Bencher) {
                    let a = make_batch_matrix(8, 64, 64);
                    let b = make_batch_matrix(8, 64, 64);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn batch16_128x128(bencher: Bencher) {
                    let a = make_batch_matrix(16, 128, 128);
                    let b = make_batch_matrix(16, 128, 128);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn batch32_64x64(bencher: Bencher) {
                    let a = make_batch_matrix(32, 64, 64);
                    let b = make_batch_matrix(32, 64, 64);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                // Transformer attention heads
                #[divan::bench]
                fn heads12_seq512_dim64(bencher: Bencher) {
                    let a = make_batch_matrix(12, 512, 64);
                    let b = make_batch_matrix(12, 64, 512);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }
            }

            // Transposed inputs (tests contiguous conversion overhead)
            #[divan::bench_group(name = "transposed")]
            mod transposed {
                use super::*;

                #[divan::bench]
                fn lhs_transposed_256x256(bencher: Bencher) {
                    let a = make_matrix(256, 256).transpose();
                    let b = make_matrix(256, 256);
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn rhs_transposed_256x256(bencher: Bencher) {
                    let a = make_matrix(256, 256);
                    let b = make_matrix(256, 256).transpose();
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn both_transposed_256x256(bencher: Bencher) {
                    let a = make_matrix(256, 256).transpose();
                    let b = make_matrix(256, 256).transpose();
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }
            }
        }
    };
}

bench_backend!(backend, "backend");

// Integer matmul benchmarks
fn make_int_matrix(rows: usize, cols: usize) -> Option<Tensor<2, burn_tensor::Int>> {
    common::try_setup(|| {
        let data: Vec<i32> = (0..rows * cols).map(|i| (i % 1000) as i32 - 500).collect();
        Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default())
    })
}

fn make_int_batch_matrix(
    batch: usize,
    rows: usize,
    cols: usize,
) -> Option<Tensor<3, burn_tensor::Int>> {
    common::try_setup(|| {
        let data: Vec<i32> = (0..batch * rows * cols)
            .map(|i| (i % 1000) as i32 - 500)
            .collect();
        Tensor::from_data(
            TensorData::new(data, [batch, rows, cols]),
            &Default::default(),
        )
    })
}

macro_rules! bench_int_backend {
    ($mod_name:ident, $backend_name:literal) => {
        #[divan::bench_group(name = $backend_name)]
        mod $mod_name {
            use super::*;

            #[divan::bench_group(name = "int_square")]
            mod square {
                use super::*;

                #[divan::bench]
                fn matmul_64x64(bencher: Bencher) {
                    let Some(a) = make_int_matrix(64, 64) else {
                        bencher.bench(|| ());
                        return;
                    };
                    let Some(b) = make_int_matrix(64, 64) else {
                        bencher.bench(|| ());
                        return;
                    };
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn matmul_128x128(bencher: Bencher) {
                    let Some(a) = make_int_matrix(128, 128) else {
                        bencher.bench(|| ());
                        return;
                    };
                    let Some(b) = make_int_matrix(128, 128) else {
                        bencher.bench(|| ());
                        return;
                    };
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn matmul_256x256(bencher: Bencher) {
                    let Some(a) = make_int_matrix(256, 256) else {
                        bencher.bench(|| ());
                        return;
                    };
                    let Some(b) = make_int_matrix(256, 256) else {
                        bencher.bench(|| ());
                        return;
                    };
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn matmul_512x512(bencher: Bencher) {
                    let Some(a) = make_int_matrix(512, 512) else {
                        bencher.bench(|| ());
                        return;
                    };
                    let Some(b) = make_int_matrix(512, 512) else {
                        bencher.bench(|| ());
                        return;
                    };
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }
            }

            #[divan::bench_group(name = "int_batched")]
            mod batched {
                use super::*;

                #[divan::bench]
                fn batch8_64x64(bencher: Bencher) {
                    let Some(a) = make_int_batch_matrix(8, 64, 64) else {
                        bencher.bench(|| ());
                        return;
                    };
                    let Some(b) = make_int_batch_matrix(8, 64, 64) else {
                        bencher.bench(|| ());
                        return;
                    };
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }

                #[divan::bench]
                fn batch16_128x128(bencher: Bencher) {
                    let Some(a) = make_int_batch_matrix(16, 128, 128) else {
                        bencher.bench(|| ());
                        return;
                    };
                    let Some(b) = make_int_batch_matrix(16, 128, 128) else {
                        bencher.bench(|| ());
                        return;
                    };
                    bencher.bench_synced(|| a.clone().matmul(b.clone()));
                }
            }
        }
    };
}

bench_int_backend!(backend_int, "backend_int");

// Broadcast matmul benchmarks
fn make_batch_matrix_4d(b1: usize, b2: usize, rows: usize, cols: usize) -> Tensor<4> {
    let data: Vec<f32> = (0..b1 * b2 * rows * cols)
        .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5)
        .collect();
    Tensor::from_data(
        TensorData::new(data, [b1, b2, rows, cols]),
        &Default::default(),
    )
}

macro_rules! bench_broadcast_backend {
    ($mod_name:ident, $backend_name:literal) => {
        #[divan::bench_group(name = $backend_name)]
        mod $mod_name {
            use super::*;

            // Broadcast: [1, M, K] x [batch, K, N] -> [batch, M, N]
            #[divan::bench]
            fn broadcast_1_vs_8_64x64(bencher: Bencher) {
                let a = make_batch_matrix(1, 64, 64);
                let b = make_batch_matrix(8, 64, 64);
                bencher.bench_synced(|| a.clone().matmul(b.clone()));
            }

            // Broadcast: [batch, M, K] x [1, K, N] -> [batch, M, N]
            #[divan::bench]
            fn broadcast_8_vs_1_64x64(bencher: Bencher) {
                let a = make_batch_matrix(8, 64, 64);
                let b = make_batch_matrix(1, 64, 64);
                bencher.bench_synced(|| a.clone().matmul(b.clone()));
            }

            // 4D broadcast: [2, 1, M, K] x [1, 4, K, N] -> [2, 4, M, N]
            #[divan::bench]
            fn broadcast_4d_2x1_vs_1x4_32x32(bencher: Bencher) {
                let a = make_batch_matrix_4d(2, 1, 32, 32);
                let b = make_batch_matrix_4d(1, 4, 32, 32);
                bencher.bench_synced(|| a.clone().matmul(b.clone()));
            }

            // Larger 4D broadcast
            #[divan::bench]
            fn broadcast_4d_4x1_vs_1x4_64x64(bencher: Bencher) {
                let a = make_batch_matrix_4d(4, 1, 64, 64);
                let b = make_batch_matrix_4d(1, 4, 64, 64);
                bencher.bench_synced(|| a.clone().matmul(b.clone()));
            }
        }
    };
}

bench_broadcast_backend!(backend_broadcast, "backend_broadcast");
