diff --git a/candle-nn/examples/cpu_benchmarks.rs b/candle-nn/examples/cpu_benchmarks.rs index 6007ff6c..9ded5f71 100644 --- a/candle-nn/examples/cpu_benchmarks.rs +++ b/candle-nn/examples/cpu_benchmarks.rs @@ -180,8 +180,25 @@ impl Benchmark for Conv2dIm2Col { const ITERS: usize = 5; } -struct Matmul; -impl Benchmark for Matmul { +struct MatMul; +impl Benchmark for MatMul { + type PreProcessData = (Tensor, Tensor); + type RunResult = Tensor; + fn preprocess() -> Result { + let lhs = Tensor::randn(0f32, 1., (1024, 1024), &Device::Cpu)?; + let rhs = Tensor::randn(0f32, 1., (1024, 1024), &Device::Cpu)?; + Ok((lhs, rhs)) + } + + fn run_one(d: &Self::PreProcessData) -> Result { + d.0.matmul(&d.1) + } + + const ITERS: usize = 100; +} + +struct MatVec; +impl Benchmark for MatVec { type PreProcessData = (Tensor, Tensor); type RunResult = Tensor; fn preprocess() -> Result { @@ -271,6 +288,7 @@ enum Task { Conv2d, Conv2dIm2Col, Matmul, + Matvec, Qmatmul, Softmax, SoftmaxLastDim, @@ -293,7 +311,8 @@ fn main() -> Result<()> { Task::Conv1d => run::(args.iters)?, Task::Conv2d => run::(args.iters)?, Task::Conv2dIm2Col => run::(args.iters)?, - Task::Matmul => run::(args.iters)?, + Task::Matmul => run::(args.iters)?, + Task::Matvec => run::(args.iters)?, Task::Softmax => run::(args.iters)?, Task::SoftmaxLastDim => run::(args.iters)?, Task::Qmatmul => run::(args.iters)?,