diff --git a/cuda-async/src/device_context.rs b/cuda-async/src/device_context.rs index a119beae28..915c2ac76d 100644 --- a/cuda-async/src/device_context.rs +++ b/cuda-async/src/device_context.rs @@ -77,6 +77,21 @@ pub struct PointerParamType { pub struct TensorParamType { pub element_type: String, pub shape: Vec, + /// Per-axis strides in elements. `-1` denotes a runtime stride; other + /// values are constants used by the compiled specialization. + pub strides: Vec, +} + +impl TensorParamType { + /// Whether the runtime strides satisfy this specialization's constants. + pub fn strides_match(&self, given: &[i32]) -> bool { + self.strides.len() == given.len() + && self + .strides + .iter() + .zip(given) + .all(|(&expected, &given)| expected == -1 || expected == given) + } } #[derive(Debug, Clone)] @@ -523,6 +538,52 @@ pub fn load_module_from_ptx(ptx_src: &str, device_id: usize) -> Result TensorParamType { + TensorParamType { + element_type: "f32".to_string(), + shape: vec![-1; strides.len()], + strides: strides.to_vec(), + } + } + + #[test] + fn static_strides_must_match() { + let param = tensor_param(&[8, 1]); + assert!(param.strides_match(&[8, 1])); + assert!(!param.strides_match(&[4, 1])); + assert!(!param.strides_match(&[8, 2])); + } + + #[test] + fn dynamic_strides_accept_different_runtime_values() { + let param = tensor_param(&[-1, 1]); + assert!(param.strides_match(&[8, 1])); + assert!(param.strides_match(&[4, 1])); + assert!(param.strides_match(&[1, 1])); + assert!(!param.strides_match(&[8, 2])); + } + + #[test] + fn stride_rank_must_match_before_comparing_axes() { + let param = tensor_param(&[-1, 1]); + assert!(!param.strides_match(&[])); + assert!(!param.strides_match(&[8])); + assert!(!param.strides_match(&[8, 1, 1])); + assert!(!tensor_param(&[-1, -1]).strides_match(&[8])); + } + + #[test] + fn rank_zero_has_no_strides() { + let param = tensor_param(&[]); + assert!(param.strides_match(&[])); + assert!(!param.strides_match(&[1])); + } +} + #[cfg(test)] mod context_map_tests { //! Host-only tests for the thread-local context map's lifecycle. They only diff --git a/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs b/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs index 33ce9dce04..846e91d2d7 100644 --- a/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs +++ b/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs @@ -503,6 +503,7 @@ pub fn generate_entry_point( final_stmnt_args.push(tensor_input.generate_arg()); fn_entry.sig.inputs.extend(tensor_input.generate_args()); fn_params_concrete_types.push(ValidParamType::Tensor(TensorParamType { + strides: stride_args[&tensor_input.var_name].clone(), element_type: tensor_input.element_type, shape: tensor_input .input_tensor_shape @@ -541,6 +542,7 @@ pub fn generate_entry_point( fn_entry.sig.inputs.extend(tensor_input.generate_args()); fn_params_concrete_types.push(ValidParamType::Tensor( TensorParamType { + strides: stride_args[&tensor_input.var_name].clone(), element_type: tensor_input.element_type, shape: tensor_input .input_tensor_shape diff --git a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs index 27dac6ab6c..ec1176c3af 100644 --- a/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs +++ b/cutile-rs/cutile-macro/src/kernel_launcher_generator.rs @@ -1855,6 +1855,8 @@ fn get_tensor_code( let given_shape: &[usize] = KernelOutputStored::partition_shape(&#var_ident); kernel_launch_assert_with(valid_shape.len() == given_shape.len(), || format!("{} rank mismatch: Expected {}, got {}", #var_name, valid_shape.len(), given_shape.len()))?; kernel_launch_assert_with(zip(valid_shape, given_shape).all(|(&expected, &given)| expected as i64 == given as i64), || format!("{} partition shape mismatch. Expected {:?}, got {:?}", #var_name, valid_shape, given_shape))?; + let given_strides = KernelOutputStored::strides(&#var_ident); + kernel_launch_assert_with(tensor_validator.strides_match(given_strides), || format!("{} strides mismatch. Expected {:?} (-1 denotes a dynamic stride), got {:?}", #var_name, tensor_validator.strides, given_strides))?; // Runtime extents for launch-time check hoisting (indexed by // signature position): root shape resolves `Dim` atoms, and the // partition slab -- this param's kernel-visible view -- resolves @@ -1894,7 +1896,8 @@ fn get_tensor_code( }).collect::>(); format!("{} partition shape mismatch. Expected {:?}, got {:?}", #var_name, valid_shape_mixed, given_shape) })?; - // TODO (hme): add validation for strides here too. + let given_strides = KernelInputStored::strides(&#var_ident); + kernel_launch_assert_with(tensor_validator.strides_match(given_strides), || format!("{} strides mismatch. Expected {:?} (-1 denotes a dynamic stride), got {:?}", #var_name, tensor_validator.strides, given_strides))?; // Runtime extents for launch-time check hoisting (indexed by // signature position). An immutable tensor is passed whole, so // its view is its root shape. diff --git a/cutile-rs/cutile/src/tensor.rs b/cutile-rs/cutile/src/tensor.rs index 68cf53b8ce..1ef3adec62 100644 --- a/cutile-rs/cutile/src/tensor.rs +++ b/cutile-rs/cutile/src/tensor.rs @@ -1881,6 +1881,8 @@ pub trait KernelOutputStored: Send { fn partition_shape_as_i32(&self) -> Vec; /// The partition shape as bound, borrowed: what launch validation reads. fn partition_shape(&self) -> &[usize]; + /// Runtime strides of the underlying tensor, in elements. + fn strides(&self) -> &[i32]; fn strides_hint(&self) -> Vec; fn spec(&self) -> &SpecializationBits; fn shape_as_i32(&self) -> Vec; @@ -1942,6 +1944,9 @@ impl KernelOutputStored for Partition> { fn partition_shape(&self) -> &[usize] { &self.partition_shape } + fn strides(&self) -> &[i32] { + self.object.strides() + } fn strides_hint(&self) -> Vec { self.object .spec @@ -2000,6 +2005,9 @@ impl KernelOutputStored for Partition<&mut Tensor> { fn partition_shape(&self) -> &[usize] { &self.partition_shape } + fn strides(&self) -> &[i32] { + self.object.strides() + } fn strides_hint(&self) -> Vec { self.object .spec @@ -2043,6 +2051,10 @@ impl KernelOutputStored for MappedLaunchPartition &[i32] { + KernelOutputStored::strides(&self.partition) + } + fn strides_hint(&self) -> Vec { self.partition.strides_hint() } @@ -2083,6 +2095,10 @@ impl KernelOutputStored for MappedLaunchPartition &[i32] { + KernelOutputStored::strides(&self.partition) + } + fn strides_hint(&self) -> Vec { self.partition.strides_hint() } @@ -2498,6 +2514,25 @@ mod tests { ); } + #[test] + fn output_bindings_report_underlying_tensor_strides() { + // The root's [8, 1] strides differ from the [4, 1] strides a + // contiguous [2, 4] tile would have. + let owned = meta_f32(&[4, 8]).partition([2, 4]); + assert_eq!(KernelOutputStored::strides(&owned), &[8, 1]); + + let mut tensor = meta_f32(&[4, 8]); + let borrowed = (&mut tensor).partition([2, 4]); + assert_eq!(KernelOutputStored::strides(&borrowed), &[8, 1]); + + let mapped_owned = meta_f32(&[4, 8]).partition([2, 4]).map([2, 2], 4); + assert_eq!(KernelOutputStored::strides(&mapped_owned), &[8, 1]); + + let mut tensor = meta_f32(&[4, 8]); + let mapped_borrowed = (&mut tensor).partition([2, 4]).map([2, 2], 4); + assert_eq!(KernelOutputStored::strides(&mapped_borrowed), &[8, 1]); + } + #[test] fn launch_grid_rejects_zero_partition_dimension_and_rank_mismatch() { // Zero partition axis: formerly a divide-by-zero panic in div_ceil. diff --git a/cutile-rs/cutile/tests/gpu/launcher_guards.rs b/cutile-rs/cutile/tests/gpu/launcher_guards.rs index 1602742217..72a7766e76 100644 --- a/cutile-rs/cutile/tests/gpu/launcher_guards.rs +++ b/cutile-rs/cutile/tests/gpu/launcher_guards.rs @@ -8,7 +8,10 @@ //! buffers it is handed, however the generics were chosen. A user //! `.generics(..)` list overrides the inferred element type, and a kernel //! specialized wider than its buffers indexes past them. +//! Stride constants must also agree with the runtime layout, including when +//! a custom input/output binding supplies stale specialization metadata. +use cuda_async::launch::AsyncKernelLaunch; use cutile::api; use cutile::half::f16; use cutile::prelude::*; @@ -25,6 +28,15 @@ mod launcher_guards_module { z.store(t); } + #[cutile::entry()] + fn copy_strided(z: &mut Tensor, x: &Tensor) { + let t: Tile = load_tile_like(x, z); + z.store(t); + } + + #[cutile::entry()] + fn mapped_stride_guard(_z: MappedPartitionMut) {} + /// Zeroes 128 elements through a raw pointer. The pointer is the only /// typed argument, so this isolates the pointer element-type check. #[cutile::entry()] @@ -55,7 +67,218 @@ mod launcher_guards_module { } } -use launcher_guards_module::{copy_elements, dump_to_missing_dir, zero_through_ptr}; +use launcher_guards_module::{ + copy_elements, copy_strided, dump_to_missing_dir, mapped_stride_guard, zero_through_ptr, +}; + +/// Inject stale specialization metadata through the public binding traits. +/// Shapes, strides, retention and kernel arguments still come from the real +/// binding. This tests the generated guard without forging a device pointer. +struct OverrideSpec { + inner: I, + spec: SpecializationBits, +} + +fn with_first_stride_one(mut spec: SpecializationBits) -> SpecializationBits { + spec.stride_one[0] = true; + spec.stride_div[0].divisor = 1; + spec +} + +impl KernelInputStored for OverrideSpec { + fn retain(&self, ctx: &ExecutionContext) -> Result<(), DeviceError> { + self.inner.retain(ctx) + } + fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + self.inner.push_kernel_args(launcher); + } + fn shape(&self) -> &[i32] { + self.inner.shape() + } + fn strides(&self) -> &[i32] { + self.inner.strides() + } + fn spec(&self) -> &SpecializationBits { + &self.spec + } + fn dtype_str(&self) -> &'static str { + self.inner.dtype_str() + } +} + +impl KernelInput for OverrideSpec { + type Stored = Self; + type Returned = Self; + fn prepare(self) -> Self { + self + } + fn recover(stored: Self) -> Self { + stored + } +} + +impl> KernelOutputStored for OverrideSpec { + fn retain(&self, ctx: &ExecutionContext) -> Result<(), DeviceError> { + self.inner.retain(ctx) + } + fn push_kernel_args(&self, launcher: &mut AsyncKernelLaunch) { + self.inner.push_kernel_args(launcher); + } + fn grid(&self) -> Result<(u32, u32, u32), Error> { + self.inner.grid() + } + fn map_shape_as_i32(&self) -> Option> { + self.inner.map_shape_as_i32() + } + fn dtype_str(&self) -> &'static str { + self.inner.dtype_str() + } + fn partition_shape_as_i32(&self) -> Vec { + self.inner.partition_shape_as_i32() + } + fn partition_shape(&self) -> &[usize] { + self.inner.partition_shape() + } + fn strides(&self) -> &[i32] { + self.inner.strides() + } + fn strides_hint(&self) -> Vec { + self.spec + .stride_one + .iter() + .map(|&is_one| if is_one { 1 } else { -1 }) + .collect() + } + fn spec(&self) -> &SpecializationBits { + &self.spec + } + fn shape_as_i32(&self) -> Vec { + self.inner.shape_as_i32() + } +} + +impl> KernelOutput for OverrideSpec { + type Stored = Self; + type Returned = Self; + fn prepare(self) -> Self { + self + } + fn recover(stored: Self) -> Self { + stored + } +} + +#[test] +fn stale_input_stride_constants_are_rejected_after_cache_fill() { + common::with_test_stack(|| { + let _guard = common::cache_test_lock(); + let x = api::arange::(256) + .reshape(&[128, 2]) + .sync() + .expect("alloc x"); + let view = x.slice(&[0..128, 0..1]).expect("strided input view"); + assert_eq!(view.shape(), &[128, 1]); + assert_eq!(view.strides(), &[2, 1]); + let spec = with_first_stride_one(view.spec().clone()); + let mut z = api::zeros::(&[128, 1]).sync().expect("alloc z"); + + // The JIT resolution is cached before argument validation. Repeating + // the failed compile-only launch exercises the cached path too. + for _ in 0..2 { + let input = OverrideSpec { + inner: &view, + spec: spec.clone(), + }; + let err = copy_strided((&mut z).partition([128, 1]), value(input)) + .compile() + .expect_err("a unit stride specialization over stride 2 must be rejected"); + let msg = err.to_string(); + assert!(msg.contains("x strides mismatch"), "{msg}"); + assert!(msg.contains("[1, 1]") && msg.contains("[2, 1]"), "{msg}"); + } + }); +} + +#[test] +fn stale_output_stride_constants_are_rejected_after_cache_fill() { + common::with_test_stack(|| { + let _guard = common::cache_test_lock(); + let x = api::ones::(&[128, 2]).sync().expect("alloc x"); + let mut z = api::zeros::(&[128, 2]).sync().expect("alloc z"); + let spec = with_first_stride_one(z.spec().clone()); + for _ in 0..2 { + let output = OverrideSpec { + inner: (&mut z).partition([128, 1]), + spec: spec.clone(), + }; + let err = copy_strided(value(output), &x) + .compile() + .expect_err("a unit stride specialization over stride 2 must be rejected"); + let msg = err.to_string(); + assert!(msg.contains("z strides mismatch"), "{msg}"); + assert!(msg.contains("[1, 1]") && msg.contains("[2, 1]"), "{msg}"); + } + + let output = OverrideSpec { + inner: (&mut z).partition([128, 1]).map([1, 2], 2), + spec, + }; + let err = mapped_stride_guard(value(output)) + .compile() + .expect_err("mapped outputs must also validate their runtime strides"); + assert!(err.to_string().contains("_z strides mismatch"), "{err}"); + }); +} + +#[test] +fn different_dynamic_input_strides_reuse_the_same_specialization() { + common::with_test_stack(|| { + let _guard = common::cache_test_lock(); + let mut previous_key = None; + for row_stride in [32usize, 48, 32] { + let x = api::arange::(128 * row_stride) + .reshape(&[128, row_stride]) + .sync() + .expect("alloc x"); + let view = x.slice(&[0..128, 0..1]).expect("strided input view"); + let mut z = api::zeros::(&[128, 1]).sync().expect("alloc z"); + // Both strides have the same divisibility hints and are dynamic. + let key = copy_strided((&mut z).partition([128, 1]), &view) + .l1_cache_key() + .expect("specialization key"); + if let Some(previous) = &previous_key { + assert_eq!(previous, &key); + } + previous_key = Some(key); + copy_strided((&mut z).partition([128, 1]), &view) + .sync() + .expect("a dynamic input stride must be accepted"); + let host = z.to_host_vec().sync().expect("copy back"); + let expected: Vec = (0..128).map(|i| (i * row_stride) as f32).collect(); + assert_eq!(host, expected); + } + }); +} + +#[test] +fn dynamic_output_strides_are_accepted() { + common::with_test_stack(|| { + for columns in [2usize, 3] { + let x = api::arange::(128 * columns) + .reshape(&[128, columns]) + .sync() + .expect("alloc x"); + let host = copy_strided(api::zeros::(&[128, columns]).partition([128, 1]), &x) + .first() + .unpartition() + .to_host_vec() + .sync() + .expect("a dynamic output stride must be accepted"); + let expected: Vec = (0..128 * columns).map(|i| i as f32).collect(); + assert_eq!(host, expected); + } + }); +} #[test] fn unwritable_dump_mlir_dir_is_an_error_not_a_panic() { diff --git a/cutile-rs/cutile/tests/specialization_bits.rs b/cutile-rs/cutile/tests/specialization_bits.rs index 9f660f788d..e720431a97 100644 --- a/cutile-rs/cutile/tests/specialization_bits.rs +++ b/cutile-rs/cutile/tests/specialization_bits.rs @@ -2,9 +2,11 @@ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +use cuda_async::device_context::{ValidParamType, Validator}; use cutile::api; use cutile::tile_kernel::{DeviceOp, TileKernel}; use cutile_compiler::compiler::utils::CompileOptions; +use cutile_compiler::compiler::{CUDATileFunctionCompiler, CUDATileModules}; use cutile_compiler::specialization::{DivHint, SpecializationBits}; use std::path::Path; use std::sync::{Arc, Mutex}; @@ -46,6 +48,16 @@ mod spec_test_module { /// scalar DivHints still lower when pointer args are present. #[cutile::entry(dump_mlir_dir = "/tmp/cutile_raw_ptr_scalar_mlir")] unsafe fn raw_ptr_scalar_kernel(_ptr: *mut f32, _n: i32) {} + + #[cutile::entry()] + fn stride_validator_kernel( + _output: &mut Tensor, + _input: &Tensor, + ) { + } + + #[cutile::entry()] + fn mapped_stride_validator_kernel(_output: MappedPartitionMut) {} } use spec_test_module::{__module_ast_self, raw_ptr_scalar_kernel}; @@ -89,6 +101,61 @@ fn compile_kernel( .expect("Failed to compile") } +// -- Compiled tensor stride metadata -- + +fn stride_validator(name: &str, strides: &[(&str, &[i32])]) -> Validator { + let modules = CUDATileModules::from_kernel(__module_ast_self()).expect("kernel module"); + let options = CompileOptions::default(); + let compiler = CUDATileFunctionCompiler::new( + &modules, + "spec_test_module", + name, + &[], + strides, + &[], + &[], + None, + "sm_120".to_string(), + &options, + ) + .expect("entry point"); + compiler.get_validator() +} + +#[test] +fn validator_preserves_input_and_output_stride_constants() { + common::with_test_stack(|| { + for (output_strides, input_strides) in [([8, 1], [-1, 1]), ([1, -1], [-1, 4])] { + let validator = stride_validator( + "stride_validator_kernel", + &[("_output", &output_strides), ("_input", &input_strides)], + ); + assert_eq!(validator.params.len(), 2); + for (param, expected) in validator.params.iter().zip([output_strides, input_strides]) { + let ValidParamType::Tensor(tensor) = param else { + panic!("expected a tensor validator"); + }; + assert_eq!(tensor.strides, expected); + } + } + }); +} + +#[test] +fn validator_preserves_mapped_output_stride_constants() { + common::with_test_stack(|| { + for expected in [[8, 1], [-1, 1]] { + let validator = + stride_validator("mapped_stride_validator_kernel", &[("_output", &expected)]); + assert_eq!(validator.params.len(), 1); + let ValidParamType::Tensor(tensor) = &validator.params[0] else { + panic!("expected a mapped tensor validator"); + }; + assert_eq!(tensor.strides, expected); + } + }); +} + // -- SpecializationBits produces correct assume_div_by in MLIR -- #[test]