# Memory Management

Efficient memory usage in RustyTorch++.

## GPU Memory Hierarchy

```
┌─────────────────────────────────────────┐
│       Global Memory (HBM) - 24-80 GB    │ ← Tensor storage
├─────────────────────────────────────────┤
│       L2 Cache - 4-50 MB                │ ← Auto-managed
├─────────────────────────────────────────┤
│       Shared Memory - 48-164 KB/SM      │ ← Kernel local
├─────────────────────────────────────────┤
│       Registers - 64K/SM                │ ← Thread local
└─────────────────────────────────────────┘
```

## Memory Monitoring

```rust
use rtx_tensor::Device;

let device = Device::cuda(0)?;

// Query memory
let total = device.total_memory()?;
let used = device.used_memory()?;
let free = device.free_memory()?;

println!("Total: {:.2} GB", total as f64 / 1e9);
println!("Used: {:.2} GB", used as f64 / 1e9);
println!("Free: {:.2} GB", free as f64 / 1e9);
```

## Tensor Pool

### How It Works

```rust
use rtx_tensor::{TensorPool, pool_stats};

// Request a buffer from the pool
let buffer = TensorPool::get([1024, 1024], &device)?;

// When buffer is dropped, memory returns to pool (not freed)
drop(buffer);

// Next allocation with same size reuses memory
let buffer2 = TensorPool::get([1024, 1024], &device)?;  // No allocation!
```

### Pool Statistics

```rust
let stats = pool_stats();
println!("Pool size: {} MB", stats.total_size / 1_000_000);
println!("In use: {} MB", stats.used_size / 1_000_000);
println!("Cached: {} MB", stats.cached_size / 1_000_000);
println!("Hit rate: {:.1}%", stats.hit_rate * 100.0);
println!("Fragmentation: {:.1}%", stats.fragmentation * 100.0);
```

### Pool Management

```rust
use rtx_tensor::clear_pool;

// Clear unused cached memory
clear_pool();

// Configure pool behavior
TensorPool::configure(PoolConfig {
    max_cached_bytes: 4 * 1024 * 1024 * 1024,  // 4 GB max cache
    max_split_size: 512 * 1024 * 1024,          // Max block to split
});
```

## Gradient Checkpointing

Trade compute for memory during training.

### Automatic Checkpointing

```rust
use rtx_autograd::{CheckpointManager, CheckpointStrategy};

// √n checkpoints - optimal memory/compute tradeoff
let checkpoint = CheckpointManager::new(CheckpointStrategy::SqrtN);

// Wrap expensive forward passes
let hidden = checkpoint.checkpoint(|| {
    let h1 = layer1.forward(&x)?;
    let h2 = layer2.forward(&h1)?;
    layer3.forward(&h2)
})?;

// During backward, activations are recomputed
hidden.backward()?;
```

### Memory Savings

| Strategy | Memory | Compute Overhead |
|----------|--------|------------------|
| None | O(n) | 1x |
| √n checkpoints | O(√n) | ~1.5x |
| Every-k | O(n/k) | ~kx |

## Memory-Efficient Attention

### Flash Attention

```rust
use rtx_flash_attention::FlashAttention;

// Standard attention: O(N²) memory
// Flash attention: O(N) memory

let flash = FlashAttention::new(config)?;
let output = flash.forward(&q, &k, &v)?;
```

### Chunked Attention

```rust
// Process attention in chunks for very long sequences
let chunk_size = 1024;
let mut outputs = vec![];

for chunk_start in (0..seq_len).step_by(chunk_size) {
    let chunk_end = (chunk_start + chunk_size).min(seq_len);
    let q_chunk = q.narrow(1, chunk_start, chunk_end - chunk_start)?;
    let out_chunk = attention.forward(&q_chunk, &k, &v)?;
    outputs.push(out_chunk);
}

let output = Tensor::cat(&outputs, 1)?;
```

## Pinned Memory

Faster CPU-GPU transfers.

```rust
use rtx_memory::PinnedBuffer;

// Allocate pinned (page-locked) host memory
let pinned = PinnedBuffer::new(size)?;

// Fill with data
pinned.copy_from_slice(&data);

// Fast async transfer to GPU
let gpu_tensor = Tensor::from_pinned_async(&pinned, shape, &device, &stream)?;
```

## Memory Debugging

### Find Memory Leaks

```rust
use rtx_memory::MemoryTracker;

// Enable tracking
MemoryTracker::enable();

// Run your code...
train_loop(&model, &data)?;

// Report allocations
let report = MemoryTracker::report();
for allocation in report.live_allocations() {
    println!("{}: {} bytes at {}",
        allocation.id,
        allocation.size,
        allocation.location
    );
}
```

### Memory Snapshots

```rust
use rtx_memory::Snapshot;

let before = Snapshot::capture(&device)?;

// Operations...
let output = model.forward(&input)?;

let after = Snapshot::capture(&device)?;

// Compare
let diff = after.diff(&before);
println!("Memory change: {} MB", diff.total_bytes / 1_000_000);
```

## Best Practices

### 1. Delete Intermediate Tensors

```rust
// Bad - keeps all intermediates
let a = model.layer1(&x)?;
let b = model.layer2(&a)?;
let c = model.layer3(&b)?;

// Good - release when not needed
let mut h = model.layer1(&x)?;
h = model.layer2(&h)?;
h = model.layer3(&h)?;
```

### 2. Use no_grad for Inference

```rust
// Training: stores activations for backward
let output = model.forward(&input)?;

// Inference: no activation storage
let output = no_grad(|| model.forward(&input))?;
```

### 3. Process in Batches

```rust
// OOM risk with large batch
let all_outputs = model.forward(&huge_batch)?;

// Safe - process in chunks
let mut outputs = vec![];
for chunk in data.chunks(safe_batch_size) {
    outputs.push(model.forward(&chunk)?);
}
```

## Next Steps

- [GPU Optimization](./gpu-optimization.md) - Compute optimization
- [Profiling](./profiling.md) - Find memory bottlenecks
- [Troubleshooting Memory](../troubleshooting/memory.md) - Fix OOM