rustytorch/crates/core/rtx-runtime/src/lib.rs

// Crate-level lint overrides (workspace lints enabled in Cargo.toml)
#![allow(unsafe_code)]
//! RustyTorch++ GPU Runtime System
//!
//! This crate provides the core runtime environment for GPU-accelerated machine learning,
//! including memory management, device abstraction, and kernel execution.
//!
//! # Architecture
//!
//! The runtime is built around several key components:
//! - **Memory Allocator**: Arena-based GPU memory management with pooling and fragmentation tracking
//! - **Device Abstraction**: Unified interface for CUDA, ROCm, and Metal devices
//! - **Stream Scheduler**: Multi-stream dependency resolution with sub-microsecond overhead
//! - **Kernel Launcher**: Type-safe parameter marshalling and execution
//!
//! # Safety
//!
//! This crate uses unsafe code in controlled contexts for GPU operations. All unsafe blocks
//! are thoroughly documented with safety invariants.

#![allow(missing_docs)]

use parking_lot::RwLock;
use std::collections::BTreeMap;
use std::sync::atomic::AtomicU64;
use tracing::info;

pub mod allocator;
pub mod device;
pub mod error;
pub mod kernel;
pub mod scheduler;
pub mod stream;

// Production CUDA backend implementation using cudarc (only when cuda feature enabled)
#[cfg(feature = "cuda")]
pub mod cuda_backend;
#[cfg(feature = "cuda")]
pub use cuda_backend::{CudaBackend, CudaEventHandle, CudaStreamHandle, DeviceManager};

// Stub types when CUDA is not available
#[cfg(not(feature = "cuda"))]
pub mod cuda_backend {
    use crate::error::Result;

    /// Stub CUDA backend when CUDA feature is disabled
    pub struct CudaBackend;
    impl CudaBackend {
        /// Create stub backend (fails - CUDA not available)
        pub fn new() -> Result<Self> {
            Err(crate::RuntimeError::BackendNotSupported(
                "CUDA feature not enabled".into(),
            ))
        }
    }
    /// Stub device manager
    pub struct DeviceManager;
    /// Stub stream handle
    #[derive(Clone)]
    pub struct CudaStreamHandle;
    impl CudaStreamHandle {
        /// Stub synchronize (no-op)
        pub fn synchronize(&self) -> Result<()> {
            Ok(())
        }
    }
    /// Stub event handle
    pub struct CudaEventHandle;
}
#[cfg(not(feature = "cuda"))]
pub use cuda_backend::{CudaBackend, CudaEventHandle, CudaStreamHandle, DeviceManager};

// Metal backend implementation for Apple Silicon (macOS only)
#[cfg(all(target_os = "macos", feature = "metal"))]
pub mod metal_backend;
#[cfg(all(target_os = "macos", feature = "metal"))]
pub use metal_backend::{MetalBackend, MetalDeviceInfo, discover_metal_devices, initialize_metal};

// Stub Metal types when not on macOS or Metal feature disabled
#[cfg(not(all(target_os = "macos", feature = "metal")))]
pub mod metal_backend {
    /// Stub Metal backend when Metal feature is disabled
    pub struct MetalBackend;
    /// Stub Metal device info
    pub struct MetalDeviceInfo;
    /// Stub discover function
    pub fn discover_metal_devices() -> Vec<MetalDeviceInfo> {
        vec![]
    }
    /// Stub initialize function
    pub fn initialize_metal() -> Result<(), &'static str> {
        Ok(())
    }
}
#[cfg(not(all(target_os = "macos", feature = "metal")))]
pub use metal_backend::{MetalBackend, MetalDeviceInfo, discover_metal_devices, initialize_metal};

// Metal fusion kernel executor
#[cfg(all(target_os = "macos", feature = "metal"))]
pub mod metal_fusion;
#[cfg(all(target_os = "macos", feature = "metal"))]
pub use metal_fusion::{MetalFusionExecutor, get_metal_executor};

// CUDA Graph capture and execution using safe APIs
#[cfg(feature = "cuda")]
pub mod cuda_graph;
#[cfg(feature = "cuda")]
pub use cuda_graph::{CleanupStatistics, CudaGraphManager, GraphManagerStatistics};

// Stub CUDA graph types
#[cfg(not(feature = "cuda"))]
pub mod cuda_graph {
    /// Stub graph manager
    pub struct CudaGraphManager;
    /// Stub statistics
    pub struct GraphManagerStatistics;
    /// Stub cleanup statistics
    pub struct CleanupStatistics;
}
#[cfg(not(feature = "cuda"))]
pub use cuda_graph::{CleanupStatistics, CudaGraphManager, GraphManagerStatistics};

// CUDA Kernel compilation and execution using safe APIs
#[cfg(feature = "cuda")]
pub mod cuda_kernels;
#[cfg(feature = "cuda")]
pub use cuda_kernels::{CudaKernelManager, KernelLaunch, KernelManagerStatistics, OccupancyInfo};

// Stub kernel types
#[cfg(not(feature = "cuda"))]
pub mod cuda_kernels {
    /// Stub kernel manager
    pub struct CudaKernelManager;
    /// Stub kernel launch
    pub struct KernelLaunch;
    /// Stub occupancy info
    pub struct OccupancyInfo;
    /// Stub statistics
    pub struct KernelManagerStatistics;
}
#[cfg(not(feature = "cuda"))]
pub use cuda_kernels::{CudaKernelManager, KernelLaunch, KernelManagerStatistics, OccupancyInfo};

// Modular kernel fusion system
pub mod fusion;
pub mod fusion_core;
pub mod fusion_execution;
pub mod fusion_patterns;

// Lazy Execution Layer for ownership-based kernel fusion
pub mod lazy;

// Tensor Core optimization system
pub mod tensor_core;

// Advanced Tensor Core optimizations
// Revolutionary tensor core module removed - contained quantum/neuromorphic features

// Memory Bandwidth Optimization System (Phase C)
pub mod memory_bandwidth;

// CUDA-specific tests (only when cuda feature enabled)
#[cfg(all(test, feature = "cuda"))]
mod cuda_backend_test;
#[cfg(all(test, feature = "cuda"))]
mod cuda_context_test;
#[cfg(all(test, feature = "cuda"))]
mod tdd_cuda_tests;

// Dynamic Compilation & Specialization Engine (Phase D)
pub mod dynamic_compilation;

// CUDA integration tests - only when cuda feature enabled
#[cfg(all(test, feature = "cuda"))]
mod cuda_integration_tests;

// rustg integration tests - basic functionality verification
#[cfg(test)]
mod rustg_integration_test;

pub use allocator::*;
pub use device::*;
pub use error::*;
pub use kernel::*;
pub use scheduler::*;
pub use stream::{Stream, StreamId};

// Stream bridge and multi-stream pool
pub mod stream_bridge;
pub use stream_bridge::{
    StreamGuard, StreamManager, StreamPool, StreamPoolConfig, StreamPoolPriority, StreamPoolStats,
    ToCudaHandle,
};
// Note: cuda_graph types are already exported above with proper feature gates
// Use specific fusion module to avoid conflicts
pub use fusion::*;
// Re-export tensor_core types explicitly to avoid conflict with device::types
pub use tensor_core::{
    engine::TensorCoreEngine, precision::PrecisionOptimizer, scheduling::SchedulingOptimizer,
    specs::TensorCoreSpecs,
};
// Revolutionary tensor core exports removed - contained quantum/neuromorphic features
pub use dynamic_compilation::*;
pub use memory_bandwidth::*;

/// Runtime result type with rich error context
pub type RuntimeResult<T> = Result<T>;

/// Global runtime instance
static RUNTIME: once_cell::sync::Lazy<Runtime> = once_cell::sync::Lazy::new(|| {
    Runtime::new().expect("Failed to initialize RustyTorch++ runtime")
});

/// Main runtime coordinator
pub struct Runtime {
    /// Device pool indexed by device ID
    devices: RwLock<BTreeMap<DeviceId, Device>>,
    /// Global allocator statistics
    stats: AtomicU64,
}

impl Runtime {
    /// Initialize the runtime system
    pub fn new() -> RuntimeResult<Self> {
        info!("Initializing RustyTorch++ Runtime");

        Ok(Self {
            devices: RwLock::new(BTreeMap::new()),
            stats: AtomicU64::new(0),
        })
    }

    /// Get the global runtime instance
    pub fn global() -> &'static Self {
        &RUNTIME
    }

    /// Discover and register all available devices
    pub fn discover_devices(&self) -> RuntimeResult<usize> {
        let mut devices = self.devices.write();
        let count = device::discover_devices(&mut devices)?;
        info!("Discovered {} devices", count);
        Ok(count)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_runtime_creation() {
        let runtime = Runtime::new().unwrap();
        assert!(runtime.discover_devices().is_ok());
    }

    #[test]
    fn test_global_runtime_access() {
        let runtime1 = Runtime::global();
        let runtime2 = Runtime::global();
        assert!(std::ptr::eq(runtime1, runtime2));
    }
}