fix: validate runtime tensor strides before kernel launch

cuda-async/src/device_context.rs

@@ -77,6 +77,21 @@ pub struct PointerParamType {
 pub struct TensorParamType {
     pub element_type: String,
     pub shape: Vec<i32>,
+    /// Per-axis strides in elements. `-1` denotes a runtime stride; other
+    /// values are constants used by the compiled specialization.
+    pub strides: Vec<i32>,
+}
+
+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<Arc<Modul
     })?
 }
 
+#[cfg(test)]
+mod tensor_stride_tests {
+    use super::TensorParamType;
+
+    fn tensor_param(strides: &[i32]) -> 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

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

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::<Vec<_>>();
                     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.

cutile-rs/cutile/src/tensor.rs

@@ -1881,6 +1881,8 @@ pub trait KernelOutputStored<T: DType>: Send {
     fn partition_shape_as_i32(&self) -> Vec<i32>;
     /// 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<i32>;
     fn spec(&self) -> &SpecializationBits;
     fn shape_as_i32(&self) -> Vec<i32>;
@@ -1942,6 +1944,9 @@ impl<T: DType> KernelOutputStored<T> for Partition<Tensor<T>> {
     fn partition_shape(&self) -> &[usize] {
         &self.partition_shape
     }
+    fn strides(&self) -> &[i32] {
+        self.object.strides()
+    }
     fn strides_hint(&self) -> Vec<i32> {
         self.object
             .spec
@@ -2000,6 +2005,9 @@ impl<T: DType> KernelOutputStored<T> for Partition<&mut Tensor<T>> {
     fn partition_shape(&self) -> &[usize] {
         &self.partition_shape
     }
+    fn strides(&self) -> &[i32] {
+        self.object.strides()
+    }
     fn strides_hint(&self) -> Vec<i32> {
         self.object
             .spec
@@ -2043,6 +2051,10 @@ impl<T: DType> KernelOutputStored<T> for MappedLaunchPartition<Partition<Tensor<
         KernelOutputStored::partition_shape(&self.partition)
     }
 
+    fn strides(&self) -> &[i32] {
+        KernelOutputStored::strides(&self.partition)
+    }
+
     fn strides_hint(&self) -> Vec<i32> {
         self.partition.strides_hint()
     }
@@ -2083,6 +2095,10 @@ impl<T: DType> KernelOutputStored<T> for MappedLaunchPartition<Partition<&mut Te
         KernelOutputStored::partition_shape(&self.partition)
     }
 
+    fn strides(&self) -> &[i32] {
+        KernelOutputStored::strides(&self.partition)
+    }
+
     fn strides_hint(&self) -> Vec<i32> {
         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.