TinyTorch/test_module_performance.py

#!/usr/bin/env python3
"""
Real Performance Testing for TinyTorch Modules
==============================================

This tests actual performance improvements in TinyTorch optimization modules.
No hallucinated numbers - only real, measured performance data.
"""

import sys
import os
import time
import tracemalloc
import numpy as np
import statistics
from typing import Dict, Tuple, Any

# Add TinyTorch to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'tinytorch'))

# Test Framework
class RealPerformanceTester:
    """Scientific performance testing with statistical rigor."""
    
    def __init__(self, runs=5):
        self.runs = runs
    
    def measure_timing(self, func, *args, **kwargs):
        """Measure execution time with multiple runs."""
        times = []
        for _ in range(self.runs):
            start = time.perf_counter()
            result = func(*args, **kwargs)
            end = time.perf_counter()
            times.append(end - start)
        
        mean_time = statistics.mean(times)
        std_time = statistics.stdev(times) if len(times) > 1 else 0
        
        return {
            'mean': mean_time,
            'std': std_time,
            'times': times,
            'result': result
        }
    
    def measure_memory(self, func, *args, **kwargs):
        """Measure memory usage."""
        tracemalloc.start()
        result = func(*args, **kwargs)
        current, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        
        return {
            'current_mb': current / 1024 / 1024,
            'peak_mb': peak / 1024 / 1024,
            'result': result
        }
    
    def compare_implementations(self, baseline_func, optimized_func, args, test_name):
        """Compare two implementations scientifically."""
        print(f"\n🧪 {test_name}")
        print("=" * 60)
        
        # Timing comparison
        baseline_timing = self.measure_timing(baseline_func, *args)
        optimized_timing = self.measure_timing(optimized_func, *args)
        
        speedup = baseline_timing['mean'] / optimized_timing['mean']
        
        print(f"  Baseline:  {baseline_timing['mean']*1000:.2f} ± {baseline_timing['std']*1000:.2f} ms")
        print(f"  Optimized: {optimized_timing['mean']*1000:.2f} ± {optimized_timing['std']*1000:.2f} ms")
        print(f"  Speedup:   {speedup:.2f}×")
        
        # Memory comparison
        baseline_memory = self.measure_memory(baseline_func, *args)
        optimized_memory = self.measure_memory(optimized_func, *args)
        
        memory_ratio = optimized_memory['peak_mb'] / baseline_memory['peak_mb']
        
        print(f"  Memory (baseline):  {baseline_memory['peak_mb']:.2f} MB")
        print(f"  Memory (optimized): {optimized_memory['peak_mb']:.2f} MB") 
        print(f"  Memory ratio: {memory_ratio:.2f}×")
        
        # Accuracy check
        baseline_result = np.array(baseline_timing['result'])
        optimized_result = np.array(optimized_timing['result'])
        
        if baseline_result.shape == optimized_result.shape:
            max_diff = np.max(np.abs(baseline_result - optimized_result))
            accuracy_ok = max_diff < 1e-5
            print(f"  Max difference: {max_diff:.2e}")
            print(f"  Accuracy: {'✅ preserved' if accuracy_ok else '❌ lost'}")
        else:
            accuracy_ok = False
            print(f"  Shapes: baseline {baseline_result.shape} vs optimized {optimized_result.shape}")
            print(f"  Accuracy: ❌ shapes don't match")
        
        success = speedup > 1.1 and accuracy_ok
        print(f"  Overall: {'✅ IMPROVEMENT' if success else '⚠️ NO IMPROVEMENT'}")
        
        return {
            'speedup': speedup,
            'memory_ratio': memory_ratio,
            'accuracy_preserved': accuracy_ok,
            'success': success
        }


def test_matrix_multiplication_optimization():
    """Test Module 16: Acceleration - Matrix multiplication optimization."""
    
    def naive_matmul(A, B):
        """Naive triple-nested loop implementation."""
        n, k = A.shape
        k2, m = B.shape
        assert k == k2, "Matrix dimensions must match"
        
        C = np.zeros((n, m), dtype=np.float32)
        for i in range(n):
            for j in range(m):
                for idx in range(k):
                    C[i, j] += A[i, idx] * B[idx, j]
        return C
    
    def blocked_matmul(A, B, block_size=32):
        """Cache-friendly blocked implementation."""
        n, k = A.shape
        k2, m = B.shape
        assert k == k2, "Matrix dimensions must match"
        
        C = np.zeros((n, m), dtype=np.float32)
        
        for i0 in range(0, n, block_size):
            for j0 in range(0, m, block_size):
                for k0 in range(0, k, block_size):
                    # Process block
                    i_end = min(i0 + block_size, n)
                    j_end = min(j0 + block_size, m)
                    k_end = min(k0 + block_size, k)
                    
                    for i in range(i0, i_end):
                        for j in range(j0, j_end):
                            for idx in range(k0, k_end):
                                C[i, j] += A[i, idx] * B[idx, j]
        return C
    
    def numpy_matmul(A, B):
        """NumPy optimized implementation."""
        return np.dot(A, B)
    
    # Create test matrices
    size = 128  # Small enough to complete quickly
    np.random.seed(42)
    A = np.random.randn(size, size).astype(np.float32)
    B = np.random.randn(size, size).astype(np.float32)
    
    tester = RealPerformanceTester(runs=3)
    
    # Test naive vs blocked
    results1 = tester.compare_implementations(
        naive_matmul, blocked_matmul, (A, B),
        "Matrix Multiplication: Naive vs Blocked"
    )
    
    # Test blocked vs numpy  
    results2 = tester.compare_implementations(
        blocked_matmul, numpy_matmul, (A, B),
        "Matrix Multiplication: Blocked vs NumPy"
    )
    
    return results1, results2


def test_attention_optimization():
    """Test Module 19: Caching - Attention mechanism optimization."""
    
    def standard_attention(Q, K, V, mask=None):
        """Standard attention computation."""
        # Compute attention scores
        scores = np.dot(Q, K.T) / np.sqrt(Q.shape[-1])
        
        # Apply mask if provided
        if mask is not None:
            scores = np.where(mask, scores, -1e9)
        
        # Softmax
        exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
        attention_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
        
        # Apply to values
        output = np.dot(attention_weights, V)
        return output, attention_weights
    
    def cached_attention_step(Q_new, K_cache, V_cache, K_new, V_new, mask=None):
        """Cached attention for incremental computation."""
        # Append new K,V to cache
        K_combined = np.concatenate([K_cache, K_new.reshape(1, -1)], axis=0)
        V_combined = np.concatenate([V_cache, V_new.reshape(1, -1)], axis=0)
        
        # Compute attention only for new query
        scores = np.dot(Q_new, K_combined.T) / np.sqrt(Q_new.shape[-1])
        
        if mask is not None:
            scores = np.where(mask, scores, -1e9)
        
        exp_scores = np.exp(scores - np.max(scores))
        attention_weights = exp_scores / np.sum(exp_scores)
        
        output = np.dot(attention_weights, V_combined)
        
        return output, K_combined, V_combined
    
    # Create test data
    seq_len = 64
    d_model = 128
    np.random.seed(42)
    
    Q = np.random.randn(seq_len, d_model).astype(np.float32)
    K = np.random.randn(seq_len, d_model).astype(np.float32)  
    V = np.random.randn(seq_len, d_model).astype(np.float32)
    
    # Causal mask
    causal_mask = np.tril(np.ones((seq_len, seq_len), dtype=bool))
    
    def standard_generation():
        """Standard attention for autoregressive generation."""
        outputs = []
        for i in range(1, seq_len):
            # Recompute attention for sequence up to position i
            Q_slice = Q[i:i+1]  # Current query
            K_slice = K[:i+1]   # All keys up to current position
            V_slice = V[:i+1]   # All values up to current position
            mask_slice = causal_mask[i:i+1, :i+1]
            
            output, _ = standard_attention(Q_slice, K_slice, V_slice, mask_slice)
            outputs.append(output[0])
        
        return np.array(outputs)
    
    def cached_generation():
        """Cached attention for autoregressive generation."""
        outputs = []
        K_cache = K[0:1]  # Initialize with first key
        V_cache = V[0:1]  # Initialize with first value
        
        for i in range(1, seq_len):
            Q_new = Q[i]    # New query
            K_new = K[i]    # New key
            V_new = V[i]    # New value
            mask_new = causal_mask[i, :i+1]
            
            output, K_cache, V_cache = cached_attention_step(
                Q_new, K_cache, V_cache, K_new, V_new, mask_new
            )
            outputs.append(output)
        
        return np.array(outputs)
    
    tester = RealPerformanceTester(runs=3)
    
    results = tester.compare_implementations(
        standard_generation, cached_generation, (),
        "Attention: Standard vs KV Cache"
    )
    
    return results


def test_quantization_performance():
    """Test Module 17: Quantization - FP32 vs INT8."""
    
    def fp32_conv(input_data, weights, bias):
        """Standard FP32 convolution."""
        # Simple convolution implementation
        batch_size, in_height, in_width, in_channels = input_data.shape
        out_channels, kernel_h, kernel_w, in_ch = weights.shape
        
        out_height = in_height - kernel_h + 1
        out_width = in_width - kernel_w + 1
        
        output = np.zeros((batch_size, out_height, out_width, out_channels), dtype=np.float32)
        
        for b in range(batch_size):
            for oh in range(out_height):
                for ow in range(out_width):
                    for oc in range(out_channels):
                        for kh in range(kernel_h):
                            for kw in range(kernel_w):
                                for ic in range(in_channels):
                                    output[b, oh, ow, oc] += (
                                        input_data[b, oh + kh, ow + kw, ic] * 
                                        weights[oc, kh, kw, ic]
                                    )
                        output[b, oh, ow, oc] += bias[oc]
        
        return output
    
    def quantized_conv(input_data, weights, bias, input_scale, weight_scale):
        """Quantized INT8 convolution simulation."""
        # Quantize inputs (simulate INT8 by using int8 data type)
        input_quantized = np.round(input_data / input_scale).astype(np.int8)
        weights_quantized = np.round(weights / weight_scale).astype(np.int8)
        
        # Run convolution in int8 (simulated - numpy doesn't have true int8 conv)
        batch_size, in_height, in_width, in_channels = input_quantized.shape
        out_channels, kernel_h, kernel_w, in_ch = weights_quantized.shape
        
        out_height = in_height - kernel_h + 1
        out_width = in_width - kernel_w + 1
        
        # Use int32 accumulator
        output = np.zeros((batch_size, out_height, out_width, out_channels), dtype=np.int32)
        
        for b in range(batch_size):
            for oh in range(out_height):
                for ow in range(out_width):
                    for oc in range(out_channels):
                        for kh in range(kernel_h):
                            for kw in range(kernel_w):
                                for ic in range(in_channels):
                                    output[b, oh, ow, oc] += (
                                        int(input_quantized[b, oh + kh, ow + kw, ic]) * 
                                        int(weights_quantized[oc, kh, kw, ic])
                                    )
                        # Add quantized bias (scaled appropriately)
                        bias_quantized = int(bias[oc] / (input_scale * weight_scale))
                        output[b, oh, ow, oc] += bias_quantized
        
        # Dequantize output
        output_scale = input_scale * weight_scale
        output_fp32 = output.astype(np.float32) * output_scale
        
        return output_fp32
    
    # Create test data
    batch_size, height, width, in_channels = 1, 28, 28, 3
    out_channels, kernel_size = 8, 3
    
    np.random.seed(42)
    input_data = np.random.randn(batch_size, height, width, in_channels).astype(np.float32)
    weights = np.random.randn(out_channels, kernel_size, kernel_size, in_channels).astype(np.float32) * 0.1
    bias = np.random.randn(out_channels).astype(np.float32) * 0.1
    
    # Quantization scales (typical values)
    input_scale = np.max(np.abs(input_data)) / 127.0
    weight_scale = np.max(np.abs(weights)) / 127.0
    
    tester = RealPerformanceTester(runs=3)
    
    results = tester.compare_implementations(
        lambda: fp32_conv(input_data, weights, bias),
        lambda: quantized_conv(input_data, weights, bias, input_scale, weight_scale),
        (),
        "Convolution: FP32 vs INT8 Quantized"
    )
    
    return results


def main():
    """Run comprehensive performance tests."""
    print("🔥 TinyTorch Real Performance Analysis")
    print("=====================================")
    print("Testing ACTUAL performance improvements in optimization modules.")
    print("No hallucinated numbers - only real, measured data.\n")
    
    all_results = {}
    
    # Test Module 16: Acceleration
    print("📊 MODULE 16: ACCELERATION TESTING")
    try:
        matmul_results = test_matrix_multiplication_optimization()
        all_results['matrix_multiplication'] = matmul_results
        print("✅ Matrix multiplication tests completed")
    except Exception as e:
        print(f"❌ Matrix multiplication tests failed: {e}")
        all_results['matrix_multiplication'] = None
    
    # Test Module 19: Caching  
    print("\n📊 MODULE 19: CACHING TESTING")
    try:
        attention_results = test_attention_optimization()
        all_results['attention_caching'] = attention_results
        print("✅ Attention caching tests completed")
    except Exception as e:
        print(f"❌ Attention caching tests failed: {e}")
        all_results['attention_caching'] = None
    
    # Test Module 17: Quantization
    print("\n📊 MODULE 17: QUANTIZATION TESTING")
    try:
        quant_results = test_quantization_performance()
        all_results['quantization'] = quant_results
        print("✅ Quantization tests completed")
    except Exception as e:
        print(f"❌ Quantization tests failed: {e}")
        all_results['quantization'] = None
    
    # Summary
    print("\n" + "="*60)
    print("📋 PERFORMANCE TESTING SUMMARY")
    print("="*60)
    
    successful_tests = 0
    total_tests = 0
    
    for test_name, results in all_results.items():
        if results is not None:
            if isinstance(results, tuple):  # Multiple sub-tests
                for i, result in enumerate(results):
                    total_tests += 1
                    if result and result.get('success', False):
                        successful_tests += 1
                        print(f"✅ {test_name}_{i}: {result['speedup']:.2f}× speedup")
                    else:
                        if result:
                            print(f"⚠️ {test_name}_{i}: {result['speedup']:.2f}× speedup (not significant)")
                        else:
                            print(f"❌ {test_name}_{i}: failed")
            else:  # Single test
                total_tests += 1
                if results.get('success', False):
                    successful_tests += 1
                    print(f"✅ {test_name}: {results['speedup']:.2f}× speedup")
                else:
                    print(f"⚠️ {test_name}: {results['speedup']:.2f}× speedup (not significant)")
        else:
            total_tests += 1
            print(f"❌ {test_name}: test failed")
    
    print(f"\n🎯 OVERALL RESULTS: {successful_tests}/{total_tests} optimizations successful")
    
    if successful_tests > 0:
        print(f"✅ TinyTorch optimization modules deliver measurable improvements!")
    else:
        print(f"⚠️ TinyTorch optimization modules need improvement - no significant speedups found")
    
    return all_results


if __name__ == "__main__":
    results = main()