perf: configurable float storage dtype (default float32, ~½ coeffs memory)#839
Draft
FBumann wants to merge 3 commits into
Draft
perf: configurable float storage dtype (default float32, ~½ coeffs memory)#839FBumann wants to merge 3 commits into
FBumann wants to merge 3 commits into
Conversation
…fault float32)
Extend the per-model dtypes mechanism (Model._dtypes) with a new "values"
key controlling the storage dtype of the model's float arrays: coefficients,
constants, right-hand sides and variable lower/upper bounds. coeffs is the
single largest array in a built model, so defaulting to float32 roughly
halves peak build memory. Solvers upcast to double internally, so solutions
are unaffected; the LP/MPS writers and direct backends handle float32 fine.
- model.py: DtypeKey gains "values"; annotations widened to type[np.number];
_resolve_dtypes validates "values" as np.float32/np.float64 (default
np.float32) and "labels" as before; add_variables casts bounds to the
values dtype; docstrings updated.
- expressions.py: LinearExpression.__init__ casts coeffs and const to the
model's values dtype (all expression construction funnels through here,
including simplify() results).
- constraints.py: Constraint.__init__ casts rhs to the model's values dtype.
Pass Model(dtypes={"values": np.float64}) to restore full float64 storage.
The float32 default is TEMPORARY, to surface the reduction in CI memory
benchmarks, and is intended to flip to opt-in (float64 default) later.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Merging this PR will improve performance by 48.67%
Performance Changes
Tip Curious why this is faster? Comment Comparing 
Footnotes
|
CI's warn_unused_ignores flags the two # type: ignore[dict-item] on the invalid-dtype rejection tests as unnecessary; np.float64/np.float16 are valid type[np.number] values (the rejection is a runtime check). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Multiplying/adding a float64 array against a float32 expression let numpy
upcast to a float64 broadcast intermediate that __init__ then downcast
back to float32 — a large transient that raised op peak memory ~17-27%
over master (CodSpeed test_op[expr_{mul,add}_array_bcast]).
Cast the other operand to the model value dtype first (extracted as
_as_value_dtype), so the elementwise op stays float32 end to end. Peak
now drops below master (mul_array_bcast 1.04MB -> 0.69MB).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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.
DO NOT MERGE
Im exploring other dtype performance gains. Exploratory for now.
In this branch the values default is float32. This needs to be opt-in, bit is default for ci to show the improvements!
Note
The following content was generated by AI.
Follow-up to the int32-labels win (#566). That change halved the integer label/
varsarrays; the float arrays were left atfloat64. In a built modelcoeffsis the single largest array (2× thevarsarray), so storing the float arrays asfloat32is a comparable memory lever.What this does. Extends the per-model
_dtypesdict (introduced in #828, previously"labels"-only) with a"values"key controlling the storage dtype of the float arrays — coefficients, constants, right-hand sides, and variable lower/upper bounds. Acceptsnp.float32ornp.float64.Default. Ships defaulting to
np.float32(opt-out) on purpose, and only for now, so CI memory benchmarks surface the reduction on this PR. The intent is to flip the default tonp.float64(opt-in) before this is considered for merge —float32coefficients lose precision and can bite badly-scaled / big-M problems, so it should not be the silent default.Model(dtypes={"values": np.float64})already restores full double-precision storage today.Implementation & verification
linopy/model.py:DtypeKey = Literal["labels", "values"]; annotations widenedsignedinteger → number;_resolve_dtypesdefaults{"labels": int32, "values": float32}and branch-validates per key;add_variablescasts bounds to the values dtype; docstrings updated.linopy/expressions.py: theBaseExpression.__init__chokepoint castscoeffs/consttomodel._dtypes["values"].linopy/constraints.py:Constraint.__init__castsrhsto the values dtype.linopy/matrices.pyleft untouched — derived/solver matrices stayfloat64(solvers want double; the win is in storage).vstackinLinearExpression.simplifyis deliberately keptfloat64: it shares a dtype with var-IDs, andfloat32would round IDs past 2²⁴. Storage still ends upfloat32via the__init__cast.Tests (
test/test_dtypes.py): default-float32 on bounds/coeffs/const/rhs,float64override, rejection offloat16, and a float32 HiGHS solve to the correct optimum.test/test_dtypes.py: 21 passedtest/test_optimization.py(solve subset across gurobi/highs/cbc/mosek/…): 340 passed, 10 skippedtest_io.py+test_matrices.py: 43 passed, 3 skippedsolvers.pyerror remains); ruff clean.Memory sanity. 100k×2-term constraint:
coeffs800000bytes vs1600000at float64 — exactly halved.🤖 Generated with Claude Code