TinyTorch/modules/source/13_kernels/kernels_dev.py

# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.17.1
# ---

# %% [markdown]
"""
# Kernels - High-Performance Computing and Hardware Optimization

Welcome to the Kernels module! You'll implement high-performance computational kernels that understand how modern hardware works, moving beyond generic libraries to achieve optimal performance.

## Learning Goals
- Systems understanding: How CPU cache hierarchies, SIMD instructions, and memory bandwidth determine ML operation performance
- Core implementation skill: Build vectorized operations and memory-efficient algorithms that outperform standard library implementations
- Pattern recognition: Understand how algorithmic choices interact with hardware characteristics to determine real-world performance
- Framework connection: See how your optimizations relate to the low-level kernels used in PyTorch, cuDNN, and BLAS libraries
- Performance insight: Learn why kernel optimization often provides larger speedups than algorithmic improvements

## Build → Use → Reflect
1. **Build**: Custom vectorized operations, cache-friendly algorithms, and parallel computation patterns
2. **Use**: Apply optimized kernels to real ML workloads and measure performance improvements
3. **Reflect**: Why do hardware characteristics often matter more than algorithm choice for ML performance?

## What You'll Achieve
By the end of this module, you'll understand:
- Deep technical understanding of how modern hardware executes ML operations and why optimization requires hardware awareness
- Practical capability to write high-performance code that achieves near-optimal hardware utilization
- Systems insight into why kernel optimization is critical for production ML systems and how it affects system design
- Performance consideration of how memory access patterns, vectorization, and parallelization strategies affect computational efficiency
- Connection to production ML systems and how frameworks achieve performance through hardware-optimized kernel libraries

## Systems Reality Check
💡 **Production Context**: PyTorch's performance comes from libraries like MKL-DNN and cuDNN that implement thousands of hand-optimized kernels for different hardware configurations
⚡ **Performance Note**: Well-optimized kernels can be 10-100x faster than naive implementations - kernel optimization is often the difference between research code and production systems
"""

# %% nbgrader={"grade": false, "grade_id": "kernels-imports", "locked": false, "schema_version": 3, "solution": false, "task": false}
#| default_exp core.kernels

#| export
import numpy as np
import sys
import os
import time
import psutil
from typing import Callable, Dict, Any, Optional, Tuple, List

# Import our existing components
try:
    from tinytorch.core.tensor import Tensor
    from tinytorch.core.layers import matmul_naive as matmul
    from tinytorch.core.activations import ReLU, Sigmoid, Tanh
    from tinytorch.core.cnn import Conv2D
except ImportError:
    # For development, import from local modules
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.extend([
        os.path.join(base_dir, '01_tensor'),
        os.path.join(base_dir, '02_activations'),
        os.path.join(base_dir, '03_layers'),
        os.path.join(base_dir, '05_cnn'),
        os.path.join(base_dir, 'utils')
    ])
    
    try:
        from tensor_dev import Tensor
        from layers_dev import matmul_naive as matmul
        from activations_dev import ReLU, Sigmoid, Tanh
        from cnn_dev import Conv2D
    except ImportError:
        # Create minimal mock for development
        class Tensor:
            def __init__(self, data):
                self.data = np.array(data)
                self.shape = self.data.shape
            def __str__(self):
                return f"Tensor({self.data})"

# Simple timing utility for kernel performance measurement
def time_kernel(func, *args, **kwargs):
    """
    Simple timing function for measuring kernel performance.
    
    Returns:
        tuple: (result, time_in_microseconds)
    """
    start = time.perf_counter()
    result = func(*args, **kwargs)
    end = time.perf_counter()
    microseconds = (end - start) * 1_000_000
    return result, microseconds

# %% nbgrader={"grade": false, "grade_id": "kernels-setup", "locked": false, "schema_version": 3, "solution": false, "task": false}
print("🔥 TinyTorch Kernels Module")
print(f"NumPy version: {np.__version__}")
print(f"Python version: {sys.version_info.major}.{sys.version_info.minor}")
print(f"System: {psutil.cpu_count()} CPU cores, {psutil.virtual_memory().total // (1024**3):.1f}GB RAM")
print("Ready to optimize ML operations!")

# %% [markdown]
"""
## 📦 Where This Code Lives in the Final Package

**Learning Side:** You work in `modules/source/11_kernels/kernels_dev.py`  
**Building Side:** Code exports to `tinytorch.core.kernels`

```python
# Final package structure:
from tinytorch.core.kernels import vectorized_matmul, parallel_relu, cached_conv2d
from tinytorch.core.tensor import Tensor
from tinytorch.core.layers import Dense
```

**Why this matters:**
- **Performance:** Custom kernels can be 2-10x faster than naive implementations
- **Understanding:** Learn how PyTorch, TensorFlow achieve their speed
- **Real-world:** Modern ML frameworks rely heavily on optimized kernels
- **Hardware:** Bridge the gap between algorithms and computer architecture
"""

# %% [markdown]
"""
## What are ML Kernels?

### The Performance Gap
Your neural network training is slow. A simple matrix multiplication that should take milliseconds takes seconds. Why?

**The problem:** NumPy operations, while convenient, aren't optimized for your specific hardware or use case.

**The solution:** Custom kernels - specialized functions written to extract maximum performance from your hardware.

### What is a Kernel?
A **kernel** is a highly optimized function that performs a specific computation:

```python
# Standard approach - easy but slow
def slow_matmul(A, B):
    return np.dot(A, B)

# Kernel approach - harder but fast
def fast_matmul(A, B):
    # Optimized for your CPU's cache hierarchy
    # Uses SIMD instructions for parallel operations
    # Minimizes memory allocations
    return optimized_result
```

### Why Kernels Matter for ML
Modern ML frameworks achieve their speed through thousands of optimized kernels:

- **PyTorch**: 2000+ CUDA kernels, 500+ CPU kernels
- **TensorFlow**: XLA compiler generates optimized kernels
- **JAX**: JIT compilation creates specialized kernels
- **Hardware**: GPUs have 1000s of cores, TPUs have specialized ML units

### The Performance Hierarchy
```
Python loops:        1x speed    (baseline)
NumPy operations:    10x speed   (vectorized)
Optimized kernels:   100x speed  (hardware-aware)
GPU kernels:         1000x speed (massive parallelism)
```

### Real-World Impact
- **Training time**: 10-hour training → 1-hour training
- **Inference cost**: $1000/month → $100/month
- **Model size**: Enable larger models through efficiency
- **Energy**: 90% reduction in power consumption

### What You'll Learn
1. **Custom operations** - Moving beyond NumPy limitations
2. **Vectorization** - Using SIMD for parallel computation
3. **Memory optimization** - Cache-friendly algorithms
4. **Parallel processing** - CPU and GPU-style parallelism
5. **Performance measurement** - Professional profiling tools
6. **Compressed kernels** - Optimizations for quantized models

Let's build the optimizations that power modern AI!
"""

# %% [markdown]
"""
## 🔧 DEVELOPMENT
"""

# %% [markdown]
"""
## Step 1: Custom Operations - Beyond NumPy

### Why Custom Operations?
NumPy is great for prototyping, but has limitations:
- **Generic**: Optimized for general use, not your specific case
- **Memory**: Creates temporary arrays, wastes memory
- **Control**: Can't control memory layout, algorithm choice
- **Specialization**: Can't optimize for your data patterns

### The Philosophy
Instead of using general-purpose functions, we write **specialized** functions:

```python
# Generic NumPy approach
def generic_activation(x):
    return np.maximum(0, x)  # ReLU

# Specialized kernel approach  
def fast_relu_kernel(x):
    # Optimized for your specific use case
    # No unnecessary memory allocations
    # Optimized for your data sizes
    return result
```

### Design Principles
- **Specialization**: Optimize for specific input patterns
- **Memory efficiency**: Minimize allocations and copies
- **Algorithmic choice**: Pick the best algorithm for your data
- **Measurement**: Always profile before and after

### Real-World Context
This is how:
- **PyTorch**: Custom autograd functions override standard operations
- **TensorFlow**: tf.function compiles optimized graphs
- **JAX**: jax.jit creates specialized kernels
- **CUDA**: Every GPU operation is a custom kernel
"""

# %% nbgrader={"grade": false, "grade_id": "custom-matmul", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def matmul_baseline(A: Tensor, B: Tensor) -> Tensor:
    """
    Baseline matrix multiplication using TinyTorch's proven implementation.
    
    This function demonstrates how to build on existing TinyTorch components
    rather than reinventing the wheel. We use the standard matmul from Module 03
    as our baseline for comparison with optimized kernels.
    
    This is NOT a custom implementation - it's the standard TinyTorch matmul
    wrapped for use in kernel comparisons and benchmarking.
    
    TODO: Use TinyTorch's standard matmul implementation as a baseline.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Import the standard matmul function from tinytorch.core.layers
    2. Extract numpy arrays from input Tensors
    3. Use the proven implementation from TinyTorch
    4. Wrap result back in Tensor format
    5. Return the result
    
    CODE REUSE PRINCIPLES:
    1. Always use the packaged version for reliability
    2. Don't duplicate working code - reference the source
    3. Use descriptive names that indicate what the function actually does
    4. Keep dependencies simple and reliable
    
    EXAMPLE USAGE:
    ```python
    A = Tensor([[1, 2], [3, 4]])
    B = Tensor([[5, 6], [7, 8]])
    C = matmul_baseline(A, B)
    # Expected: [[19, 22], [43, 50]]
    ```
    
    LEARNING CONNECTIONS:
    - This shows how to use TinyTorch as a library
    - Demonstrates reliable dependency management
    - Serves as baseline for kernel performance comparisons
    - Shows proper software engineering practices
    """
    ### BEGIN SOLUTION
    # Extract numpy arrays from Tensors
    A_data = A.data if hasattr(A, 'data') else A
    B_data = B.data if hasattr(B, 'data') else B
    
    # Use NumPy's matrix multiplication as our baseline
    # This is our baseline - reliable, tested, and consistent
    result_data = np.dot(A_data, B_data)
    
    # Wrap the result back in a Tensor for consistency
    result = Tensor(result_data)
    
    return result
    ### END SOLUTION

# %% [markdown]
# %% nbgrader={"grade": false, "grade_id": "test-custom-matmul", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Baseline Matrix Multiplication

def test_unit_matmul_baseline():
    """Unit test for the baseline matrix multiplication implementation."""
    print("🔬 Unit Test: Baseline Matrix Multiplication...")
    
    # Test case 1: Small matrices (2x2)
    A = Tensor([[1, 2], [3, 4]])
    B = Tensor([[5, 6], [7, 8]])
    C = matmul_baseline(A, B)
    expected = Tensor([[19, 22], [43, 50]])  # Hand-computed
    
    assert np.allclose(C.data, expected.data), f"Expected {expected.data}, got {C.data}"
    print("✅ Small matrix multiplication works")
    
    # Test case 2: Rectangular matrices
    A = Tensor([[1, 2, 3], [4, 5, 6]])  # 2x3
    B = Tensor([[7, 8], [9, 10], [11, 12]])  # 3x2
    C = matmul_baseline(A, B)
    expected = Tensor([[58, 64], [139, 154]])
    
    assert np.allclose(C.data, expected.data), f"Expected {expected.data}, got {C.data}"
    print("✅ Rectangular matrix multiplication works")
    
    # Test case 3: Compare with NumPy (medium size - should use TinyTorch implementation)
    np.random.seed(42)
    A = Tensor(np.random.randn(32, 32))
    B = Tensor(np.random.randn(32, 32))
    
    C_baseline = matmul_baseline(A, B)
    C_numpy = Tensor(np.dot(A.data, B.data))
    
    assert np.allclose(C_baseline.data, C_numpy.data, rtol=1e-10), "Baseline implementation differs from NumPy"
    print("✅ Baseline implementation matches NumPy")
    
    # Test case 4: Large matrix
    A = Tensor(np.random.randn(100, 100))
    B = Tensor(np.random.randn(100, 100))
    C = matmul_baseline(A, B)
    
    assert C.shape == (100, 100), f"Expected shape (100, 100), got {C.shape}"
    print("✅ Large matrix multiplication works")
    
    print("📈 Progress: Baseline Matrix Multiplication ✓")

# %% [markdown]
"""
## Step 2: Vectorized Operations - SIMD Principles

### What is Vectorization?
**Vectorization** means processing multiple data elements in parallel using SIMD (Single Instruction, Multiple Data) operations.

### The Problem with Loops
```python
# Scalar processing - one element at a time
def slow_relu(x):
    result = np.zeros_like(x)
    for i in range(len(x)):
        result[i] = max(0, x[i])  # One operation per cycle
    return result
```

### The Vectorization Solution
```python
# Vector processing - multiple elements at once
def fast_relu(x):
    return np.maximum(0, x)  # Many operations per cycle
```

### Why Vectorization Matters
- **CPU SIMD**: Modern CPUs can process 4-8 floats simultaneously
- **GPU parallelism**: GPUs have thousands of cores for parallel processing
- **Memory bandwidth**: Better utilization of memory transfers
- **Compiler optimization**: Enables automatic vectorization

### SIMD Principles
1. **Data parallelism**: Same operation on multiple data elements
2. **Memory alignment**: Aligned data enables faster SIMD instructions
3. **Batch processing**: Process data in chunks that fit SIMD registers
4. **Avoid branches**: Conditional operations break SIMD efficiency

### Real-World Context
- **NumPy**: All operations are vectorized using BLAS/LAPACK
- **PyTorch**: Vectorized operations compile to SIMD instructions
- **GPU kernels**: Thousands of parallel threads process data
- **AVX-512**: Intel's latest SIMD can process 16 floats at once
"""

# %% nbgrader={"grade": false, "grade_id": "vectorized-relu", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def vectorized_relu(x: Tensor) -> Tensor:
    """
    Vectorized ReLU implementation demonstrating SIMD principles.
    
    This function shows how to write operations that take advantage of
    CPU vectorization capabilities for better performance.
    
    TODO: Implement a vectorized ReLU that's optimized for performance.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Extract numpy array from Tensor
    2. Use NumPy's vectorized operations (these compile to SIMD instructions)
    3. Apply ReLU: f(x) = max(0, x) for all elements simultaneously
    4. Return result as Tensor
    
    VECTORIZATION TECHNIQUES:
    1. Use np.maximum instead of loops - this is vectorized
    2. Ensure input is contiguous in memory for better SIMD performance
    3. Consider using specific dtypes (float32 vs float64) for SIMD alignment
    4. Avoid conditional operations that break vectorization
    
    EXAMPLE USAGE:
    ```python
    x = Tensor([-2, -1, 0, 1, 2])
    y = vectorized_relu(x)
    # Expected: [0, 0, 0, 1, 2]
    ```
    
    PERFORMANCE CONSIDERATIONS:
    - np.maximum is vectorized and uses SIMD instructions
    - Memory layout matters: contiguous arrays are faster
    - Data type matters: float32 allows more SIMD parallelism than float64
    - Avoid Python loops - they can't be vectorized
    
    LEARNING CONNECTIONS:
    - This is how PyTorch's ReLU is implemented under the hood
    - GPU kernels use similar principles with thousands of parallel threads
    - Modern CPUs can process 4-16 floats simultaneously with SIMD
    """
    ### BEGIN SOLUTION
    # Extract numpy array
    x_data = x.data if hasattr(x, 'data') else x
    
    # Ensure contiguous memory layout for better SIMD performance
    if not x_data.flags.c_contiguous:
        x_data = np.ascontiguousarray(x_data)
    
    # Vectorized ReLU using NumPy's maximum function
    # This compiles to SIMD instructions on modern CPUs
    result = np.maximum(0, x_data)
    
    return Tensor(result)
    ### END SOLUTION

# %% nbgrader={"grade": false, "grade_id": "vectorized-operations", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def vectorized_operations(x: Tensor, y: Tensor) -> Dict[str, Tensor]:
    """
    Demonstration of various vectorized operations.
    
    Shows how multiple operations can be vectorized for better performance.
    
    TODO: Implement a collection of vectorized operations.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Extract numpy arrays from input Tensors
    2. Implement vectorized versions of common operations
    3. Use NumPy's built-in vectorized functions
    4. Return dictionary of results
    
    OPERATIONS TO IMPLEMENT:
    - element_wise_multiply: x * y (element-wise)
    - element_wise_add: x + y (element-wise)
    - squared_difference: (x - y)^2
    - euclidean_distance: sqrt(sum((x - y)^2))
    - dot_product: sum(x * y)
    
    VECTORIZATION PRINCIPLES:
    - Use NumPy operations instead of Python loops
    - Combine operations when possible: (x - y)**2 instead of subtract then square
    - Consider memory layout and data types
    - Measure performance improvements
    
    EXAMPLE USAGE:
    ```python
    x = Tensor([1, 2, 3, 4])
    y = Tensor([2, 3, 4, 5])
    results = vectorized_operations(x, y)
    # Returns dict with all vectorized operation results
    ```
    """
    ### BEGIN SOLUTION
    # Extract numpy arrays
    x_data = x.data if hasattr(x, 'data') else x
    y_data = y.data if hasattr(y, 'data') else y
    
    # Ensure arrays are the same shape for element-wise operations
    assert x_data.shape == y_data.shape, f"Shape mismatch: {x_data.shape} vs {y_data.shape}"
    
    # Vectorized operations
    results = {
        'element_wise_multiply': Tensor(x_data * y_data),
        'element_wise_add': Tensor(x_data + y_data),
        'squared_difference': Tensor((x_data - y_data) ** 2),
        'euclidean_distance': Tensor(np.sqrt(np.sum((x_data - y_data) ** 2))),
        'dot_product': Tensor(np.dot(x_data.flatten(), y_data.flatten()))
    }
    
    return results
    ### END SOLUTION

# %% nbgrader={"grade": false, "grade_id": "test-vectorized-operations", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Vectorized Operations

def test_unit_vectorized_operations():
    """Unit test for the vectorized operations implementation."""
    print("🔬 Unit Test: Vectorized Operations...")
    
    # Test vectorized ReLU
    x = Tensor([-2, -1, 0, 1, 2])
    y = vectorized_relu(x)
    expected = [0, 0, 0, 1, 2]
    
    assert np.allclose(y.data, expected), f"Expected {expected}, got {y.data}"
    print("✅ Vectorized ReLU works")
    
    # Test vectorized operations
    x = Tensor([1, 2, 3, 4])
    y = Tensor([2, 3, 4, 5])
    results = vectorized_operations(x, y)
    
    # Check element-wise multiply
    expected_mul = [2, 6, 12, 20]
    assert np.allclose(results['element_wise_multiply'].data, expected_mul), \
        f"Expected {expected_mul}, got {results['element_wise_multiply'].data}"
    print("✅ Element-wise multiply works")
    
    # Check element-wise add
    expected_add = [3, 5, 7, 9]
    assert np.allclose(results['element_wise_add'].data, expected_add), \
        f"Expected {expected_add}, got {results['element_wise_add'].data}"
    print("✅ Element-wise add works")
    
    # Check squared difference
    expected_sq_diff = [1, 1, 1, 1]  # (1-2)^2, (2-3)^2, etc.
    assert np.allclose(results['squared_difference'].data, expected_sq_diff), \
        f"Expected {expected_sq_diff}, got {results['squared_difference'].data}"
    print("✅ Squared difference works")
    
    # Check dot product
    expected_dot = 40  # 1*2 + 2*3 + 3*4 + 4*5 = 2 + 6 + 12 + 20 = 40
    assert np.allclose(results['dot_product'].data, expected_dot), \
        f"Expected {expected_dot}, got {results['dot_product'].data}"
    print("✅ Dot product works")
    
    print("📈 Progress: Vectorized Operations ✓")

# %% [markdown]
"""
## Step 3: Memory Layout Optimization - Cache-Friendly Algorithms

### Why Memory Layout Matters
Modern CPUs are **memory-bound**, not compute-bound. The bottleneck isn't how fast you can multiply numbers—it's how fast you can get data from memory.

### The Memory Hierarchy
```
CPU Registers:    1 cycle     (fastest, tiny)
L1 Cache:         3 cycles    (fast, small)
L2 Cache:         10 cycles   (medium, medium)
L3 Cache:         40 cycles   (slow, large)
Main Memory:      200+ cycles (slowest, huge)
```

### Cache-Friendly Principles
1. **Spatial locality**: Access nearby memory locations
2. **Temporal locality**: Reuse recently accessed data
3. **Cache lines**: Memory is loaded in 64-byte chunks
4. **Cache blocking**: Process data in cache-sized chunks

### Real-World Impact
- **Matrix multiplication**: Cache-friendly algorithms are 10x faster
- **Image processing**: Row-major vs column-major access patterns
- **Neural networks**: Memory layout affects training speed significantly

### The Problem with Naive Algorithms
```python
# Cache-unfriendly: jumps around memory
def slow_transpose(A):
    for i in range(rows):
        for j in range(cols):
            B[j, i] = A[i, j]  # Poor cache locality
```

### Cache-Friendly Solution
```python
# Cache-friendly: processes data in blocks
def fast_transpose(A):
    # Process in cache-sized blocks
    for block_i in range(0, rows, BLOCK_SIZE):
        for block_j in range(0, cols, BLOCK_SIZE):
            # Process block - good cache locality
            for i in range(block_i, min(block_i + BLOCK_SIZE, rows)):
                for j in range(block_j, min(block_j + BLOCK_SIZE, cols)):
                    B[j, i] = A[i, j]
```
"""

# %% nbgrader={"grade": false, "grade_id": "cache-friendly-matmul", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def cache_friendly_matmul(A: Tensor, B: Tensor, block_size: int = 32) -> Tensor:
    """
    Cache-friendly matrix multiplication using blocking technique.
    
    This implementation uses cache blocking to improve memory access patterns
    and achieve better performance on modern CPUs.
    
    TODO: Implement cache-friendly matrix multiplication using blocking.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Extract numpy arrays and get dimensions
    2. Pre-allocate output matrix
    3. Use three nested loops for blocks: block_i, block_j, block_k
    4. Within each block, use three nested loops for elements: i, j, k
    5. Process data in cache-sized blocks for better locality
    
    BLOCKING ALGORITHM:
    1. Divide matrices into blocks of size block_size x block_size
    2. For each block of C, compute contribution from corresponding A and B blocks
    3. This keeps data in cache longer, reducing memory access time
    
    CACHE OPTIMIZATION PRINCIPLES:
    - Process data in small blocks that fit in cache
    - Reuse data as much as possible while it's in cache
    - Access memory in predictable patterns
    - Minimize cache misses
    
    EXAMPLE USAGE:
    ```python
    A = Tensor([[1, 2], [3, 4]])
    B = Tensor([[5, 6], [7, 8]])
    C = cache_friendly_matmul(A, B, block_size=2)
    # Expected: [[19, 22], [43, 50]]
    ```
    
    PERFORMANCE HINTS:
    - block_size should be chosen based on cache size
    - Typical L1 cache: 32KB, so block_size=32 for float32 matrices
    - Experiment with different block sizes for your hardware
    - This algorithm is O(n^3) but with much better constants
    
    LEARNING CONNECTIONS:
    - This is how BLAS libraries achieve high performance
    - GPUs use similar tiling strategies for shared memory
    - Modern compilers can sometimes do this automatically
    """
    ### BEGIN SOLUTION
    # Extract numpy arrays
    A_data = A.data if hasattr(A, 'data') else A
    B_data = B.data if hasattr(B, 'data') else B
    
    # Get dimensions
    m, k = A_data.shape
    k2, n = B_data.shape
    assert k == k2, f"Cannot multiply {A_data.shape} and {B_data.shape}"
    
    # Pre-allocate output matrix
    C = np.zeros((m, n), dtype=A_data.dtype)
    
    # Cache-friendly blocked matrix multiplication
    for block_i in range(0, m, block_size):
        for block_j in range(0, n, block_size):
            for block_k in range(0, k, block_size):
                # Define block boundaries
                end_i = min(block_i + block_size, m)
                end_j = min(block_j + block_size, n)
                end_k = min(block_k + block_size, k)
                
                # Process block - good cache locality
                for i in range(block_i, end_i):
                    for j in range(block_j, end_j):
                        for k_idx in range(block_k, end_k):
                            C[i, j] += A_data[i, k_idx] * B_data[k_idx, j]
    
    return Tensor(C)
    ### END SOLUTION

# %% nbgrader={"grade": false, "grade_id": "test-cache-friendly", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Cache-Friendly Matrix Multiplication

def test_unit_cache_friendly_matmul():
    """Unit test for the cache-friendly matrix multiplication implementation."""
    print("🔬 Unit Test: Cache-Friendly Matrix Multiplication...")
    
    # Test case 1: Small matrices
    A = Tensor([[1, 2], [3, 4]])
    B = Tensor([[5, 6], [7, 8]])
    C = cache_friendly_matmul(A, B, block_size=2)
    expected = [[19, 22], [43, 50]]
    
    assert np.allclose(C.data, expected), f"Expected {expected}, got {C.data}"
    print("✅ Small matrix cache-friendly multiplication works")
    
    # Test case 2: Larger matrices with different block sizes
    np.random.seed(42)
    A = Tensor(np.random.randn(64, 64))
    B = Tensor(np.random.randn(64, 64))
    
    C_blocked = cache_friendly_matmul(A, B, block_size=16)
    C_numpy = Tensor(np.dot(A.data, B.data))
    
    assert np.allclose(C_blocked.data, C_numpy.data, rtol=1e-4), \
        "Cache-friendly implementation differs from NumPy"
    print("✅ Cache-friendly implementation matches NumPy")
    
    # Test case 3: Non-square matrices
    A = Tensor(np.random.randn(48, 32))
    B = Tensor(np.random.randn(32, 48))
    
    C_blocked = cache_friendly_matmul(A, B, block_size=8)
    C_numpy = Tensor(np.dot(A.data, B.data))
    
    assert np.allclose(C_blocked.data, C_numpy.data, rtol=1e-4), \
        "Non-square cache-friendly implementation differs from NumPy"
    print("✅ Non-square matrix cache-friendly multiplication works")
    
    print("📈 Progress: Cache-Friendly Algorithms ✓")

# %% [markdown]
"""
## Step 4: Parallel Processing - CPU and GPU-Style Computing

### Why Parallel Processing?
Modern hardware has multiple cores, and ML workloads are inherently parallel. We need to use all available compute resources.

### Types of Parallelism
1. **Data parallelism**: Split data across processors
2. **Task parallelism**: Split operations across processors
3. **Pipeline parallelism**: Different stages on different processors
4. **Model parallelism**: Split model across processors

### CPU vs GPU Parallelism
- **CPU**: Few cores (4-64), complex operations, low latency
- **GPU**: Many cores (1000s), simple operations, high throughput

### Parallel Processing Patterns
```python
# Sequential processing
for i in range(n):
    result[i] = expensive_operation(data[i])

# Parallel processing
with ThreadPoolExecutor() as executor:
    futures = [executor.submit(expensive_operation, data[i]) for i in range(n)]
    results = [f.result() for f in futures]
```

### Real-World Context
- **PyTorch**: Parallel data loading, distributed training
- **TensorFlow**: tf.data for parallel preprocessing
- **NumPy**: Multithreaded BLAS operations
- **GPU kernels**: Thousands of parallel threads
"""

# %% nbgrader={"grade": false, "grade_id": "parallel-relu", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def parallel_relu(x: Tensor, num_workers: int = 4) -> Tensor:
    """
    Parallel ReLU implementation using multiple CPU cores.
    
    This function demonstrates data parallelism by splitting the input
    across multiple worker processes.
    
    TODO: Implement parallel ReLU using multiprocessing or threading.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Extract numpy array from Tensor
    2. Split array into chunks for parallel processing
    3. Define worker function that applies ReLU to a chunk
    4. Use ThreadPoolExecutor to process chunks in parallel
    5. Combine results from all workers
    6. Return result as Tensor
    
    PARALLELIZATION STRATEGY:
    1. Split input into num_workers chunks
    2. Each worker processes its chunk independently
    3. Apply ReLU: max(0, x) to each chunk
    4. Combine results preserving original order
    
    EXAMPLE USAGE:
    ```python
    x = Tensor(np.random.randn(1000))
    y = parallel_relu(x, num_workers=4)
    # Processes data using 4 parallel workers
    ```
    
    PERFORMANCE CONSIDERATIONS:
    - Overhead of parallel processing may not be worth it for small arrays
    - Threading vs multiprocessing trade-offs
    - Chunk size should be large enough to amortize overhead
    - Consider memory bandwidth limitations
    
    LEARNING CONNECTIONS:
    - This is how PyTorch processes batches in parallel
    - GPUs naturally do this with thousands of parallel threads
    - Modern deep learning frameworks heavily use parallelism
    """
    ### BEGIN SOLUTION
    from concurrent.futures import ThreadPoolExecutor
    
    # Extract numpy array
    x_data = x.data if hasattr(x, 'data') else x
    
    # For small arrays, parallel processing isn't worth the overhead
    if x_data.size < 1000:
        return Tensor(np.maximum(0, x_data))
    
    # Split array into chunks
    chunk_size = max(1, x_data.size // num_workers)
    chunks = []
    flat_data = x_data.flatten()
    
    for i in range(0, len(flat_data), chunk_size):
        chunks.append(flat_data[i:i + chunk_size])
    
    # Worker function
    def relu_chunk(chunk):
        return np.maximum(0, chunk)
    
    # Process chunks in parallel
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        future_to_chunk = {executor.submit(relu_chunk, chunk): i for i, chunk in enumerate(chunks)}
        results = [None] * len(chunks)
        
        for future in future_to_chunk:
            chunk_idx = future_to_chunk[future]
            results[chunk_idx] = future.result()
    
    # Combine results
    combined_result = np.concatenate(results)
    
    # Reshape back to original shape
    result = combined_result.reshape(x_data.shape)
    
    return Tensor(result)
    ### END SOLUTION

# %% nbgrader={"grade": false, "grade_id": "parallel-batch-processing", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def parallel_batch_processing(batch_data: List[Tensor], operation: Callable, num_workers: int = 4) -> List[Tensor]:
    """
    Process a batch of tensors in parallel using multiple workers.
    
    This function demonstrates how to parallelize operations across
    multiple data samples, similar to how modern ML frameworks work.
    
    TODO: Implement parallel batch processing.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Take a list of Tensors and an operation function
    2. Use ThreadPoolExecutor to process multiple tensors simultaneously
    3. Apply the operation to each tensor in parallel
    4. Return list of results in original order
    
    PARALLELIZATION STRATEGY:
    1. Each worker processes one tensor at a time
    2. Multiple workers can process different tensors simultaneously
    3. Preserve order of results to match input order
    
    EXAMPLE USAGE:
    ```python
    batch = [Tensor(np.random.randn(100, 100)) for _ in range(8)]
    relu_op = lambda x: vectorized_relu(x)
    results = parallel_batch_processing(batch, relu_op, num_workers=4)
    # Processes 8 tensors using 4 parallel workers
    ```
    
    PERFORMANCE CONSIDERATIONS:
    - Each tensor should be large enough to justify parallel overhead
    - Balance number of workers with available CPU cores
    - Consider memory usage with multiple workers
    - Thread vs process pool trade-offs
    
    LEARNING CONNECTIONS:
    - This is how PyTorch's DataLoader processes batches
    - Similar to how GPUs process multiple samples simultaneously
    - Foundation for distributed training across multiple nodes
    """
    ### BEGIN SOLUTION
    from concurrent.futures import ThreadPoolExecutor
    
    # For small batches, parallel processing might not be worth it
    if len(batch_data) < num_workers:
        return [operation(tensor) for tensor in batch_data]
    
    # Process batch in parallel
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        # Submit all tasks
        future_to_index = {executor.submit(operation, tensor): i for i, tensor in enumerate(batch_data)}
        
        # Collect results in original order
        results = [None] * len(batch_data)
        for future in future_to_index:
            index = future_to_index[future]
            results[index] = future.result()
    
    return results
    ### END SOLUTION

# %% nbgrader={"grade": false, "grade_id": "test-parallel-processing", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Parallel Processing

def test_unit_parallel_processing():
    """Unit test for the parallel processing implementations."""
    print("🔬 Unit Test: Parallel Processing...")
    
    # Test parallel ReLU
    x = Tensor(np.array([-2, -1, 0, 1, 2]))
    y = parallel_relu(x, num_workers=2)
    expected = [0, 0, 0, 1, 2]
    
    assert np.allclose(y.data, expected), f"Expected {expected}, got {y.data}"
    print("✅ Parallel ReLU works")
    
    # Test parallel ReLU with larger data
    x_large = Tensor(np.random.randn(2000))
    y_large = parallel_relu(x_large, num_workers=4)
    y_sequential = vectorized_relu(x_large)
    
    assert np.allclose(y_large.data, y_sequential.data), \
        "Parallel ReLU differs from sequential version"
    print("✅ Parallel ReLU matches sequential version")
    
    # Test parallel batch processing
    batch = [Tensor(np.random.randn(100)) for _ in range(8)]
    relu_op = lambda x: vectorized_relu(x)
    
    results_parallel = parallel_batch_processing(batch, relu_op, num_workers=4)
    results_sequential = [relu_op(tensor) for tensor in batch]
    
    assert len(results_parallel) == len(results_sequential), \
        f"Expected {len(results_sequential)} results, got {len(results_parallel)}"
    
    for i, (parallel, sequential) in enumerate(zip(results_parallel, results_sequential)):
        assert np.allclose(parallel.data, sequential.data), \
            f"Batch item {i}: parallel differs from sequential"
    
    print("✅ Parallel batch processing works")
    print("📈 Progress: Parallel Processing ✓")

# Test will be run in main block

# %% [markdown]
"""
## Step 5: Simple Performance Measurement - Timing Your Kernels

### Why Timing Matters
> "Premature optimization is the root of all evil" - Donald Knuth

But **measured optimization** based on simple timing is essential for understanding kernel performance.

### What We'll Measure
1. **Execution time**: How long does each kernel take?
2. **Relative performance**: Which implementation is faster?
3. **Scale effects**: How does performance change with data size?
4. **Optimization impact**: Did our changes actually help?

### The Simple Timing Process
1. **Measure baseline**: Time the standard implementation
2. **Time optimizations**: Measure your improved versions
3. **Compare results**: See which is faster
4. **Verify correctness**: Ensure optimized code produces correct results

### Our Simple Timing Tool
We use `time.perf_counter()` for microsecond-precision timing:
- **Precise**: Measures actual execution time
- **Simple**: Easy to understand and use
- **Realistic**: Shows kernel performance at the right scale
- **Educational**: Immediate feedback on optimization impact

### Real-World Context
- **Kernel operations**: Typically take 10-1000 microseconds
- **Optimization impact**: Good kernels are 2-10x faster
- **Professional tools**: Production systems use sophisticated profilers
- **Foundation**: Simple timing teaches measurement principles
"""

# %% nbgrader={"grade": false, "grade_id": "test-profiling", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Simple Kernel Timing

def test_unit_simple_kernel_timing():
    """Unit test for the simple kernel timing capabilities."""
    print("🔬 Unit Test: Simple Kernel Timing...")
    
    # Test timing different matrix multiplication methods
    np.random.seed(42)
    A = Tensor(np.random.randn(100, 100))
    B = Tensor(np.random.randn(100, 100))
    
    # Time NumPy matmul
    result_numpy, time_numpy = time_kernel(lambda: Tensor(np.dot(A.data, B.data)))
    print(f"🔍 NumPy matmul: {time_numpy:.1f} μs")
    
    # Time baseline matmul  
    result_baseline, time_baseline = time_kernel(matmul_baseline, A, B)
    print(f"🔍 Baseline matmul: {time_baseline:.1f} μs")
    
    # Time cache-friendly matmul
    result_cache, time_cache = time_kernel(cache_friendly_matmul, A, B, 16)
    print(f"🔍 Cache-friendly matmul: {time_cache:.1f} μs")
    
    # Verify results are similar
    assert np.allclose(result_numpy.data, result_baseline.data, rtol=1e-4), \
        "NumPy and baseline results differ"
    assert np.allclose(result_numpy.data, result_cache.data, rtol=1e-2), \
        "NumPy and cache-friendly results differ"
    
    print("✅ All matrix multiplication methods produce correct results")
    
    # Test timing parallel vs sequential ReLU
    x_large = Tensor(np.random.randn(10000))
    
    result_seq, time_seq = time_kernel(vectorized_relu, x_large)
    result_par, time_par = time_kernel(parallel_relu, x_large, 4)
    
    print(f"🔍 Sequential ReLU: {time_seq:.1f} μs")
    print(f"🔍 Parallel ReLU: {time_par:.1f} μs")
    
    # Verify results are the same
    assert np.allclose(result_seq.data, result_par.data), \
        "Sequential and parallel ReLU results differ"
    
    print("✅ Simple timing works correctly")
    print("📈 Progress: Simple Kernel Timing ✓")

# Test will be run in main block

# %% [markdown]
"""
## Step 6: Compressed Model Kernels - Optimizing Quantized Operations

### Why Compressed Model Kernels?
Modern deployment requires smaller, faster models:
- **Mobile devices**: Limited compute and memory
- **Edge computing**: Real-time inference requirements
- **Cloud costs**: Reduce computational expenses
- **Energy efficiency**: Lower power consumption

### Types of Model Compression
1. **Quantization**: Reduce precision (float32 → int8)
2. **Pruning**: Remove unimportant weights
3. **Knowledge distillation**: Train smaller models
4. **Low-rank approximation**: Factorize weight matrices

### Quantization Fundamentals
```python
# Original: 32-bit floating point
weights_fp32 = np.array([1.234, -0.567, 2.891])

# Quantized: 8-bit integer
scale = max(weights_fp32) / 127
weights_int8 = np.round(weights_fp32 / scale).astype(np.int8)

# Dequantized for computation
weights_dequant = weights_int8 * scale
```

### Why Custom Kernels for Compression?
- **Integer arithmetic**: Faster than floating-point on many devices
- **Memory bandwidth**: 4x less data to transfer
- **Specialized instructions**: CPUs have optimized int8 operations
- **Accumulation**: Need to handle precision carefully

### Real-World Context
- **TensorFlow Lite**: Quantized inference kernels
- **PyTorch Mobile**: Optimized int8 operations
- **ONNX Runtime**: Hardware-specific quantized kernels
- **Hardware accelerators**: TPUs, Neural Processing Units
"""

# %% nbgrader={"grade": false, "grade_id": "quantized-matmul", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def quantized_matmul(A: Tensor, B: Tensor, scale_A: float = 1.0, scale_B: float = 1.0) -> Tensor:
    """
    Quantized matrix multiplication kernel for compressed models.
    
    This function demonstrates how to perform matrix multiplication
    with quantized (int8) weights while maintaining numerical accuracy.
    
    TODO: Implement quantized matrix multiplication.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Extract numpy arrays from Tensors
    2. Quantize inputs to int8 using provided scales
    3. Perform integer matrix multiplication
    4. Rescale result back to appropriate range
    5. Return result as Tensor
    
    QUANTIZATION PROCESS:
    1. Quantize: int8_value = round(float_value / scale)
    2. Compute: int8_result = int8_A @ int8_B
    3. Rescale: float_result = int8_result * scale_A * scale_B
    
    EXAMPLE USAGE:
    ```python
    A = Tensor([[1.0, 2.0], [3.0, 4.0]])
    B = Tensor([[0.5, 1.5], [2.5, 3.5]])
    C = quantized_matmul(A, B, scale_A=1.0/127, scale_B=1.0/127)
    # Should approximate regular matrix multiplication
    ```
    
    PERFORMANCE CONSIDERATIONS:
    - int8 operations are often faster than float32
    - Memory usage is 4x lower
    - Accumulation in int32 to prevent overflow
    - Careful handling of scales to maintain precision
    
    LEARNING CONNECTIONS:
    - This is how TensorFlow Lite performs quantized inference
    - Similar to how mobile ML accelerators work
    - Foundation for edge deployment of neural networks
    """
    ### BEGIN SOLUTION
    # Extract numpy arrays
    A_data = A.data if hasattr(A, 'data') else A
    B_data = B.data if hasattr(B, 'data') else B
    
    # Quantize inputs to int8
    A_int8 = np.round(A_data / scale_A).astype(np.int8)
    B_int8 = np.round(B_data / scale_B).astype(np.int8)
    
    # Perform integer matrix multiplication
    # Use int32 for accumulation to prevent overflow
    C_int32 = np.dot(A_int8.astype(np.int32), B_int8.astype(np.int32))
    
    # Rescale result back to float
    C_float = C_int32 * scale_A * scale_B
    
    return Tensor(C_float)
    ### END SOLUTION

# %% nbgrader={"grade": false, "grade_id": "quantized-relu", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
def quantized_relu(x: Tensor, scale: float = 1.0) -> Tensor:
    """
    Quantized ReLU implementation for compressed models.
    
    This function shows how to apply ReLU activation to quantized values
    while maintaining the quantization format.
    
    TODO: Implement quantized ReLU activation.
    
    STEP-BY-STEP IMPLEMENTATION:
    1. Extract numpy array from Tensor
    2. Quantize input to int8 using provided scale
    3. Apply ReLU in integer domain: max(0, x)
    4. Keep result in int8 format (no rescaling needed for ReLU)
    5. Convert back to float using scale
    6. Return result as Tensor
    
    QUANTIZED RELU PROCESS:
    1. Quantize: int8_value = round(float_value / scale)
    2. Apply ReLU: int8_result = max(0, int8_value)
    3. Dequantize: float_result = int8_result * scale
    
    EXAMPLE USAGE:
    ```python
    x = Tensor([-1.0, 0.0, 1.0, 2.0])
    y = quantized_relu(x, scale=1.0/127)
    # Should produce [0.0, 0.0, 1.0, 2.0] (approximately)
    ```
    
    OPTIMIZATION NOTES:
    - ReLU in int8 is just max(0, x) - very fast
    - No floating-point operations needed during activation
    - Maintains quantization format throughout
    - Can be vectorized efficiently
    
    LEARNING CONNECTIONS:
    - This is how quantized neural networks maintain speed
    - Similar to how mobile processors optimize ML inference
    - Foundation for real-time edge computing applications
    """
    ### BEGIN SOLUTION
    # Extract numpy array
    x_data = x.data if hasattr(x, 'data') else x
    
    # Quantize input to int8
    x_int8 = np.round(x_data / scale).astype(np.int8)
    
    # Apply ReLU in integer domain
    x_relu_int8 = np.maximum(0, x_int8)
    
    # Convert back to float
    x_relu_float = x_relu_int8 * scale
    
    return Tensor(x_relu_float)
    ### END SOLUTION

# %% nbgrader={"grade": false, "grade_id": "test-compressed-kernels", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Compressed Model Kernels

def test_unit_compressed_kernels():
    """Unit test for the compressed model kernel implementations."""
    print("🔬 Unit Test: Compressed Model Kernels...")
    
    # Test quantized matrix multiplication
    A = Tensor([[1.0, 2.0], [3.0, 4.0]])
    B = Tensor([[0.5, 1.5], [2.5, 3.5]])
    
    # Regular matrix multiplication
    C_regular = matmul_baseline(A, B)
    
    # Quantized matrix multiplication
    # Use larger scales to prevent int8 overflow
    scale_A = 1.0 / 20  # Max value 4.0 / (1/20) = 80, fits in int8
    scale_B = 1.0 / 20  # Max value 3.5 / (1/20) = 70, fits in int8
    C_quantized = quantized_matmul(A, B, scale_A, scale_B)
    
    # Should be approximately equal (some quantization error expected)
    assert np.allclose(C_regular.data, C_quantized.data, rtol=0.1), \
        f"Regular: {C_regular.data}, Quantized: {C_quantized.data}"
    print("✅ Quantized matrix multiplication works")
    
    # Test quantized ReLU
    x = Tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
    
    # Regular ReLU
    y_regular = vectorized_relu(x)
    
    # Quantized ReLU
    # Use larger scale to prevent int8 overflow
    scale = 1.0 / 50  # Max value 2.0 / (1/50) = 100, fits in int8
    y_quantized = quantized_relu(x, scale)
    
    # Should be approximately equal
    assert np.allclose(y_regular.data, y_quantized.data, rtol=0.1), \
        f"Regular: {y_regular.data}, Quantized: {y_quantized.data}"
    print("✅ Quantized ReLU works")
    
    # Test that quantized operations can be timed
    # This shows the performance characteristics of quantized vs regular operations
    x_large = Tensor(np.random.randn(1000))
    
    # Time regular ReLU
    _, time_regular = time_kernel(vectorized_relu, x_large)
    
    # Time quantized ReLU
    _, time_quantized = time_kernel(quantized_relu, x_large, 1.0/127)
    
    print(f"🔍 Regular ReLU: {time_regular:.1f} μs")
    print(f"🔍 Quantized ReLU: {time_quantized:.1f} μs")
    
    print("✅ Quantized operations timing works")
    print("📈 Progress: Compressed Model Kernels ✓")

# Test will be run in main block

# %% nbgrader={"grade": false, "grade_id": "final-performance-test", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Comprehensive Kernel Performance Comparison

def final_performance_test():
    """Comprehensive performance test of all implemented kernels."""
    print("🔬 Final Performance Test: Comprehensive Kernel Comparison")
    print("=" * 60)
    
    # Create test data
    np.random.seed(42)
    A = Tensor(np.random.randn(256, 256))
    B = Tensor(np.random.randn(256, 256))
    x = Tensor(np.random.randn(10000))
    
    print("\n📊 Matrix Multiplication Performance:")
    print("-" * 40)
    
    # Test different matrix multiplication methods
    methods = [
        ("NumPy", lambda: Tensor(np.dot(A.data, B.data))),
        ("Baseline", lambda: matmul_baseline(A, B)),
        ("Cache-friendly", lambda: cache_friendly_matmul(A, B, 32)),
        ("Quantized", lambda: quantized_matmul(A, B, 1.0/127, 1.0/127))
    ]
    
    results = {}
    for name, method in methods:
        result, time_us = time_kernel(method)
        results[name] = (result, time_us)
        print(f"{name:15}: {time_us:.1f} μs")
    
    print("\n📊 ReLU Activation Performance:")
    print("-" * 40)
    
    # Test different ReLU methods
    relu_methods = [
        ("Vectorized", lambda: vectorized_relu(x)),
        ("Parallel", lambda: parallel_relu(x, 4)),
        ("Quantized", lambda: quantized_relu(x, 1.0/127))
    ]
    
    relu_results = {}
    for name, method in relu_methods:
        result, time_us = time_kernel(method)
        relu_results[name] = (result, time_us)
        print(f"{name:15}: {time_us:.1f} μs")
    
    print("\n✅ All kernels implemented successfully!")
    print("📈 Progress: Complete Kernels Module ✓")
    
    # Verify correctness
    print("\n🔍 Correctness Verification:")
    print("-" * 40)
    
    # Check that all matrix multiplication methods produce similar results
    base_result = results["NumPy"][0]
    for name, (result, _) in results.items():
        if name != "NumPy":
            if name == "Quantized":
                # Skip quantized comparison in final test - already validated individually
                print(f"⚠️  Skipping {name} comparison (quantization errors expected)")
            else:
                assert np.allclose(base_result.data, result.data, rtol=1e-2), \
                    f"{name} differs from NumPy"
    
    # Check that all ReLU methods produce similar results
    base_relu = relu_results["Vectorized"][0]
    for name, (result, _) in relu_results.items():
        if name != "Vectorized":
            if name == "Quantized":
                # Skip quantized ReLU comparison - already validated individually
                print(f"⚠️  Skipping {name} ReLU comparison (quantization errors expected)")
            else:
                assert np.allclose(base_relu.data, result.data, rtol=1e-4), \
                    f"{name} ReLU differs from vectorized"
    
    print("✅ All implementations produce correct results!")
    
    print("\n🎉 CONGRATULATIONS! 🎉")
    print("You've successfully implemented hardware-optimized ML kernels!")
    print("You now understand the performance optimizations that power modern AI frameworks.")

# Run the final test
if __name__ == "__main__":
    # Run individual kernel tests
    test_unit_matmul_baseline()
    test_unit_vectorized_operations()
    test_unit_cache_friendly_matmul()
    
    # Run final performance test
    final_performance_test()

# %% [markdown]
"""
## Step 7: ML Systems - Production Kernel Optimization Profiler

### GPU Architecture and Custom Kernels in Production ML

In production ML systems, kernel optimization becomes critical for performance and cost efficiency. Modern ML frameworks rely on thousands of specialized kernels that are optimized for specific hardware architectures and use cases.

### The Production Reality
Real ML deployments face:
- **Inference latency**: Sub-millisecond requirements for real-time applications
- **Throughput demands**: Processing millions of requests per second
- **Hardware diversity**: CPUs, GPUs, TPUs, custom ASICs
- **Memory constraints**: Limited bandwidth and capacity
- **Energy efficiency**: Power consumption in data centers and edge devices

### GPU Kernel Optimization Patterns
Modern GPUs require specialized optimization techniques:
- **Memory coalescing**: Optimizing memory access patterns for GPU memory hierarchy
- **Warp divergence analysis**: Ensuring efficient execution across GPU thread warps
- **Shared memory optimization**: Leveraging fast on-chip memory for data reuse
- **Tensor core utilization**: Maximizing mixed-precision compute throughput
- **Kernel fusion**: Combining multiple operations to reduce memory overhead
- **Multi-GPU scaling**: Coordinating computation across multiple devices

### Real-World Context
- **NVIDIA cuDNN**: Thousands of optimized GPU kernels for deep learning
- **Intel oneDNN**: CPU-optimized kernels for inference
- **Triton**: Python-like language for writing GPU kernels
- **TensorRT**: Runtime optimization for NVIDIA GPUs
- **Custom silicon**: TPUs, AWS Inferentia, Apple Neural Engine
"""

# %% nbgrader={"grade": false, "grade_id": "kernel-optimization-profiler", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
class KernelOptimizationProfiler:
    """
    Production-grade kernel optimization profiler for ML systems.
    
    This class provides comprehensive analysis tools for optimizing ML kernels
    across different hardware architectures, focusing on GPU optimization patterns
    and production deployment scenarios.
    
    Key Features:
    - CUDA kernel performance analysis
    - Memory coalescing pattern detection
    - Warp divergence analysis
    - Shared memory optimization
    - Tensor core utilization metrics
    - Kernel fusion opportunities
    - Multi-GPU scaling analysis
    """
    
    def __init__(self, hardware_config: Optional[Dict[str, Any]] = None):
        """
        Initialize the kernel optimization profiler.
        
        Args:
            hardware_config: Dictionary containing hardware specifications
        """
        self.hardware_config = hardware_config or self._detect_hardware()
        self.profile_results = {}
        self.optimization_recommendations = []
        
    def _detect_hardware(self) -> Dict[str, Any]:
        """Detect current hardware configuration."""
        return {
            'cpu_cores': psutil.cpu_count(),
            'memory_gb': psutil.virtual_memory().total // (1024**3),
            'cache_sizes': {
                'l1': 32768,  # Typical L1 cache size in bytes
                'l2': 262144,  # Typical L2 cache size in bytes  
                'l3': 8388608  # Typical L3 cache size in bytes
            },
            'gpu_available': False,  # Would check for CUDA/OpenCL in real implementation
            'gpu_memory_gb': 0,
            'tensor_cores': False,
            'warp_size': 32  # NVIDIA GPU warp size
        }
    
    def analyze_cuda_kernel_performance(self, kernel_func: Callable, input_data: Tensor, 
                                      iterations: int = 100) -> Dict[str, Any]:
        """
        Analyze CUDA kernel performance characteristics.
        
        In a real implementation, this would interface with CUDA profiling tools
        to measure actual GPU kernel performance metrics.
        """
        # Simulate CUDA kernel analysis
        total_time = 0
        memory_bandwidth = 0
        compute_utilization = 0
        
        for _ in range(iterations):
            result, execution_time = time_kernel(kernel_func, input_data)
            total_time += execution_time
            
            # Simulate GPU metrics calculation
            data_size = input_data.data.nbytes
            memory_bandwidth += (data_size * 2) / (execution_time / 1_000_000)  # Read + Write
            compute_utilization += np.random.uniform(0.3, 0.9)  # Simulated utilization
        
        avg_time = total_time / iterations
        avg_bandwidth = memory_bandwidth / iterations
        avg_utilization = compute_utilization / iterations
        
        analysis = {
            'avg_execution_time_us': avg_time,
            'memory_bandwidth_gb_s': avg_bandwidth / (1024**3),
            'compute_utilization': avg_utilization,
            'theoretical_peak_bandwidth': 900,  # GB/s for high-end GPU
            'bandwidth_efficiency': min(100, (avg_bandwidth / (1024**3)) / 900 * 100),
            'bottleneck_analysis': self._identify_bottlenecks(avg_bandwidth / (1024**3), avg_utilization)
        }
        
        self.profile_results['cuda_analysis'] = analysis
        return analysis
    
    def analyze_memory_coalescing(self, access_pattern: str, data_shape: Tuple[int, ...]) -> Dict[str, Any]:
        """
        Analyze memory access patterns for GPU coalescing efficiency.
        
        Memory coalescing is critical for GPU performance - threads in a warp
        should access contiguous memory locations.
        """
        coalescing_efficiency = 1.0
        
        if access_pattern == 'row_major':
            # Good coalescing for row-major access
            coalescing_efficiency = 0.95
        elif access_pattern == 'column_major':
            # Poor coalescing for column-major access
            coalescing_efficiency = 0.3
        elif access_pattern == 'strided':
            # Moderate coalescing for strided access
            stride = data_shape[1] if len(data_shape) > 1 else 1
            coalescing_efficiency = max(0.1, 1.0 / stride)
        elif access_pattern == 'random':
            # Very poor coalescing for random access
            coalescing_efficiency = 0.1
        
        analysis = {
            'access_pattern': access_pattern,
            'data_shape': data_shape,
            'coalescing_efficiency': coalescing_efficiency,
            'memory_transactions': self._calculate_memory_transactions(data_shape, coalescing_efficiency),
            'optimization_potential': 1.0 - coalescing_efficiency
        }
        
        self.profile_results['memory_coalescing'] = analysis
        return analysis
    
    def analyze_warp_divergence(self, conditional_operations: int, total_operations: int) -> Dict[str, Any]:
        """
        Analyze warp divergence patterns in kernel execution.
        
        Warp divergence occurs when threads in a warp take different execution paths,
        reducing parallelism efficiency.
        """
        divergence_ratio = conditional_operations / total_operations
        efficiency_loss = divergence_ratio * 0.5  # Simplified model
        
        analysis = {
            'conditional_operations': conditional_operations,
            'total_operations': total_operations,
            'divergence_ratio': divergence_ratio,
            'efficiency_loss': efficiency_loss,
            'warp_efficiency': 1.0 - efficiency_loss,
            'optimization_suggestions': self._generate_divergence_optimizations(divergence_ratio)
        }
        
        self.profile_results['warp_divergence'] = analysis
        return analysis
    
    def analyze_shared_memory_usage(self, kernel_data_size: int, reuse_factor: float) -> Dict[str, Any]:
        """
        Analyze shared memory optimization opportunities.
        
        Shared memory is fast on-chip memory that can dramatically improve
        performance when used effectively for data reuse.
        """
        shared_memory_size = 48 * 1024  # 48KB typical shared memory per SM
        bank_conflicts = self._estimate_bank_conflicts(kernel_data_size)
        
        analysis = {
            'data_size_bytes': kernel_data_size,
            'shared_memory_available': shared_memory_size,
            'utilization_ratio': min(1.0, kernel_data_size / shared_memory_size),
            'reuse_factor': reuse_factor,
            'bank_conflicts': bank_conflicts,
            'performance_gain': min(10.0, reuse_factor * (1.0 - bank_conflicts)),
            'optimization_opportunities': self._identify_shared_memory_optimizations(kernel_data_size, reuse_factor)
        }
        
        self.profile_results['shared_memory'] = analysis
        return analysis
    
    def analyze_tensor_core_utilization(self, operation_type: str, data_types: List[str]) -> Dict[str, Any]:
        """
        Analyze tensor core utilization for mixed-precision operations.
        
        Tensor cores provide massive acceleration for mixed-precision matrix operations
        when data shapes and types are optimized correctly.
        """
        tensor_core_compatible = (
            operation_type in ['matmul', 'conv2d'] and
            any(dtype in ['float16', 'bfloat16', 'int8'] for dtype in data_types)
        )
        
        if tensor_core_compatible:
            theoretical_speedup = 4.0  # Typical tensor core speedup
            actual_utilization = 0.7   # Realistic utilization
        else:
            theoretical_speedup = 1.0
            actual_utilization = 0.0
        
        analysis = {
            'operation_type': operation_type,
            'data_types': data_types,
            'tensor_core_compatible': tensor_core_compatible,
            'theoretical_speedup': theoretical_speedup,
            'actual_utilization': actual_utilization,
            'performance_gain': theoretical_speedup * actual_utilization,
            'optimization_requirements': self._get_tensor_core_requirements()
        }
        
        self.profile_results['tensor_core'] = analysis
        return analysis
    
    def analyze_kernel_fusion_opportunities(self, operation_sequence: List[str]) -> Dict[str, Any]:
        """
        Analyze opportunities for kernel fusion to reduce memory overhead.
        
        Kernel fusion combines multiple operations into a single kernel,
        reducing memory bandwidth requirements and improving performance.
        """
        fusable_patterns = [
            ['matmul', 'relu'],
            ['conv2d', 'batchnorm', 'relu'],
            ['add', 'relu'],
            ['mul', 'add']
        ]
        
        fusion_opportunities = []
        memory_savings = 0
        
        for pattern in fusable_patterns:
            if self._sequence_contains_pattern(operation_sequence, pattern):
                fusion_opportunities.append(pattern)
                memory_savings += len(pattern) - 1  # Save intermediate results
        
        analysis = {
            'operation_sequence': operation_sequence,
            'fusion_opportunities': fusion_opportunities,
            'memory_savings_factor': memory_savings,
            'performance_improvement': min(2.0, 1 + memory_savings * 0.3),
            'implementation_complexity': len(fusion_opportunities) * 2
        }
        
        self.profile_results['kernel_fusion'] = analysis
        return analysis
    
    def analyze_multi_gpu_scaling(self, data_size: int, num_gpus: int) -> Dict[str, Any]:
        """
        Analyze multi-GPU scaling patterns and communication overhead.
        
        Multi-GPU deployments require careful optimization of data distribution
        and communication patterns to achieve good scaling efficiency.
        """
        communication_overhead = self._calculate_communication_overhead(data_size, num_gpus)
        compute_scaling = min(num_gpus, data_size / 1000)  # Simplified scaling model
        
        analysis = {
            'data_size': data_size,
            'num_gpus': num_gpus,
            'communication_overhead': communication_overhead,
            'compute_scaling': compute_scaling,
            'scaling_efficiency': compute_scaling / num_gpus,
            'bottleneck_type': 'communication' if communication_overhead > 0.3 else 'compute',
            'optimization_strategies': self._get_multi_gpu_optimizations(communication_overhead)
        }
        
        self.profile_results['multi_gpu'] = analysis
        return analysis
    
    def generate_optimization_report(self) -> str:
        """Generate comprehensive optimization report with recommendations."""
        report = ["🚀 Kernel Optimization Analysis Report", "=" * 50, ""]
        
        for analysis_type, results in self.profile_results.items():
            report.append(f"📊 {analysis_type.replace('_', ' ').title()} Analysis:")
            report.append("-" * 30)
            
            for key, value in results.items():
                if isinstance(value, float):
                    report.append(f"  {key}: {value:.3f}")
                elif isinstance(value, list):
                    report.append(f"  {key}: {', '.join(map(str, value))}")
                else:
                    report.append(f"  {key}: {value}")
            report.append("")
        
        # Add optimization recommendations
        report.append("🎯 Optimization Recommendations:")
        report.append("-" * 30)
        for rec in self.optimization_recommendations:
            report.append(f"  • {rec}")
        
        return "\n".join(report)
    
    # Helper methods
    def _identify_bottlenecks(self, bandwidth_gb_s: float, utilization: float) -> str:
        """Identify performance bottlenecks."""
        if bandwidth_gb_s < 100:
            return "Memory bandwidth limited"
        elif utilization < 0.5:
            return "Compute utilization limited"
        else:
            return "Well balanced"
    
    def _calculate_memory_transactions(self, shape: Tuple[int, ...], efficiency: float) -> int:
        """Calculate memory transaction count."""
        total_elements = np.prod(shape)
        return int(total_elements / (32 * efficiency))  # 32 threads per warp
    
    def _generate_divergence_optimizations(self, divergence_ratio: float) -> List[str]:
        """Generate warp divergence optimization suggestions."""
        suggestions = []
        if divergence_ratio > 0.3:
            suggestions.append("Reduce conditional operations in inner loops")
            suggestions.append("Use predicated execution instead of branching")
        if divergence_ratio > 0.5:
            suggestions.append("Restructure algorithm to minimize thread divergence")
        return suggestions
    
    def _estimate_bank_conflicts(self, data_size: int) -> float:
        """Estimate shared memory bank conflicts."""
        # Simplified model - assumes some degree of bank conflicts
        return min(0.5, data_size / (32 * 4))  # 32 banks, 4 bytes per bank
    
    def _identify_shared_memory_optimizations(self, size: int, reuse: float) -> List[str]:
        """Identify shared memory optimization opportunities."""
        optimizations = []
        if reuse > 2.0:
            optimizations.append("High reuse factor - shared memory beneficial")
        if size < 16384:  # 16KB
            optimizations.append("Data fits in shared memory - implement tiling")
        return optimizations
    
    def _get_tensor_core_requirements(self) -> List[str]:
        """Get tensor core optimization requirements."""
        return [
            "Use mixed precision (float16/bfloat16)",
            "Ensure matrix dimensions are multiples of 8",
            "Use proper memory layout (NHWC for convolutions)"
        ]
    
    def _sequence_contains_pattern(self, sequence: List[str], pattern: List[str]) -> bool:
        """Check if operation sequence contains fusable pattern."""
        for i in range(len(sequence) - len(pattern) + 1):
            if sequence[i:i+len(pattern)] == pattern:
                return True
        return False
    
    def _calculate_communication_overhead(self, data_size: int, num_gpus: int) -> float:
        """Calculate multi-GPU communication overhead."""
        # Simplified model based on data size and GPU count
        return min(0.8, (data_size / 1000) / num_gpus + 0.1)
    
    def _get_multi_gpu_optimizations(self, overhead: float) -> List[str]:
        """Get multi-GPU optimization strategies."""
        strategies = []
        if overhead > 0.3:
            strategies.append("Implement gradient compression")
            strategies.append("Use asynchronous communication")
        if overhead > 0.5:
            strategies.append("Increase batch size to amortize communication")
        return strategies

# %% nbgrader={"grade": false, "grade_id": "test-kernel-profiler", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Kernel Optimization Profiler

def test_unit_kernel_optimization_profiler():
    """Unit test for the kernel optimization profiler."""
    print("🔬 Unit Test: Kernel Optimization Profiler...")
    
    # Create profiler instance
    profiler = KernelOptimizationProfiler()
    
    # Test CUDA kernel analysis
    x = Tensor(np.random.randn(1000))
    cuda_analysis = profiler.analyze_cuda_kernel_performance(vectorized_relu, x, iterations=10)
    
    assert 'avg_execution_time_us' in cuda_analysis
    assert 'memory_bandwidth_gb_s' in cuda_analysis
    assert 'compute_utilization' in cuda_analysis
    print("✅ CUDA kernel analysis works")
    
    # Test memory coalescing analysis
    memory_analysis = profiler.analyze_memory_coalescing('row_major', (1024, 1024))
    
    assert memory_analysis['coalescing_efficiency'] > 0.9
    assert 'optimization_potential' in memory_analysis
    print("✅ Memory coalescing analysis works")
    
    # Test warp divergence analysis
    warp_analysis = profiler.analyze_warp_divergence(100, 1000)
    
    assert warp_analysis['divergence_ratio'] == 0.1
    assert 'warp_efficiency' in warp_analysis
    print("✅ Warp divergence analysis works")
    
    # Test shared memory analysis
    shared_analysis = profiler.analyze_shared_memory_usage(16384, 3.0)
    
    assert 'performance_gain' in shared_analysis
    assert shared_analysis['reuse_factor'] == 3.0
    print("✅ Shared memory analysis works")
    
    # Test tensor core analysis
    tensor_analysis = profiler.analyze_tensor_core_utilization('matmul', ['float16'])
    
    assert tensor_analysis['tensor_core_compatible'] == True
    assert tensor_analysis['theoretical_speedup'] > 1.0
    print("✅ Tensor core analysis works")
    
    # Test kernel fusion analysis
    fusion_analysis = profiler.analyze_kernel_fusion_opportunities(['matmul', 'relu', 'add'])
    
    assert len(fusion_analysis['fusion_opportunities']) > 0
    assert 'performance_improvement' in fusion_analysis
    print("✅ Kernel fusion analysis works")
    
    # Test multi-GPU analysis
    gpu_analysis = profiler.analyze_multi_gpu_scaling(10000, 4)
    
    assert gpu_analysis['num_gpus'] == 4
    assert 'scaling_efficiency' in gpu_analysis
    print("✅ Multi-GPU analysis works")
    
    # Test report generation
    report = profiler.generate_optimization_report()
    
    assert "Kernel Optimization Analysis Report" in report
    assert len(report) > 100  # Should be a substantial report
    print("✅ Optimization report generation works")
    
    print("📈 Progress: Kernel Optimization Profiler ✓")

# Run the test
test_unit_kernel_optimization_profiler()

# %% [markdown]
"""
## Step 7: ML Systems - Production Kernel Optimization Profiler

### GPU Architecture and Custom Kernels in Production ML

In production ML systems, kernel optimization becomes critical for performance and cost efficiency. Modern ML frameworks rely on thousands of specialized kernels that are optimized for specific hardware architectures and use cases.

### The Production Reality
Real ML deployments face:
- **Inference latency**: Sub-millisecond requirements for real-time applications
- **Throughput demands**: Processing millions of requests per second
- **Hardware diversity**: CPUs, GPUs, TPUs, custom ASICs
- **Memory constraints**: Limited bandwidth and capacity
- **Energy efficiency**: Power consumption in data centers and edge devices

### GPU Kernel Optimization Patterns
Modern GPUs require specialized optimization techniques:
- **Memory coalescing**: Optimizing memory access patterns for GPU memory hierarchy
- **Warp divergence analysis**: Ensuring efficient execution across GPU thread warps
- **Shared memory optimization**: Leveraging fast on-chip memory for data reuse
- **Tensor core utilization**: Maximizing mixed-precision compute throughput
- **Kernel fusion**: Combining multiple operations to reduce memory overhead
- **Multi-GPU scaling**: Coordinating computation across multiple devices

### Real-World Context
- **NVIDIA cuDNN**: Thousands of optimized GPU kernels for deep learning
- **Intel oneDNN**: CPU-optimized kernels for inference
- **Triton**: Python-like language for writing GPU kernels
- **TensorRT**: Runtime optimization for NVIDIA GPUs
- **Custom silicon**: TPUs, AWS Inferentia, Apple Neural Engine
"""

# %% nbgrader={"grade": false, "grade_id": "kernel-optimization-profiler", "locked": false, "schema_version": 3, "solution": true, "task": false}
#| export
class KernelOptimizationProfiler:
    """
    Production-grade kernel optimization profiler for ML systems.
    
    This class provides comprehensive analysis tools for optimizing ML kernels
    across different hardware architectures, focusing on GPU optimization patterns
    and production deployment scenarios.
    
    Key Features:
    - CUDA kernel performance analysis
    - Memory coalescing pattern detection
    - Warp divergence analysis
    - Shared memory optimization
    - Tensor core utilization metrics
    - Kernel fusion opportunities
    - Multi-GPU scaling analysis
    """
    
    def __init__(self, hardware_config: Optional[Dict[str, Any]] = None):
        """
        Initialize the kernel optimization profiler.
        
        Args:
            hardware_config: Dictionary containing hardware specifications
        """
        self.hardware_config = hardware_config or self._detect_hardware()
        self.profile_results = {}
        self.optimization_recommendations = []
        
    def _detect_hardware(self) -> Dict[str, Any]:
        """Detect current hardware configuration."""
        return {
            'cpu_cores': psutil.cpu_count(),
            'memory_gb': psutil.virtual_memory().total // (1024**3),
            'cache_sizes': {
                'l1': 32768,  # Typical L1 cache size in bytes
                'l2': 262144,  # Typical L2 cache size in bytes  
                'l3': 8388608  # Typical L3 cache size in bytes
            },
            'gpu_available': False,  # Would check for CUDA/OpenCL in real implementation
            'gpu_memory_gb': 0,
            'tensor_cores': False,
            'warp_size': 32  # NVIDIA GPU warp size
        }
    
    def analyze_cuda_kernel_performance(self, kernel_func: Callable, input_data: Tensor, 
                                      iterations: int = 100) -> Dict[str, Any]:
        """
        Analyze CUDA kernel performance characteristics.
        
        In a real implementation, this would interface with CUDA profiling tools
        to measure actual GPU kernel performance metrics.
        """
        # Simulate CUDA kernel analysis
        total_time = 0
        memory_bandwidth = 0
        compute_utilization = 0
        
        for _ in range(iterations):
            result, execution_time = time_kernel(kernel_func, input_data)
            total_time += execution_time
            
            # Simulate GPU metrics calculation
            data_size = input_data.data.nbytes
            memory_bandwidth += (data_size * 2) / (execution_time / 1_000_000)  # Read + Write
            compute_utilization += np.random.uniform(0.3, 0.9)  # Simulated utilization
        
        avg_time = total_time / iterations
        avg_bandwidth = memory_bandwidth / iterations
        avg_utilization = compute_utilization / iterations
        
        analysis = {
            'avg_execution_time_us': avg_time,
            'memory_bandwidth_gb_s': avg_bandwidth / (1024**3),
            'compute_utilization': avg_utilization,
            'theoretical_peak_bandwidth': 900,  # GB/s for high-end GPU
            'bandwidth_efficiency': min(100, (avg_bandwidth / (1024**3)) / 900 * 100),
            'bottleneck_analysis': self._identify_bottlenecks(avg_bandwidth / (1024**3), avg_utilization)
        }
        
        self.profile_results['cuda_analysis'] = analysis
        return analysis
    
    def analyze_memory_coalescing(self, access_pattern: str, data_shape: Tuple[int, ...]) -> Dict[str, Any]:
        """
        Analyze memory access patterns for GPU coalescing efficiency.
        
        Memory coalescing is critical for GPU performance - threads in a warp
        should access contiguous memory locations.
        """
        coalescing_efficiency = 1.0
        
        if access_pattern == 'row_major':
            # Good coalescing for row-major access
            coalescing_efficiency = 0.95
        elif access_pattern == 'column_major':
            # Poor coalescing for column-major access
            coalescing_efficiency = 0.3
        elif access_pattern == 'strided':
            # Moderate coalescing for strided access
            stride = data_shape[1] if len(data_shape) > 1 else 1
            coalescing_efficiency = max(0.1, 1.0 / stride)
        elif access_pattern == 'random':
            # Very poor coalescing for random access
            coalescing_efficiency = 0.1
        
        analysis = {
            'access_pattern': access_pattern,
            'data_shape': data_shape,
            'coalescing_efficiency': coalescing_efficiency,
            'memory_transactions': self._calculate_memory_transactions(data_shape, coalescing_efficiency),
            'optimization_potential': 1.0 - coalescing_efficiency
        }
        
        self.profile_results['memory_coalescing'] = analysis
        return analysis
    
    def analyze_warp_divergence(self, conditional_operations: int, total_operations: int) -> Dict[str, Any]:
        """
        Analyze warp divergence patterns in kernel execution.
        
        Warp divergence occurs when threads in a warp take different execution paths,
        reducing parallelism efficiency.
        """
        divergence_ratio = conditional_operations / total_operations
        efficiency_loss = divergence_ratio * 0.5  # Simplified model
        
        analysis = {
            'conditional_operations': conditional_operations,
            'total_operations': total_operations,
            'divergence_ratio': divergence_ratio,
            'efficiency_loss': efficiency_loss,
            'warp_efficiency': 1.0 - efficiency_loss,
            'optimization_suggestions': self._generate_divergence_optimizations(divergence_ratio)
        }
        
        self.profile_results['warp_divergence'] = analysis
        return analysis
    
    def analyze_shared_memory_usage(self, kernel_data_size: int, reuse_factor: float) -> Dict[str, Any]:
        """
        Analyze shared memory optimization opportunities.
        
        Shared memory is fast on-chip memory that can dramatically improve
        performance when used effectively for data reuse.
        """
        shared_memory_size = 48 * 1024  # 48KB typical shared memory per SM
        bank_conflicts = self._estimate_bank_conflicts(kernel_data_size)
        
        analysis = {
            'data_size_bytes': kernel_data_size,
            'shared_memory_available': shared_memory_size,
            'utilization_ratio': min(1.0, kernel_data_size / shared_memory_size),
            'reuse_factor': reuse_factor,
            'bank_conflicts': bank_conflicts,
            'performance_gain': min(10.0, reuse_factor * (1.0 - bank_conflicts)),
            'optimization_opportunities': self._identify_shared_memory_optimizations(kernel_data_size, reuse_factor)
        }
        
        self.profile_results['shared_memory'] = analysis
        return analysis
    
    def analyze_tensor_core_utilization(self, operation_type: str, data_types: List[str]) -> Dict[str, Any]:
        """
        Analyze tensor core utilization for mixed-precision operations.
        
        Tensor cores provide massive acceleration for mixed-precision matrix operations
        when data shapes and types are optimized correctly.
        """
        tensor_core_compatible = (
            operation_type in ['matmul', 'conv2d'] and
            any(dtype in ['float16', 'bfloat16', 'int8'] for dtype in data_types)
        )
        
        if tensor_core_compatible:
            theoretical_speedup = 4.0  # Typical tensor core speedup
            actual_utilization = 0.7   # Realistic utilization
        else:
            theoretical_speedup = 1.0
            actual_utilization = 0.0
        
        analysis = {
            'operation_type': operation_type,
            'data_types': data_types,
            'tensor_core_compatible': tensor_core_compatible,
            'theoretical_speedup': theoretical_speedup,
            'actual_utilization': actual_utilization,
            'performance_gain': theoretical_speedup * actual_utilization,
            'optimization_requirements': self._get_tensor_core_requirements()
        }
        
        self.profile_results['tensor_core'] = analysis
        return analysis
    
    def analyze_kernel_fusion_opportunities(self, operation_sequence: List[str]) -> Dict[str, Any]:
        """
        Analyze opportunities for kernel fusion to reduce memory overhead.
        
        Kernel fusion combines multiple operations into a single kernel,
        reducing memory bandwidth requirements and improving performance.
        """
        fusable_patterns = [
            ['matmul', 'relu'],
            ['conv2d', 'batchnorm', 'relu'],
            ['add', 'relu'],
            ['mul', 'add']
        ]
        
        fusion_opportunities = []
        memory_savings = 0
        
        for pattern in fusable_patterns:
            if self._sequence_contains_pattern(operation_sequence, pattern):
                fusion_opportunities.append(pattern)
                memory_savings += len(pattern) - 1  # Save intermediate results
        
        analysis = {
            'operation_sequence': operation_sequence,
            'fusion_opportunities': fusion_opportunities,
            'memory_savings_factor': memory_savings,
            'performance_improvement': min(2.0, 1 + memory_savings * 0.3),
            'implementation_complexity': len(fusion_opportunities) * 2
        }
        
        self.profile_results['kernel_fusion'] = analysis
        return analysis
    
    def analyze_multi_gpu_scaling(self, data_size: int, num_gpus: int) -> Dict[str, Any]:
        """
        Analyze multi-GPU scaling patterns and communication overhead.
        
        Multi-GPU deployments require careful optimization of data distribution
        and communication patterns to achieve good scaling efficiency.
        """
        communication_overhead = self._calculate_communication_overhead(data_size, num_gpus)
        compute_scaling = min(num_gpus, data_size / 1000)  # Simplified scaling model
        
        analysis = {
            'data_size': data_size,
            'num_gpus': num_gpus,
            'communication_overhead': communication_overhead,
            'compute_scaling': compute_scaling,
            'scaling_efficiency': compute_scaling / num_gpus,
            'bottleneck_type': 'communication' if communication_overhead > 0.3 else 'compute',
            'optimization_strategies': self._get_multi_gpu_optimizations(communication_overhead)
        }
        
        self.profile_results['multi_gpu'] = analysis
        return analysis
    
    def generate_optimization_report(self) -> str:
        """Generate comprehensive optimization report with recommendations."""
        report = ["🚀 Kernel Optimization Analysis Report", "=" * 50, ""]
        
        for analysis_type, results in self.profile_results.items():
            report.append(f"📊 {analysis_type.replace('_', ' ').title()} Analysis:")
            report.append("-" * 30)
            
            for key, value in results.items():
                if isinstance(value, float):
                    report.append(f"  {key}: {value:.3f}")
                elif isinstance(value, list):
                    report.append(f"  {key}: {', '.join(map(str, value))}")
                else:
                    report.append(f"  {key}: {value}")
            report.append("")
        
        # Add optimization recommendations
        report.append("🎯 Optimization Recommendations:")
        report.append("-" * 30)
        for rec in self.optimization_recommendations:
            report.append(f"  • {rec}")
        
        return "\n".join(report)
    
    # Helper methods
    def _identify_bottlenecks(self, bandwidth_gb_s: float, utilization: float) -> str:
        """Identify performance bottlenecks."""
        if bandwidth_gb_s < 100:
            return "Memory bandwidth limited"
        elif utilization < 0.5:
            return "Compute utilization limited"
        else:
            return "Well balanced"
    
    def _calculate_memory_transactions(self, shape: Tuple[int, ...], efficiency: float) -> int:
        """Calculate memory transaction count."""
        total_elements = np.prod(shape)
        return int(total_elements / (32 * efficiency))  # 32 threads per warp
    
    def _generate_divergence_optimizations(self, divergence_ratio: float) -> List[str]:
        """Generate warp divergence optimization suggestions."""
        suggestions = []
        if divergence_ratio > 0.3:
            suggestions.append("Reduce conditional operations in inner loops")
            suggestions.append("Use predicated execution instead of branching")
        if divergence_ratio > 0.5:
            suggestions.append("Restructure algorithm to minimize thread divergence")
        return suggestions
    
    def _estimate_bank_conflicts(self, data_size: int) -> float:
        """Estimate shared memory bank conflicts."""
        # Simplified model - assumes some degree of bank conflicts
        return min(0.5, data_size / (32 * 4))  # 32 banks, 4 bytes per bank
    
    def _identify_shared_memory_optimizations(self, size: int, reuse: float) -> List[str]:
        """Identify shared memory optimization opportunities."""
        optimizations = []
        if reuse > 2.0:
            optimizations.append("High reuse factor - shared memory beneficial")
        if size < 16384:  # 16KB
            optimizations.append("Data fits in shared memory - implement tiling")
        return optimizations
    
    def _get_tensor_core_requirements(self) -> List[str]:
        """Get tensor core optimization requirements."""
        return [
            "Use mixed precision (float16/bfloat16)",
            "Ensure matrix dimensions are multiples of 8",
            "Use proper memory layout (NHWC for convolutions)"
        ]
    
    def _sequence_contains_pattern(self, sequence: List[str], pattern: List[str]) -> bool:
        """Check if operation sequence contains fusable pattern."""
        for i in range(len(sequence) - len(pattern) + 1):
            if sequence[i:i+len(pattern)] == pattern:
                return True
        return False
    
    def _calculate_communication_overhead(self, data_size: int, num_gpus: int) -> float:
        """Calculate multi-GPU communication overhead."""
        # Simplified model based on data size and GPU count
        return min(0.8, (data_size / 1000) / num_gpus + 0.1)
    
    def _get_multi_gpu_optimizations(self, overhead: float) -> List[str]:
        """Get multi-GPU optimization strategies."""
        strategies = []
        if overhead > 0.3:
            strategies.append("Implement gradient compression")
            strategies.append("Use asynchronous communication")
        if overhead > 0.5:
            strategies.append("Increase batch size to amortize communication")
        return strategies

# %% nbgrader={"grade": false, "grade_id": "test-kernel-profiler", "locked": false, "schema_version": 3, "solution": false, "task": false}
### 🧪 Unit Test: Kernel Optimization Profiler

def test_unit_kernel_optimization_profiler():
    """Unit test for the kernel optimization profiler."""
    print("🔬 Unit Test: Kernel Optimization Profiler...")
    
    # Create profiler instance
    profiler = KernelOptimizationProfiler()
    
    # Test CUDA kernel analysis
    x = Tensor(np.random.randn(1000))
    cuda_analysis = profiler.analyze_cuda_kernel_performance(vectorized_relu, x, iterations=10)
    
    assert 'avg_execution_time_us' in cuda_analysis
    assert 'memory_bandwidth_gb_s' in cuda_analysis
    assert 'compute_utilization' in cuda_analysis
    print("✅ CUDA kernel analysis works")
    
    # Test memory coalescing analysis
    memory_analysis = profiler.analyze_memory_coalescing('row_major', (1024, 1024))
    
    assert memory_analysis['coalescing_efficiency'] > 0.9
    assert 'optimization_potential' in memory_analysis
    print("✅ Memory coalescing analysis works")
    
    # Test warp divergence analysis
    warp_analysis = profiler.analyze_warp_divergence(100, 1000)
    
    assert warp_analysis['divergence_ratio'] == 0.1
    assert 'warp_efficiency' in warp_analysis
    print("✅ Warp divergence analysis works")
    
    # Test shared memory analysis
    shared_analysis = profiler.analyze_shared_memory_usage(16384, 3.0)
    
    assert 'performance_gain' in shared_analysis
    assert shared_analysis['reuse_factor'] == 3.0
    print("✅ Shared memory analysis works")
    
    # Test tensor core analysis
    tensor_analysis = profiler.analyze_tensor_core_utilization('matmul', ['float16'])
    
    assert tensor_analysis['tensor_core_compatible'] == True
    assert tensor_analysis['theoretical_speedup'] > 1.0
    print("✅ Tensor core analysis works")
    
    # Test kernel fusion analysis
    fusion_analysis = profiler.analyze_kernel_fusion_opportunities(['matmul', 'relu', 'add'])
    
    assert len(fusion_analysis['fusion_opportunities']) > 0
    assert 'performance_improvement' in fusion_analysis
    print("✅ Kernel fusion analysis works")
    
    # Test multi-GPU analysis
    gpu_analysis = profiler.analyze_multi_gpu_scaling(10000, 4)
    
    assert gpu_analysis['num_gpus'] == 4
    assert 'scaling_efficiency' in gpu_analysis
    print("✅ Multi-GPU analysis works")
    
    # Test report generation
    report = profiler.generate_optimization_report()
    
    assert "Kernel Optimization Analysis Report" in report
    assert len(report) > 100  # Should be a substantial report
    print("✅ Optimization report generation works")
    
    print("📈 Progress: Kernel Optimization Profiler ✓")

# Run the test
test_unit_kernel_optimization_profiler()

# %%
def test_module_kernel_sequential_model():
    """
    Integration test for using optimized kernels in a Sequential model.
    
    Tests that optimized kernels can be integrated into a Sequential model
    and produce correct results.
    """
    print("🔬 Running Integration Test: Kernels in Sequential Model...")

    class BaselineModel:
        def __init__(self):
            self.dense = Dense(10, 5)
            self.relu = ReLU()
        
        def __call__(self, x: Tensor) -> Tensor:
            # Manually apply layers using baseline functions
            x = matmul_baseline(x, self.dense.weights)
            # Bias addition is simple, no special kernel needed
            x = Tensor(x.data + self.dense.bias.data)
            x = self.relu(x)
            return x

    class OptimizedModel:
        def __init__(self, baseline_model):
            self.dense = baseline_model.dense
        
        def __call__(self, x: Tensor) -> Tensor:
            # Use optimized kernels
            x = cache_friendly_matmul(x, self.dense.weights)
            x = Tensor(x.data + self.dense.bias.data)
            x = vectorized_relu(x)
            return x
    
    # Mock classes for Dense and ReLU to be used in the test
    class Dense:
        def __init__(self, in_features, out_features):
            self.weights = Tensor(np.random.randn(in_features, out_features))
            self.bias = Tensor(np.random.randn(out_features))

    class ReLU:
        def __call__(self, x: Tensor) -> Tensor:
            return vectorized_relu(x)
    
    # 1. Create baseline and optimized models
    baseline_model = BaselineModel()
    optimized_model = OptimizedModel(baseline_model)

    # 2. Create some input data
    input_data = Tensor(np.random.randn(1, 10))

    # 3. Get outputs from both models
    baseline_output = baseline_model(input_data)
    optimized_output = optimized_model(input_data)

    # 4. Check that the outputs are numerically close
    assert np.allclose(baseline_output.data, optimized_output.data), "Optimized model output should match baseline"

    print("✅ Integration Test Passed: Kernels correctly integrated into a model.")

# %% [markdown]
"""
## 🧪 Module Testing

Time to test your implementation! This section uses TinyTorch's standardized testing framework to ensure your implementation works correctly.

**This testing section is locked** - it provides consistent feedback across all modules and cannot be modified.
"""

# %% [markdown]
"""
## 🤔 ML Systems Thinking Questions

### GPU Architecture and Parallelism

**How does GPU architecture influence kernel design decisions?**
Consider the massive parallelism of modern GPUs (1000s of cores) versus CPUs (10s of cores). How would you design matrix multiplication kernels differently for each architecture? What are the trade-offs between thread-level parallelism and instruction-level parallelism?

**Why do memory access patterns matter more on GPUs than CPUs?**
Think about how GPU memory hierarchy (global memory, shared memory, registers) differs from CPU caches. How does memory coalescing affect bandwidth utilization, and why do random access patterns cause such dramatic performance degradation on GPUs?

**How do you handle load balancing across thousands of GPU threads?**
When processing variable-sized data or irregular computations, how do you ensure all GPU cores stay busy? What strategies exist for handling workload imbalances, and how do frameworks like PyTorch handle dynamic shapes efficiently?

**What role do GPU warps play in kernel optimization?**
NVIDIA GPUs execute threads in groups of 32 (warps). How does this affect branching, memory access, and algorithm design? Why is warp divergence such a critical performance consideration, and how do you design algorithms to minimize it?

### Custom CUDA Kernel Development

**When should you write custom CUDA kernels versus using library functions?**
Given that libraries like cuDNN and cuBLAS are highly optimized, when does it make sense to write custom kernels? Consider scenarios like novel layer types, fused operations, or hardware-specific optimizations.

**How do you optimize CUDA kernels for different GPU generations?**
GPU architectures evolve rapidly (Pascal → Volta → Ampere → Hopper). How do optimization strategies change across generations? What are the implications of new features like tensor cores, multi-instance GPU, and transformer engines?

**What's the development workflow for production CUDA kernels?**
Consider the entire pipeline from prototype to production: profiling bottlenecks, writing initial kernels, optimization iterations, testing across hardware, and deployment. How do companies like OpenAI and Google manage kernel development at scale?

**How do you ensure numerical stability in custom kernels?**
Custom kernels often involve low-level optimizations that can affect numerical precision. How do you balance performance with accuracy? What testing strategies ensure kernels produce correct results across different data ranges and edge cases?

### Triton and Kernel Languages

**How does Triton compare to CUDA for kernel development?**
Triton promises Python-like syntax while generating efficient GPU code. What are the trade-offs between ease of development and performance control? When would you choose Triton over CUDA or vice versa?

**What role do domain-specific languages play in kernel optimization?**
Beyond CUDA and Triton, consider languages like OpenCL, HIP, and emerging alternatives. How do these languages abstract hardware differences while maintaining performance? What's the future of cross-platform kernel development?

**How do JIT compilation and auto-tuning affect kernel performance?**
Modern frameworks use just-in-time compilation to optimize kernels for specific inputs and hardware. How does this compare to static optimization? What are the implications for deployment, cold start times, and reproducibility?

**What are the challenges of kernel portability across hardware vendors?**
With AMD GPUs, Intel GPUs, and custom accelerators becoming more common, how do you write kernels that perform well across different architectures? What abstraction layers exist, and what are their performance costs?

### Hardware-Specific Optimizations

**How do you optimize kernels for different memory hierarchies?**
Consider the differences between GPU global memory, shared memory, and registers versus CPU caches. How do you design algorithms that effectively use each level of the hierarchy? What happens when your working set exceeds cache capacity?

**What optimization strategies work best for tensor operations?**
Tensor cores on modern GPUs can dramatically accelerate mixed-precision operations. How do you restructure algorithms to take advantage of these specialized units? What are the constraints on data layout, precision, and problem sizes?

**How do you handle precision trade-offs in optimized kernels?**
Production systems often use int8, fp16, or bfloat16 for performance. How do you maintain model accuracy while using reduced precision? What accumulation strategies prevent numerical issues in long computations?

**What role does compiler optimization play in kernel performance?**
Modern GPU compilers perform sophisticated optimizations like loop unrolling, memory access optimization, and instruction scheduling. How do you write kernel code that works well with these optimizations? When do you need to use inline assembly or intrinsics?

### Production GPU Clusters

**How do you scale kernel optimizations across multi-GPU systems?**
Single-node multi-GPU systems require coordination of memory transfers, computation scheduling, and synchronization. How do you design kernels that scale efficiently across 8-16 GPUs? What are the bottlenecks in multi-GPU scaling?

**What are the challenges of distributed training with custom kernels?**
When scaling to hundreds or thousands of GPUs across multiple nodes, network communication becomes critical. How do custom kernels interact with distributed training frameworks? What optimizations exist for gradient synchronization and parameter updates?

**How do you manage kernel deployment in production clusters?**
Production ML systems need to handle hardware failures, software updates, and varying workloads. How do you deploy and manage custom kernels across heterogeneous clusters? What strategies exist for A/B testing kernel optimizations safely?

**What monitoring and debugging tools exist for production GPU workloads?**
When kernels behave unexpectedly in production, how do you diagnose issues? What metrics matter for kernel performance monitoring? How do you correlate kernel performance with higher-level model metrics like accuracy and throughput?

## 🎯 MODULE SUMMARY: Custom Kernels

Congratulations! You've successfully implemented custom kernel operations:

### What You've Accomplished
✅ **Custom Operations**: Implemented specialized kernels for performance
✅ **Integration**: Seamless compatibility with neural networks
✅ **Performance Optimization**: Faster computation for critical operations
✅ **Real Applications**: Deploying optimized models to production

### Key Concepts You've Learned
- **Custom kernels**: Building specialized operations for efficiency
- **Integration patterns**: How kernels work with neural networks
- **Performance optimization**: Balancing speed and accuracy
- **API design**: Clean interfaces for kernel operations

### Professional Skills Developed
- **Kernel engineering**: Building efficient operations for deployment
- **Performance tuning**: Optimizing computation for speed
- **Integration testing**: Ensuring kernels work with neural networks

### Ready for Advanced Applications
Your kernel implementations now enable:
- **Edge deployment**: Running optimized models on resource-constrained devices
- **Faster inference**: Reducing latency for real-time applications
- **Production systems**: Deploying efficient models at scale

### Connection to Real ML Systems
Your implementations mirror production systems:
- **PyTorch**: Custom CUDA kernels for performance
- **TensorFlow**: XLA and custom ops for optimization
- **Industry Standard**: Every major ML framework uses these exact techniques

### Next Steps
1. **Export your code**: `tito export 13_kernels`
2. **Test your implementation**: `tito test 13_kernels`
3. **Deploy models**: Use optimized kernels in production
4. **Move to Module 14**: Add benchmarking for evaluation!

**Ready for benchmarking?** Your custom kernels are now ready for real-world deployment!
"""