From 0437c812fa3ee318298f920fa375e6acdef965d1 Mon Sep 17 00:00:00 2001 From: Raul Estrada Date: Thu, 8 Oct 2026 01:48:11 -0600 Subject: [PATCH] fix(cutile): validate full tensor shapes before allocation Signed-off-by: Raul Estrada --- cutile-rs/cutile/src/api.rs | 107 ++++++++++++++++++-- cutile-rs/cutile/tests/gpu/tensor_guards.rs | 67 ++++++++++++ 2 files changed, 166 insertions(+), 8 deletions(-) diff --git a/cutile-rs/cutile/src/api.rs b/cutile-rs/cutile/src/api.rs index 6e70bc4ad4..4d56232ea1 100644 --- a/cutile-rs/cutile/src/api.rs +++ b/cutile-rs/cutile/src/api.rs @@ -530,6 +530,12 @@ pub fn ones(shape: &[usize]) -> impl DeviceOp> { /// Allocates GPU memory and fills it with the specified value. This uses a GPU kernel /// to initialize the memory efficiently. /// +/// ## Errors +/// +/// Returns an error on execution if a dimension is zero, the element count +/// overflows `usize`, or the flattened element count exceeds `i32::MAX`. +/// Invalid shapes are rejected before GPU allocation or kernel execution. +/// /// ## Examples /// /// ```rust,ignore @@ -540,15 +546,40 @@ pub fn ones(shape: &[usize]) -> impl DeviceOp> { /// let matrix = api::full(-1, &[128, 128]).await; /// ``` pub fn full(val: T, shape: &[usize]) -> impl DeviceOp> { + let len = match checked_full_len(shape) { + Ok(len) => len, + Err(message) => return fail::>(message).boxed(), + }; let shape = shape.to_vec(); - let len = shape.iter().product::(); - Tensor::::uninitialized(len).then(move |t| { - // TODO (hme): It's awkward to assume_init this before actually initializing it. - let partition_size = 128; - let result = unsafe { t.assume_init() }.partition([partition_size]); - let (_, res) = value((val, result)).then(full_apply).unzip(); - res.unpartition().reshape(&shape) - }) + Tensor::::uninitialized(len) + .then(move |t| { + // TODO (hme): It's awkward to assume_init this before actually initializing it. + let partition_size = 128; + let result = unsafe { t.assume_init() }.partition([partition_size]); + let (_, res) = value((val, result)).then(full_apply).unzip(); + res.unpartition().reshape(&shape) + }) + .boxed() +} + +// `full` initializes a flat tensor before reshaping it. Its element count +// must fit that tensor's i32 dimension, which also bounds every positive +// target dimension and contiguous stride. An empty shape remains a scalar. +fn checked_full_len(shape: &[usize]) -> Result { + let len = shape + .iter() + .try_fold(1usize, |len, &dim| len.checked_mul(dim)) + .ok_or_else(|| format!("full: shape {shape:?} element count overflows usize"))?; + if len == 0 { + return Err(format!("full: shape {shape:?} contains a zero dimension")); + } + if i32::try_from(len).is_err() { + return Err(format!( + "full: shape {shape:?} element count {len} exceeds i32::MAX ({})", + i32::MAX + )); + } + Ok(len) } pub fn fill(tensor: Tensor, val: T) -> impl DeviceOp> { @@ -963,3 +994,63 @@ pub trait DeviceOpReshapeShared: } impl>>> DeviceOpReshapeShared for DI {} + +#[cfg(test)] +mod full_shape_tests { + use super::{checked_full_len, full, ones, zeros}; + + #[test] + fn accepts_valid_shapes_and_i32_boundary() { + assert_eq!(checked_full_len(&[]), Ok(1)); + assert_eq!(checked_full_len(&[1024]), Ok(1024)); + assert_eq!(checked_full_len(&[3, 5]), Ok(15)); + assert_eq!(checked_full_len(&[2, 3, 4]), Ok(24)); + assert_eq!( + checked_full_len(&[i32::MAX as usize]), + Ok(i32::MAX as usize) + ); + } + + #[test] + fn rejects_zero_dimensions() { + for shape in [&[0][..], &[3, 0][..], &[0, usize::MAX][..]] { + let error = checked_full_len(shape).expect_err("zero dimension must fail"); + assert!(error.contains("zero dimension"), "{error}"); + } + } + + #[test] + fn rejects_usize_product_overflow() { + // The first product wraps to 1 without overflow checks. The second + // overflows even though every individual dimension fits in i32. + for shape in [&[usize::MAX, usize::MAX][..], &[i32::MAX as usize; 3][..]] { + let error = checked_full_len(shape).expect_err("shape product must fail"); + assert!(error.contains("overflows usize"), "{error}"); + } + } + + #[test] + fn rejects_flattened_length_above_i32_max() { + // Both products fit usize on 32-bit and 64-bit hosts. + for shape in [&[i32::MAX as usize + 1][..], &[32768, 65536][..]] { + let error = checked_full_len(shape).expect_err("flat dimension must fail"); + assert!(error.contains("exceeds i32::MAX"), "{error}"); + } + } + + #[test] + fn invalid_shapes_do_not_panic_during_construction() { + // Only construct the lazy operations: this must not need a CUDA + // context, panic on arithmetic, or reach the zero-length assertion. + for shape in [ + &[usize::MAX, usize::MAX][..], + &[i32::MAX as usize + 1][..], + &[32768, 65536][..], + &[0][..], + ] { + let _ = full(7.0f32, shape); + let _ = zeros::(shape); + let _ = ones::(shape); + } + } +} diff --git a/cutile-rs/cutile/tests/gpu/tensor_guards.rs b/cutile-rs/cutile/tests/gpu/tensor_guards.rs index 381a949780..12754ab099 100644 --- a/cutile-rs/cutile/tests/gpu/tensor_guards.rs +++ b/cutile-rs/cutile/tests/gpu/tensor_guards.rs @@ -9,6 +9,7 @@ use std::panic::{self, AssertUnwindSafe}; +use cuda_async::error::DeviceError; use cutile::api; use cutile::prelude::*; @@ -104,3 +105,69 @@ fn oversized_allocation_is_an_error_not_a_panic() { .expect("allocation after OOM"); assert_eq!(t.shape(), &[16]); } + +const INVALID_FULL_SHAPES: &[(&[usize], &str)] = &[ + (&[usize::MAX, usize::MAX], "overflows usize"), + (&[i32::MAX as usize; 3], "overflows usize"), + (&[i32::MAX as usize + 1], "exceeds i32::MAX"), + (&[32768, 65536], "exceeds i32::MAX"), + (&[0], "zero dimension"), + (&[3, 0], "zero dimension"), +]; + +fn assert_full_shape_error(error: DeviceError, expected: &str) { + match error { + DeviceError::Internal(message) => { + assert!(message.starts_with("full: "), "{message}"); + assert!(message.contains(expected), "{message}"); + } + other => panic!("expected full shape validation error, got: {other:?}"), + } +} + +#[test] +fn full_reports_invalid_shapes_as_errors() { + // Check the specific error, not just is_err(): in release builds the + // old product could wrap, allocate and fill, then fail during reshape. + for &(shape, expected) in INVALID_FULL_SHAPES { + let error = api::full(7.0f32, shape) + .sync() + .expect_err("invalid full shape must fail"); + assert_full_shape_error(error, expected); + } +} + +#[test] +fn full_wrappers_report_invalid_shapes_as_errors() { + for &(shape, expected) in INVALID_FULL_SHAPES { + let zeros_error = api::zeros::(shape) + .sync() + .expect_err("invalid zeros shape must fail"); + assert_full_shape_error(zeros_error, expected); + + let ones_error = api::ones::(shape) + .sync() + .expect_err("invalid ones shape must fail"); + assert_full_shape_error(ones_error, expected); + } +} + +#[test] +fn full_and_wrappers_preserve_shape_and_values() { + // Exercise a matrix reshape and a partial edge tile of the fill kernel. + for (shape, expected_shape) in [ + (&[2usize, 3][..], &[2i32, 3][..]), + (&[129usize][..], &[129i32][..]), + ] { + let tensors = [ + (api::full(3.25f32, shape).sync().expect("valid full"), 3.25), + (api::zeros::(shape).sync().expect("valid zeros"), 0.0), + (api::ones::(shape).sync().expect("valid ones"), 1.0), + ]; + for (tensor, expected) in tensors { + assert_eq!(tensor.shape(), expected_shape); + let host: Vec = tensor.to_host_vec().sync().expect("copy"); + assert_eq!(host, vec![expected; shape.iter().product()]); + } + } +}