1 files changed, 19 insertions, 0 deletions
diff --git a/candle-nn/examples/cpu_benchmarks.rs b/candle-nn/examples/cpu_benchmarks.rs
index 001be116..430316b8 100644
--- a/candle-nn/examples/cpu_benchmarks.rs
+++ b/candle-nn/examples/cpu_benchmarks.rs
@@ -238,6 +238,23 @@ impl Benchmark for QMatMul {
     const ITERS: usize = 100;
 }
 
+struct Cat;
+impl Benchmark for Cat {
+    type PreProcessData = (Tensor, Tensor);
+    type RunResult = Tensor;
+    fn preprocess() -> Result<Self::PreProcessData> {
+        let lhs = Tensor::randn(0f32, 1., (1, 32, 2000, 128), &Device::Cpu)?;
+        let rhs = Tensor::randn(0f32, 1., (1, 32, 1, 128), &Device::Cpu)?;
+        Ok((lhs, rhs))
+    }
+
+    fn run_one(d: &Self::PreProcessData) -> Result<Self::RunResult> {
+        Tensor::cat(&[&d.0, &d.1], 2)
+    }
+
+    const ITERS: usize = 1000;
+}
+
 struct Softmax;
 impl Benchmark for Softmax {
     type PreProcessData = Tensor;
@@ -295,6 +312,7 @@ enum Task {
     Qmatmul,
     Softmax,
     SoftmaxLastDim,
+    Cat,
 }
 
 #[derive(Parser, Debug)]
@@ -319,6 +337,7 @@ fn main() -> Result<()> {
         Task::Softmax => run::<Softmax>(args.iters)?,
         Task::SoftmaxLastDim => run::<SoftmaxLastDim>(args.iters)?,
         Task::Qmatmul => run::<QMatMul>(args.iters)?,
+        Task::Cat => run::<Cat>(args.iters)?,
     }
     Ok(())
 }