core(tensor): avoid redundant checks in hot paths; use debug_asserts and clarify try_* docs

Author: ariawisp

crates/cubecl-core/src/frontend/container/tensor/launch.rs

@@ -206,6 +206,16 @@ impl<R: Runtime> ArgSettings<R> for TensorArg<'_, R> {
 impl<'a, R: Runtime> TensorHandleRef<'a, R> {
     /// Convert the handle into a [tensor argument](TensorArg).
     pub fn as_tensor_arg(&'a self, vectorization: u8) -> TensorArg<'a, R> {
+        // In debug builds, assert that the requested vectorization is supported
+        // by the runtime. Validation of the chosen factor should normally be
+        // performed upstream (at selection time) to avoid redundant checks in
+        // hot paths.
+        debug_assert!(
+            R::supported_line_sizes().contains(&vectorization),
+            "unsupported vectorization {} (supported: {:?})",
+            vectorization,
+            R::supported_line_sizes()
+        );
         unsafe {
             TensorArg::from_raw_parts_and_size(
                 self.handle,
@@ -218,10 +228,14 @@ impl<'a, R: Runtime> TensorHandleRef<'a, R> {
     }
     /// Convert the handle into a [tensor argument](TensorArg) with basic safety checks
     /// for vectorization compatibility.
-    /// Try to convert the handle into a tensor argument, validating that the
-    /// requested vectorization factor is supported by the runtime. This does not
-    /// enforce inner-most contiguity or alignment requirements as kernels may
-    /// legally vectorize along axes other than the innermost.
+    ///
+    /// Note: This convenience is primarily intended for host wrappers / FFI
+    /// ingestion paths. In internal code, prefer validating the chosen
+    /// vectorization factor at selection time and then calling
+    /// [`as_tensor_arg`], to avoid redundant work in hot paths.
+    ///
+    /// This does not enforce inner‑most contiguity or alignment requirements as
+    /// kernels may vectorize along axes other than the innermost.
     pub fn try_as_tensor_arg(
         &'a self,
         vectorization: u8,
@@ -244,6 +258,24 @@ impl<'a, R: Runtime> TensorHandleRef<'a, R> {
         shape: &'a [usize],
         elem_size: usize,
     ) -> Self {
+        // Basic invariants for debug builds only; upstream layers are expected
+        // to ensure correctness in release builds.
+        debug_assert_eq!(
+            shape.len(),
+            strides.len(),
+            "rank mismatch (shape={}, strides={})",
+            shape.len(),
+            strides.len()
+        );
+        debug_assert!(elem_size > 0, "element size must be > 0");
+        // Disallow zero strides when corresponding dimension extent > 1
+        for (i, (&s, &d)) in strides.iter().zip(shape.iter()).enumerate() {
+            debug_assert!(
+                !(s == 0 && d > 1),
+                "zero stride on axis {} with extent > 1",
+                i
+            );
+        }
         Self {
             handle,
             strides,
@@ -254,6 +286,11 @@ impl<'a, R: Runtime> TensorHandleRef<'a, R> {
     }
 
     /// Safely create a tensor handle from raw parts with basic shape/stride validation.
+    ///
+    /// Note: This is mainly useful for host / FFI boundaries to surface clear
+    /// errors early. Internal code should ensure these invariants when
+    /// constructing handles and may use the `unsafe` constructor directly in
+    /// performance‑critical paths.
     pub fn try_from_parts(
         handle: &'a cubecl_runtime::server::Handle,
         strides: &'a [usize],