torsh_utils/
cpp_extension.rs

1//! C++ Extension utilities for ToRSh
2//!
3//! This module provides utilities for building C++ extensions that integrate
4//! with the ToRSh framework, similar to PyTorch's cpp_extension module.
5
6use std::collections::HashMap;
7use std::env;
8use std::fs;
9use std::path::{Path, PathBuf};
10use std::process::Command;
11
12/// JIT compilation configuration
13#[derive(Debug, Clone, Default)]
14pub struct JitCompilationConfig {
15    /// Enable just-in-time compilation
16    pub enabled: bool,
17    /// Cache compiled kernels
18    pub cache_enabled: bool,
19    /// Cache directory
20    pub cache_dir: Option<PathBuf>,
21    /// Optimization level for JIT (0-3)
22    pub optimization_level: u8,
23    /// Enable CUDA JIT compilation
24    pub cuda_jit: bool,
25    /// CUDA JIT cache size (in MB)
26    pub cuda_cache_size: usize,
27    /// Maximum number of registers for CUDA kernels
28    pub cuda_max_registers: Option<u32>,
29}
30
31/// Custom operation definition
32#[derive(Debug, Clone)]
33pub struct CustomOpDefinition {
34    /// Operation name
35    pub name: String,
36    /// Operation type (forward, backward, both)
37    pub op_type: CustomOpType,
38    /// Input tensor shapes (None means dynamic)
39    pub input_shapes: Vec<Option<Vec<usize>>>,
40    /// Output tensor shapes (None means dynamic)
41    pub output_shapes: Vec<Option<Vec<usize>>>,
42    /// CPU implementation source
43    pub cpu_source: Option<String>,
44    /// CUDA implementation source
45    pub cuda_source: Option<String>,
46    /// Custom compile flags for this operation
47    pub compile_flags: Vec<String>,
48    /// Operation schema for validation
49    pub schema: OpSchema,
50}
51
52/// Custom operation type
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub enum CustomOpType {
55    Forward,
56    Backward,
57    ForwardBackward,
58}
59
60/// Operation schema for validation and optimization
61#[derive(Debug, Clone, Default)]
62pub struct OpSchema {
63    /// Input tensor types
64    pub input_types: Vec<TensorType>,
65    /// Output tensor types
66    pub output_types: Vec<TensorType>,
67    /// Whether the operation is elementwise
68    pub is_elementwise: bool,
69    /// Whether the operation is deterministic
70    pub is_deterministic: bool,
71    /// Memory requirement estimation
72    pub memory_requirement: MemoryRequirement,
73}
74
75/// Tensor type information
76#[derive(Debug, Clone)]
77pub struct TensorType {
78    /// Data type (f32, f64, i32, etc.)
79    pub dtype: String,
80    /// Minimum number of dimensions
81    pub min_dims: usize,
82    /// Maximum number of dimensions (None means unlimited)
83    pub max_dims: Option<usize>,
84    /// Whether the tensor can be sparse
85    pub supports_sparse: bool,
86}
87
88/// Memory requirement estimation
89#[derive(Debug, Clone, Default)]
90pub enum MemoryRequirement {
91    #[default]
92    Unknown,
93    /// O(1) memory
94    Constant,
95    /// O(n) memory where n is input size
96    Linear,
97    /// O(n²) memory
98    Quadratic,
99    /// Custom memory formula
100    Custom(String),
101}
102
103/// Cross-platform build configuration
104#[derive(Debug, Clone, Default)]
105pub struct CrossPlatformConfig {
106    /// Target platforms to build for
107    pub target_platforms: Vec<TargetPlatform>,
108    /// Windows-specific settings
109    pub windows: WindowsConfig,
110    /// macOS-specific settings
111    pub macos: MacOsConfig,
112    /// Linux-specific settings
113    pub linux: LinuxConfig,
114    /// Enable cross-compilation
115    pub cross_compile: bool,
116    /// Docker-based building
117    pub use_docker: bool,
118}
119
120/// Target platform specification
121#[derive(Debug, Clone, PartialEq, Eq, Hash)]
122pub enum TargetPlatform {
123    WindowsX64,
124    WindowsX86,
125    MacOsX64,
126    MacOsArm64,
127    LinuxX64,
128    LinuxArm64,
129    LinuxAarch64,
130}
131
132/// Windows-specific build configuration
133#[derive(Debug, Clone, Default)]
134pub struct WindowsConfig {
135    /// Visual Studio version to use
136    pub vs_version: Option<String>,
137    /// Windows SDK version
138    pub sdk_version: Option<String>,
139    /// Use clang instead of MSVC
140    pub use_clang: bool,
141    /// Enable Windows-specific optimizations
142    pub enable_simd: bool,
143}
144
145/// macOS-specific build configuration
146#[derive(Debug, Clone, Default)]
147pub struct MacOsConfig {
148    /// Minimum macOS version
149    pub min_version: Option<String>,
150    /// Xcode version to use
151    pub xcode_version: Option<String>,
152    /// Enable Metal Performance Shaders
153    pub enable_mps: bool,
154    /// Universal binary (x64 + ARM64)
155    pub universal_binary: bool,
156}
157
158/// Linux-specific build configuration
159#[derive(Debug, Clone, Default)]
160pub struct LinuxConfig {
161    /// GCC/Clang version preference
162    pub compiler_preference: CompilerPreference,
163    /// Enable Intel MKL
164    pub enable_mkl: bool,
165    /// Enable OpenMP
166    pub enable_openmp: bool,
167    /// Distribution-specific packages
168    pub distro_packages: Vec<String>,
169}
170
171/// Compiler preference on Linux
172#[derive(Debug, Clone, Default, PartialEq, Eq)]
173pub enum CompilerPreference {
174    #[default]
175    Auto,
176    Gcc,
177    Clang,
178    Intel,
179}
180
181/// Configuration for building a C++ extension
182#[derive(Debug, Clone)]
183pub struct CppExtensionConfig {
184    /// Name of the extension module
185    pub name: String,
186    /// Source files to compile
187    pub sources: Vec<PathBuf>,
188    /// Include directories
189    pub include_dirs: Vec<PathBuf>,
190    /// Library directories
191    pub library_dirs: Vec<PathBuf>,
192    /// Libraries to link
193    pub libraries: Vec<String>,
194    /// Extra compiler flags
195    pub extra_compile_args: Vec<String>,
196    /// Extra linker flags
197    pub extra_link_args: Vec<String>,
198    /// Whether to build with CUDA support
199    pub with_cuda: bool,
200    /// CUDA architectures to target
201    pub cuda_archs: Vec<String>,
202    /// Whether to enable debug symbols
203    pub debug: bool,
204    /// Output directory
205    pub build_dir: PathBuf,
206    /// JIT compilation settings
207    pub jit_config: JitCompilationConfig,
208    /// Custom operation definitions
209    pub custom_ops: Vec<CustomOpDefinition>,
210    /// Cross-platform build settings
211    pub cross_platform: CrossPlatformConfig,
212}
213
214impl CppExtensionConfig {
215    /// Create a new C++ extension configuration
216    pub fn new(name: impl Into<String>, sources: Vec<PathBuf>) -> Self {
217        let name = name.into();
218        let build_dir = env::temp_dir().join("torsh_cpp_extensions").join(&name);
219
220        Self {
221            name,
222            sources,
223            include_dirs: vec![],
224            library_dirs: vec![],
225            libraries: vec![],
226            extra_compile_args: vec![],
227            extra_link_args: vec![],
228            with_cuda: false,
229            cuda_archs: vec![
230                "sm_70".to_string(),
231                "sm_75".to_string(),
232                "sm_80".to_string(),
233                "sm_86".to_string(),
234                "sm_89".to_string(),
235            ],
236            debug: false,
237            build_dir,
238            jit_config: JitCompilationConfig::default(),
239            custom_ops: vec![],
240            cross_platform: CrossPlatformConfig::default(),
241        }
242    }
243
244    /// Add include directory
245    pub fn include_dir(mut self, dir: impl AsRef<Path>) -> Self {
246        self.include_dirs.push(dir.as_ref().to_path_buf());
247        self
248    }
249
250    /// Add library directory
251    pub fn library_dir(mut self, dir: impl AsRef<Path>) -> Self {
252        self.library_dirs.push(dir.as_ref().to_path_buf());
253        self
254    }
255
256    /// Add library to link
257    pub fn library(mut self, lib: impl Into<String>) -> Self {
258        self.libraries.push(lib.into());
259        self
260    }
261
262    /// Add extra compile arguments
263    pub fn extra_compile_arg(mut self, arg: impl Into<String>) -> Self {
264        self.extra_compile_args.push(arg.into());
265        self
266    }
267
268    /// Add extra link arguments
269    pub fn extra_link_arg(mut self, arg: impl Into<String>) -> Self {
270        self.extra_link_args.push(arg.into());
271        self
272    }
273
274    /// Enable CUDA support
275    pub fn cuda(mut self, cuda_archs: Vec<String>) -> Self {
276        self.with_cuda = true;
277        self.cuda_archs = cuda_archs;
278        self
279    }
280
281    /// Enable debug symbols
282    pub fn debug(mut self) -> Self {
283        self.debug = true;
284        self
285    }
286
287    /// Set build directory
288    pub fn build_dir(mut self, dir: impl AsRef<Path>) -> Self {
289        self.build_dir = dir.as_ref().to_path_buf();
290        self
291    }
292
293    /// Enable JIT compilation
294    pub fn jit(mut self, config: JitCompilationConfig) -> Self {
295        self.jit_config = config;
296        self
297    }
298
299    /// Add custom operation
300    pub fn custom_op(mut self, op: CustomOpDefinition) -> Self {
301        self.custom_ops.push(op);
302        self
303    }
304
305    /// Set cross-platform build configuration
306    pub fn cross_platform(mut self, config: CrossPlatformConfig) -> Self {
307        self.cross_platform = config;
308        self
309    }
310
311    /// Enable JIT compilation with default settings
312    pub fn enable_jit(mut self) -> Self {
313        self.jit_config.enabled = true;
314        self.jit_config.cache_enabled = true;
315        self.jit_config.optimization_level = 2;
316        self
317    }
318
319    /// Enable CUDA JIT compilation
320    pub fn enable_cuda_jit(mut self) -> Self {
321        self.jit_config.cuda_jit = true;
322        self.jit_config.cuda_cache_size = 256; // 256 MB default
323        self
324    }
325}
326
327/// Build result containing the path to the compiled extension
328#[derive(Debug)]
329pub struct BuildResult {
330    /// Path to the compiled shared library
331    pub library_path: PathBuf,
332    /// Include directories for using the extension
333    pub include_dirs: Vec<PathBuf>,
334    /// JIT compilation results
335    pub jit_info: Option<JitBuildInfo>,
336    /// Custom operations that were compiled
337    pub compiled_ops: Vec<String>,
338    /// Cross-platform build artifacts
339    pub platform_artifacts: HashMap<TargetPlatform, PathBuf>,
340}
341
342/// JIT compilation build information
343#[derive(Debug)]
344pub struct JitBuildInfo {
345    /// JIT cache directory
346    pub cache_dir: PathBuf,
347    /// Number of kernels compiled
348    pub kernel_count: usize,
349    /// CUDA JIT compilation info
350    pub cuda_info: Option<CudaJitInfo>,
351}
352
353/// CUDA JIT compilation information
354#[derive(Debug)]
355pub struct CudaJitInfo {
356    /// PTX cache size in bytes
357    pub ptx_cache_size: usize,
358    /// Number of CUDA kernels
359    pub kernel_count: usize,
360    /// GPU compute capability used
361    pub compute_capability: Vec<String>,
362    /// Runtime compilation cache hits
363    pub cache_hits: usize,
364    /// Runtime compilation cache misses
365    pub cache_misses: usize,
366    /// JIT compilation time in milliseconds
367    pub compilation_time_ms: f64,
368}
369
370/// CUDA device information
371#[derive(Debug, Clone)]
372pub struct CudaDeviceInfo {
373    /// Device index
374    pub device_id: u32,
375    /// Device name
376    pub name: String,
377    /// Compute capability (e.g., "8.0")
378    pub compute_capability: String,
379    /// Total global memory in bytes
380    pub total_memory: usize,
381    /// Maximum threads per block
382    pub max_threads_per_block: u32,
383    /// Maximum grid dimensions
384    pub max_grid_size: [u32; 3],
385    /// Maximum block dimensions
386    pub max_block_size: [u32; 3],
387    /// Warp size
388    pub warp_size: u32,
389    /// Number of multiprocessors
390    pub multiprocessor_count: u32,
391    /// Maximum shared memory per block
392    pub shared_memory_per_block: usize,
393}
394
395/// Advanced CUDA kernel compilation options
396#[derive(Debug, Clone)]
397pub struct CudaKernelCompilationOptions {
398    /// Optimization level (0-3)
399    pub optimization_level: u8,
400    /// Enable fast math operations
401    pub fast_math: bool,
402    /// Maximum register count per thread
403    pub max_registers: Option<u32>,
404    /// Use cache for global memory loads
405    pub use_cache: bool,
406    /// Generate debug information
407    pub debug_info: bool,
408    /// Compile for specific GPU architecture
409    pub target_arch: Option<String>,
410    /// Custom compiler flags
411    pub custom_flags: Vec<String>,
412}
413
414impl Default for CudaKernelCompilationOptions {
415    fn default() -> Self {
416        Self {
417            optimization_level: 2,
418            fast_math: false,
419            max_registers: None,
420            use_cache: true,
421            debug_info: false,
422            target_arch: None,
423            custom_flags: vec![],
424        }
425    }
426}
427
428/// Runtime CUDA kernel management
429#[derive(Debug)]
430pub struct RuntimeCudaKernel {
431    /// Kernel name
432    pub name: String,
433    /// PTX source code
434    pub ptx_source: String,
435    /// Compiled module handle (would be CUmodule in real implementation)
436    pub module_handle: Option<usize>,
437    /// Kernel function handle (would be CUfunction in real implementation)
438    pub function_handle: Option<usize>,
439    /// Compilation options used
440    pub compilation_options: CudaKernelCompilationOptions,
441    /// Grid and block configuration
442    pub launch_config: CudaLaunchConfig,
443}
444
445/// CUDA kernel launch configuration
446#[derive(Debug, Clone)]
447pub struct CudaLaunchConfig {
448    /// Grid dimensions
449    pub grid_size: [u32; 3],
450    /// Block dimensions
451    pub block_size: [u32; 3],
452    /// Shared memory size in bytes
453    pub shared_memory_size: usize,
454    /// CUDA stream handle
455    pub stream: Option<usize>,
456}
457
458/// Build a C++ extension
459pub fn build_cpp_extension(config: &CppExtensionConfig) -> Result<BuildResult, String> {
460    // Create build directory
461    fs::create_dir_all(&config.build_dir)
462        .map_err(|e| format!("Failed to create build directory: {}", e))?;
463
464    // Setup JIT compilation if enabled
465    let jit_info = if config.jit_config.enabled {
466        Some(setup_jit_compilation(config)?)
467    } else {
468        None
469    };
470
471    // Generate custom operation sources
472    let mut generated_sources = vec![];
473    let mut compiled_ops = vec![];
474
475    for custom_op in &config.custom_ops {
476        let generated_source = generate_custom_op_source(custom_op)?;
477        generated_sources.push(generated_source);
478        compiled_ops.push(custom_op.name.clone());
479    }
480
481    // Build for each target platform
482    let mut platform_artifacts = HashMap::new();
483
484    if config.cross_platform.target_platforms.is_empty() {
485        // Build for current platform
486        let artifact = build_for_platform(config, None, &generated_sources, &jit_info)?;
487        platform_artifacts.insert(detect_current_platform(), artifact);
488    } else {
489        // Build for specified platforms
490        for platform in &config.cross_platform.target_platforms {
491            let artifact =
492                build_for_platform(config, Some(platform), &generated_sources, &jit_info)?;
493            platform_artifacts.insert(platform.clone(), artifact);
494        }
495    }
496
497    // Get the main artifact (current platform or first specified)
498    let main_artifact = platform_artifacts
499        .get(&detect_current_platform())
500        .or_else(|| platform_artifacts.values().next())
501        .ok_or("No artifacts built")?
502        .clone();
503
504    Ok(BuildResult {
505        library_path: main_artifact,
506        include_dirs: config.include_dirs.clone(),
507        jit_info,
508        compiled_ops,
509        platform_artifacts,
510    })
511}
512
513/// Setup JIT compilation
514fn setup_jit_compilation(config: &CppExtensionConfig) -> Result<JitBuildInfo, String> {
515    let cache_dir = config
516        .jit_config
517        .cache_dir
518        .clone()
519        .unwrap_or_else(|| config.build_dir.join("jit_cache"));
520
521    fs::create_dir_all(&cache_dir)
522        .map_err(|e| format!("Failed to create JIT cache directory: {}", e))?;
523
524    let cuda_info = if config.jit_config.cuda_jit && config.with_cuda {
525        Some(setup_cuda_jit(config, &cache_dir)?)
526    } else {
527        None
528    };
529
530    Ok(JitBuildInfo {
531        cache_dir,
532        kernel_count: config.custom_ops.len(),
533        cuda_info,
534    })
535}
536
537/// Setup CUDA JIT compilation
538fn setup_cuda_jit(config: &CppExtensionConfig, cache_dir: &Path) -> Result<CudaJitInfo, String> {
539    let cuda_cache_dir = cache_dir.join("cuda");
540    fs::create_dir_all(&cuda_cache_dir)
541        .map_err(|e| format!("Failed to create CUDA cache directory: {}", e))?;
542
543    // Initialize CUDA runtime and query device capabilities
544    let device_info = query_cuda_devices()?;
545    let available_archs = device_info
546        .iter()
547        .map(|dev| format!("sm_{}", dev.compute_capability.replace(".", "")))
548        .collect::<Vec<_>>();
549
550    // Setup PTX cache structure
551    let ptx_cache_dir = cuda_cache_dir.join("ptx");
552    let cubin_cache_dir = cuda_cache_dir.join("cubin");
553    fs::create_dir_all(&ptx_cache_dir)
554        .map_err(|e| format!("Failed to create PTX cache directory: {}", e))?;
555    fs::create_dir_all(&cubin_cache_dir)
556        .map_err(|e| format!("Failed to create CUBIN cache directory: {}", e))?;
557
558    // Configure JIT compilation options
559    configure_cuda_jit_options(config)?;
560
561    // Validate CUDA kernel sources for syntax
562    for op in &config.custom_ops {
563        if let Some(cuda_source) = &op.cuda_source {
564            validate_cuda_kernel_syntax(cuda_source, &op.name)?;
565        }
566    }
567
568    Ok(CudaJitInfo {
569        ptx_cache_size: config.jit_config.cuda_cache_size * 1024 * 1024, // Convert MB to bytes
570        kernel_count: config
571            .custom_ops
572            .iter()
573            .filter(|op| op.cuda_source.is_some())
574            .count(),
575        compute_capability: available_archs,
576        cache_hits: 0,
577        cache_misses: 0,
578        compilation_time_ms: 0.0,
579    })
580}
581
582/// Generate custom operation source code
583fn generate_custom_op_source(op: &CustomOpDefinition) -> Result<PathBuf, String> {
584    // Generate C++ source code for the custom operation
585    let source_content = match &op.op_type {
586        CustomOpType::Forward => generate_forward_op(&op.name, &op.cpu_source, &op.cuda_source)?,
587        CustomOpType::Backward => generate_backward_op(&op.name, &op.cpu_source, &op.cuda_source)?,
588        CustomOpType::ForwardBackward => {
589            generate_forward_backward_op(&op.name, &op.cpu_source, &op.cuda_source)?
590        }
591    };
592
593    // Write to temporary file
594    let temp_file = env::temp_dir().join(format!("{}_custom_op.cpp", op.name));
595    fs::write(&temp_file, source_content)
596        .map_err(|e| format!("Failed to write custom op source: {}", e))?;
597
598    Ok(temp_file)
599}
600
601/// Generate forward operation source
602fn generate_forward_op(
603    name: &str,
604    cpu_source: &Option<String>,
605    cuda_source: &Option<String>,
606) -> Result<String, String> {
607    let mut source = format!(
608        r#"// Generated custom operation: {}
609#include <torsh/tensor.h>
610#include <torsh/autograd.h>
611
612namespace torsh {{
613namespace ops {{
614
615"#,
616        name
617    );
618
619    // Add CPU implementation
620    if let Some(cpu_impl) = cpu_source {
621        source.push_str(&format!(
622            r#"
623// CPU implementation
624Tensor {}_cpu_forward(const std::vector<Tensor>& inputs) {{
625    {}
626}}
627"#,
628            name, cpu_impl
629        ));
630    }
631
632    // Add CUDA implementation
633    if let Some(cuda_impl) = cuda_source {
634        source.push_str(&format!(
635            r#"
636#ifdef TORSH_USE_CUDA
637// CUDA implementation
638Tensor {}_cuda_forward(const std::vector<Tensor>& inputs) {{
639    {}
640}}
641#endif
642"#,
643            name, cuda_impl
644        ));
645    }
646
647    // Add dispatcher
648    source.push_str(&format!(
649        r#"
650// Operation dispatcher
651Tensor {}_forward(const std::vector<Tensor>& inputs) {{
652#ifdef TORSH_USE_CUDA
653    if (inputs[0].is_cuda()) {{
654        return {}_cuda_forward(inputs);
655    }}
656#endif
657    return {}_cpu_forward(inputs);
658}}
659
660// Register operation
661TORSH_REGISTER_OP("{}", {}_forward);
662
663}} // namespace ops
664}} // namespace torsh
665"#,
666        name, name, name, name, name
667    ));
668
669    Ok(source)
670}
671
672/// Generate backward operation source
673fn generate_backward_op(
674    name: &str,
675    cpu_source: &Option<String>,
676    cuda_source: &Option<String>,
677) -> Result<String, String> {
678    // Similar structure to forward op but for backward pass
679    let mut source = format!(
680        r#"// Generated custom backward operation: {}
681#include <torsh/tensor.h>
682#include <torsh/autograd.h>
683
684namespace torsh {{
685namespace ops {{
686"#,
687        name
688    );
689
690    if let Some(cpu_impl) = cpu_source {
691        source.push_str(&format!(
692            r#"
693std::vector<Tensor> {}_cpu_backward(const std::vector<Tensor>& grad_outputs, const std::vector<Tensor>& inputs) {{
694    {}
695}}
696"#,
697            name, cpu_impl
698        ));
699    }
700
701    if let Some(cuda_impl) = cuda_source {
702        source.push_str(&format!(
703            r#"
704#ifdef TORSH_USE_CUDA
705std::vector<Tensor> {}_cuda_backward(const std::vector<Tensor>& grad_outputs, const std::vector<Tensor>& inputs) {{
706    {}
707}}
708#endif
709"#,
710            name, cuda_impl
711        ));
712    }
713
714    source.push_str(&format!(
715        r#"
716std::vector<Tensor> {}_backward(const std::vector<Tensor>& grad_outputs, const std::vector<Tensor>& inputs) {{
717#ifdef TORSH_USE_CUDA
718    if (inputs[0].is_cuda()) {{
719        return {}_cuda_backward(grad_outputs, inputs);
720    }}
721#endif
722    return {}_cpu_backward(grad_outputs, inputs);
723}}
724
725TORSH_REGISTER_BACKWARD_OP("{}", {}_backward);
726
727}} // namespace ops
728}} // namespace torsh
729"#,
730        name, name, name, name, name
731    ));
732
733    Ok(source)
734}
735
736/// Generate forward and backward operation source
737fn generate_forward_backward_op(
738    name: &str,
739    cpu_source: &Option<String>,
740    cuda_source: &Option<String>,
741) -> Result<String, String> {
742    // Combine forward and backward generation
743    let forward_source = generate_forward_op(name, cpu_source, cuda_source)?;
744    let backward_source =
745        generate_backward_op(&format!("{}_backward", name), cpu_source, cuda_source)?;
746
747    Ok(format!("{}\n\n{}", forward_source, backward_source))
748}
749
750/// Build for a specific platform
751fn build_for_platform(
752    config: &CppExtensionConfig,
753    target_platform: Option<&TargetPlatform>,
754    generated_sources: &[PathBuf],
755    _jit_info: &Option<JitBuildInfo>,
756) -> Result<PathBuf, String> {
757    // Determine compiler based on platform
758    let (compiler, extra_flags) =
759        match target_platform {
760            Some(TargetPlatform::WindowsX64) | Some(TargetPlatform::WindowsX86) => {
761                if config.cross_platform.windows.use_clang {
762                    (
763                        "clang++".to_string(),
764                        vec![
765                            "-target".to_string(),
766                            get_windows_target(target_platform.expect(
767                                "target_platform should be Some for Windows platform branch",
768                            )),
769                        ],
770                    )
771                } else {
772                    ("cl.exe".to_string(), vec!["/std:c++17".to_string()])
773                }
774            }
775            Some(TargetPlatform::MacOsX64) | Some(TargetPlatform::MacOsArm64) => {
776                let target = match target_platform
777                    .expect("target_platform should be Some for macOS platform branch")
778                {
779                    TargetPlatform::MacOsX64 => "x86_64-apple-darwin",
780                    TargetPlatform::MacOsArm64 => "arm64-apple-darwin",
781                    _ => unreachable!(),
782                };
783                (
784                    "clang++".to_string(),
785                    vec!["-target".to_string(), target.to_string()],
786                )
787            }
788            Some(TargetPlatform::LinuxX64)
789            | Some(TargetPlatform::LinuxArm64)
790            | Some(TargetPlatform::LinuxAarch64) => {
791                match config.cross_platform.linux.compiler_preference {
792                    CompilerPreference::Clang => ("clang++".to_string(), vec![]),
793                    CompilerPreference::Gcc => ("g++".to_string(), vec![]),
794                    CompilerPreference::Intel => ("icpc".to_string(), vec![]),
795                    CompilerPreference::Auto => (
796                        env::var("CXX").unwrap_or_else(|_| "g++".to_string()),
797                        vec![],
798                    ),
799                }
800            }
801            None => {
802                // Current platform
803                if config.with_cuda {
804                    ("nvcc".to_string(), vec![])
805                } else {
806                    (
807                        env::var("CXX").unwrap_or_else(|_| "c++".to_string()),
808                        vec![],
809                    )
810                }
811            }
812        };
813
814    // Build compile command
815    let mut cmd = Command::new(&compiler);
816
817    // Add platform-specific flags
818    cmd.args(&extra_flags);
819
820    // Add include directories
821    for include_dir in &config.include_dirs {
822        cmd.arg(format!("-I{}", include_dir.display()));
823    }
824
825    // Add ToRSh include directory
826    if let Ok(torsh_include) = env::var("TORSH_INCLUDE_DIR") {
827        cmd.arg(format!("-I{}", torsh_include));
828    }
829
830    // Add standard flags (platform-specific)
831    if compiler.contains("cl.exe") {
832        // MSVC flags
833        cmd.arg("/std:c++17");
834        if !config.debug {
835            cmd.arg("/O2");
836            cmd.arg("/DNDEBUG");
837        } else {
838            cmd.arg("/Od");
839            cmd.arg("/Zi");
840        }
841    } else {
842        // GCC/Clang flags
843        cmd.arg("-std=c++17");
844        cmd.arg("-fPIC");
845        if !config.debug {
846            cmd.arg("-O3");
847            cmd.arg("-DNDEBUG");
848        } else {
849            cmd.arg("-g");
850            cmd.arg("-O0");
851        }
852    }
853
854    // Add CUDA specific flags
855    if config.with_cuda && compiler.contains("nvcc") {
856        for arch in &config.cuda_archs {
857            cmd.arg(format!(
858                "-gencode=arch=compute_{},code={}",
859                &arch[3..],
860                arch
861            ));
862        }
863        cmd.arg("-x").arg("cu");
864    }
865
866    // Add extra compile args
867    for arg in &config.extra_compile_args {
868        cmd.arg(arg);
869    }
870
871    // Add source files (original + generated)
872    for source in &config.sources {
873        cmd.arg(source);
874    }
875    for source in generated_sources {
876        cmd.arg(source);
877    }
878
879    // Output file
880    let platform_suffix = target_platform
881        .map(|p| format!("_{:?}", p))
882        .unwrap_or_default();
883    let output_file = config
884        .build_dir
885        .join(format!("lib{}{}.so", config.name, platform_suffix));
886
887    if compiler.contains("cl.exe") {
888        cmd.arg("/Fe:").arg(&output_file);
889        cmd.arg("/LD"); // Create DLL
890    } else {
891        cmd.arg("-shared");
892        cmd.arg("-o").arg(&output_file);
893    }
894
895    // Add library directories
896    for lib_dir in &config.library_dirs {
897        if compiler.contains("cl.exe") {
898            cmd.arg(format!("/LIBPATH:{}", lib_dir.display()));
899        } else {
900            cmd.arg(format!("-L{}", lib_dir.display()));
901        }
902    }
903
904    // Add libraries
905    for lib in &config.libraries {
906        if compiler.contains("cl.exe") {
907            cmd.arg(format!("{}.lib", lib));
908        } else {
909            cmd.arg(format!("-l{}", lib));
910        }
911    }
912
913    // Add extra link args
914    for arg in &config.extra_link_args {
915        cmd.arg(arg);
916    }
917
918    // Execute build
919    let output = cmd
920        .output()
921        .map_err(|e| format!("Failed to execute compiler {}: {}", compiler, e))?;
922
923    if !output.status.success() {
924        let stderr = String::from_utf8_lossy(&output.stderr);
925        return Err(format!(
926            "Compilation failed for platform {:?}:\n{}",
927            target_platform, stderr
928        ));
929    }
930
931    Ok(output_file)
932}
933
934/// Detect current platform
935fn detect_current_platform() -> TargetPlatform {
936    match env::consts::OS {
937        "windows" => match env::consts::ARCH {
938            "x86_64" => TargetPlatform::WindowsX64,
939            "x86" => TargetPlatform::WindowsX86,
940            _ => TargetPlatform::WindowsX64, // Default
941        },
942        "macos" => match env::consts::ARCH {
943            "aarch64" => TargetPlatform::MacOsArm64,
944            _ => TargetPlatform::MacOsX64,
945        },
946        "linux" => match env::consts::ARCH {
947            "aarch64" => TargetPlatform::LinuxAarch64,
948            "arm64" => TargetPlatform::LinuxArm64,
949            _ => TargetPlatform::LinuxX64,
950        },
951        _ => TargetPlatform::LinuxX64, // Default fallback
952    }
953}
954
955/// Get Windows target string
956fn get_windows_target(platform: &TargetPlatform) -> String {
957    match platform {
958        TargetPlatform::WindowsX64 => "x86_64-pc-windows-msvc".to_string(),
959        TargetPlatform::WindowsX86 => "i686-pc-windows-msvc".to_string(),
960        _ => "x86_64-pc-windows-msvc".to_string(), // Default
961    }
962}
963
964/// Load a C++ extension from a shared library
965pub fn load_cpp_extension(library_path: &Path) -> Result<(), String> {
966    // This would typically use libloading or similar to dynamically load the library
967    // For now, we just verify the file exists
968    if !library_path.exists() {
969        return Err(format!("Library not found: {}", library_path.display()));
970    }
971
972    // In a real implementation, we would:
973    // 1. Load the shared library
974    // 2. Register any custom operators
975    // 3. Initialize any global state
976
977    Ok(())
978}
979
980/// Generate a simple C++ extension template
981pub fn generate_extension_template(name: &str, output_dir: &Path) -> Result<(), String> {
982    fs::create_dir_all(output_dir)
983        .map_err(|e| format!("Failed to create output directory: {}", e))?;
984
985    // Generate header file
986    let header_content = format!(
987        r#"#pragma once
988
989#include <torsh/tensor.h>
990#include <torsh/module.h>
991
992namespace torsh {{
993namespace ops {{
994
995// Example custom operation
996Tensor {}_forward(const Tensor& input);
997
998}} // namespace ops
999}} // namespace torsh
1000"#,
1001        name
1002    );
1003
1004    let header_path = output_dir.join(format!("{}.h", name));
1005    fs::write(&header_path, header_content)
1006        .map_err(|e| format!("Failed to write header file: {}", e))?;
1007
1008    // Generate source file
1009    let source_content = format!(
1010        r#"#include "{}.h"
1011#include <torsh/autograd.h>
1012#include <iostream>
1013
1014namespace torsh {{
1015namespace ops {{
1016
1017Tensor {}_forward(const Tensor& input) {{
1018    // Example implementation
1019    auto output = input.clone();
1020    
1021    // Perform custom operation
1022    // This is where you would implement your custom logic
1023    
1024    return output;
1025}}
1026
1027// Register the operation
1028TORSH_LIBRARY(TORCH_EXTENSION_NAME, m) {{
1029    m.def("{}_forward", &{}_forward);
1030}}
1031
1032}} // namespace ops
1033}} // namespace torsh
1034"#,
1035        name, name, name, name
1036    );
1037
1038    let source_path = output_dir.join(format!("{}.cpp", name));
1039    fs::write(&source_path, source_content)
1040        .map_err(|e| format!("Failed to write source file: {}", e))?;
1041
1042    // Generate setup script
1043    let setup_content = format!(
1044        r#"use torsh_utils::cpp_extension::{{CppExtensionConfig, build_cpp_extension}};
1045use std::path::PathBuf;
1046
1047fn main() {{
1048    let config = CppExtensionConfig::new("{}", vec![
1049        PathBuf::from("{}.cpp"),
1050    ])
1051    .include_dir(".")
1052    .extra_compile_arg("-Wall")
1053    .extra_compile_arg("-Wextra");
1054
1055    match build_cpp_extension(&config) {{
1056        Ok(result) => {{
1057            println!("Extension built successfully!");
1058            println!("Library: {{:?}}", result.library_path);
1059        }}
1060        Err(e) => {{
1061            eprintln!("Build failed: {{}}", e);
1062            std::process::exit(1);
1063        }}
1064    }}
1065}}
1066"#,
1067        name, name
1068    );
1069
1070    let setup_path = output_dir.join("build.rs");
1071    fs::write(&setup_path, setup_content)
1072        .map_err(|e| format!("Failed to write setup script: {}", e))?;
1073
1074    Ok(())
1075}
1076
1077/// Check if CUDA is available for building extensions
1078pub fn cuda_is_available() -> bool {
1079    Command::new("nvcc")
1080        .arg("--version")
1081        .output()
1082        .map(|output| output.status.success())
1083        .unwrap_or(false)
1084}
1085
1086/// Get CUDA architectures available on the system.
1087///
1088/// Queries `nvidia-smi` to enumerate the compute capabilities of GPUs actually present.
1089/// Returns an empty `Vec` when `nvidia-smi` is unavailable or fails — never fabricates
1090/// a list of architectures.
1091pub fn get_cuda_arch_list() -> Vec<String> {
1092    let output = Command::new("nvidia-smi")
1093        .args(["--query-gpu=compute_cap", "--format=csv,noheader,nounits"])
1094        .output();
1095
1096    match output {
1097        Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
1098            .lines()
1099            .filter_map(|line| {
1100                let cap = line.trim().replace('.', "_");
1101                if cap.is_empty() {
1102                    None
1103                } else {
1104                    Some(format!("sm_{}", cap))
1105                }
1106            })
1107            .collect(),
1108        // nvidia-smi unavailable or failed — honest empty list
1109        _ => vec![],
1110    }
1111}
1112
1113/// Query CUDA devices on the system.
1114///
1115/// Attempts to enumerate real GPUs via `nvidia-smi`. Returns an empty `Vec` when the CUDA
1116/// toolchain is absent or `nvidia-smi` cannot be executed — never fabricates device info.
1117///
1118/// Fields that require the CUDA driver API to query (`multiprocessor_count`,
1119/// `shared_memory_per_block`) are set to `0` (honest "unknown") rather than invented values.
1120/// `max_threads_per_block`, `max_grid_size`, and `max_block_size` are set to the CUDA
1121/// specification minimum guaranteed values, which are safe conservative lower bounds.
1122fn query_cuda_devices() -> Result<Vec<CudaDeviceInfo>, String> {
1123    if !cuda_is_available() {
1124        // CUDA toolchain not detected — return empty list (no devices confirmed)
1125        return Ok(vec![]);
1126    }
1127
1128    // Try nvidia-smi to get the real device list
1129    let output = Command::new("nvidia-smi")
1130        .args([
1131            "--query-gpu=index,name,memory.total,compute_cap",
1132            "--format=csv,noheader,nounits",
1133        ])
1134        .output();
1135
1136    let output = match output {
1137        Ok(o) if o.status.success() => o,
1138        // nvidia-smi not available or failed — cannot enumerate devices
1139        _ => return Ok(vec![]),
1140    };
1141
1142    let stdout = String::from_utf8_lossy(&output.stdout);
1143    let mut devices = Vec::new();
1144
1145    for line in stdout.lines() {
1146        let line = line.trim();
1147        if line.is_empty() {
1148            continue;
1149        }
1150        // Format: "index, name, memory_MiB, compute_cap"
1151        let parts: Vec<&str> = line.splitn(4, ',').map(str::trim).collect();
1152        if parts.len() < 4 {
1153            continue;
1154        }
1155        let device_id: u32 = parts[0].parse().unwrap_or(0);
1156        let name = parts[1].to_string();
1157        let total_memory_mib: u64 = parts[2].parse().unwrap_or(0);
1158        let compute_capability = parts[3].to_string();
1159
1160        devices.push(CudaDeviceInfo {
1161            device_id,
1162            name,
1163            compute_capability,
1164            total_memory: (total_memory_mib * 1024 * 1024) as usize, // MiB → bytes
1165            // CUDA spec minimum guaranteed values (safe conservative lower bounds, not invented)
1166            max_threads_per_block: 1024,
1167            max_grid_size: [65535, 65535, 65535],
1168            max_block_size: [1024, 1024, 64],
1169            warp_size: 32,
1170            // Unknown without CUDA driver API — honest zero
1171            multiprocessor_count: 0,
1172            shared_memory_per_block: 0,
1173        });
1174    }
1175
1176    Ok(devices)
1177}
1178
1179/// Configure CUDA JIT compilation options
1180fn configure_cuda_jit_options(config: &CppExtensionConfig) -> Result<(), String> {
1181    // In a real implementation, this would configure CUDA driver JIT options
1182    // Such as:
1183    // - cuLinkCreate with JIT options
1184    // - Setting optimization level
1185    // - Configuring cache behavior
1186    // - Setting debug/profiling options
1187
1188    if config.jit_config.cuda_jit {
1189        // Validate JIT configuration
1190        if config.jit_config.cuda_cache_size == 0 {
1191            return Err("CUDA JIT cache size must be greater than 0".to_string());
1192        }
1193
1194        if config.jit_config.optimization_level > 3 {
1195            return Err("CUDA JIT optimization level must be 0-3".to_string());
1196        }
1197
1198        // Configure JIT options based on config
1199        // This is where we would set:
1200        // - CU_JIT_OPTIMIZATION_LEVEL
1201        // - CU_JIT_CACHE_MODE
1202        // - CU_JIT_MAX_REGISTERS
1203        // - CU_JIT_THREADS_PER_BLOCK
1204    }
1205
1206    Ok(())
1207}
1208
1209/// Validate CUDA kernel syntax
1210fn validate_cuda_kernel_syntax(cuda_source: &str, op_name: &str) -> Result<(), String> {
1211    // Basic syntax validation for CUDA kernel source
1212    let required_patterns = [
1213        "__global__", // Kernel function marker
1214        "__device__", // Or device function marker
1215        "__host__",   // Or host function marker
1216    ];
1217
1218    // Check if at least one CUDA pattern is present
1219    let has_cuda_pattern = required_patterns
1220        .iter()
1221        .any(|pattern| cuda_source.contains(pattern));
1222
1223    if !has_cuda_pattern {
1224        return Err(format!(
1225            "CUDA source for operation '{}' does not contain valid CUDA kernel markers (__global__, __device__, or __host__)",
1226            op_name
1227        ));
1228    }
1229
1230    // Check for common syntax errors
1231    let brackets_open = cuda_source.chars().filter(|&c| c == '{').count();
1232    let brackets_close = cuda_source.chars().filter(|&c| c == '}').count();
1233
1234    if brackets_open != brackets_close {
1235        return Err(format!(
1236            "CUDA source for operation '{}' has mismatched braces ({{ and }})",
1237            op_name
1238        ));
1239    }
1240
1241    // Check for semicolon at end of statements (basic check)
1242    let lines: Vec<&str> = cuda_source.lines().collect();
1243    for (i, line) in lines.iter().enumerate() {
1244        let trimmed = line.trim();
1245        if !trimmed.is_empty()
1246            && !trimmed.starts_with("//")
1247            && !trimmed.starts_with("/*")
1248            && !trimmed.ends_with('{')
1249            && !trimmed.ends_with('}')
1250            && !trimmed.ends_with(';')
1251            && !trimmed.starts_with('#')
1252        {
1253            return Err(format!(
1254                "CUDA source for operation '{}' line {} may be missing semicolon: '{}'",
1255                op_name,
1256                i + 1,
1257                trimmed
1258            ));
1259        }
1260    }
1261
1262    Ok(())
1263}
1264
1265/// Compile CUDA kernel at runtime
1266pub fn compile_cuda_kernel_runtime(
1267    kernel_source: &str,
1268    kernel_name: &str,
1269    options: &CudaKernelCompilationOptions,
1270) -> Result<RuntimeCudaKernel, String> {
1271    // Validate CUDA availability
1272    if !cuda_is_available() {
1273        return Err("CUDA is not available for runtime compilation".to_string());
1274    }
1275
1276    // Validate kernel source
1277    validate_cuda_kernel_syntax(kernel_source, kernel_name)?;
1278
1279    // In a real implementation, this would:
1280    // 1. Use CUDA Driver API to compile PTX from source
1281    // 2. Load the compiled module
1282    // 3. Get kernel function handle
1283    // 4. Configure launch parameters
1284
1285    // Generate PTX source (mock)
1286    let ptx_source = format!(
1287        r#"
1288.version 8.0
1289.target sm_80
1290.address_size 64
1291
1292.visible .entry {}(
1293    .param .u64 param_0
1294)
1295{{
1296    // Generated PTX code would go here
1297    ret;
1298}}
1299"#,
1300        kernel_name
1301    );
1302
1303    // Mock launch configuration
1304    let launch_config = CudaLaunchConfig {
1305        grid_size: [1, 1, 1],
1306        block_size: [256, 1, 1],
1307        shared_memory_size: 0,
1308        stream: None,
1309    };
1310
1311    Ok(RuntimeCudaKernel {
1312        name: kernel_name.to_string(),
1313        ptx_source,
1314        module_handle: Some(1),   // Mock handle
1315        function_handle: Some(1), // Mock handle
1316        compilation_options: options.clone(),
1317        launch_config,
1318    })
1319}
1320
1321/// Launch a runtime-compiled CUDA kernel
1322pub fn launch_cuda_kernel(
1323    kernel: &RuntimeCudaKernel,
1324    args: &[*mut std::ffi::c_void],
1325) -> Result<(), String> {
1326    // In a real implementation, this would:
1327    // 1. Validate kernel is loaded
1328    // 2. Set kernel parameters
1329    // 3. Launch kernel with configured grid/block dimensions
1330    // 4. Handle synchronization if needed
1331
1332    if kernel.module_handle.is_none() || kernel.function_handle.is_none() {
1333        return Err(format!("Kernel '{}' is not properly loaded", kernel.name));
1334    }
1335
1336    // Validate launch configuration
1337    if kernel.launch_config.grid_size[0] == 0 || kernel.launch_config.block_size[0] == 0 {
1338        return Err(format!(
1339            "Invalid launch configuration for kernel '{}'",
1340            kernel.name
1341        ));
1342    }
1343
1344    // Mock kernel launch validation
1345    println!(
1346        "Launching CUDA kernel '{}' with grid {:?} and block {:?}",
1347        kernel.name, kernel.launch_config.grid_size, kernel.launch_config.block_size
1348    );
1349
1350    // Validate argument count (basic check)
1351    if args.is_empty() {
1352        return Err(format!(
1353            "No arguments provided for kernel '{}'",
1354            kernel.name
1355        ));
1356    }
1357
1358    Ok(())
1359}
1360
1361/// Auto-tune CUDA kernel launch parameters
1362pub fn auto_tune_cuda_kernel(
1363    kernel: &mut RuntimeCudaKernel,
1364    input_sizes: &[usize],
1365) -> Result<CudaLaunchConfig, String> {
1366    // Query device properties for optimal configuration
1367    let devices = query_cuda_devices()?;
1368    let device = devices
1369        .first()
1370        .ok_or("No CUDA devices available for auto-tuning")?;
1371
1372    // Calculate optimal block size based on kernel complexity and device properties
1373    let optimal_block_size = if input_sizes.iter().any(|&size| size > 10000) {
1374        // Large inputs: use larger blocks for better memory coalescing
1375        device.max_threads_per_block.min(512)
1376    } else {
1377        // Small inputs: use smaller blocks to avoid warp underutilization
1378        device.max_threads_per_block.min(256)
1379    };
1380
1381    // Calculate grid size based on input size and block size
1382    let total_elements = input_sizes.iter().max().copied().unwrap_or(1);
1383    let optimal_grid_size =
1384        (total_elements + optimal_block_size as usize - 1) / optimal_block_size as usize;
1385
1386    // Limit grid size to device maximum
1387    let clamped_grid_size = (optimal_grid_size as u32).min(device.max_grid_size[0]);
1388
1389    let optimized_config = CudaLaunchConfig {
1390        grid_size: [clamped_grid_size, 1, 1],
1391        block_size: [optimal_block_size, 1, 1],
1392        shared_memory_size: 0, // Auto-tune shared memory based on kernel requirements
1393        stream: kernel.launch_config.stream,
1394    };
1395
1396    // Update kernel configuration
1397    kernel.launch_config = optimized_config.clone();
1398
1399    Ok(optimized_config)
1400}
1401
1402#[cfg(test)]
1403mod tests {
1404    use super::*;
1405    use std::env;
1406
1407    #[test]
1408    fn test_cpp_extension_config() {
1409        let config = CppExtensionConfig::new("test_ext", vec![PathBuf::from("test.cpp")])
1410            .include_dir("/usr/include")
1411            .library("torsh")
1412            .extra_compile_arg("-std=c++17");
1413
1414        assert_eq!(config.name, "test_ext");
1415        assert_eq!(config.sources.len(), 1);
1416        assert_eq!(config.include_dirs.len(), 1);
1417        assert_eq!(config.libraries.len(), 1);
1418    }
1419
1420    #[test]
1421    fn test_generate_template() {
1422        let temp_dir = env::temp_dir().join("torsh_test_template");
1423        let result = generate_extension_template("test_op", &temp_dir);
1424
1425        assert!(result.is_ok());
1426        assert!(temp_dir.join("test_op.h").exists());
1427        assert!(temp_dir.join("test_op.cpp").exists());
1428        assert!(temp_dir.join("build.rs").exists());
1429
1430        // Cleanup
1431        let _ = fs::remove_dir_all(temp_dir);
1432    }
1433
1434    #[test]
1435    fn test_cuda_detection() {
1436        // cuda_is_available() checks for nvcc in PATH — passes or fails depending on system
1437        let available = cuda_is_available();
1438        println!("CUDA available: {}", available);
1439
1440        // get_cuda_arch_list queries nvidia-smi; may return empty even when nvcc is present
1441        let archs = get_cuda_arch_list();
1442        println!("CUDA arch list: {:?}", archs);
1443        // All returned entries must be well-formed sm_XX strings
1444        for arch in &archs {
1445            assert!(
1446                arch.starts_with("sm_"),
1447                "arch should start with sm_, got: {}",
1448                arch
1449            );
1450        }
1451    }
1452
1453    #[test]
1454    fn test_query_cuda_devices_no_panic() {
1455        // Should not panic even when CUDA is absent
1456        let result = query_cuda_devices();
1457        assert!(
1458            result.is_ok(),
1459            "query_cuda_devices should return Ok on all platforms"
1460        );
1461        let devices = result.unwrap();
1462        // Entries sourced from nvidia-smi: names must be non-empty, memory reasonable
1463        for dev in &devices {
1464            assert!(!dev.name.is_empty(), "device name should not be empty");
1465            assert!(
1466                !dev.compute_capability.is_empty(),
1467                "compute_capability should not be empty"
1468            );
1469        }
1470    }
1471
1472    #[test]
1473    fn test_get_cuda_arch_list_no_panic() {
1474        // Should not panic even when nvidia-smi is absent; just returns empty vec
1475        let archs = get_cuda_arch_list();
1476        // If non-empty, all entries should start with "sm_"
1477        for arch in &archs {
1478            assert!(
1479                arch.starts_with("sm_"),
1480                "arch should start with sm_, got: {}",
1481                arch
1482            );
1483        }
1484    }
1485
1486    #[test]
1487    fn test_cuda_kernel_compilation_options() {
1488        let default_options = CudaKernelCompilationOptions::default();
1489        assert_eq!(default_options.optimization_level, 2);
1490        assert!(!default_options.fast_math);
1491        assert!(default_options.use_cache);
1492        assert!(!default_options.debug_info);
1493
1494        let custom_options = CudaKernelCompilationOptions {
1495            optimization_level: 3,
1496            fast_math: true,
1497            max_registers: Some(64),
1498            debug_info: true,
1499            target_arch: Some("sm_80".to_string()),
1500            ..Default::default()
1501        };
1502
1503        assert_eq!(custom_options.optimization_level, 3);
1504        assert!(custom_options.fast_math);
1505        assert_eq!(custom_options.max_registers, Some(64));
1506        assert!(custom_options.debug_info);
1507    }
1508
1509    #[test]
1510    fn test_cuda_kernel_syntax_validation() {
1511        // Valid CUDA kernel
1512        let valid_kernel = r#"
1513        __global__ void test_kernel(float* input, float* output) {
1514            int idx = blockIdx.x * blockDim.x + threadIdx.x;
1515            output[idx] = input[idx] * 2.0f;
1516        }
1517        "#;
1518
1519        assert!(validate_cuda_kernel_syntax(valid_kernel, "test_kernel").is_ok());
1520
1521        // Invalid kernel (missing __global__)
1522        let invalid_kernel = r#"
1523        void test_kernel(float* input, float* output) {
1524            int idx = blockIdx.x * blockDim.x + threadIdx.x;
1525            output[idx] = input[idx] * 2.0f;
1526        }
1527        "#;
1528
1529        assert!(validate_cuda_kernel_syntax(invalid_kernel, "test_kernel").is_err());
1530
1531        // Invalid kernel (mismatched braces)
1532        let invalid_braces = r#"
1533        __global__ void test_kernel(float* input, float* output) {
1534            int idx = blockIdx.x * blockDim.x + threadIdx.x;
1535            output[idx] = input[idx] * 2.0f;
1536        // Missing closing brace
1537        "#;
1538
1539        assert!(validate_cuda_kernel_syntax(invalid_braces, "test_kernel").is_err());
1540    }
1541
1542    #[test]
1543    fn test_runtime_cuda_kernel_compilation() {
1544        let kernel_source = r#"
1545        __global__ void vector_add(float* a, float* b, float* c, int n) {
1546            int idx = blockIdx.x * blockDim.x + threadIdx.x;
1547            if (idx < n) {
1548                c[idx] = a[idx] + b[idx];
1549            }
1550        }
1551        "#;
1552
1553        let options = CudaKernelCompilationOptions::default();
1554
1555        // This should work even without CUDA (returns mock result)
1556        if cuda_is_available() {
1557            let result = compile_cuda_kernel_runtime(kernel_source, "vector_add", &options);
1558            if let Ok(kernel) = result {
1559                assert_eq!(kernel.name, "vector_add");
1560                assert!(!kernel.ptx_source.is_empty());
1561                assert!(kernel.module_handle.is_some());
1562                assert!(kernel.function_handle.is_some());
1563            }
1564        }
1565    }
1566
1567    #[test]
1568    fn test_cuda_launch_config_auto_tuning() {
1569        if cuda_is_available() {
1570            let kernel_source = r#"
1571            __global__ void simple_kernel(float* data) {
1572                int idx = blockIdx.x * blockDim.x + threadIdx.x;
1573                data[idx] *= 2.0f;
1574            }
1575            "#;
1576
1577            let options = CudaKernelCompilationOptions::default();
1578            let result = compile_cuda_kernel_runtime(kernel_source, "simple_kernel", &options);
1579
1580            if let Ok(mut kernel) = result {
1581                let input_sizes = vec![1024, 2048, 4096];
1582                let tuned_config = auto_tune_cuda_kernel(&mut kernel, &input_sizes);
1583
1584                if let Ok(config) = tuned_config {
1585                    assert!(config.grid_size[0] > 0);
1586                    assert!(config.block_size[0] > 0);
1587                    assert!(config.block_size[0] <= 1024); // Max threads per block
1588                }
1589            }
1590        }
1591    }
1592
1593    #[test]
1594    fn test_custom_op_with_cuda_jit() {
1595        let custom_op = CustomOpDefinition {
1596            name: "custom_relu".to_string(),
1597            op_type: CustomOpType::Forward,
1598            input_shapes: vec![None], // Dynamic shape
1599            output_shapes: vec![None],
1600            cpu_source: Some("return torch::relu(inputs[0]);".to_string()),
1601            cuda_source: Some(
1602                r#"
1603            __global__ void relu_kernel(float* input, float* output, int size) {
1604                int idx = blockIdx.x * blockDim.x + threadIdx.x;
1605                if (idx < size) {
1606                    output[idx] = fmaxf(0.0f, input[idx]);
1607                }
1608            }
1609            "#
1610                .to_string(),
1611            ),
1612            compile_flags: vec!["-O3".to_string()],
1613            schema: OpSchema {
1614                input_types: vec![TensorType {
1615                    dtype: "float32".to_string(),
1616                    min_dims: 1,
1617                    max_dims: None,
1618                    supports_sparse: false,
1619                }],
1620                output_types: vec![TensorType {
1621                    dtype: "float32".to_string(),
1622                    min_dims: 1,
1623                    max_dims: None,
1624                    supports_sparse: false,
1625                }],
1626                is_elementwise: true,
1627                is_deterministic: true,
1628                memory_requirement: MemoryRequirement::Linear,
1629            },
1630        };
1631
1632        let config = CppExtensionConfig::new("custom_relu_ext", vec![])
1633            .enable_cuda_jit()
1634            .custom_op(custom_op);
1635
1636        assert!(config.jit_config.cuda_jit);
1637        assert_eq!(config.custom_ops.len(), 1);
1638        assert_eq!(config.custom_ops[0].name, "custom_relu");
1639    }
1640}
torsh_utils/cpp_extension.rs

torsh_utils/
cpp_extension.rs
