feat(billing): dtype-aware cost accounting#144
Open
spMohanty wants to merge 36 commits into
Open
Conversation
Adds a dtype_rates table alongside the existing per-op weights table: packaged defaults, custom-file override, unit-mode fallback, and a fail-closed lookup for dtypes outside the loaded table. Also syncs the new exception class into the generated client error module.
Thread optional dtypes/complex_factor_override through deduct(), deduct_after(), and the shared _charge_op() choke point, changing the formula from int(flop_cost * weight) to int(flop_cost * dtype_rate * complex_factor * weight). Resolution happens before the budget check so an unsupported dtype fails closed without charging. dtypes stays optional (None => dtype-neutral, identical to prior billing) so all existing call sites are unaffected until they migrate. Mirrors the same formula in accounting._weight_cost for public cost estimates.
Threads dtypes=/billing_operand through every unary, binary, and reduction deduct() site in _pointwise.py, plus the custom ops (clip, around/round, sort_complex, isclose, average, count_nonzero, the median/percentile/quantile family, ptp, outer, kron, cross, diff, gradient, ediff1d, convolve, correlate, corrcoef, cov, trapezoid/trapz, interp), so complex and fp64 inputs finally bill their registry complex_factor and dtype_rate instead of the pre-migration flat rate. Binary sites bill from the original (pre-coercion) operands via billing_operand so NEP 50 weak scalar promotion is preserved. A few ops needed numpy's actual (non-obvious) dtype-resolution behavior verified before billing: diff/ediff1d's prepend/to_begin route through concatenate (no NEP 50); cov/corrcoef always compute at a float64 floor; interp's search/blend is always float64 but the interpolated value keeps fp's dtype. Left unmigrated (dtypes=None, unchanged pre-migration billing): - The five generic ufunc-method fallbacks (_counted_ufunc_outer/ reduce_generic/accumulate_generic/reduceat/at) bill under a dynamic op_name (e.g. "subtract.reduce") that has no registry entry, so adding dtypes= would make any complex-dtype call raise UnsupportedDtypeError instead of billing correctly. Flagged separately for follow-up. - _einsum_routed_binary, tensordot, and vdot route through the einsum accumulation model and carry an "exact" complex_factor requiring a computed override; out of scope here. tests/test_cost_convention.py::test_conformance now has one expected failure: sort_complex bills 600 instead of 300 (its complex_factor of 2.0 is now applied) — to be reconciled in a follow-up task.
The five generic ufunc-method fallbacks (_counted_ufunc_outer/ reduce_generic/accumulate_generic/reduceat/at) bill under a dynamic op_name like "subtract.reduce" that has no direct registry entry. Adding dtypes= there naively would make complex_factor_for fail closed and raise UnsupportedDtypeError on every complex-dtype call through these sites, rather than billing correctly. Add a base-ufunc-name fallback to complex_factor_for: when the exact op_name misses the registry, strip a known ufunc-method suffix (.reduce/.accumulate/.reduceat/.outer/.at) and look up the base ufunc's factor instead — the per-element arithmetic of these methods is the base ufunc's (a reduce over multiply is a chain of complex multiplies). The fallback only fires when the exact name isn't already a registry key, so dotted keys like fft.fft, linalg.svd, and linalg.outer always resolve on the direct lookup and are never stripped. A base op classified "illegal" for complex (e.g. logaddexp) still raises through the fallback, matching numpy's own behavior. Declare dtypes= at all five sites now that the fallback makes it safe: .outer bills both operands' dtypes, the rest bill the primary operand's dtype (plus out= where applicable); index/selector operands (reduceat's indices, at's index/value args) are excluded. No test_cost_convention.py shifts: the file has no np.<ufunc>.reduce/ .accumulate/.outer/.reduceat/.at assertions.
Wires complex_factor_override into every "exact"-classified contraction
deduct site (einsum, the shared dot/matmul/inner/vecdot/matvec/vecmat
router, both tensordot deduct sites, vdot) using complex_real_total's
mu/adds decomposition, so complex-dtype contractions bill their true
real-FLOP cost instead of silently falling back to dtype-neutral
billing (which previously let einsum('i,i->i', z, z) undercut an
equivalent multiply(z, z) call by ~12x). einsum_path stays declared
dtype-neutral (dtypes=()) since it does no value arithmetic.
tensordot's general fallback path has branches (oversized symmetry, or
rank>52) that compute a plain int with no accumulation decomposition;
those keep complex_factor_override=None and rely on the registry's
"exact" classification to fail closed via RuntimeError if a complex
dtype ever reaches them, rather than inventing an unverified ratio.
Real-dtype billing is unchanged: complex_factor_for short-circuits to
1.0 for real dtypes regardless of the override.
Threads dtypes= through every budget.deduct()/deduct_after() site in _sorting_ops.py, _counting_ops.py, _polynomial.py, _window.py, _unwrap.py, and _symmetric.py (53 sites), so these ops finally bill their registry complex_factor and dtype_rate instead of the pre-migration flat rate. Rule applied per site: declare the operand dtype(s) (both operands for two-array ops like polyadd/searchsorted/allclose/digitize; the single operand otherwise); include an explicit dtype=/out=/w= argument's dtype where the op has one (trace, polyfit's weights). logspace/ geomspace follow the existing linspace/arange pattern (dtype= kwarg, else result_type(start, stop)). Window ops (bartlett/blackman/ hamming/hanning/kaiser) take an int length with no array operand, so they declare dtypes=() (dtype-neutral). tests/test_fma2_cost_fixes.py::test_geomspace_logspace_cost_broadcast_output_times_transcendental now fails: geomspace/logspace resolve to float64 and correctly bill 2x under production rates (1600 instead of 800, 160000 instead of 80000) instead of the pre-migration flat rate. Reconciliation needed alongside the existing sort_complex shift from the pointwise migration.
Threads dtype-aware billing through the remaining fft.*/linalg.*/ random.*/stats.* deduct sites (56 total). FFT declares the input array dtype (registry complex_factor is priced-in at 1.0, no override needed); linalg declares all array operands (registry factor 4.0 applies uniformly); random declares the actual output dtype, reading it off the already-computed result where the same factory backs many distributions with different dtypes; stats declares the forced-float64 coerced input. random.permutation/shuffle/choice/symmetric (module-level) and the Generator/RandomState choice/permutation/permuted/shuffle methods intentionally keep dtypes=() -- these pass through arbitrary (possibly complex) caller data that numpy handles fine, but the registry's complex_factor="illegal" classification for them would incorrectly reject legitimate complex input if declared.
deduct()/deduct_after() no longer default dtypes to None; both now raise TypeError up front (before _charge_op) if dtypes is omitted or explicitly None, with a message pointing callers at dtypes=() for dtype-neutral ops. _charge_op keeps its own dtypes=None default since its only two callers (both now guarded) always pass a tuple. Adds a permanent AST guard (test_deduct_declares_dtypes.py) asserting every deduct/deduct_after call site under src/flopscope names dtypes=. Renames the Task 5 migration-era test to test_dtype_neutral_and_none_rejected: dtypes=() still bills as dtype-neutral, dtypes=None now raises.
np.real/np.imag return a view (or constant-fill) of a complex value's component -- no floating-point arithmetic -- so charging numel FLOPs at weight 1.0 was an overcount. Reclassify both to the free tier: flop_cost=0 hardcoded at the call site via a new _free_unary binding, registry category counted_unary -> free, default weight 1.0 -> 0.0 (mirrors iscomplexobj/ isrealobj). Functional behavior (out=, symmetry, NaN/Inf checks) unchanged.
Pre-migration tests hardcoded the old (no dtype_rate/complex_factor) billing model. Reconciles each shift against the new charged = flop_cost * dtype_rate * complex_factor * weight formula, with a derivation comment at every re-pinned assertion: - float64-default inputs under production weights now bill 2x (dtype_rate) in test_cost_constant_unification.py, test_production_weight_billing.py, test_fma2_cost_fixes.py, and test_ufunc_alias_parity.py. - genuinely-complex inputs pick up their op's registry complex_factor in test_cost_convention.py, test_cost_formula_vs_code.py, and test_pointwise_coverage.py (sort_complex, real_if_close). - test_no_bare_numpy_in_counted_ops.py: allowlist the bare `_np.dtype(...)` calls the migration added to normalize billing dtypes -- same O(1) metadata class as the already-allowlisted result_type/promote_types/can_cast, not a compute call. - drop the dead "imag" case from test_cost_formula_vs_code.py's _unary_input (imag was moved to the free tier and removed from the numel-parametrized list).
Document the four-factor billing formula charged = int(flop_cost x dtype_rate x complex_factor x weight) as the public single source of truth. - Rewrite the billing-model intro from two layers to four, with the extended separation invariant (width policy in dtype_rate; complex structure in complex_factor, computed per call for the contraction family). - Add a Dtype and precision section: why width and structure are priced, the 32-bit-class billing unit + 14-dtype rate table + fail-closed rule, complex arithmetic from first principles (atom table, per-op classification, exact contraction rule, FFT priced-in derivation), the resolved-dtype rule with worked cases, a worked billing table, and the packing stance. - Sweep the rest of the doc: FMA=2 unit cross-ref; two new non-exploitability rows (complex/width packing); one complex-classification line per family section; move real/imag to the free tier. - Extend tests/test_cost_model_coverage.py: assert the section heading, the verbatim four-factor formula, and one rate-table row per supported dtype.
Two dtype-billing follow-ups tracked during the four-factor work. Random movement ops (shuffle/permutation/choice + Generator.permuted) relocate caller-supplied data rather than sampling a distribution, so they are movement ops, not value generators. Their registry complex_factor was "illegal" -- valid only for the real-only samplers, and a footgun here: numpy permutes/shuffles/selects complex arrays fine, so any future dtype declaration at these sites would wrongly raise on a legitimate complex operand. It is now 1.0. Billing stays dtype-neutral (dtypes=()): the shape[axis] cost counts the dtype-independent Fisher-Yates swap / selection work, so a complex or fp64 population costs the same as fp32 -- unlike a genuine sampler, which bills the width of the values it synthesizes. Stripping a FlopscopeArray operand to a base ndarray before numpy's C path also fixes a crash in the in-place shuffle (numpy's internal empty_like tripped the fnp-reentry guard). asarray with an explicit value-changing dtype= did the same work as astype but billed 0 (weight 0.0), making it a free alternative to astype for downcasting. asarray now mirrors astype: numel at the heavier of source/target rate on a value-changing cast, 0 for a no-op or lossless widening (weight 0.0 -> 1.0, relabel free -> counted_custom). Regression tests pin dtype-neutral billing for the movement ops (including the shuffle crash and registry hygiene) and astype-parity for asarray casts.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 732a401f32
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+1027
to
+1029
| billing_dtypes: tuple = (a.dtype,) | ||
| if kwargs.get("dtype") is not None: | ||
| billing_dtypes += (_np.dtype(kwargs["dtype"]),) |
There was a problem hiding this comment.
When dtype is omitted, several NumPy reductions do not actually compute in a.dtype: on 64-bit platforms sum/prod/cumsum of an int32 input promote to an int64 accumulator, and the mean/variance factories have the same issue for integer inputs promoting to float64. Declaring only a.dtype here means calls like fnp.sum(np.ones(1000, dtype=np.int32)) are charged at the 32-bit rate even though NumPy executes the reduction in a 64-bit dtype, leaving a width-discount path in the new dtype-aware accounting.
Useful? React with 👍 / 👎.
Collaborator
Author
There was a problem hiding this comment.
Addressed in 4896aec — reductions now bill numpy's widened accumulator dtype (int64/uint64 for sum/prod/cumsum/cumprod, float64 for integer mean/var/std); extremum/index/boolean reductions intentionally keep the input dtype.
Comment on lines
+73
to
+74
| else: | ||
| billing_dtypes = () |
There was a problem hiding this comment.
For Generator/RandomState sampler methods with size=None, NumPy returns a scalar (for example rng.standard_normal() returns a float64 scalar), so this else path declares the operation dtype-neutral. Under production dtype rates that undercharges scalar draws such as fnp.random.default_rng(0).standard_normal() or integers(dtype=np.int64) by skipping the float64/int64 rate, even though the module-level sampler wrapper already handles scalar result dtypes.
Useful? React with 👍 / 👎.
Collaborator
Author
There was a problem hiding this comment.
Addressed in 4896aec — size=None sampler draws now bill their scalar output dtype (float64/int64) instead of dtype-neutral.
…lling Two CI jobs outside the root `make ci` (client-server-integration runs the flopscope-server suite; website-tests runs the node regression tests) exercised billing paths the dtype-aware change touched: - Three flopscope-server tests call BudgetContext.deduct() directly to stage a synthetic op for the namespace-breakdown assertions. deduct() now requires the dtypes= keyword, so they declare dtypes=() (dtype-neutral), matching the plumbing-site convention. - The symmetry-guide regression test pins two einsum costs computed from fnp.random.randn (float64) data. Dtype-aware billing bills float64 at rate 2.0, so the base 54/81 become 108/162; the expected values and the explaining comment are updated.
Two width-discount paths surfaced in review, both undercounting integer or scalar work under the new dtype-aware accounting: - sum/prod/cumsum/cumprod (and their nan* variants) widen a narrow-integer or boolean input to numpy's 64-bit accumulator (int64/uint64), and mean/var/std of an integer input compute in float64. Billing the input dtype charged these at the 32-bit rate while numpy runs at 64-bit. They now bill the accumulator / compute dtype. Extremum (max/min), index (argmax/argmin) and boolean (all/any) reductions are unchanged: they do not widen, and an int64 index output does not reflect the input-dtype comparison work. - A size=None Generator/RandomState sampler draw returns a scalar, which the dispatcher billed dtype-neutral; it now bills the scalar's output dtype (float64/int64), matching the module-level sampler wrapper and the array draw's per-element rate. Regression tests in tests/test_dtype_cost.py.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
We identified a complex-packing and a differential dtype bit-packing issue that
cause FLOP undercounting: complex operations billed the same as their real
counterparts (letting two real payloads ride one complex op), and all dtypes
billed at one width (letting two lower-precision multiplies ride one float64
multiply). This PR closes both by making billing dtype-aware.
The model
dtype_rate— policy, configurable like weights (dtype_ratesinsrc/flopscope/data/default_weights.json). Anchor: 1 FLOP = one 32-bit-classreal op. All ≤32-bit dtypes 1.0;
int64/uint64/float642.0; extendedprecision priced by width class (
float963.0,float1284.0 — packingnarrower payloads through their mantissas still loses at these rates);
complex entries carry component width (
complex641.0,complex1282.0,complex1923.0,complex2564.0). Only genuinely unknown future numericdtypes fail closed (
UnsupportedDtypeError); non-numeric kinds(object/str/datetime/…) bill neutrally at 1.0 (no floating-point arithmetic
is possible through them; keeps the numpy-compat guarantee intact). The
evaluation wire codec still accepts only the 14 standard dtypes.
complex_factor— math, per-op registry field: the textbook real-FLOPcount per billed unit on complex inputs (add 2, multiply 6 = 4M+2A, divide
11, abs 4, sqrt 10, lexicographic compare 2, transcendentals from their
closed forms, linalg 4). The contraction family (einsum/matmul/dot/…) bills
exactly — 6·multiplies + 2·adds computed from the einsum accumulation
decomposition per call — so
einsum('i,i->i', …)cannot dodgemultiply'sfactor. FFT is priced-in at 1.0 (its
5n·log₂nformula already counts thecomplex butterflies). Every charged op carries an explicit classification,
enforced by a completeness guard.
np.result_typeover the declared operands plusany explicit
dtype=/out=— declared via a now-requireddtypes=keywordat every deduct site (263 sites; an AST guard keeps future ops honest). NEP
50 weak-scalar promotion is preserved (
f32_array * 2.0bills float32).astypebills the heavier of source/destination rate.real/imagarereclassified to the free tier (component extraction is a view).
Reference tables
dtype rates — policy, the full
dtype_ratestable (18 entries), grouped byrate. Complex dtypes sit at their component precision; the complex arithmetic
is billed separately by
complex_factor.bool,int8,int16,int32,uint8,uint16,uint32,float16,float32,complex64int64,uint64,float64,complex128float96,complex192float128,complex256Only the 14 standard dtypes cross the eval wire; the four extended-precision
rows are priced for in-process use (defense-in-depth — the wire codec rejects
them anyway).
complex factors — math, per-op registry field (real dtypes are always 1.0).
Every charged op is classified (a completeness guard enforces it); the factor
is the real-arithmetic decomposition of one billed unit, with
z = a + bi.conjnegatesb;fft.*'s5N·log₂Nis already a real-FLOP count)conj,angle,fft.fft,concatenate,take,shuffle(a+bi) + (c+di) = (a+c) + (b+d)i→ 2 real adds; ordering compares lexicographically (real, then imaginary) → also 2add,subtract,sum,mean,cumsum,sort,less,equallog(z) = ln(abs(z)) + i·atan2(b, a)— a modulus, one reallog, oneatan2log,log2,log10,log1pabs(z−μ)²accumulated) — mostly real adds, some multipliesvar,std,nanvar,nanstdstats.*pdf/cdf/ppfforce-cast their input to float64, so the complex path is unreachable;interpis a genuine 3stats.norm.pdf,interpexp(a+bi) = exp(a)·(cos b + i·sin b)— one realexp, onesin, onecos, two multsexp,exp2,expm1sin(a+bi) = sin a·cosh b + i·cos a·sinh b— four real transcendentals + two mults (tan= 3.6 needs a little more)sin,cos,tan,sinh,cosh,tanhlog/sqrtclosed form, e.g.arcsin z = −i·log(iz + sqrt(1 − z²))arcsin,arccos,arctan,arcsinhabs(z) = sqrt(a² + b²)→ 2 mul + 1 add + 1 sqrt = 4; dense linalg is complex-matmul-dominated (≈4× the real arithmetic)abs,linalg.inv,linalg.svd,linalg.eig,covcrossz² = (a² − b²) + (2ab)ireusesa,b→ 3 mul + 1 sub + 1 add = 5 (cheaper than a general multiply)squarez**w = exp(w·log z)— a complexlog, a complexmultiply, a complexexp(blended)power,float_power,geomspace,logspace(a+bi)(c+di) = (ac − bd) + (ad + bc)i→ 4 real mul + 2 real add = 6;reciprocalcosts the same (a special-cased divide)multiply,prod,cumprod,outer,kron,reciprocalsign = z / abs(z)— the modulus (4) plus the real divides and a zero-guardsignabs= 4) plus further real roots and a sign/halving combinationsqrtz / w = z·conj(w) / abs(w)²— numerator 4 mul + 2 add, denominator 2 mul + 1 add, then 2 real divides → 11divide,true_dividesinc(z) = sin(pi·z) / (pi·z)— a complexsin(≈3.4) composed with a complexdivide(11)sincexact6·mults + 2·adds. A length-Kdot is8K − 2real FLOPs (ratio 6 atK=1, →4 asK→∞), billed per call — which blocks alias-shopping:einsum('i,i->i', z, z)bills 6, same asmultiplyeinsum,matmul,dot,inner,vdotillegalatan2, real-only math/samplers) → fail closedbitwise_*,floor,ceil,trunc,mod,arctan2The two axes compose: a
complex128multiply billsflop_cost × 2.0 (component = float64) × 6.0 (complex multiply) × weight.What participants see
payloads into a complex op now loses (elementwise 3×, matmul ~2×), and
packing two float32 multiplies into one float64 multiply is break-even
before packing overhead — both locked by invariant tests
(
tests/test_dtype_cost.py).pricier side. Fixed-float64-output ops with no array operand (window
functions,
fft.fftfreq, the random samplers) bill the float64 rate.astypeand a value-changingasarray(x, dtype=...)are billed identically(numel at the heavier rate) —
asarrayis no longer a free way to downcast.docs/reference/cost-model.mdgained a Dtype and precision section — thepolicy-vs-math split, the rate table, the per-op complex derivations (with the
FFT butterfly proof), the resolved-dtype rule with worked examples, and the
stance on packing (not banned; priced to break-even-or-losing at 32/64-bit
widths) — plus a full self-consistency sweep and an extended doc coverage
guard.
Verification
make ciexit 0): lint, lockfile,commit lint, typecheck, full suite (4808 tests, 0 failed), NumPy
compatibility harness (numpy's own test suite under monkeypatch — 0 failed),
client/server sync gates, docs build.
flopscope-clientsuite green (1353 tests). Registry data andops.jsonregenerated (
--checkclean); the per-opcomplex_factoris publiclyauditable in
ops.json.completeness guards all in
tests/test_dtype_cost.py,tests/test_dtype_billing.py,tests/test_complex_factor_completeness.py,tests/test_deduct_declares_dtypes.py.Rollout notes
eval redeploy. Existing submissions re-price on re-eval (float64-default
submissions ≈2×; complex-packing submissions substantially more).
dtype_ratestable is intentionally tunable policy: revisiting values(e.g. sub-32-bit discounts) needs only a data change.
Follow-ups closed in this PR
The items tracked during the four-factor work are resolved here:
shuffle/choice/permutation/permuted,module-level and
Generator/RandomState) are reclassified from"illegal"tocomplex_factor1.0. They keep billing dtype-neutral —their
shape[axis]cost counts the dtype-independent Fisher-Yatesswap/selection work, so a complex or fp64 population costs the same as fp32
(a genuine sampler, which synthesizes values, still bills its output width).
The
"illegal"classification was a footgun, not a live bug: numpy handlescomplex here, so it only mattered once a dtype was declared at these sites.
fnp.random.shufflecrash onFlopscopeArrayinputs: fixed by strippingto a base ndarray before numpy's C path (its internal
empty_liketrippedthe reentry guard on the in-place shuffle).
asarraywith an explicit value-changingdtype=no longer converts forfree — it now mirrors
astype(numel at the heavier source/target rate);a no-op or lossless-widening
asarraystays free.Genuinely deferred (documented, no billing impact): the
complex_factoron thestats.*pdf/cdf/ppffamily is dead metadata — those ops force-cast theirinput to float64, so the complex path is unreachable and the value is never
consulted.