feat(pool): add AdaptiveAvgPool3d kernels

crates/cubek-pool/src/definition/base.rs

@@ -36,6 +36,17 @@ pub struct PoolBackwardProblem<const N: usize> {
     pub mode: PoolMode<N>,
 }
 
+#[cfg(feature = "benchmarks")]
+impl<const N: usize> PoolBackwardProblem<N> {
+    /// Reconstruct the channels-last input shape from the stored problem dimensions.
+    pub(crate) fn input_shape(&self) -> Shape {
+        let mut shape = vec![self.out_grad_shape[0]];
+        shape.extend_from_slice(&self.input_size);
+        shape.push(self.out_grad_shape[N + 1]);
+        Shape::from(shape)
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum PoolMode<const N: usize> {
     Max(MaxPoolOptions<N>),

crates/cubek-pool/src/definition/error.rs

@@ -13,4 +13,22 @@ pub enum PoolError {
 
     #[error("Channel count mismatch: input has {input} but output has {output}")]
     ChannelMismatch { input: usize, output: usize },
+
+    #[error("{tensor} spatial dimensions must be non-zero, got {actual:?}")]
+    InvalidSpatialSize {
+        tensor: &'static str,
+        actual: Vec<usize>,
+    },
+
+    #[error("Output spatial shape mismatch: expected {expected:?} but got {actual:?}")]
+    OutputSizeMismatch {
+        expected: Vec<usize>,
+        actual: Vec<usize>,
+    },
+
+    #[error("Input gradient shape mismatch: expected {expected:?} but got {actual:?}")]
+    InputGradientShapeMismatch {
+        expected: Vec<usize>,
+        actual: Vec<usize>,
+    },
 }

crates/cubek-pool/src/eval/cpu_reference/backward/adaptive_avg_pool.rs

@@ -1,58 +1,58 @@
-use crate::{definition::AdaptiveAvgPoolOptions, eval::cpu_reference::decode_index};
+use crate::eval::cpu_reference::decode_index_simple;
 use cubek_test_utils::HostData;
 
 pub fn run_adaptive_avg_pool_backward<const N: usize>(
     grad_output: &HostData,
-    _opts: &AdaptiveAvgPoolOptions<N>,
     grad_input_dims: &[usize],
     grad_output_dims: &[usize],
     grad_input_strides: &[usize],
 ) -> Vec<f32> {
     let total: usize = grad_input_dims.iter().product();
     let mut grad_input = vec![0.0; total];
+    let batch_size = grad_output_dims[0];
+    let channels = grad_output_dims[N + 1];
+    let spatial_output = &grad_output_dims[1..N + 1];
+    let total_spatial_output: usize = spatial_output.iter().product();
 
-    if N != 2 {
-        return grad_input;
-    }
-
-    let out_h = grad_input_dims[1];
-    let out_w = grad_input_dims[2];
-    let grad_h = grad_output_dims[1];
-    let grad_w = grad_output_dims[2];
-
-    for (i, grad_val) in grad_input.iter_mut().enumerate().take(total) {
-        let coords = decode_index(i, grad_input_dims, grad_input_strides);
-        let batch = coords[0];
-        let ih = coords[1];
-        let iw = coords[2];
-        let channel = coords[3];
-
-        let oh_start = start_index(ih, out_h, grad_h);
-        let oh_end = end_index(ih, out_h, grad_h);
-        let ow_start = start_index(iw, out_w, grad_w);
-        let ow_end = end_index(iw, out_w, grad_w);
-
-        let mut grad_acc = 0.0f32;
+    for batch in 0..batch_size {
+        for output_linear in 0..total_spatial_output {
+            let output_coords = decode_index_simple(output_linear, spatial_output);
+            let mut starts = [0; N];
+            let mut ends = [0; N];
+            for d in 0..N {
+                starts[d] = start_index(
+                    output_coords[d],
+                    grad_output_dims[d + 1],
+                    grad_input_dims[d + 1],
+                );
+                ends[d] = end_index(
+                    output_coords[d],
+                    grad_output_dims[d + 1],
+                    grad_input_dims[d + 1],
+                );
+            }
 
-        for oh in oh_start..oh_end {
-            let ih_start = start_index(oh, grad_h, out_h);
-            let ih_end = end_index(oh, grad_h, out_h);
+            let window_shape: [usize; N] = core::array::from_fn(|d| ends[d] - starts[d]);
+            let window_volume: usize = window_shape.iter().product();
 
-            if ih >= ih_start && ih < ih_end {
-                for ow in ow_start..ow_end {
-                    let iw_start = start_index(ow, grad_w, out_w);
-                    let iw_end = end_index(ow, grad_w, out_w);
+            for channel in 0..channels {
+                let mut grad_coords = Vec::with_capacity(N + 2);
+                grad_coords.push(batch);
+                grad_coords.extend_from_slice(&output_coords);
+                grad_coords.push(channel);
+                let contribution = grad_output.get_f32(&grad_coords) / window_volume as f32;
 
-                    if iw >= iw_start && iw < iw_end {
-                        let count = (ih_end - ih_start) * (iw_end - iw_start);
-                        let out_coords = vec![batch, oh, ow, channel];
-                        grad_acc += grad_output.get_f32(&out_coords) / count as f32;
+                for window_linear in 0..window_volume {
+                    let window_coords = decode_index_simple(window_linear, &window_shape);
+                    let mut input_offset = batch * grad_input_strides[0];
+                    for d in 0..N {
+                        input_offset += (starts[d] + window_coords[d]) * grad_input_strides[d + 1];
                     }
+                    input_offset += channel * grad_input_strides[N + 1];
+                    grad_input[input_offset] += contribution;
                 }
             }
         }
-
-        *grad_val = grad_acc;
     }
 
     grad_input

crates/cubek-pool/src/eval/cpu_reference/backward/mod.rs

@@ -31,12 +31,7 @@ pub fn strategy_result(
     let dtype = f32_elem_type();
     let indices_dtype = i32_elem_type();
     let out_grad_shape = problem.out_grad_shape.to_vec();
-    let input_shape = vec![
-        out_grad_shape[0],
-        problem.input_size[0],
-        problem.input_size[1],
-        out_grad_shape[3],
-    ];
+    let input_shape = problem.input_shape().to_vec();
 
     let (input_handle, _input_host) = make_random_f32_host(&client, input_shape.clone(), seed);
     let (out_grad_handle, _out_grad_host) =
@@ -114,12 +109,7 @@ pub fn cpu_reference_result(
     }
 
     let out_grad_shape = problem.out_grad_shape.to_vec();
-    let input_shape = vec![
-        out_grad_shape[0],
-        problem.input_size[0],
-        problem.input_size[1],
-        out_grad_shape[3],
-    ];
+    let input_shape = problem.input_shape().to_vec();
 
     if let Some(p) = progress {
         let total: usize = input_shape.iter().product();

crates/cubek-pool/src/eval/cpu_reference/forward/adaptive_avg_pool.rs

@@ -1,6 +1,6 @@
 use crate::{
     definition::AdaptiveAvgPoolOptions,
-    eval::cpu_reference::{decode_index, forward::decode_index_simple},
+    eval::cpu_reference::{decode_index, decode_index_simple},
 };
 use cubek_test_utils::HostData;
 

crates/cubek-pool/src/eval/cpu_reference/forward/avg_pool.rs

@@ -1,6 +1,6 @@
 use crate::{
     definition::AvgPoolOptions,
-    eval::cpu_reference::{decode_index, forward::decode_index_simple},
+    eval::cpu_reference::{decode_index, decode_index_simple},
 };
 use cubek_test_utils::HostData;
 

crates/cubek-pool/src/eval/cpu_reference/forward/max_pool.rs

@@ -1,9 +1,6 @@
 use crate::{
     definition::MaxPoolOptions,
-    eval::cpu_reference::{
-        decode_index,
-        forward::{decode_index_simple, get_window_coords},
-    },
+    eval::cpu_reference::{decode_index, decode_index_simple, forward::get_window_coords},
 };
 use cubek_test_utils::HostData;
 

crates/cubek-pool/src/eval/cpu_reference/forward/mod.rs

@@ -8,7 +8,7 @@ pub use max_pool::{run_max_pool, run_max_pool_with_indices};
 
 use super::{f32_elem_type, i32_elem_type, make_random_f32_host, make_zero_handle};
 use crate::definition::{PoolForwardProblem, PoolMode};
-use crate::eval::cpu_reference::{cpu_reference_pool, decode_index, geometry::PoolGeometry};
+use crate::eval::cpu_reference::{cpu_reference_pool, geometry::PoolGeometry};
 use crate::{pool2d, pool2d_with_indices};
 use cubecl::{TestRuntime, client::ComputeClient};
 use cubek_test_utils::{
@@ -36,19 +36,6 @@ pub(crate) fn get_window_coords<const N: usize>(
     Some(in_coords)
 }
 
-pub(crate) fn decode_index_simple(index: usize, shape: &[usize]) -> Vec<usize> {
-    let strides = row_major_strides_vec(shape);
-    decode_index(index, shape, &strides)
-}
-
-pub(crate) fn row_major_strides_vec(shape: &[usize]) -> Vec<usize> {
-    let mut strides = vec![1; shape.len()];
-    for i in (0..shape.len() - 1).rev() {
-        strides[i] = strides[i + 1] * shape[i + 1];
-    }
-    strides
-}
-
 pub fn strategy_result(
     client: ComputeClient<TestRuntime>,
     problem: PoolForwardProblem<2>,

crates/cubek-pool/src/eval/cpu_reference/mod.rs

@@ -9,10 +9,7 @@ use crate::{
     },
     eval::cpu_reference::{
         backward::{run_adaptive_avg_pool_backward, run_avg_pool_backward, run_max_pool_backward},
-        forward::{
-            row_major_strides_vec, run_adaptive_avg_pool, run_avg_pool, run_max_pool,
-            run_max_pool_with_indices,
-        },
+        forward::{run_adaptive_avg_pool, run_avg_pool, run_max_pool, run_max_pool_with_indices},
         geometry::PoolGeometry,
     },
 };
@@ -118,12 +115,7 @@ pub fn cpu_reference_pool_backward<const N: usize>(
     problem: PoolBackwardProblem<N>,
 ) -> HostData {
     let out_dims = grad_output.shape.to_vec();
-    let input_shape = Shape::from(vec![
-        problem.out_grad_shape[0],
-        problem.input_size[0],
-        problem.input_size[1],
-        problem.out_grad_shape[3],
-    ]);
+    let input_shape = problem.input_shape();
     let in_dims = input_shape.to_vec();
     let in_strides = row_major_strides_vec(&in_dims);
 
@@ -138,8 +130,13 @@ pub fn cpu_reference_pool_backward<const N: usize>(
         PoolMode::Avg(_opts) => {
             run_avg_pool_backward(grad_output, _opts, &in_dims, &out_dims, &in_strides)
         }
-        PoolMode::AdaptiveAvg(_opts) => {
-            run_adaptive_avg_pool_backward(grad_output, _opts, &in_dims, &out_dims, &in_strides)
+        PoolMode::AdaptiveAvg(opts) => {
+            assert_eq!(
+                &out_dims[1..N + 1],
+                opts.output_size.as_slice(),
+                "adaptive output-gradient shape must match options"
+            );
+            run_adaptive_avg_pool_backward::<N>(grad_output, &in_dims, &out_dims, &in_strides)
         }
     };
 
@@ -175,3 +172,16 @@ pub(crate) fn decode_index(mut index: usize, shape: &[usize], strides: &[usize])
     }
     coords
 }
+
+pub(crate) fn decode_index_simple(index: usize, shape: &[usize]) -> Vec<usize> {
+    let strides = row_major_strides_vec(shape);
+    decode_index(index, shape, &strides)
+}
+
+pub(crate) fn row_major_strides_vec(shape: &[usize]) -> Vec<usize> {
+    let mut strides = vec![1; shape.len()];
+    for i in (0..shape.len() - 1).rev() {
+        strides[i] = strides[i + 1] * shape[i + 1];
+    }
+    strides
+}

crates/cubek-pool/src/kernel/backward/adaptive_avg_pool2d_backward.rs

@@ -1,4 +1,7 @@
-use super::super::{decompose_linear, shape_divmod};
+use super::super::{
+    adaptive_end_index as end_index, adaptive_start_index as start_index, decompose_linear,
+    shape_divmod,
+};
 use crate::definition::{AdaptiveAvgPoolOptions, PoolError};
 use crate::kernel::forward::{Position, view4d};
 use cubecl::{
@@ -61,23 +64,6 @@ fn adaptive_avg_pool2d_backward_direct<E: Numeric, N: Size>(
     output.write((b, ih, iw, c), grad_acc);
 }
 
-#[cube]
-fn start_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize {
-    (output_size_index * input_size) / output_size
-}
-
-#[cube]
-fn end_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize {
-    let index = (output_size_index + 1) * input_size;
-    let index = index.div_ceil(output_size);
-
-    if input_size < index {
-        input_size
-    } else {
-        index
-    }
-}
-
 pub(crate) fn adaptive_avg_pool2d_backward_launch<R: Runtime>(
     client: &ComputeClient<R>,
     input: TensorBinding<R>,