[WIP] batching - #80
Draft
simeonschaub wants to merge 25 commits into
Draft
Conversation
Solutions, scratch space and MILP data go from Vector to Matrix, while the quantities derived from them go from Scalar to Vector, one entry per column of the batch. `BatchedNumber` covers both cases, so unbatched problems keep scalar errors and step sizes. Restart candidates are selected column by column, but the decisions to restart or terminate are reduced over the batch, since the whole batch shares a single state. The per-iteration loop stays allocation free in both cases: `Scratch` gains two per-column buffers and `prog_showvalues` is only evaluated when the progress bar prints it. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
`KKTErrors` becomes mutable and is filled in place by `kkt_errors!`, which now takes its destination as an argument: with batching the per-column buffers are reused, without it the scalar fields are simply overwritten. Column-wise reductions write into a `1 × nbatch` alias of a scratch buffer, since `sum!` only avoids allocating with a preallocated destination of the reduced shape. `RestartStats` keeps the absolute errors it compares, which also removes the repeated `absolute` calls in `should_restart`, and holds the buffer for the candidate that gets discarded. Termination and restart checks are now allocation free as well, for both batched and unbatched problems, which a new test guards against. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Instead of waiting for every column to satisfy the restart criteria, apply the three usual conditions to the mean of the per-column absolute KKT errors, so a single decision covers the whole batch. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Use `instance` for anything addressing a single instance of the batch (`instance(x, i)`, `EachInstance`, `instance_vec`, `instance_num`, ...) and `batched` for operations spanning all of them (`batched_apply!`, `batched_all`, `batched_mean`, `batched_select!`, ...). The instance count becomes `nbinstances`, in line with `nbvar` and `nbcons`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add a backend-parameterized `test_batching(matrix_type, backend)` next to `test_moi`, checking instance iteration, per-instance KKT errors, iterates and full solves of a batch against the same problems solved one at a time. Run it in the CUDA group for both `GPUSparseMatrixCSR` and `cuSPARSE.CuSparseMatrixCSR`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Collaborator
Author
|
Hmm, allocation tests are always finicky, especially with coverage involved. I'll see whether it's real or whether we can work around it somehow. OpenCL failure is genuinely interesting. It only fails on 1.12 and not on CUDA, I wonder whether there's any miscompilation going on there? Will try to track it down |
Member
|
Thanks, I'll take a look next week! |
The CSR kernels compute `c[i] = α * s + β * c[i]`, so they read the destination even when β is zero. `kkt_errors!` passed a float `zero(T)` into a `scratch.x` freshly obtained from `similar`, and `0.0 * NaN` is NaN, so uninitialized memory leaked into `c_At_y`. Under pocl the recycled host pages regularly do contain NaN bit patterns, which made the batched OpenCL tests fail on Julia 1.12. The Bool `false` is a strong zero, so it discards the garbage. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
simeonschaub
force-pushed
the
sds/batching
branch
from
July 24, 2026 15:07
37ef004 to
c3b7351
Compare
Collaborator
Author
No, just boring old reading from uninitialized memory... 😞 We should probably have separate |
simeonschaub
commented
Jul 27, 2026
written by claude
written by claude
Cover every subset of the four field groups that can vary across a batch (objective, variable bounds, constraint matrix, constraint bounds). Combinations with a batched constraint matrix are marked broken where they currently fail: slicing a `BatchedGPUSparseMatrixCSR` on the CPU, and preconditioning inside `solve`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`PrimalDualSolution(milp)` copied the shape of `lv` and `lc`, so a MILP batched only through its objective (or only through its constraint matrix) produced an unbatched starting point. Take the number of columns from `nbinstances` instead, and reuse the constructor in the two-argument `solve`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`view(A, :, :, k)` hands the k-th column of `nzval` to `GPUSparseMatrixCSR`, which required a `DenseVector`. GPU array types return one of their own from a contiguous view, but on the CPU the slice is a `SubArray`, so `instance`, `EachInstance` and `spectral_norm` all failed on a batched constraint matrix. Accept any `AbstractVector` of nonzeros and relax the matching `mul!` methods. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A `MILP` carries a single pair of scalings `D1` and `D2` for the whole batch, so `pdlp_preconditioner` cannot handle a batched constraint matrix. Say so instead of failing with a `MethodError` from the `ConstraintMatrix` constructor. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`batched_zeros` picked between a vector and a matrix from the number of instances, so `PrimalDualSolution(milp)` returned a `Union`. Whether a MILP is batched is a property of its type, unlike how many instances it holds: expose that as `isbatched` and pass it down as a `Val`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The `isbatched` docstring cross-references `nbinstances`, which had no docstring of its own, so Documenter could not resolve the `@ref` and the build failed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Only the nonzeros of `GPUSparseMatrixCSR` had to accept an `AbstractVector`, for the sake of slicing a batched matrix. Widening the operands as well made the matrix method ambiguous with the `AbstractTriangular` and `Bidiagonal` ones from LinearAlgebra, which are gone on 1.12 but not on 1.10. Also loosen the batched spectral norm comparison there, where the power method lands further from `opnorm` on the matrix the 1.10 RNG draws. Co-Authored-By: Claude Opus 5 <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
2 participants
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.
An initial implementation of batched linear solves, as described in https://openreview.net/pdf?id=A4qADJulVa
@gdalle, I know we discussed just making all scalars vectors and all vectors matrices, but I ended up going with a parametric approach, since at least for
MILP, we do need to actually be able to tell what is batched and what stays the same across batches and it made sense to go with this dual approach for things like the errors and restart stats as well. My hope is that this also keeps downstream breakage to a minimum, since other than additional type parameters, a lot of the structs actually stay the same.