Hardware Constraints

memory_budgets_mb: [0.5, 1.0, 2.0, 5.0]  # Reporting thresholds
active_memory_budget_mb: 1.0              # Hard acceptance threshold

def model_memory_mb(model: nn.Module) -> float:
    """
    Calculate total model memory footprint in megabytes.
    
    Args:
        model: PyTorch model instance
    
    Returns:
        Total memory in MB (parameters + buffers)
    """
    total_bytes = 0
    for tensor in model.state_dict().values():
        if isinstance(tensor, torch.Tensor):
            total_bytes += tensor.numel() * tensor.element_size()
    return total_bytes / (1024 ** 2)

def memory_violations(memory_mb: float, budgets_mb: list[float]) -> dict[str, bool]:
    """
    Generate per-budget violation flags.
    
    Args:
        memory_mb: Model memory footprint
        budgets_mb: List of budget thresholds
    
    Returns:
        Dictionary mapping "violates_{budget}mb" -> bool
    """
    return {f"violates_{budget}mb": memory_mb > budget for budget in budgets_mb}

{
    "violates_0.5mb": True,   # 1.2 > 0.5
    "violates_1.0mb": True,   # 1.2 > 1.0
    "violates_2.0mb": False,  # 1.2 < 2.0
    "violates_5.0mb": False   # 1.2 < 5.0
}

# From experiments.run_sweep (src/edge_opt/experiments.py:76)
metrics: PerfMetrics = collect_metrics(
    variant, val_loader, device,
    power_watts=power_watts,
    precision=metric_precision,
    latency_multiplier=latency_multiplier,
    benchmark_repeats=benchmark_repeats,
)

# Check violations for all budgets (soft constraints)
violations = memory_violations(metrics.memory_mb, memory_budgets_mb)

# Check active budget (hard constraint)
rejected = metrics.memory_mb > active_memory_budget_mb

row = {
    "pruning_level": pruning,
    "precision": precision,
    "accepted": not rejected,        # Pass/fail flag
    "active_budget_mb": active_memory_budget_mb,
    **asdict(metrics),               # All performance metrics
    **violations,                    # Per-budget violation flags
}

def estimate_layerwise_stats(
    model: nn.Module, 
    batch_size: int, 
    input_shape: tuple[int, int, int] = (1, 28, 28)
) -> pd.DataFrame:
    """
    Estimate layer-wise activation memory and compute (MACs).
    
    Args:
        model: SmallCNN instance
        batch_size: Batch size for activation calculations
        input_shape: (C, H, W) input tensor shape
    
    Returns:
        DataFrame with columns:
        - layer: Layer name
        - output_elements: Number of output activation elements
        - parameter_bytes: Layer parameter size (weights + biases)
        - activation_bytes: Output activation memory
        - macs: Multiply-accumulate operations
    """
    channels, height, width = input_shape
    bytes_per_value = 4  # FP32 assumption
    
    # Conv1 analysis
    conv1 = model.conv1
    h1, w1 = _conv2d_output_shape(height, width, kernel=3, padding=1)
    h1_pool, w1_pool = h1 // 2, w1 // 2  # MaxPool2d(2)
    
    conv1_elements = batch_size * conv1.out_channels * h1 * w1
    conv1_macs = batch_size * conv1.out_channels * h1 * w1 * conv1.in_channels * 3 * 3
    
    # Similar calculations for conv2 and classifier...

def _conv2d_output_shape(
    height: int, width: int, 
    kernel: int, padding: int, stride: int = 1
) -> tuple[int, int]:
    """
    Calculate Conv2d output spatial dimensions.
    
    Formula: out = (in + 2*padding - kernel) // stride + 1
    """
    out_h = (height + (2 * padding) - kernel) // stride + 1
    out_w = (width + (2 * padding) - kernel) // stride + 1
    return out_h, out_w

def summarize_hardware(
    layerwise_df: pd.DataFrame,
    latency_ms: float,
    memory_bandwidth_gbps: float,
) -> dict[str, float]:
    """
    Estimate bandwidth utilization from layer-wise statistics.
    
    Args:
        layerwise_df: Output from estimate_layerwise_stats
        latency_ms: Measured inference latency
        memory_bandwidth_gbps: Configured memory bandwidth (GB/s)
    
    Returns:
        Dictionary with bandwidth and compute metrics
    """
    total_bytes = float(
        layerwise_df["parameter_bytes"].sum() + 
        layerwise_df["activation_bytes"].sum()
    )
    total_macs = float(layerwise_df["macs"].sum())
    
    latency_s = max(latency_ms / 1000.0, 1e-9)  # Avoid division by zero
    
    # Achieved bandwidth: bytes moved per second
    achieved_bandwidth_gbps = (total_bytes / latency_s) / 1e9
    
    # Utilization: fraction of configured bandwidth
    bandwidth_utilization = achieved_bandwidth_gbps / max(memory_bandwidth_gbps, 1e-9)
    
    # Achieved throughput: MACs per second
    achieved_gmacs = (total_macs / latency_s) / 1e9
    
    return {
        "estimated_total_bytes": total_bytes,
        "estimated_total_macs": total_macs,
        "achieved_bandwidth_gbps": achieved_bandwidth_gbps,
        "configured_memory_bandwidth_gbps": memory_bandwidth_gbps,
        "bandwidth_utilization": bandwidth_utilization,
        "achieved_gmacs": achieved_gmacs,
    }

cpu_frequency_scale: 2.0  # Simulate 50% frequency (2x latency)

# From metrics.collect_metrics (src/edge_opt/metrics.py:86)
latency_mean, latency_std, latency_p95 = measure_latency_distribution(
    model, sample_input, repeats=benchmark_repeats
)

# Apply frequency scaling
latency = latency_mean * latency_multiplier  # latency_multiplier = cpu_frequency_scale
throughput = sample_input.shape[0] / (latency / 1000.0)
energy_proxy = (latency / 1000.0) * power_watts

return PerfMetrics(
    accuracy=accuracy,
    latency_ms=latency,
    latency_std_ms=latency_std * latency_multiplier,
    latency_p95_ms=latency_p95 * latency_multiplier,
    throughput_sps=throughput,  # Inversely affected by latency
    memory_mb=memory,
    energy_proxy_j=energy_proxy,  # Linearly affected by latency
)

def measure_latency_distribution(
    model: nn.Module, 
    sample_input: torch.Tensor, 
    repeats: int = 5, 
    num_runs: int = 100, 
    warmup: int = 10
) -> tuple[float, float, float]:
    """
    Measure latency statistics across multiple measurement windows.
    
    Args:
        model: Model in evaluation mode
        sample_input: Example input tensor
        repeats: Number of measurement windows
        num_runs: Iterations per measurement window
        warmup: Warmup iterations before each window
    
    Returns:
        (mean_ms, std_ms, p95_ms)
    """
    latencies = [
        measure_latency(model, sample_input, num_runs=num_runs, warmup=warmup) 
        for _ in range(repeats)
    ]
    latency_tensor = torch.tensor(latencies, dtype=torch.float32)
    return (
        float(latency_tensor.mean()), 
        float(latency_tensor.std(unbiased=False)), 
        float(torch.quantile(latency_tensor, 0.95))
    )

def save_hardware_artifacts(
    output_dir: Path,
    layerwise_df: pd.DataFrame,
    precision_df: pd.DataFrame,
    summary: dict[str, float],
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # CSV reports
    layerwise_df.to_csv(output_dir / "layerwise_breakdown.csv", index=False)
    precision_df.to_csv(output_dir / "precision_tradeoffs.csv", index=False)
    pd.DataFrame([summary]).to_csv(output_dir / "hardware_summary.csv", index=False)
    
    # Visualization plots (see below)
    ...