diff --git a/crates/cubek-pool/src/definition/base.rs b/crates/cubek-pool/src/definition/base.rs index 21c3309b6..2657d2365 100644 --- a/crates/cubek-pool/src/definition/base.rs +++ b/crates/cubek-pool/src/definition/base.rs @@ -36,6 +36,17 @@ pub struct PoolBackwardProblem { pub mode: PoolMode, } +#[cfg(feature = "benchmarks")] +impl PoolBackwardProblem { + /// 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 { Max(MaxPoolOptions), diff --git a/crates/cubek-pool/src/definition/error.rs b/crates/cubek-pool/src/definition/error.rs index 2b4fe335b..00b634f6f 100644 --- a/crates/cubek-pool/src/definition/error.rs +++ b/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, + }, + + #[error("Output spatial shape mismatch: expected {expected:?} but got {actual:?}")] + OutputSizeMismatch { + expected: Vec, + actual: Vec, + }, + + #[error("Input gradient shape mismatch: expected {expected:?} but got {actual:?}")] + InputGradientShapeMismatch { + expected: Vec, + actual: Vec, + }, } diff --git a/crates/cubek-pool/src/eval/cpu_reference/backward/adaptive_avg_pool.rs b/crates/cubek-pool/src/eval/cpu_reference/backward/adaptive_avg_pool.rs index feb0f1add..8b79fdb98 100644 --- a/crates/cubek-pool/src/eval/cpu_reference/backward/adaptive_avg_pool.rs +++ b/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( grad_output: &HostData, - _opts: &AdaptiveAvgPoolOptions, grad_input_dims: &[usize], grad_output_dims: &[usize], grad_input_strides: &[usize], ) -> Vec { 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 diff --git a/crates/cubek-pool/src/eval/cpu_reference/backward/mod.rs b/crates/cubek-pool/src/eval/cpu_reference/backward/mod.rs index 8e40c911d..568243714 100644 --- a/crates/cubek-pool/src/eval/cpu_reference/backward/mod.rs +++ b/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(); diff --git a/crates/cubek-pool/src/eval/cpu_reference/forward/adaptive_avg_pool.rs b/crates/cubek-pool/src/eval/cpu_reference/forward/adaptive_avg_pool.rs index c1f483ff6..5ef67ea87 100644 --- a/crates/cubek-pool/src/eval/cpu_reference/forward/adaptive_avg_pool.rs +++ b/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; diff --git a/crates/cubek-pool/src/eval/cpu_reference/forward/avg_pool.rs b/crates/cubek-pool/src/eval/cpu_reference/forward/avg_pool.rs index 951d044e6..0047002a3 100644 --- a/crates/cubek-pool/src/eval/cpu_reference/forward/avg_pool.rs +++ b/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; diff --git a/crates/cubek-pool/src/eval/cpu_reference/forward/max_pool.rs b/crates/cubek-pool/src/eval/cpu_reference/forward/max_pool.rs index 175ad3bba..7579e9059 100644 --- a/crates/cubek-pool/src/eval/cpu_reference/forward/max_pool.rs +++ b/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; diff --git a/crates/cubek-pool/src/eval/cpu_reference/forward/mod.rs b/crates/cubek-pool/src/eval/cpu_reference/forward/mod.rs index 074b842f9..b5f4a9e80 100644 --- a/crates/cubek-pool/src/eval/cpu_reference/forward/mod.rs +++ b/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( Some(in_coords) } -pub(crate) fn decode_index_simple(index: usize, shape: &[usize]) -> Vec { - let strides = row_major_strides_vec(shape); - decode_index(index, shape, &strides) -} - -pub(crate) fn row_major_strides_vec(shape: &[usize]) -> Vec { - 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, problem: PoolForwardProblem<2>, diff --git a/crates/cubek-pool/src/eval/cpu_reference/mod.rs b/crates/cubek-pool/src/eval/cpu_reference/mod.rs index d2311a5a1..46a7fa3f3 100644 --- a/crates/cubek-pool/src/eval/cpu_reference/mod.rs +++ b/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( problem: PoolBackwardProblem, ) -> 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( 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::(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 { + let strides = row_major_strides_vec(shape); + decode_index(index, shape, &strides) +} + +pub(crate) fn row_major_strides_vec(shape: &[usize]) -> Vec { + let mut strides = vec![1; shape.len()]; + for i in (0..shape.len() - 1).rev() { + strides[i] = strides[i + 1] * shape[i + 1]; + } + strides +} diff --git a/crates/cubek-pool/src/kernel/backward/adaptive_avg_pool2d_backward.rs b/crates/cubek-pool/src/kernel/backward/adaptive_avg_pool2d_backward.rs index a31c9392e..1c254788c 100644 --- a/crates/cubek-pool/src/kernel/backward/adaptive_avg_pool2d_backward.rs +++ b/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( 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( client: &ComputeClient, input: TensorBinding, diff --git a/crates/cubek-pool/src/kernel/backward/adaptive_avg_pool3d_backward.rs b/crates/cubek-pool/src/kernel/backward/adaptive_avg_pool3d_backward.rs new file mode 100644 index 000000000..319685f63 --- /dev/null +++ b/crates/cubek-pool/src/kernel/backward/adaptive_avg_pool3d_backward.rs @@ -0,0 +1,123 @@ +use super::super::{ + accumulator_dtype, adaptive_end_index as end_index, adaptive_start_index as start_index, + adaptive_window_address_type, decompose_linear_5d, shape_divmod, +}; +use crate::definition::PoolError; +use crate::kernel::forward::{Position3d, view5d}; +use cubecl::{ + CubeDim, Runtime, calculate_cube_count_elemwise, + num_traits::Zero, + prelude::{TensorBinding, *}, + std::{FastDivmod, tensor::ViewMut}, + tensor_vector_size_parallel, +}; + +#[cube(launch, address_type = "dynamic")] +fn adaptive_avg_pool3d_backward_direct( + grad: &Tensor>, + mut output: ViewMut<'_, Vector, Position3d>, + out_shape: Sequence>, + working_units: usize, + #[define(EI)] _dtype: ElemType, + #[define(EA)] _acc_dtype: ElemType, +) { + if ABSOLUTE_POS >= working_units { + terminate!(); + } + + let (_, in_d, in_h, in_w, _) = output.shape(); + let (grad_d, grad_h, grad_w) = (grad.shape(1), grad.shape(2), grad.shape(3)); + let (b, id, ih, iw, c) = decompose_linear_5d(ABSOLUTE_POS * output.vector_size(), &out_shape); + + let od_start = start_index(id, in_d, grad_d); + let od_end = end_index(id, in_d, grad_d); + let oh_start = start_index(ih, in_h, grad_h); + let oh_end = end_index(ih, in_h, grad_h); + let ow_start = start_index(iw, in_w, grad_w); + let ow_end = end_index(iw, in_w, grad_w); + + // Each input position gathers every output window that can contain it, so work items never + // race to update the same gradient. The inverse bounds are conservative for uneven windows; + // the containment checks below discard candidates outside the exact forward window. + let mut grad_acc = Vector::::zero(); + let index_base = b * grad.stride(0) + c * grad.stride(4); + + for od in od_start..od_end { + let id_start = start_index(od, grad_d, in_d); + let id_end = end_index(od, grad_d, in_d); + if id >= id_start && id < id_end { + for oh in oh_start..oh_end { + let ih_start = start_index(oh, grad_h, in_h); + let ih_end = end_index(oh, grad_h, in_h); + if ih >= ih_start && ih < ih_end { + for ow in ow_start..ow_end { + let iw_start = start_index(ow, grad_w, in_w); + let iw_end = end_index(ow, grad_w, in_w); + if iw >= iw_start && iw < iw_end { + let volume = + (id_end - id_start) * (ih_end - ih_start) * (iw_end - iw_start); + let index = index_base + + od * grad.stride(1) + + oh * grad.stride(2) + + ow * grad.stride(3); + grad_acc += + Vector::::cast_from(grad[index / grad.vector_size()]) + / Vector::cast_from(volume); + } + } + } + } + } + } + + output.write((b, id, ih, iw, c), Vector::cast_from(grad_acc)); +} + +pub(crate) fn adaptive_avg_pool3d_backward_launch( + client: &ComputeClient, + out_grad: TensorBinding, + output: TensorBinding, + dtype: ElemType, +) -> Result<(), PoolError> { + let acc_dtype = accumulator_dtype(dtype); + let grad_vector_size = tensor_vector_size_parallel( + client.io_optimized_vector_sizes(dtype.size()), + &out_grad.shape, + &out_grad.strides, + out_grad.shape.len() - 1, + ); + let output_vector_size = tensor_vector_size_parallel( + client.io_optimized_vector_sizes(dtype.size()), + &output.shape, + &output.strides, + output.shape.len() - 1, + ); + let vector_size = grad_vector_size.min(output_vector_size); + + let working_units = output.shape.iter().product::() / vector_size as usize; + let cube_dim = CubeDim::new(client, working_units); + let cube_count = calculate_cube_count_elemwise(client, working_units, cube_dim); + let address_type = out_grad + .required_address_type(dtype.size()) + .max(output.required_address_type(dtype.size())) + .max(adaptive_window_address_type( + &output.shape[1..4], + &out_grad.shape[1..4], + )); + + adaptive_avg_pool3d_backward_direct::launch( + client, + cube_count, + cube_dim, + address_type, + vector_size, + out_grad.into_tensor_arg(), + view5d(output.clone(), vector_size), + shape_divmod(&output), + working_units, + dtype, + acc_dtype, + ); + + Ok(()) +} diff --git a/crates/cubek-pool/src/kernel/backward/mod.rs b/crates/cubek-pool/src/kernel/backward/mod.rs index 9311dfdeb..96bbb0817 100644 --- a/crates/cubek-pool/src/kernel/backward/mod.rs +++ b/crates/cubek-pool/src/kernel/backward/mod.rs @@ -1,7 +1,9 @@ mod adaptive_avg_pool2d_backward; +mod adaptive_avg_pool3d_backward; mod avg_pool2d_backward; mod max_pool2d_backward; pub(crate) use adaptive_avg_pool2d_backward::*; +pub(crate) use adaptive_avg_pool3d_backward::*; pub(crate) use avg_pool2d_backward::*; pub(crate) use max_pool2d_backward::*; diff --git a/crates/cubek-pool/src/kernel/forward/adaptive_avg_pool2d.rs b/crates/cubek-pool/src/kernel/forward/adaptive_avg_pool2d.rs index 40754b521..c9abf310f 100644 --- a/crates/cubek-pool/src/kernel/forward/adaptive_avg_pool2d.rs +++ b/crates/cubek-pool/src/kernel/forward/adaptive_avg_pool2d.rs @@ -1,6 +1,8 @@ use super::{ - super::decompose_linear, - super::shape_divmod, + super::{ + adaptive_end_index as end_index, adaptive_start_index as start_index, decompose_linear, + shape_divmod, + }, pool2d::{Position, view4d}, }; use crate::definition::{AdaptiveAvgPoolOptions, PoolError}; @@ -57,23 +59,6 @@ fn adaptive_avg_pool2d_direct( output.write((b, oh, ow, c), sum / Vector::cast_from(num_ih * num_iw)); } -#[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_launch( client: &ComputeClient, input: TensorBinding, diff --git a/crates/cubek-pool/src/kernel/forward/adaptive_avg_pool3d.rs b/crates/cubek-pool/src/kernel/forward/adaptive_avg_pool3d.rs new file mode 100644 index 000000000..7a500d68c --- /dev/null +++ b/crates/cubek-pool/src/kernel/forward/adaptive_avg_pool3d.rs @@ -0,0 +1,113 @@ +use super::pool3d::{Position3d, view5d}; +use crate::{ + definition::PoolError, + kernel::{ + accumulator_dtype, adaptive_end_index as end_index, adaptive_start_index as start_index, + adaptive_window_address_type, decompose_linear_5d, shape_divmod, + }, +}; +use cubecl::{ + CubeDim, Runtime, calculate_cube_count_elemwise, + num_traits::Zero, + prelude::{TensorBinding, *}, + std::{FastDivmod, tensor::ViewMut}, + tensor_vector_size_parallel, +}; + +#[cube(launch, address_type = "dynamic")] +fn adaptive_avg_pool3d_direct( + input: &Tensor>, + mut output: ViewMut<'_, Vector, Position3d>, + out_shape: Sequence>, + working_units: usize, + #[define(EI)] _dtype: ElemType, + #[define(EA)] _acc_dtype: ElemType, +) { + if ABSOLUTE_POS >= working_units { + terminate!(); + } + + let (b, od, oh, ow, c) = decompose_linear_5d(ABSOLUTE_POS * output.vector_size(), &out_shape); + let (_, out_d, out_h, out_w, _) = output.shape(); + let (in_d, in_h, in_w) = (input.shape(1), input.shape(2), input.shape(3)); + + let id_start = start_index(od, out_d, in_d); + let id_end = end_index(od, out_d, in_d); + let ih_start = start_index(oh, out_h, in_h); + let ih_end = end_index(oh, out_h, in_h); + let iw_start = start_index(ow, out_w, in_w); + let iw_end = end_index(ow, out_w, in_w); + + let mut sum = Vector::::zero(); + let index_input_base = b * input.stride(0) + c * input.stride(4); + + for id in id_start..id_end { + let index_input_d = id * input.stride(1); + for ih in ih_start..ih_end { + let index_input_h = ih * input.stride(2); + for iw in iw_start..iw_end { + let index_input = + index_input_base + index_input_d + index_input_h + iw * input.stride(3); + sum += Vector::cast_from(input[index_input / input.vector_size()]); + } + } + } + + let volume = (id_end - id_start) * (ih_end - ih_start) * (iw_end - iw_start); + let average = sum / Vector::cast_from(volume); + output.write((b, od, oh, ow, c), Vector::cast_from(average)); +} + +pub(crate) fn adaptive_avg_pool3d_launch( + client: &ComputeClient, + input: TensorBinding, + output: TensorBinding, + dtype: ElemType, +) -> Result<(), PoolError> { + let acc_dtype = accumulator_dtype(dtype); + let input_vector_size = tensor_vector_size_parallel( + client.io_optimized_vector_sizes(dtype.size()), + &input.shape, + &input.strides, + input.shape.len() - 1, + ); + let output_vector_size = tensor_vector_size_parallel( + client.io_optimized_vector_sizes(dtype.size()), + &output.shape, + &output.strides, + output.shape.len() - 1, + ); + let vector_size = input_vector_size.min(output_vector_size); + let working_units = output.shape.iter().product::() / vector_size as usize; + + if working_units == 0 { + return Ok(()); + } + + let cube_dim = CubeDim::new(client, working_units); + let cube_count = calculate_cube_count_elemwise(client, working_units, cube_dim); + + let address_type = input + .required_address_type(dtype.size()) + .max(output.required_address_type(dtype.size())) + .max(adaptive_window_address_type( + &input.shape[1..4], + &output.shape[1..4], + )); + + adaptive_avg_pool3d_direct::launch( + client, + cube_count, + cube_dim, + address_type, + vector_size, + input.into_tensor_arg(), + view5d(output.clone(), vector_size), + shape_divmod(&output), + working_units, + dtype, + acc_dtype, + ); + + Ok(()) +} diff --git a/crates/cubek-pool/src/kernel/forward/mod.rs b/crates/cubek-pool/src/kernel/forward/mod.rs index f0cd704c1..2365b600b 100644 --- a/crates/cubek-pool/src/kernel/forward/mod.rs +++ b/crates/cubek-pool/src/kernel/forward/mod.rs @@ -1,9 +1,13 @@ mod adaptive_avg_pool2d; +mod adaptive_avg_pool3d; mod avg_pool2d; mod max_pool2d; mod pool2d; +mod pool3d; pub(crate) use adaptive_avg_pool2d::*; +pub(crate) use adaptive_avg_pool3d::*; pub(crate) use avg_pool2d::*; pub(crate) use max_pool2d::*; pub(crate) use pool2d::*; +pub(crate) use pool3d::*; diff --git a/crates/cubek-pool/src/kernel/forward/pool3d.rs b/crates/cubek-pool/src/kernel/forward/pool3d.rs new file mode 100644 index 000000000..a9a77a1b8 --- /dev/null +++ b/crates/cubek-pool/src/kernel/forward/pool3d.rs @@ -0,0 +1,30 @@ +use cubecl::{ + Runtime, + prelude::{TensorBinding, *}, + std::tensor::{ + launch::ViewArg, + layout::fixed_dim::{FixedDimLayout, FixedDimLayoutLaunch}, + }, +}; + +pub(crate) type Position3d = (usize, usize, usize, usize, usize); + +pub(crate) fn view5d( + tensor: TensorBinding, + vector_size: VectorSize, +) -> ViewArg { + let shape = ( + tensor.shape[0], + tensor.shape[1], + tensor.shape[2], + tensor.shape[3], + tensor.shape[4], + ); + let layout = FixedDimLayoutLaunch::::from_shape_handle_unchecked( + &tensor, + shape, + vector_size, + ); + let buffer = tensor.into_tensor_arg(); + ViewArg::new_tensor::>(buffer, layout) +} diff --git a/crates/cubek-pool/src/kernel/mod.rs b/crates/cubek-pool/src/kernel/mod.rs index eb0860e8a..39dffdccb 100644 --- a/crates/cubek-pool/src/kernel/mod.rs +++ b/crates/cubek-pool/src/kernel/mod.rs @@ -1,7 +1,11 @@ pub(crate) mod backward; pub(crate) mod forward; -use cubecl::{prelude::*, std::FastDivmod}; +use cubecl::{ + ir::{AddressType, ElemType, FloatKind}, + prelude::*, + std::FastDivmod, +}; pub(crate) fn shape_divmod( binding: &TensorBinding, @@ -13,6 +17,26 @@ pub(crate) fn shape_divmod( out_seq } +#[cube] +pub(crate) fn adaptive_start_index( + output_size_index: usize, + output_size: usize, + input_size: usize, +) -> usize { + (output_size_index * input_size) / output_size +} + +#[cube] +pub(crate) fn adaptive_end_index( + output_size_index: usize, + output_size: usize, + input_size: usize, +) -> usize { + ((output_size_index + 1) * input_size) + .div_ceil(output_size) + .min(input_size) +} + #[cube] pub(crate) fn decompose_linear( index: usize, @@ -25,3 +49,42 @@ pub(crate) fn decompose_linear( (b, oh, ow, c) } + +#[cube] +pub(crate) fn decompose_linear_5d( + index: usize, + shape: &Sequence>, +) -> (usize, usize, usize, usize, usize) { + let (remainder, c) = shape[4].div_mod(index); + let (remainder, ow) = shape[3].div_mod(remainder); + let (remainder, oh) = shape[2].div_mod(remainder); + let (remainder, od) = shape[1].div_mod(remainder); + let (_, b) = shape[0].div_mod(remainder); + + (b, od, oh, ow, c) +} + +pub(crate) fn accumulator_dtype(input: ElemType) -> ElemType { + match input { + ElemType::Float(FloatKind::F16) + | ElemType::Float(FloatKind::BF16) + | ElemType::Float(FloatKind::Flex32) => ElemType::Float(FloatKind::F32), + _ => input, + } +} + +/// Account for intermediate products in adaptive window bounds, not just buffer addresses. +pub(crate) fn adaptive_window_address_type( + input_size: &[usize], + output_size: &[usize], +) -> AddressType { + if input_size.iter().zip(output_size).any(|(input, output)| { + input + .checked_mul(*output) + .is_none_or(|product| product > u32::MAX as usize) + }) { + AddressType::U64 + } else { + AddressType::U32 + } +} diff --git a/crates/cubek-pool/src/lib.rs b/crates/cubek-pool/src/lib.rs index 40f332496..2a80e43c2 100644 --- a/crates/cubek-pool/src/lib.rs +++ b/crates/cubek-pool/src/lib.rs @@ -11,12 +11,12 @@ mod kernel; use crate::definition::{PoolError, PoolMode}; use crate::kernel::{ backward::{ - adaptive_avg_pool2d_backward_launch, avg_pool2d_backward_launch, - max_pool2d_with_indices_backward_launch, + adaptive_avg_pool2d_backward_launch, adaptive_avg_pool3d_backward_launch, + avg_pool2d_backward_launch, max_pool2d_with_indices_backward_launch, }, forward::{ - adaptive_avg_pool2d_launch, avg_pool2d_launch, max_pool2d_launch, - max_pool2d_with_indices_launch, + adaptive_avg_pool2d_launch, adaptive_avg_pool3d_launch, avg_pool2d_launch, + max_pool2d_launch, max_pool2d_with_indices_launch, }, }; @@ -30,8 +30,8 @@ pub fn pool2d( mode: PoolMode<2>, dtype: ElemType, ) -> Result<(), PoolError> { - validate_rank(input.shape.len(), output.shape.len())?; - validate_nhwc_consistency(&input.shape, &output.shape)?; + validate_rank(input.shape.len(), output.shape.len(), 4)?; + validate_batch_channel_consistency(&input.shape, &output.shape)?; match mode { PoolMode::Max(max_options) => max_pool2d_launch(client, input, output, max_options, dtype), @@ -42,6 +42,42 @@ pub fn pool2d( } } +/// Pool3d public wrapper. +/// +/// Expects matching floating-point input and output tensors in NDHWC layout. Only +/// [`PoolMode::AdaptiveAvg`] is supported, and its configured output size must match the output +/// tensor's spatial shape. +/// +/// # Errors +/// +/// Returns [`PoolError::InvalidRank`] unless both tensors have rank five, +/// [`PoolError::BatchMismatch`] or [`PoolError::ChannelMismatch`] when their outer dimensions do +/// not agree, [`PoolError::InvalidSpatialSize`] when either tensor has a zero spatial dimension, +/// [`PoolError::OutputSizeMismatch`] when the configured spatial size differs from the output +/// tensor, and [`PoolError::UnsupportedMode`] for any other pooling mode. +pub fn pool3d( + client: &ComputeClient, + input: TensorBinding, + output: TensorBinding, + mode: PoolMode<3>, + dtype: ElemType, +) -> Result<(), PoolError> { + validate_rank(input.shape.len(), output.shape.len(), 5)?; + validate_batch_channel_consistency(&input.shape, &output.shape)?; + + match mode { + PoolMode::AdaptiveAvg(options) => { + validate_spatial_size(&input.shape, "input")?; + validate_spatial_size(&output.shape, "output")?; + validate_output_size(&output.shape, &options.output_size)?; + adaptive_avg_pool3d_launch(client, input, output, dtype) + } + _ => Err(PoolError::UnsupportedMode { + mode: format!("{0:?}", mode), + }), + } +} + /// Pool2d with indices public wrapper /// /// Expects input in NHWC layout. Output indices are expected to be in the same layout as well. @@ -53,10 +89,10 @@ pub fn pool2d_with_indices( mode: PoolMode<2>, dtype: ElemType, ) -> Result<(), PoolError> { - validate_rank(input.shape.len(), output.shape.len())?; - validate_rank(input.shape.len(), indices.shape.len())?; - validate_nhwc_consistency(&input.shape, &output.shape)?; - validate_nhwc_consistency(&input.shape, &indices.shape)?; + validate_rank(input.shape.len(), output.shape.len(), 4)?; + validate_rank(input.shape.len(), indices.shape.len(), 4)?; + validate_batch_channel_consistency(&input.shape, &output.shape)?; + validate_batch_channel_consistency(&input.shape, &indices.shape)?; match mode { PoolMode::Max(max_options) => { @@ -79,10 +115,10 @@ pub fn pool2d_backward( mode: PoolMode<2>, dtype: ElemType, ) -> Result<(), PoolError> { - validate_rank(input.shape.len(), out_grad.shape.len())?; - validate_rank(input.shape.len(), in_grad.shape.len())?; - validate_nhwc_consistency(&input.shape, &out_grad.shape)?; - validate_nhwc_consistency(&input.shape, &in_grad.shape)?; + validate_rank(input.shape.len(), out_grad.shape.len(), 4)?; + validate_rank(input.shape.len(), in_grad.shape.len(), 4)?; + validate_batch_channel_consistency(&input.shape, &out_grad.shape)?; + validate_batch_channel_consistency(&input.shape, &in_grad.shape)?; match mode { PoolMode::Avg(avg_options) => { @@ -102,6 +138,53 @@ pub fn pool2d_backward( } } +/// Pool3d backward public wrapper. +/// +/// Expects matching floating-point input, output-gradient, and input-gradient tensors in NDHWC +/// layout. Only [`PoolMode::AdaptiveAvg`] is supported. The configured output size must match the +/// output gradient, and the input gradient must have exactly the input shape. +/// +/// # Errors +/// +/// Returns [`PoolError::InvalidRank`] unless every tensor has rank five, +/// [`PoolError::BatchMismatch`] or [`PoolError::ChannelMismatch`] when their outer dimensions do +/// not agree, [`PoolError::InputGradientShapeMismatch`] when the input-gradient shape differs from +/// the input, [`PoolError::InvalidSpatialSize`] when the input or output gradient has a zero spatial +/// dimension, [`PoolError::OutputSizeMismatch`] when the configured spatial size differs from the +/// output gradient, and [`PoolError::UnsupportedMode`] for any other pooling mode. +pub fn pool3d_backward( + client: &ComputeClient, + input: TensorBinding, + out_grad: TensorBinding, + in_grad: TensorBinding, + mode: PoolMode<3>, + dtype: ElemType, +) -> Result<(), PoolError> { + validate_rank(input.shape.len(), out_grad.shape.len(), 5)?; + validate_rank(input.shape.len(), in_grad.shape.len(), 5)?; + validate_batch_channel_consistency(&input.shape, &out_grad.shape)?; + validate_batch_channel_consistency(&input.shape, &in_grad.shape)?; + + if input.shape != in_grad.shape { + return Err(PoolError::InputGradientShapeMismatch { + expected: input.shape.to_vec(), + actual: in_grad.shape.to_vec(), + }); + } + + match mode { + PoolMode::AdaptiveAvg(options) => { + validate_spatial_size(&input.shape, "input")?; + validate_spatial_size(&out_grad.shape, "output gradient")?; + validate_output_size(&out_grad.shape, &options.output_size)?; + adaptive_avg_pool3d_backward_launch(client, out_grad, in_grad, dtype) + } + _ => Err(PoolError::UnsupportedMode { + mode: format!("{0:?}", mode), + }), + } +} + /// Pool2d backward with indices public wrapper /// /// Expects input and output gradients in NHWC layout. Output indices are expected to be in the same layout as well. @@ -116,12 +199,12 @@ pub fn pool2d_with_indices_backward( dtype: ElemType, indices_dtype: ElemType, ) -> Result<(), PoolError> { - validate_rank(input.shape.len(), out_grad.shape.len())?; - validate_rank(input.shape.len(), in_grad.shape.len())?; - validate_rank(input.shape.len(), indices.shape.len())?; - validate_nhwc_consistency(&input.shape, &out_grad.shape)?; - validate_nhwc_consistency(&input.shape, &in_grad.shape)?; - validate_nhwc_consistency(&input.shape, &indices.shape)?; + validate_rank(input.shape.len(), out_grad.shape.len(), 4)?; + validate_rank(input.shape.len(), in_grad.shape.len(), 4)?; + validate_rank(input.shape.len(), indices.shape.len(), 4)?; + validate_batch_channel_consistency(&input.shape, &out_grad.shape)?; + validate_batch_channel_consistency(&input.shape, &in_grad.shape)?; + validate_batch_channel_consistency(&input.shape, &indices.shape)?; match mode { PoolMode::Max(max_options) => max_pool2d_with_indices_backward_launch( @@ -140,9 +223,12 @@ pub fn pool2d_with_indices_backward( } } -/// Check that both tensors are 4D (Batch, Height, Width, Channels). -fn validate_rank(input_rank: usize, output_rank: usize) -> Result<(), PoolError> { - if input_rank != 4 || output_rank != 4 { +fn validate_rank( + input_rank: usize, + output_rank: usize, + expected_rank: usize, +) -> Result<(), PoolError> { + if input_rank != expected_rank || output_rank != expected_rank { return Err(PoolError::InvalidRank { input: input_rank, output: output_rank, @@ -151,9 +237,8 @@ fn validate_rank(input_rank: usize, output_rank: usize) -> Result<(), PoolError> Ok(()) } -/// Check that Batch (0) and Channel (3) dimensions match. -/// Height (1) and Width (2) are allowed to differ for resizing. -fn validate_nhwc_consistency( +/// Check that the batch and final channel dimensions match. +fn validate_batch_channel_consistency( input_shape: &[usize], output_shape: &[usize], ) -> Result<(), PoolError> { @@ -164,12 +249,36 @@ fn validate_nhwc_consistency( }); } - if input_shape[3] != output_shape[3] { + let input_channel = input_shape[input_shape.len() - 1]; + let output_channel = output_shape[output_shape.len() - 1]; + if input_channel != output_channel { return Err(PoolError::ChannelMismatch { - input: input_shape[3], - output: output_shape[3], + input: input_channel, + output: output_channel, + }); + } + + Ok(()) +} + +fn validate_output_size(output_shape: &[usize], expected: &[usize; 3]) -> Result<(), PoolError> { + let actual = &output_shape[1..4]; + if actual != expected { + return Err(PoolError::OutputSizeMismatch { + expected: expected.to_vec(), + actual: actual.to_vec(), }); } + Ok(()) +} +fn validate_spatial_size(shape: &[usize], tensor: &'static str) -> Result<(), PoolError> { + let actual = &shape[1..4]; + if actual.contains(&0) { + return Err(PoolError::InvalidSpatialSize { + tensor, + actual: actual.to_vec(), + }); + } Ok(()) } diff --git a/crates/cubek-pool/tests/pool/backward/adaptive_avg_pool2d.rs b/crates/cubek-pool/tests/pool/backward/adaptive_avg_pool2d.rs index 741cf7d95..7ce41b6f2 100644 --- a/crates/cubek-pool/tests/pool/backward/adaptive_avg_pool2d.rs +++ b/crates/cubek-pool/tests/pool/backward/adaptive_avg_pool2d.rs @@ -9,7 +9,7 @@ fn test_adaptive_avg_pool2d_backward_global() { let client = TestRuntime::client(&Default::default()); let problem = make_problem( [8, 8], - Shape::from([2, 4, 1, 1]), + Shape::from([2, 1, 1, 4]), false, AdaptiveAvgPoolOptions { output_size: [1, 1], @@ -30,7 +30,7 @@ fn test_adaptive_avg_pool2d_backward_square() { let client = TestRuntime::client(&Default::default()); let problem = make_problem( [7, 7], - Shape::from([1, 2, 3, 3]), + Shape::from([1, 3, 3, 2]), false, AdaptiveAvgPoolOptions { output_size: [3, 3], @@ -51,7 +51,7 @@ fn test_adaptive_avg_pool2d_backward_non_square() { let client = TestRuntime::client(&Default::default()); let problem = make_problem( [10, 10], - Shape::from([2, 3, 3, 5]), + Shape::from([2, 3, 5, 3]), false, AdaptiveAvgPoolOptions { output_size: [3, 5], @@ -72,7 +72,7 @@ fn test_adaptive_avg_pool2d_backward_large_input() { let client = TestRuntime::client(&Default::default()); let problem = make_problem( [14, 14], - Shape::from([1, 8, 7, 7]), + Shape::from([1, 7, 7, 8]), false, AdaptiveAvgPoolOptions { output_size: [7, 7], diff --git a/crates/cubek-pool/tests/pool/backward/adaptive_avg_pool3d.rs b/crates/cubek-pool/tests/pool/backward/adaptive_avg_pool3d.rs new file mode 100644 index 000000000..62c622b69 --- /dev/null +++ b/crates/cubek-pool/tests/pool/backward/adaptive_avg_pool3d.rs @@ -0,0 +1,293 @@ +use crate::pool::{build_output_tensor, output_host_f32, validate_test}; +use cubecl::{ + Runtime, TestRuntime, + ir::{ElemType, FloatKind}, + prelude::*, + zspace::Shape, +}; +use cubek_pool::{ + definition::{ + AdaptiveAvgPoolOptions, AvgPoolOptions, MaxPoolOptions, PoolBackwardProblem, PoolError, + PoolMode, + }, + eval::cpu_reference::cpu_reference_pool_backward, + pool3d_backward, +}; +use cubek_test_utils::TestInput; + +const TOLERANCE: f32 = 1e-5; + +#[test] +fn test_adaptive_avg_pool3d_backward_global() { + run_case([5, 7, 4], [2, 1, 1, 1, 3], [1, 1, 1], 2001); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_divisible_vectorized_channels() { + run_case([6, 8, 10], [2, 3, 4, 5, 8], [3, 4, 5], 2002); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_non_divisible_overlap_odd_channels() { + run_case([5, 7, 4], [2, 3, 4, 3, 7], [3, 4, 3], 2003); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_output_larger_than_input() { + run_case([2, 3, 2], [1, 4, 5, 3, 5], [4, 5, 3], 2004); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_large_dimension_index_math() { + run_case( + [100_000, 1, 1], + [1, 100_000, 1, 1, 1], + [100_000, 1, 1], + 2005, + ); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_conserves_gradient_sum() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + let input = build_output_tensor(&client, vec![2, 5, 7, 4, 3], dtype); + let (out_grad, out_grad_host) = TestInput::builder(client.clone(), [2, 3, 4, 3, 3]) + .uniform(2112, -1.0, 1.0) + .generate_with_f32_host_data(); + let in_grad = build_output_tensor(&client, vec![2, 5, 7, 4, 3], dtype); + + pool3d_backward( + &client, + input.clone().binding(), + out_grad.binding(), + in_grad.clone().binding(), + AdaptiveAvgPoolOptions::new([3, 4, 3]).into(), + dtype, + ) + .expect("adaptive average pool 3d backward should launch"); + + let in_grad_host = output_host_f32(&client, in_grad); + let expected_sum: f32 = out_grad_host + .iter_indexed_f32() + .map(|(_, value)| value) + .sum(); + let actual_sum: f32 = in_grad_host + .iter_indexed_f32() + .map(|(_, value)| value) + .sum(); + let tolerance = 1e-4 * expected_sum.abs().max(1.0); + assert!( + (actual_sum - expected_sum).abs() <= tolerance, + "input-gradient sum {actual_sum} differs from output-gradient sum {expected_sum}" + ); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_validates_output_and_input_gradient_shapes() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + let input = build_output_tensor(&client, vec![2, 5, 7, 4, 3], dtype); + let options = AdaptiveAvgPoolOptions::new([3, 4, 3]); + + let out_grad_mismatch = build_output_tensor(&client, vec![2, 3, 4, 2, 3], dtype); + let in_grad = build_output_tensor(&client, vec![2, 5, 7, 4, 3], dtype); + assert!(matches!( + pool3d_backward( + &client, + input.clone().binding(), + out_grad_mismatch.binding(), + in_grad.clone().binding(), + options.clone().into(), + dtype, + ), + Err(PoolError::OutputSizeMismatch { .. }) + )); + + let out_grad = build_output_tensor(&client, vec![2, 3, 4, 3, 3], dtype); + let in_grad_mismatch = build_output_tensor(&client, vec![2, 5, 7, 5, 3], dtype); + assert!(matches!( + pool3d_backward( + &client, + input.binding(), + out_grad.binding(), + in_grad_mismatch.binding(), + options.into(), + dtype, + ), + Err(PoolError::InputGradientShapeMismatch { .. }) + )); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_rejects_zero_spatial_dimensions() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + + let input = build_output_tensor(&client, vec![1, 0, 2, 2, 1], dtype); + let out_grad = build_output_tensor(&client, vec![1, 1, 1, 1, 1], dtype); + let in_grad = build_output_tensor(&client, vec![1, 0, 2, 2, 1], dtype); + assert!(matches!( + pool3d_backward( + &client, + input.binding(), + out_grad.binding(), + in_grad.binding(), + AdaptiveAvgPoolOptions::new([1, 1, 1]).into(), + dtype, + ), + Err(PoolError::InvalidSpatialSize { .. }) + )); + + let input = build_output_tensor(&client, vec![1, 2, 2, 2, 1], dtype); + let out_grad = build_output_tensor(&client, vec![1, 0, 1, 1, 1], dtype); + let in_grad = build_output_tensor(&client, vec![1, 2, 2, 2, 1], dtype); + assert!(matches!( + pool3d_backward( + &client, + input.binding(), + out_grad.binding(), + in_grad.binding(), + AdaptiveAvgPoolOptions::new([0, 1, 1]).into(), + dtype, + ), + Err(PoolError::InvalidSpatialSize { .. }) + )); +} + +#[test] +fn test_adaptive_avg_pool3d_backward_accepts_empty_batch_and_channels() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + + for (input_shape, out_grad_shape) in [ + (vec![0, 2, 2, 2, 1], vec![0, 1, 1, 1, 1]), + (vec![1, 2, 2, 2, 0], vec![1, 1, 1, 1, 0]), + ] { + let input = build_output_tensor(&client, input_shape.clone(), dtype); + let out_grad = build_output_tensor(&client, out_grad_shape, dtype); + let in_grad = build_output_tensor(&client, input_shape, dtype); + pool3d_backward( + &client, + input.binding(), + out_grad.binding(), + in_grad.binding(), + AdaptiveAvgPoolOptions::new([1, 1, 1]).into(), + dtype, + ) + .expect("empty batch and channel dimensions should produce empty input gradients"); + } +} + +#[test] +fn test_adaptive_avg_pool3d_backward_rejects_other_pool_modes() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + let input = build_output_tensor(&client, vec![1, 4, 4, 4, 1], dtype); + let out_grad = build_output_tensor(&client, vec![1, 2, 2, 2, 1], dtype); + let in_grad = build_output_tensor(&client, vec![1, 4, 4, 4, 1], dtype); + + for mode in [ + PoolMode::Avg(AvgPoolOptions::new( + [2, 2, 2], + [2, 2, 2], + [0, 0, 0], + false, + false, + )), + PoolMode::Max(MaxPoolOptions::new( + [2, 2, 2], + [2, 2, 2], + [0, 0, 0], + [1, 1, 1], + false, + )), + ] { + assert!(matches!( + pool3d_backward( + &client, + input.clone().binding(), + out_grad.clone().binding(), + in_grad.clone().binding(), + mode, + dtype, + ), + Err(PoolError::UnsupportedMode { .. }) + )); + } +} + +#[test] +fn test_adaptive_avg_pool3d_backward_f16_large_global_accumulates_in_f32() { + let client = TestRuntime::client(&Default::default()); + let dtype = ElemType::Float(FloatKind::F16); + let input_shape = vec![1, 41, 41, 41, 1]; + let input = build_output_tensor(&client, input_shape.clone(), dtype); + let out_grad = TestInput::builder(client.clone(), [1, 1, 1, 1, 1]) + .dtype(dtype) + .custom(vec![1.0]) + .generate_without_host_data(); + let in_grad = build_output_tensor(&client, input_shape, dtype); + + pool3d_backward( + &client, + input.binding(), + out_grad.binding(), + in_grad.clone().binding(), + AdaptiveAvgPoolOptions::new([1, 1, 1]).into(), + dtype, + ) + .expect("f16 adaptive average pool 3d backward should launch"); + + let actual = output_host_f32(&client, in_grad); + let mut sum = 0.0; + for (_, value) in actual.iter_indexed_f32() { + assert!(value.is_finite()); + assert!(value > 0.0, "expected a non-zero input gradient"); + sum += value; + } + assert!((sum - 1.0).abs() <= 1e-2, "expected sum 1, got {sum}"); +} + +fn run_case( + input_size: [usize; 3], + out_grad_shape: [usize; 5], + output_size: [usize; 3], + seed: u64, +) { + let client = TestRuntime::client(&Default::default()); + let problem = PoolBackwardProblem { + input_size, + out_grad_shape: Shape::from(out_grad_shape), + with_indices: false, + mode: AdaptiveAvgPoolOptions::new(output_size).into(), + }; + let input_shape = vec![ + problem.out_grad_shape[0], + problem.input_size[0], + problem.input_size[1], + problem.input_size[2], + problem.out_grad_shape[4], + ]; + let (input, input_data) = TestInput::builder(client.clone(), input_shape.clone()) + .uniform(seed, -1.0, 1.0) + .generate_with_f32_host_data(); + let (out_grad, out_grad_data) = + TestInput::builder(client.clone(), problem.out_grad_shape.to_vec()) + .uniform(seed + 1, -1.0, 1.0) + .generate_with_f32_host_data(); + let in_grad = build_output_tensor(&client, input_shape, input.dtype); + + let result = pool3d_backward( + &client, + input.clone().binding(), + out_grad.binding(), + in_grad.clone().binding(), + problem.mode.clone(), + input.dtype, + ); + let actual = output_host_f32(&client, in_grad); + let expected = cpu_reference_pool_backward(&out_grad_data, &input_data, problem); + + validate_test(result, actual, expected, TOLERANCE); +} diff --git a/crates/cubek-pool/tests/pool/backward/mod.rs b/crates/cubek-pool/tests/pool/backward/mod.rs index c3e794b98..3783112b0 100644 --- a/crates/cubek-pool/tests/pool/backward/mod.rs +++ b/crates/cubek-pool/tests/pool/backward/mod.rs @@ -1,4 +1,5 @@ mod adaptive_avg_pool2d; +mod adaptive_avg_pool3d; mod avg_pool2d; mod max_pool2d; diff --git a/crates/cubek-pool/tests/pool/forward/adaptive_avg_pool3d.rs b/crates/cubek-pool/tests/pool/forward/adaptive_avg_pool3d.rs new file mode 100644 index 000000000..1786e8342 --- /dev/null +++ b/crates/cubek-pool/tests/pool/forward/adaptive_avg_pool3d.rs @@ -0,0 +1,246 @@ +use crate::pool::{build_output_tensor, output_host_f32, validate_test}; +use cubecl::{ + Runtime, TestRuntime, + ir::{ElemType, FloatKind}, + prelude::*, + zspace::Shape, +}; +use cubek_pool::{ + definition::{ + AdaptiveAvgPoolOptions, AvgPoolOptions, MaxPoolOptions, PoolError, PoolForwardProblem, + PoolMode, + }, + eval::cpu_reference::{cpu_reference_pool, geometry::PoolGeometry}, + pool3d, +}; +use cubek_test_utils::TestInput; + +const TOLERANCE: f32 = 1e-5; + +#[test] +fn test_adaptive_avg_pool3d_global() { + run_case([2, 9, 11, 7, 8], [1, 1, 1], 1001); +} + +#[test] +fn test_adaptive_avg_pool3d_divisible_downsample_vectorized_channels() { + run_case([2, 6, 8, 10, 8], [3, 4, 5], 1002); +} + +#[test] +fn test_adaptive_avg_pool3d_non_divisible_asymmetric_odd_channels() { + run_case([2, 5, 7, 4, 7], [3, 4, 3], 1003); +} + +#[test] +fn test_adaptive_avg_pool3d_non_divisible_large_windows() { + run_case([1, 19, 20, 15, 8], [2, 3, 2], 1005); +} + +#[test] +fn test_adaptive_avg_pool3d_output_larger_than_input() { + run_case([1, 2, 3, 2, 5], [4, 5, 3], 1004); +} + +#[test] +fn test_adaptive_avg_pool3d_large_dimension_index_math() { + run_case([1, 100_000, 1, 1, 1], [100_000, 1, 1], 1006); +} + +#[test] +fn test_adaptive_avg_pool3d_validates_rank_batch_channel_and_output_size() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + let options = AdaptiveAvgPoolOptions::new([3, 4, 3]); + + let invalid_rank_input = build_output_tensor(&client, vec![2, 5, 7, 3], dtype); + let valid_output = build_output_tensor(&client, vec![2, 3, 4, 3, 3], dtype); + assert!(matches!( + pool3d( + &client, + invalid_rank_input.binding(), + valid_output.clone().binding(), + options.clone().into(), + dtype, + ), + Err(PoolError::InvalidRank { .. }) + )); + + let input = build_output_tensor(&client, vec![2, 5, 7, 4, 3], dtype); + let batch_mismatch = build_output_tensor(&client, vec![1, 3, 4, 3, 3], dtype); + assert!(matches!( + pool3d( + &client, + input.clone().binding(), + batch_mismatch.binding(), + options.clone().into(), + dtype, + ), + Err(PoolError::BatchMismatch { .. }) + )); + + let channel_mismatch = build_output_tensor(&client, vec![2, 3, 4, 3, 4], dtype); + assert!(matches!( + pool3d( + &client, + input.clone().binding(), + channel_mismatch.binding(), + options.clone().into(), + dtype, + ), + Err(PoolError::ChannelMismatch { .. }) + )); + + let output_size_mismatch = build_output_tensor(&client, vec![2, 3, 4, 2, 3], dtype); + assert!(matches!( + pool3d( + &client, + input.binding(), + output_size_mismatch.binding(), + options.into(), + dtype, + ), + Err(PoolError::OutputSizeMismatch { .. }) + )); +} + +#[test] +fn test_adaptive_avg_pool3d_rejects_zero_spatial_dimensions() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + + let input = build_output_tensor(&client, vec![1, 0, 2, 2, 1], dtype); + let output = build_output_tensor(&client, vec![1, 1, 1, 1, 1], dtype); + assert!(matches!( + pool3d( + &client, + input.binding(), + output.binding(), + AdaptiveAvgPoolOptions::new([1, 1, 1]).into(), + dtype, + ), + Err(PoolError::InvalidSpatialSize { .. }) + )); + + let input = build_output_tensor(&client, vec![1, 2, 2, 2, 1], dtype); + let output = build_output_tensor(&client, vec![1, 0, 1, 1, 1], dtype); + assert!(matches!( + pool3d( + &client, + input.binding(), + output.binding(), + AdaptiveAvgPoolOptions::new([0, 1, 1]).into(), + dtype, + ), + Err(PoolError::InvalidSpatialSize { .. }) + )); +} + +#[test] +fn test_adaptive_avg_pool3d_accepts_empty_batch_and_channels() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + + for (input_shape, output_shape) in [ + (vec![0, 2, 2, 2, 1], vec![0, 1, 1, 1, 1]), + (vec![1, 2, 2, 2, 0], vec![1, 1, 1, 1, 0]), + ] { + let input = build_output_tensor(&client, input_shape, dtype); + let output = build_output_tensor(&client, output_shape, dtype); + pool3d( + &client, + input.binding(), + output.binding(), + AdaptiveAvgPoolOptions::new([1, 1, 1]).into(), + dtype, + ) + .expect("empty batch and channel dimensions should produce empty output"); + } +} + +#[test] +fn test_adaptive_avg_pool3d_rejects_other_pool_modes() { + let client = TestRuntime::client(&Default::default()); + let dtype = f32::elem_type_native(); + let input = build_output_tensor(&client, vec![1, 4, 4, 4, 1], dtype); + let output = build_output_tensor(&client, vec![1, 2, 2, 2, 1], dtype); + + for mode in [ + PoolMode::Avg(AvgPoolOptions::new( + [2, 2, 2], + [2, 2, 2], + [0, 0, 0], + false, + false, + )), + PoolMode::Max(MaxPoolOptions::new( + [2, 2, 2], + [2, 2, 2], + [0, 0, 0], + [1, 1, 1], + false, + )), + ] { + assert!(matches!( + pool3d( + &client, + input.clone().binding(), + output.clone().binding(), + mode, + dtype, + ), + Err(PoolError::UnsupportedMode { .. }) + )); + } +} + +#[test] +fn test_adaptive_avg_pool3d_f16_large_global_accumulates_in_f32() { + let client = TestRuntime::client(&Default::default()); + let dtype = ElemType::Float(FloatKind::F16); + let input_shape = vec![1, 41, 41, 41, 1]; + let input = TestInput::builder(client.clone(), input_shape.clone()) + .dtype(dtype) + .custom(vec![1.0; input_shape.iter().product()]) + .generate_without_host_data(); + let output = build_output_tensor(&client, vec![1, 1, 1, 1, 1], dtype); + pool3d( + &client, + input.binding(), + output.clone().binding(), + AdaptiveAvgPoolOptions::new([1, 1, 1]).into(), + dtype, + ) + .expect("f16 adaptive average pool 3d should launch"); + + let actual = output_host_f32(&client, output).get_f32(&[0, 0, 0, 0, 0]); + assert!(actual.is_finite()); + assert!((actual - 1.0).abs() <= 1e-3, "expected 1, got {actual}"); +} + +fn run_case(input_shape: [usize; 5], output_size: [usize; 3], seed: u64) { + let client = TestRuntime::client(&Default::default()); + let problem = PoolForwardProblem { + input_shape: Shape::from(input_shape), + with_indices: false, + mode: AdaptiveAvgPoolOptions::new(output_size).into(), + }; + let (input, input_data) = TestInput::builder(client.clone(), problem.input_shape.to_vec()) + .uniform(seed, -1.0, 1.0) + .generate_with_f32_host_data(); + let output_shape = problem.output_shape(&problem.input_shape).to_vec(); + let expected = cpu_reference_pool(&input_data, problem.clone()); + let dtype = input.dtype; + + let output = build_output_tensor(&client, output_shape, dtype); + let result = pool3d( + &client, + input.binding(), + output.clone().binding(), + problem.mode, + dtype, + ); + let actual = output_host_f32(&client, output); + + validate_test(result, actual, expected, TOLERANCE); +} diff --git a/crates/cubek-pool/tests/pool/forward/mod.rs b/crates/cubek-pool/tests/pool/forward/mod.rs index 85cb61054..188bff45f 100644 --- a/crates/cubek-pool/tests/pool/forward/mod.rs +++ b/crates/cubek-pool/tests/pool/forward/mod.rs @@ -1,4 +1,5 @@ mod adaptive_avg_pool2d; +mod adaptive_avg_pool3d; mod avg_pool2d; mod max_pool2d;