# RustyTorch++ System Patterns

## Architecture Overview (Phase-Aligned)

```mermaid
graph TB
    subgraph "User Layer"
        Python[Python API]
        Rust[Rust API]
        CPP[C++ API]
        WASM[WebAssembly]
    end
    
    subgraph "Phase 2: Tensor & Autograd"
        API[API Layer]
        TE[rtx-tensor]
        AG[rtx-autograd]
        IR[rtx-ir Graph IR]
    end
    
    subgraph "Phase 1: Core Compiler & Runtime"
        COMP[rtx-compiler]
        RT[rtx-runtime]
        KERNEL[rtx-kernel]
        RUSTG[rustg GPU Compiler]
    end
    
    subgraph "Phase 3-4: Advanced Execution"
        DIST[rtx-dist]
        SYNTH[rtx-synth]
        TUNE[Auto-Tuner]
    end
    
    subgraph "Phase 5: Serving"
        SERVE[rtx-serve]
        SCHED[Inference Scheduler]
        KV[Paged KV Cache]
    end
    
    subgraph "Backend Layer"
        CUDA[CUDA sm_120]
        ROCM[ROCm]
        METAL[Metal]
    end
    
    subgraph "Infrastructure"
        EVOLVE[rtx-evolve Phase 8]
        GOV[rtx-governance]
        PROF[rtx-profiler]
        BENCH[rtx-bench]
    end
    
    Python --> API
    API --> TE
    TE --> AG
    AG --> IR
    IR --> COMP
    COMP --> RUSTG
    RUSTG --> KERNEL
    KERNEL --> RT
    
    RT --> CUDA
    RT --> ROCM
    RT --> METAL
    
    DIST --> RT
    SYNTH --> KERNEL
    SERVE --> RT
    EVOLVE --> TUNE
```

## Core Design Patterns

### 1. Trait-Based Abstraction
```rust
// Core tensor trait enabling backend flexibility
pub trait Tensor: Send + Sync {
    type Elem: TensorElement;
    type Shape: Shape;
    type Backend: Backend;
    
    fn shape(&self) -> &Self::Shape;
    fn data(&self) -> &[Self::Elem];
    fn grad(&self) -> Option<&Self>;
}
```

**Rationale**: Enables compile-time polymorphism and zero-cost abstractions.

### 2. Builder Pattern for Complex Operations
```rust
// Flexible configuration without breaking changes
Conv2d::builder()
    .in_channels(3)
    .out_channels(64)
    .kernel_size((3, 3))
    .stride(1)
    .padding(1)
    .build()
```

**Rationale**: Maintains API stability while allowing extensibility.

### 3. Arena Allocation for Tensors
```rust
// Memory efficiency through region-based management
pub struct TensorArena {
    regions: Vec<MemoryRegion>,
    allocator: BumpAllocator,
}
```

**Rationale**: Reduces fragmentation and improves cache locality.

### 4. Lazy Evaluation with Graph Construction
```rust
// Build computation graph before execution
let graph = ComputeGraph::new();
let result = graph.build(|| {
    let x = tensor!([1, 2, 3]);
    let y = tensor!([4, 5, 6]);
    x.matmul(&y)
});
graph.execute()
```

**Rationale**: Enables optimization passes and kernel fusion.

## Component Architecture (Phase-Specific)

### Phase 1: Core Compiler & Runtime
- **rtx-compiler**: IR passes, rustg lowering, golden tests
- **rtx-runtime**: Device/Stream/Event/Graph abstractions, pooled allocator
- **rtx-kernel**: Fused kernels (MLP, LayerNorm, RoPE), kernel cache
- **Key Features**: CUDA Graphs capture/replay, AMP with loss scaler

### Phase 2: Tensor API & Autograd
- **rtx-tensor**: Shape/dtype/device, views/strides, broadcasting
- **rtx-autograd**: Tape engine, backward registry, checkpointing hooks
- **rtx-ir**: Mid-level Graph IR with cost/shape metadata, serialization
- **Key Features**: DLPack interop, deterministic mode

### Phase 3: Distributed Training
- **rtx-dist**: Process groups, NCCL/RCCL collectives, rendezvous
- **Parallelism**: DP/TP/PP/SP with auto-hybrid planner
- **Sharding**: FSDP/ZeRO-style with offload hooks
- **Key Features**: Elastic recovery, WAL checkpoints

### Phase 4: Auto-Kernel Synthesis
- **rtx-synth**: Pattern library, template emitters, synthesis flow
- **Auto-Tuner**: Bandit/grid search, hardware profile DB
- **AOT Compiler**: Graph bundles with compatibility checks
- **Key Features**: Per-SKU optimization, persistent kernel cache

### Phase 5: Inference Runtime
- **rtx-serve**: Scheduler with lanes, paged KV cache, quantization
- **Decoding**: Speculative/assisted, early-exit strategies
- **APIs**: gRPC/HTTP streaming, Python/Rust SDKs
- **Key Features**: Continuous batching, vLLM-class performance

### Phase 6-10: Platform Evolution
- **Unified Graph**: ETL + Model ops in same IR
- **Self-Optimizing**: Telemetry-driven auto-tuning
- **Governance**: SBOM, signatures, provenance
- **Evolution**: Agent-in-the-loop improvements
- **Multi-Tenant**: Global routing, quotas, federation

## Critical Implementation Paths

### Forward Pass Execution
```
1. User API Call
2. Tensor Validation
3. Graph Construction
4. Optimization Pass
5. Backend Selection
6. Kernel Dispatch
7. Memory Allocation
8. Computation
9. Result Return
```

### Backward Pass Execution
```
1. Loss Computation
2. Gradient Tape Replay
3. Chain Rule Application
4. Gradient Accumulation
5. Optimizer Step
6. Parameter Update
7. Gradient Clear
```

### Distributed Training Flow
```
1. Model Replication
2. Data Sharding
3. Forward Pass (Local)
4. Gradient Computation (Local)
5. Gradient Aggregation (All-Reduce)
6. Parameter Update (Local)
7. Synchronization Barrier
```

## Memory Management Strategy

### Tensor Lifecycle
1. **Allocation**: Via arena allocator with size classes
2. **Ownership**: Rust ownership for automatic deallocation
3. **Sharing**: Arc for multi-threaded access
4. **Gradient**: Separate allocation with weak references

### GPU Memory Hierarchy
```
Host Memory (RAM)
    ↓ (Async Transfer)
Device Memory (VRAM)
    ↓ (Kernel Launch)
Shared Memory (SM)
    ↓ (Thread Access)
Registers
```

### Memory Optimization Techniques
- **Gradient Checkpointing**: Trade compute for memory
- **Memory Pooling**: Reuse allocations across iterations
- **Operator Fusion**: Reduce intermediate tensor allocation
- **Quantization**: Reduce precision for memory savings

## Error Handling Philosophy

### Principle: Fail Fast, Recover Gracefully
```rust
pub type Result<T> = std::result::Result<T, RustyTorchError>;

#[derive(Error, Debug)]
pub enum RustyTorchError {
    #[error("Shape mismatch: expected {expected:?}, got {got:?}")]
    ShapeMismatch { expected: Shape, got: Shape },
    
    #[error("Out of memory: requested {requested} bytes")]
    OutOfMemory { requested: usize },
    
    #[error("Backend error: {0}")]
    BackendError(String),
}
```

### Error Categories
1. **Compile-Time**: Shape mismatches, type errors
2. **Runtime**: OOM, device errors, numerical instability
3. **Distributed**: Network failures, node crashes
4. **User**: Invalid configurations, API misuse

## Concurrency Model

### Thread Safety Guarantees
- **Tensors**: Send + Sync for parallel processing
- **Modules**: Arc<Mutex<>> for shared state
- **Autograd**: Thread-local tape with merge capability

### Parallelism Levels
1. **Data Parallelism**: Batch dimension splitting
2. **Model Parallelism**: Layer distribution
3. **Pipeline Parallelism**: Micro-batch processing
4. **Tensor Parallelism**: Operation splitting

## Extension Points

### Custom Operators
```rust
pub trait CustomOp: Send + Sync {
    fn forward(&self, inputs: &[Tensor]) -> Result<Tensor>;
    fn backward(&self, grad: &Tensor) -> Result<Vec<Tensor>>;
}
```

### Backend Plugins
```rust
pub trait Backend: Send + Sync {
    fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;
    fn conv2d(&self, input: &Tensor, kernel: &Tensor) -> Result<Tensor>;
    // ... other operations
}
```

### Optimization Passes
```rust
pub trait OptimizationPass {
    fn optimize(&self, graph: &mut ComputeGraph) -> Result<()>;
}
```

## Performance Patterns (Phase-Aligned)

### Phase 1: Foundation Performance
- **CUDA Graphs**: Capture/replay for reduced launch overhead
- **Pooled Allocator**: Arena-based with deterministic ordering
- **Stream Scheduling**: Multi-stream with dependency DAG
- **Initial Fusions**: MLP, LayerNorm, RoPE kernels

### Phase 2: Operator Performance
- **Broadcasting**: Efficient view-based operations
- **Contiguous Checks**: Non-contiguous tensor handling
- **Stable Reductions**: Deterministic sum/mean operations
- **AMP Integration**: fp16/bf16 with loss scaling

### Phase 3: Distributed Performance
- **Overlap**: Communication/compute scheduling
- **Bucketing**: Gradient aggregation optimization
- **Topology-Aware**: PCIe/NVLink/IB bandwidth optimization
- **Sharding**: Memory reduction via FSDP/ZeRO

### Phase 4: Auto-Optimization
- **Hardware Profiles**: Per-SKU latency/bandwidth DB
- **Kernel Synthesis**: Generate specialized kernels
- **Auto-Tuning**: Bandit search for optimal params
- **AOT Compilation**: Pre-compiled graph bundles

### Phase 5: Inference Optimization
- **Continuous Batching**: Dynamic batch merging
- **Paged KV Cache**: GPU/CPU/NVMe tiering
- **Quantization**: INT8/INT4/FP8 for latency
- **Speculative Decoding**: Draft-verify acceleration

## Testing Strategy

### Test Levels
1. **Unit Tests**: Individual operations
2. **Integration Tests**: Module combinations
3. **Property Tests**: Invariant verification
4. **Benchmark Tests**: Performance regression
5. **Distributed Tests**: Multi-node scenarios

### Test Patterns
```rust
#[test]
fn test_operation() {
    // Arrange
    let tensor = Tensor::randn([32, 64]);
    
    // Act
    let result = tensor.relu();
    
    // Assert
    assert!(result.all(|x| x >= 0.0));
}
```

## Monitoring & Observability

### Metrics Collection
- Operation latencies
- Memory usage
- GPU utilization
- Network throughput
- Cache hit rates

### Tracing Integration
```rust
#[instrument]
pub fn matmul(a: &Tensor, b: &Tensor) -> Result<Tensor> {
    span!(Level::DEBUG, "matmul", shape_a = ?a.shape(), shape_b = ?b.shape());
    // Implementation
}
```

### Profiling Hooks
- Pre/post operation callbacks
- Memory allocation tracking
- Kernel execution timing
- Gradient flow visualization

## Rust 2024 Edition Patterns (December 2024)

### Float Comparison Safety Pattern
**Requirement**: Rust 2024 edition requires NaN-safe float comparisons

```rust
// BEFORE (Rust 2021 - panics on NaN)
values.sort_by(|a, b| a.partial_cmp(b).unwrap());

// AFTER (Rust 2024 - NaN-safe)
values.sort_by(|a, b| a.total_cmp(b));
```

**Applied to**: 200+ files across the workspace
**Rationale**: `total_cmp()` provides a total ordering that handles NaN, -0.0, and +0.0 correctly

### Module Organization Pattern (rtx-nlg)
**Problem**: rtx-nlg had 245+ compilation errors due to missing dependencies
**Solution**: Create local modules instead of modifying shared crates

```
crates/models/rtx-nlg/src/
├── dialogue/
│   └── mod.rs          # Conversational AI module (LOCAL)
├── tensor_helpers.rs   # Local tensor operations (LOCAL)
└── lib.rs              # Updated exports
```

**Rationale**: Keeps core crates (rtx-tensor) stable while allowing model-specific helpers

### Dependency Hygiene Pattern
**Problem**: Legacy nom 3.2.1 pulled in via unused transitive dependency
**Solution**: Regular dependency auditing

```bash
# Check for unused dependencies
cargo machete

# Check dependency tree for specific versions
cargo tree -i nom

# Remove unused dependencies
cargo rm npy  # from rtx-vision-advanced
```

**Rationale**: Prevents compatibility issues with new Rust editions

### Excluded Crates Pattern
**Rationale**: Some crates require special handling

| Crate | Pattern | Reason |
|-------|---------|--------|
| `integration_tests` | Workspace exclude | References unimplemented APIs |
| `rtx-flash-metal-attention` | Workspace exclude | Platform-specific (macOS only) |
| `demos/ui/src-tauri` | Workspace exclude | Different MSRV requirements |

```toml
# Cargo.toml workspace configuration
[workspace]
exclude = [
    "demos/ui/src-tauri",
    "crates/training/rtx-flash-metal-attention",
    "integration_tests",
]
```

---

*System Patterns Last Updated: 2025-12-16*
*Rust Edition: 2024 (Rust 1.92+)*