Add more conv2d bench cases to candle-nn benches (#3131)

Author: slckl

candle-nn/benches/benchmarks/conv.rs

@@ -6,23 +6,33 @@ use std::time::Instant;
 
 const B: usize = 1;
 const C: usize = 1;
-const M: usize = 128;
-const K: usize = 128;
-const K_SIZE: usize = 3;
 
-fn run(input: Tensor, weight: Tensor, bias: Tensor, config: Conv2dConfig) {
-    Conv2d::new(weight, Some(bias), config)
-        .forward(&input)
-        .unwrap();
+fn run(input: Tensor, weight: Tensor, bias: Option<Tensor>, config: Conv2dConfig) {
+    Conv2d::new(weight, bias, config).forward(&input).unwrap();
 }
 
-fn run_conv2d_benchmark(c: &mut Criterion, device: &Device, dtype: DType, name: &str) {
-    let weight = Tensor::ones((1, 1, K_SIZE, K_SIZE), dtype, device)
+fn run_conv2d_benchmark(
+    c: &mut Criterion,
+    device: &Device,
+    dtype: DType,
+    k_size: usize,
+    m: usize,
+    bias: bool,
+) {
+    let weight = Tensor::ones((1, C, k_size, k_size), dtype, device)
         .unwrap()
         .to_dtype(dtype)
         .unwrap();
-    let bias = Tensor::zeros(K, dtype, device).unwrap();
-    let input = Tensor::ones((B, C, M, K), dtype, device).unwrap();
+    let bias_t = if bias {
+        Some(Tensor::zeros(m, dtype, device).unwrap())
+    } else {
+        None
+    };
+    let input = Tensor::ones((B, C, m, m), dtype, device).unwrap();
+    let name = format!(
+        "conv2d_{dtype:?}_i{m}_k{k_size}x{k_size}_{}",
+        if bias { "b" } else { "nb" }
+    );
 
     let mut group = c.benchmark_group(device.bench_name(name));
     group.bench_function("iter", move |b| {
@@ -32,7 +42,7 @@ fn run_conv2d_benchmark(c: &mut Criterion, device: &Device, dtype: DType, name:
                 run(
                     black_box(input.clone()),
                     black_box(weight.clone()),
-                    black_box(bias.clone()),
+                    black_box(bias_t.clone()),
                     Default::default(),
                 );
             }
@@ -46,8 +56,14 @@ fn run_conv2d_benchmark(c: &mut Criterion, device: &Device, dtype: DType, name:
 fn criterion_benchmark(c: &mut Criterion) {
     let device = BenchDeviceHandler::new().unwrap();
     for d in device.devices {
-        run_conv2d_benchmark(c, &d, DType::F32, "conv2d_f32");
-        run_conv2d_benchmark(c, &d, DType::F16, "conv2d_f16");
+        run_conv2d_benchmark(c, &d, DType::F32, 3, 128, true);
+        run_conv2d_benchmark(c, &d, DType::F32, 1, 128, false);
+        run_conv2d_benchmark(c, &d, DType::F32, 5, 128, false);
+        run_conv2d_benchmark(c, &d, DType::F32, 3, 512, false);
+        run_conv2d_benchmark(c, &d, DType::F16, 3, 128, true);
+        run_conv2d_benchmark(c, &d, DType::F16, 1, 128, false);
+        run_conv2d_benchmark(c, &d, DType::F16, 5, 128, false);
+        run_conv2d_benchmark(c, &d, DType::F16, 5, 512, false);
     }
 }