Autogenerate a lot of strides

Author: LegNeato

crates/cuda_std/build.rs

@@ -0,0 +1,76 @@
+use std::env;
+use std::fs;
+use std::io::Write;
+use std::path::Path;
+
+/// Maximum stride value to support
+const MAX_STRIDE: usize = 8192;
+
+fn main() {
+    // Tell cargo to rerun if build.rs changes
+    println!("cargo:rerun-if-changed=build.rs");
+    
+    // Generate stride implementations
+    generate_stride_impls();
+}
+
+fn generate_stride_impls() {
+    let out_dir = env::var_os("OUT_DIR").unwrap();
+    let out_path = Path::new(&out_dir);
+    
+    // Create stride subdirectory in OUT_DIR
+    let stride_dir = out_path.join("stride");
+    fs::create_dir_all(&stride_dir).unwrap();
+    
+    // Define stride requirements for each type
+    // (type_name, file_name, start, step)
+    let stride_configs = vec![
+        // 16-byte alignment types (multiple of element count to fill 16 bytes)
+        ("f16", "f16", 8, 8),        // f16: 2 bytes × 8 = 16 bytes
+        ("bf16", "bf16", 8, 8),      // bf16: 2 bytes × 8 = 16 bytes
+        ("i16", "i16", 8, 8),        // i16: 2 bytes × 8 = 16 bytes
+        ("u16", "u16", 8, 8),        // u16: 2 bytes × 8 = 16 bytes
+        ("f32", "f32", 4, 4),        // f32: 4 bytes × 4 = 16 bytes
+        ("i32", "i32", 4, 4),        // i32: 4 bytes × 4 = 16 bytes
+        ("u32", "u32", 4, 4),        // u32: 4 bytes × 4 = 16 bytes
+        ("f64", "f64", 2, 2),        // f64: 8 bytes × 2 = 16 bytes
+        ("i64", "i64", 2, 2),        // i64: 8 bytes × 2 = 16 bytes
+        ("u64", "u64", 2, 2),        // u64: 8 bytes × 2 = 16 bytes
+        ("i8", "i8", 16, 16),        // i8: 1 byte × 16 = 16 bytes
+        ("u8", "u8", 16, 16),        // u8: 1 byte × 16 = 16 bytes
+        ("bool", "bool", 16, 16),    // bool: 1 byte × 16 = 16 bytes
+    ];
+    
+    for (type_name, file_name, start, step) in stride_configs {
+        generate_type_strides(&stride_dir, type_name, file_name, start, step);
+    }
+}
+
+fn generate_type_strides(
+    out_dir: &Path,
+    type_name: &str,
+    file_name: &str,
+    start: usize,
+    step: usize,
+) {
+    let dest_path = out_dir.join(format!("{}_impls.rs", file_name));
+    let mut file = fs::File::create(&dest_path).unwrap();
+    
+    // Write header
+    writeln!(file, "// @generated").unwrap();
+    writeln!(file, "// Auto-generated stride implementations for {}", type_name).unwrap();
+    writeln!(file, "// DO NOT EDIT THIS FILE MANUALLY").unwrap();
+    writeln!(file).unwrap();
+    writeln!(file, "// {} requires stride to be multiple of {} (up to {})", 
+             type_name, step, MAX_STRIDE).unwrap();
+    writeln!(file).unwrap();
+    
+    // Generate implementations
+    // These files are included in modules that already have the necessary imports
+    let mut stride = start;
+    while stride <= MAX_STRIDE {
+        writeln!(file, "impl ValidStride for StrideValidator<{}, {}> {{}}", 
+                 type_name, stride).unwrap();
+        stride += step;
+    }
+}

crates/cuda_std/src/warp/matrix/mod.rs

@@ -17,6 +17,11 @@ pub mod ops;
 mod intrinsics;
 use intrinsics::*;
 
+// Stride validation module
+pub mod stride;
+// Re-export for convenience
+pub use stride::{StrideValidator, ValidStride};
+
 // ============================================================================
 // Shape Types with Compile-Time Validation
 // ============================================================================
@@ -1749,80 +1754,7 @@ impl MmaWithShapeAndLayout<f64, f64, f64, dims::Shape<8, 8, 4>, layout::Row, lay
     }
 }
 
-// ============================================================================
-// Stride Validation
-// ============================================================================
-
-/// Compile-time stride validation
-pub struct StrideValidator<T, const STRIDE: usize>(PhantomData<T>);
-
-/// Trait for valid strides
-#[diagnostic::on_unimplemented(
-    message = "`{Self}` is not a valid stride configuration",
-    label = "invalid stride for tensor core operations",
-    note = "f16/bf16 require stride to be a multiple of 8",
-    note = "f32/i32 require stride to be a multiple of 4",
-    note = "i8/u8/bool require stride to be a multiple of 16",
-    note = "f64 requires stride to be a multiple of 2"
-)]
-pub trait ValidStride: sealed::Sealed {}
-
-// f16 requires stride to be multiple of 8
-impl ValidStride for StrideValidator<f16, 8> {}
-impl ValidStride for StrideValidator<f16, 16> {}
-impl ValidStride for StrideValidator<f16, 24> {}
-impl ValidStride for StrideValidator<f16, 32> {}
-impl ValidStride for StrideValidator<f16, 40> {}
-impl ValidStride for StrideValidator<f16, 48> {}
-impl ValidStride for StrideValidator<f16, 56> {}
-impl ValidStride for StrideValidator<f16, 64> {}
-
-// f32 requires stride to be multiple of 4
-impl ValidStride for StrideValidator<f32, 4> {}
-impl ValidStride for StrideValidator<f32, 8> {}
-impl ValidStride for StrideValidator<f32, 12> {}
-impl ValidStride for StrideValidator<f32, 16> {}
-impl ValidStride for StrideValidator<f32, 20> {}
-impl ValidStride for StrideValidator<f32, 24> {}
-impl ValidStride for StrideValidator<f32, 28> {}
-impl ValidStride for StrideValidator<f32, 32> {}
-
-// bf16 requires stride to be multiple of 8 (same as f16)
-impl ValidStride for StrideValidator<bf16, 8> {}
-impl ValidStride for StrideValidator<bf16, 16> {}
-impl ValidStride for StrideValidator<bf16, 24> {}
-impl ValidStride for StrideValidator<bf16, 32> {}
-impl ValidStride for StrideValidator<bf16, 40> {}
-impl ValidStride for StrideValidator<bf16, 48> {}
-impl ValidStride for StrideValidator<bf16, 56> {}
-impl ValidStride for StrideValidator<bf16, 64> {}
-
-// i8/u8 require stride to be multiple of 16
-impl ValidStride for StrideValidator<i8, 16> {}
-impl ValidStride for StrideValidator<i8, 32> {}
-impl ValidStride for StrideValidator<i8, 48> {}
-impl ValidStride for StrideValidator<i8, 64> {}
-
-impl ValidStride for StrideValidator<u8, 16> {}
-impl ValidStride for StrideValidator<u8, 32> {}
-impl ValidStride for StrideValidator<u8, 48> {}
-impl ValidStride for StrideValidator<u8, 64> {}
-
-// i32 requires stride to be multiple of 4 (same as f32)
-impl ValidStride for StrideValidator<i32, 4> {}
-impl ValidStride for StrideValidator<i32, 8> {}
-impl ValidStride for StrideValidator<i32, 12> {}
-impl ValidStride for StrideValidator<i32, 16> {}
-impl ValidStride for StrideValidator<i32, 20> {}
-impl ValidStride for StrideValidator<i32, 24> {}
-impl ValidStride for StrideValidator<i32, 28> {}
-impl ValidStride for StrideValidator<i32, 32> {}
-
-// bool uses u8 stride requirements (multiple of 16)
-impl ValidStride for StrideValidator<bool, 16> {}
-impl ValidStride for StrideValidator<bool, 32> {}
-impl ValidStride for StrideValidator<bool, 48> {}
-impl ValidStride for StrideValidator<bool, 64> {}
+// Stride validation types and implementations are now in the stride module
 
 // ============================================================================
 // Ergonomic Builder API
@@ -1990,5 +1922,5 @@ mod sealed {
     impl Sealed for bool {}
 
     // Seal stride validators
-    impl<T, const S: usize> Sealed for StrideValidator<T, S> {}
+    // StrideValidator sealing is now in stride module
 }

crates/cuda_std/src/warp/matrix/stride/bf16.rs

@@ -0,0 +1,7 @@
+//! Stride implementations for bf16 type
+
+use super::{StrideValidator, ValidStride};
+use crate::bf16;
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/bf16_impls.rs"));

crates/cuda_std/src/warp/matrix/stride/bool.rs

@@ -0,0 +1,6 @@
+//! Stride implementations for bool type
+
+use super::{StrideValidator, ValidStride};
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/bool_impls.rs"));

crates/cuda_std/src/warp/matrix/stride/f16.rs

@@ -0,0 +1,7 @@
+//! Stride implementations for f16 type
+
+use super::{StrideValidator, ValidStride};
+use crate::f16;
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/f16_impls.rs"));

crates/cuda_std/src/warp/matrix/stride/f32.rs

@@ -0,0 +1,6 @@
+//! Stride implementations for f32 type
+
+use super::{StrideValidator, ValidStride};
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/f32_impls.rs"));

crates/cuda_std/src/warp/matrix/stride/f64.rs

@@ -0,0 +1,6 @@
+//! Stride implementations for f64 type
+
+use super::{StrideValidator, ValidStride};
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/f64_impls.rs"));

crates/cuda_std/src/warp/matrix/stride/i16.rs

@@ -0,0 +1,6 @@
+//! Stride implementations for i16 type
+
+use super::{StrideValidator, ValidStride};
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/i16_impls.rs"));

crates/cuda_std/src/warp/matrix/stride/i32.rs

@@ -0,0 +1,6 @@
+//! Stride implementations for i32 type
+
+use super::{StrideValidator, ValidStride};
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/i32_impls.rs"));

crates/cuda_std/src/warp/matrix/stride/i64.rs

@@ -0,0 +1,6 @@
+//! Stride implementations for i64 type
+
+use super::{StrideValidator, ValidStride};
+
+// Include the auto-generated implementations
+include!(concat!(env!("OUT_DIR"), "/stride/i64_impls.rs"));