Sourced from whest-starterkit @ aaa3882.

Manage Your FLOP Budget

← Documentation

🎯 When to use this page

Use this page to understand how FLOP budgets work and how to optimize your estimator to stay within budget.

Why FLOPs, not wall-clock time

This challenge scores estimators by analytical FLOP count, not execution time. Every mathematical operation your estimator performs is tracked by flopscope — a NumPy-compatible library that counts floating-point operations deterministically from tensor shapes.

This means your score is hardware-independent: the same estimator produces the same FLOP count on a laptop and a GPU cluster. You can focus on algorithmic efficiency rather than hardware tuning.

For the full flopscope API and cost model, see the flopscope documentation.

Which operations cost FLOPs

Key insight: fnp.matmul on (n, n) matrices costs O(n^3) FLOPs. For width-256 networks, a single matmul costs ~33M FLOPs. Most of your budget goes to matrix operations.

Check your budget usage

Wrap your estimator logic in a BudgetContext to see how many FLOPs it consumes:

import flopscope as flops

with flops.BudgetContext(flop_budget=68_000_000_000) as budget:
    result = estimator.predict(mlp, budget=68_000_000_000)

print(f"FLOPs used: {budget.flops_used:,}")
print(f"FLOPs remaining: {budget.flops_remaining:,}")

If you also want a wall-clock guardrail while debugging locally, set wall_time_limit_s on the same BudgetContext:

with flops.BudgetContext(
    flop_budget=68_000_000_000,
    wall_time_limit_s=2.0,
) as budget:
    result = estimator.predict(mlp, budget=68_000_000_000)

Get a per-operation breakdown

Use budget.summary() for the current explicit context or flops.budget_summary() for the session/global view to see which operations consume the most FLOPs:

import flopscope as flops

with flops.BudgetContext(flop_budget=68_000_000_000) as budget:
    result = estimator.predict(mlp, budget=68_000_000_000)
    print(budget.summary())

flops.budget_summary()

This prints a table showing each operation's name, call count, and cumulative FLOP cost — letting you identify the expensive operations to optimize.

The same summaries also show timing data:

• wall_time_s: total elapsed time for the context
• flopscope_backend_time_s: time spent inside counted flopscope backend calls
• flopscope_overhead_time_s: time spent inside flopscope dispatch and bookkeeping
• residual_wall_time_s: time spent outside flopscope backend and dispatch work

In whest run, the CLI flags map to these concepts as follows:

• --wall-time-limit: forwards a wall-clock limit into the estimator's BudgetContext
• --residual-wall-time-limit: adds a WhestBench scoring check on the reported residual_wall_time_s

Interpret whest run output

When you run your estimator with whest run, the per-MLP report includes:

• flops_used: total FLOPs your estimator consumed for that MLP.
• budget_exhausted: true if your estimator exceeded the FLOP budget — predictions were zeroed.
• final_layer_mse / all_layers_mse: your prediction accuracy (lower is better).

If budget_exhausted is true, your predictions were discarded. You need to reduce FLOP usage.

Worked walkthrough: mean propagation, line by line

The table below profiles examples/02_mean_propagation.py on the default Stage 1 MLP (width=256, depth=8 — same shape as the Stage-3 grader). Numbers are aggregated across all 8 layers; per-layer cost is roughly the row total divided by 8. Reproduce with ctx.summary() inside a flopscope.BudgetContext after a single predict() call (profiled under flopscope 0.5.0).

The full ~2.7 M FLOPs spends only ~0.004% of the default 6.8e10 grader budget, so mean propagation lands well below the multiplier floor at this shape — see Scoring Model.

Two takeaways:

• matmul dominates. ~77% of predict() cost is the two matmuls per layer (the pointwise ReLU-moment terms — multiply — are the visible ~20% remainder). Halving the matmul count (e.g., switching to a diagonal-only formulation, or fusing into a single einsum like examples/03_covariance_propagation.py does for the symmetric cov-update) buys you most of that back.
• Reductions, sqrt, and divides are free in practice. Don't twist your code to avoid them; the cost is in the tens of FLOPs per layer.

The same pattern holds for examples/03_covariance_propagation.py, where the O(width³) symmetry-aware einsum lands at ~404 M FLOPs per predict() (~0.6% of the grader budget) — ~150× more expensive than mean propagation (its full covariance is genuinely heavier than mean propagation's diagonal variance), but still leaving plenty of headroom.

Optimization tips

1. Matmul dominates. Each fnp.matmul(W.T, mu) on a (width, width) matrix costs O(width^2) FLOPs per layer. Reducing the number of matmuls (or their dimensions) has the biggest impact.

2. Diagonal approximations save FLOPs. Mean propagation uses diagonal variance (O(width^2) per layer) instead of full covariance propagation (O(width^3) per layer). Choose the right level of approximation for your budget.

3. Array creation is free. fnp.array(), fnp.zeros(), fnp.ones(), fnp.eye() cost 0 FLOPs. Precompute and store intermediate values freely.

4. Pick one strategy per estimator. Use either mean propagation or full covariance as your default implementation, then optimize it for the fixed budget.

➡️ Next step

• Write an Estimator
• Scoring Model
• Profile Simulation
• Estimator Contract