From a2d8dd936f8e75cfa0f9b24580bcf757b1013baf Mon Sep 17 00:00:00 2001 From: Rachit2323 Date: Sat, 3 Oct 2026 23:15:58 +0530 Subject: [PATCH] fix(cutile): reject oversized partition dimensions Signed-off-by: Rachit2323 --- cutile-rs/CHANGELOG.md | 6 + .../src/kernel_launcher_generator.rs | 8 +- cutile-rs/cutile/src/tensor.rs | 119 ++++++++++++++++-- cutile-rs/cutile/src/tile_kernel.rs | 20 ++- 4 files changed, 132 insertions(+), 21 deletions(-) diff --git a/cutile-rs/CHANGELOG.md b/cutile-rs/CHANGELOG.md index 2b03a2524f..de7ec3c759 100644 --- a/cutile-rs/CHANGELOG.md +++ b/cutile-rs/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ## [Unreleased] +### Fixed + +- Partition and mapped-partition dimensions above `i32::MAX` now fail launch + validation instead of wrapping to negative compiler and kernel metadata + (#1382). + ## [0.4.0] - 2026-09-25 Debug builds support cuda-gdb and Nsight. DGX Spark is now supported. diff --git a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs index 27dac6ab6c..e0a9c9f804 100644 --- a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs +++ b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs @@ -1931,7 +1931,7 @@ fn infer_mapped_partition_map_generics( ) -> Result<(), Error> { let mapped_shape_expr = |dim: usize| { format!( - "vec![KernelOutputStored::map_shape_as_i32(&{var_name}).expect(\"MappedPartitionMut missing map shape\")[{dim}].to_string()]" + "vec![KernelOutputStored::try_map_shape_as_i32(&{var_name})?.expect(\"MappedPartitionMut missing map shape\")[{dim}].to_string()]" ) }; match map_shape { @@ -1946,7 +1946,7 @@ fn infer_mapped_partition_map_generics( required_generics.expressions.insert( ident, Some(format!( - "KernelOutputStored::map_shape_as_i32(&{var_name}).expect(\"MappedPartitionMut missing map shape\").iter().map(|x| x.to_string()).collect::>()" + "KernelOutputStored::try_map_shape_as_i32(&{var_name})?.expect(\"MappedPartitionMut missing map shape\").iter().map(|x| x.to_string()).collect::>()" )), ); Ok(()) @@ -2042,7 +2042,7 @@ pub fn infer_shape_params_from_tensor_type( SupportedGenericType::ConstArray => { // This is a CGA type. if is_mutable { - required_generics.expressions.insert(last_ident.clone(), Some(format!("KernelOutputStored::partition_shape_as_i32(&{var_name}).iter().map(|x| x.to_string()).collect::>()"))); + required_generics.expressions.insert(last_ident.clone(), Some(format!("KernelOutputStored::try_partition_shape_as_i32(&{var_name})?.iter().map(|x| x.to_string()).collect::>()"))); } else { // This might make sense for a small tensor. required_generics.expressions.insert(last_ident.clone(), Some(format!("{var_name}.shape().iter().map(|x| x.to_string()).collect::>()"))); @@ -2097,7 +2097,7 @@ pub fn infer_shape_params_from_tensor_type( } SupportedGenericType::ConstScalar => { if is_mutable { - required_generics.expressions.insert(ident.clone(), Some(format!("vec![KernelOutputStored::partition_shape_as_i32(&{var_name})[{i}].to_string()]"))); + required_generics.expressions.insert(ident.clone(), Some(format!("vec![KernelOutputStored::try_partition_shape_as_i32(&{var_name})?[{i}].to_string()]"))); } else { required_generics.expressions.insert( ident.clone(), diff --git a/cutile-rs/cutile/src/tensor.rs b/cutile-rs/cutile/src/tensor.rs index 68cf53b8ce..ab920b1cc3 100644 --- a/cutile-rs/cutile/src/tensor.rs +++ b/cutile-rs/cutile/src/tensor.rs @@ -361,6 +361,7 @@ impl

MappedLaunchPartition

{ partition_grid: (u32, u32, u32), num_tile_blocks: u32, ) -> Result<(u32, u32, u32), Error> { + checked_dimensions_i32("Mapped partition", &self.map_shape)?; let map_rank = self.map_shape.len(); if map_rank == 0 || map_rank > 3 { return tensor_error_result( @@ -832,6 +833,23 @@ fn checked_num_bytes_i32(shape: &[i32]) -> Result { .ok_or_else(|| crate::error::tensor_error("Tensor byte size overflowed usize.")) } +// Converts host-side shape metadata to the i32 representation used by the +// compiler and kernel ABI. An unchecked cast would turn dimensions above +// i32::MAX negative while grid validation still accepted them as u32. +fn checked_dimensions_i32(kind: &str, dimensions: &[usize]) -> Result, Error> { + dimensions + .iter() + .map(|&dimension| { + i32::try_from(dimension).map_err(|_| { + crate::error::tensor_error(&format!( + "{kind} dimension {dimension} exceeds i32::MAX ({})", + i32::MAX + )) + }) + }) + .collect() +} + // Launch grid for a partition binding: the ceiling division of every tensor axis by the // matching partition axis, so a partial edge tile still gets a block. Both shapes are // caller-supplied, so every step is checked: a zero partition axis would divide by zero, @@ -859,14 +877,16 @@ fn partition_launch_grid( )); } let axis = |i: usize| -> Result { - // `shape[i] > 0` was checked above, so the cast is lossless. + // Both tensor and partition dimensions cross the kernel ABI as i32. + // `shape[i] > 0` was checked above, so its cast is lossless. let extent = shape[i] as u32; - let tile = u32::try_from(partition_shape[i]).map_err(|_| { + let tile = i32::try_from(partition_shape[i]).map_err(|_| { crate::error::tensor_error(&format!( - "Partition dimension {} exceeds u32::MAX.", - partition_shape[i] + "Partition dimension {} exceeds i32::MAX ({}).", + partition_shape[i], + i32::MAX )) - })?; + })? as u32; Ok(extent.div_ceil(tile)) }; match shape.len() { @@ -1877,8 +1897,16 @@ pub trait KernelOutputStored: Send { fn map_shape_as_i32(&self) -> Option> { None } + /// Converts mapped-partition dimensions to the compiler ABI representation. + fn try_map_shape_as_i32(&self) -> Result>, Error> { + Ok(self.map_shape_as_i32()) + } fn dtype_str(&self) -> &'static str; fn partition_shape_as_i32(&self) -> Vec; + /// Converts partition dimensions to the compiler ABI representation. + fn try_partition_shape_as_i32(&self) -> Result, Error> { + checked_dimensions_i32("Partition", self.partition_shape()) + } /// The partition shape as bound, borrowed: what launch validation reads. fn partition_shape(&self) -> &[usize]; fn strides_hint(&self) -> Vec; @@ -1924,10 +1952,16 @@ impl KernelOutputStored for Partition> { launcher.push_arg(*stride); } for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); + launcher.push_arg( + i32::try_from(*dim) + .expect("partition dimension must be validated before kernel launch"), + ); } for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + launcher.push_arg( + i32::try_from(*stride) + .expect("partition stride originates from valid i32 tensor metadata"), + ); } } fn grid(&self) -> Result<(u32, u32, u32), Error> { @@ -1937,7 +1971,8 @@ impl KernelOutputStored for Partition> { T::DTYPE.as_str() } fn partition_shape_as_i32(&self) -> Vec { - self.partition_shape.iter().map(|&x| x as i32).collect() + self.try_partition_shape_as_i32() + .expect("partition dimensions must fit the compiler ABI") } fn partition_shape(&self) -> &[usize] { &self.partition_shape @@ -1982,10 +2017,16 @@ impl KernelOutputStored for Partition<&mut Tensor> { launcher.push_arg(*stride); } for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); + launcher.push_arg( + i32::try_from(*dim) + .expect("partition dimension must be validated before kernel launch"), + ); } for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + launcher.push_arg( + i32::try_from(*stride) + .expect("partition stride originates from valid i32 tensor metadata"), + ); } } fn grid(&self) -> Result<(u32, u32, u32), Error> { @@ -1995,7 +2036,8 @@ impl KernelOutputStored for Partition<&mut Tensor> { T::DTYPE.as_str() } fn partition_shape_as_i32(&self) -> Vec { - self.partition_shape.iter().map(|&x| x as i32).collect() + self.try_partition_shape_as_i32() + .expect("partition dimensions must fit the compiler ABI") } fn partition_shape(&self) -> &[usize] { &self.partition_shape @@ -2029,7 +2071,12 @@ impl KernelOutputStored for MappedLaunchPartition Option> { - Some(self.map_shape.iter().map(|&dim| dim as i32).collect()) + self.try_map_shape_as_i32() + .expect("mapped partition dimensions must fit the compiler ABI") + } + + fn try_map_shape_as_i32(&self) -> Result>, Error> { + checked_dimensions_i32("Mapped partition", &self.map_shape).map(Some) } fn dtype_str(&self) -> &'static str { @@ -2069,7 +2116,12 @@ impl KernelOutputStored for MappedLaunchPartition Option> { - Some(self.map_shape.iter().map(|&dim| dim as i32).collect()) + self.try_map_shape_as_i32() + .expect("mapped partition dimensions must fit the compiler ABI") + } + + fn try_map_shape_as_i32(&self) -> Result>, Error> { + checked_dimensions_i32("Mapped partition", &self.map_shape).map(Some) } fn dtype_str(&self) -> &'static str { @@ -2517,6 +2569,47 @@ mod tests { .is_err()); } + #[test] + fn partition_metadata_rejects_dimensions_above_i32_max() { + let largest_valid = meta_f32(&[8]).partition([i32::MAX as usize]); + assert_eq!(KernelOutputStored::grid(&largest_valid).unwrap(), (1, 1, 1)); + assert_eq!( + largest_valid.try_partition_shape_as_i32().unwrap(), + vec![i32::MAX] + ); + + let oversized = i32::MAX as usize + 1; + let partition = meta_f32(&[8]).partition([oversized]); + + let grid_error = KernelOutputStored::grid(&partition).unwrap_err(); + assert!( + grid_error.to_string().contains("exceeds i32::MAX"), + "{grid_error}" + ); + let metadata_error = partition.try_partition_shape_as_i32().unwrap_err(); + assert!( + metadata_error.to_string().contains("exceeds i32::MAX"), + "{metadata_error}" + ); + } + + #[test] + fn mapped_partition_metadata_rejects_dimensions_above_i32_max() { + let oversized = i32::MAX as usize + 1; + let partition = meta_f32(&[8]).partition([8]).map([oversized], 1); + + let grid_error = KernelOutputStored::grid(&partition).unwrap_err(); + assert!( + grid_error.to_string().contains("exceeds i32::MAX"), + "{grid_error}" + ); + let metadata_error = partition.try_map_shape_as_i32().unwrap_err(); + assert!( + metadata_error.to_string().contains("exceeds i32::MAX"), + "{metadata_error}" + ); + } + #[test] fn borrowed_mutable_partition_requires_unique_storage() { let mut t = meta_f32(&[8]); diff --git a/cutile-rs/cutile/src/tile_kernel.rs b/cutile-rs/cutile/src/tile_kernel.rs index 503c826e95..2c29ce370c 100644 --- a/cutile-rs/cutile/src/tile_kernel.rs +++ b/cutile-rs/cutile/src/tile_kernel.rs @@ -1461,10 +1461,16 @@ impl KernelArgument for &Partition> { launcher.push_arg(*stride); } for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); + launcher.push_arg( + i32::try_from(*dim) + .expect("partition dimension must be validated before kernel launch"), + ); } for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + launcher.push_arg( + i32::try_from(*stride) + .expect("partition stride originates from valid i32 tensor metadata"), + ); } } } @@ -1482,10 +1488,16 @@ impl KernelArgument for &Partition<&mut Tensor> { launcher.push_arg(*stride); } for dim in self.partition_shape.iter() { - launcher.push_arg(*dim as i32); + launcher.push_arg( + i32::try_from(*dim) + .expect("partition dimension must be validated before kernel launch"), + ); } for stride in self.partition_strides.iter() { - launcher.push_arg(*stride as i32); + launcher.push_arg( + i32::try_from(*stride) + .expect("partition stride originates from valid i32 tensor metadata"), + ); } } }