implement batching - #80
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>
|
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 |
|
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>
37ef004 to
c3b7351
Compare
This comment was marked as resolved.
This comment was marked as resolved.
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>
Upstream moved the GPU test files under `test/gpu/<backend>/`, so the CUDA batching testset and the OpenCL one follow the CUDA and OpenCL runners there. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Upstream gave the sparse matrix-vector kernels a strong zero, which the batched matrix-matrix ones need just as much: `kkt_errors!` multiplies into the uninitialized `scratch.x`, so a float zero would read whatever NaN garbage lives there. With every CSR path honoring it, the `false` that stood in for `zero(T)` can go. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Metal only supports single precision, so `test_batching` takes the float type to solve in. Two iterates of a batch drift apart much faster in `Float32`, hence the looser tolerance on the comparison against the single solves. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`runtests.jl` includes every test file into `Main`, so the eight files including `fixtures.jl` each redefined the same handful of methods and buried the CI log under overwritten-method warnings. Include them once, up front. `fixtures.jl` loads CoolPDLP, so it has to come after every `set_preferences!` call: the per-group setup moves above it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Only the batch and the solver called it, always with a one α and a zero β, so the three other branches of the batched matrix-vector method went untested. Walk through all four, and refuse a batch whose instances disagree on the pattern. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Ok, I think this PR is ready! I looked into the missing coverage codecov is explaining about, but it seems to be an upstream issue with Julia. The function bodies all show up as hit, only the first line of the function definition is sometimes marked as uncovered Edit: coverage issue should be resolved by #85 |
|
I'll take a detailed look on Wednesday, thanks! You could use a break from batching anyway, to look at more flow-specific stuff |
gdalle
left a comment
gdalle
left a comment
There was a problem hiding this comment.
I mostly reviewed the code in src, didn't look at the tests in detail. For me, the main question is whether we want to allow other batched matrix types or force users to pick our own batched CSR.
| kwargs... | ||
| ) | ||
| return map(axes(K, 3)) do k | ||
| spectral_norm(view(K, :, :, k), view(Kᵀ, :, :, k); kwargs...) |
There was a problem hiding this comment.
This will be really slow unless the batched sparse matrices have a specialized view method.
I'm not sure it's safe to assume that users will always rely on our own batched matrix implementations?
There was a problem hiding this comment.
I think it's not unreasonable to assume that batched constraint matrices should be nicely splittable along batches, but that probably falls under the larger discussion to be had whether to keep batched matrix support for now and if so, what we want to impose upon them
| Base.eltype(::Type{<:BatchedDiagonal{T}}) where {T} = T | ||
| Base.size(D::BatchedDiagonal) = (size(D.diag, 1), size(D.diag, 1), size(D.diag, 2)) | ||
| Base.size(D::BatchedDiagonal, i::Integer) = i <= 3 ? size(D)[i] : 1 | ||
| LinearAlgebra.diag(D::BatchedDiagonal) = D.diag |
There was a problem hiding this comment.
Is this coherent with the definition of diag?
There was a problem hiding this comment.
Not sure, a lot of these overloads are a little bit iffy, since Julia doesn't normally define multiplication for 3d arrays
| @@ -20,29 +20,35 @@ $(TYPEDFIELDS) | |||
| """ | |||
| struct MILP{ | |||
There was a problem hiding this comment.
The docstring needs to be updated too
| M <: AbstractArray{T}, | ||
| Mt <: AbstractArray{T}, |
There was a problem hiding this comment.
This is a bit too vague, we could impose Union{AbstractMatrix{T},AbstractArray{T,3}} for readability
|
|
||
| function is_feasible( | ||
| x, milp::MILP; | ||
| x::AbstractVector, milp::MILP; |
There was a problem hiding this comment.
Batching at this level would probably be more efficient, even though we don't call this function often
There was a problem hiding this comment.
Yeah, you're right. I wasn't sure whether it was worth the extra complexity since it's not really performance critical. I think we can leave it as a later TODO for now, or do you think it's important to fix before merging?
There was a problem hiding this comment.
a good way to track this would be to try to eliminate all usages of instance
| """ | ||
| BatchedGPUSparseMatrixCSR(A, nbinstances) | ||
|
|
||
| Repeat `A` into a batch of `nbinstances` identical matrices. |
There was a problem hiding this comment.
Why would this ever be necessary? In that case, can't we just keep A as a matrix?
There was a problem hiding this comment.
I think it's mostly used for testing
| @inline batchval(v::AbstractVector, k, ::Integer) = v[k] | ||
| @inline batchval(m::AbstractMatrix, k, batch_idx::Integer) = m[k, batch_idx] | ||
|
|
||
| @kernel function spmm_csr!( |
There was a problem hiding this comment.
So this is the one which is much slower than cuSPARSE's SpMM?
There was a problem hiding this comment.
It would be worth reading up on SpMM kernels
There was a problem hiding this comment.
I didn't benchmark this one in detail, I was mostly talking about my other matrix zoo PR before
| if α_is_one && β_is_zero | ||
| kernel!(c, A.rowptr, A.colval, A.nzval, b, One(), Zero(); ndrange = size(c)) | ||
| elseif α_is_one | ||
| kernel!(c, A.rowptr, A.colval, A.nzval, b, One(), β; ndrange = size(c)) | ||
| elseif β_is_zero | ||
| kernel!(c, A.rowptr, A.colval, A.nzval, b, α, Zero(); ndrange = size(c)) | ||
| else | ||
| kernel!(c, A.rowptr, A.colval, A.nzval, b, α, β; ndrange = size(c)) | ||
| end |
There was a problem hiding this comment.
There's a lot of duplication for this 4-way split, can we avoid it?
There was a problem hiding this comment.
Yes, agree. We should probably just use LinearAlgebra's MulAddMul, I ended up reinventing the wheel here a little
|
Since it will be verified in most of our applications and since it would allow considerable code simplifications, what would you think of stashing the batched matrix additions on a branch (to avoid losing them) and trimming the current PR by assuming that A is constant across batches? It would remove a lot of complexity and edge cases, so I think it might allow us to get this merged before JuliaCon |
|
Yes, you're probably right. It bugs me a little to not get that in as well, but I agree it's probably the right call for now. There are too many open questions for how to handle 3d arrays representing batched matrices and operations like matmul and inv properly. |
| necessary = absolute(err_candidate, ω) <= necessary_decay * absolute(err_restart, ω) | ||
| no_local_progress = absolute(err_candidate, ω) > absolute(err_candidate_last, ω) | ||
| long_inner_loop = inner >= artificial_decay * total | ||
| candidate = batched_mean(abs_candidate) |
There was a problem hiding this comment.
Parametrizable aggregation function here instead of mean?
| end | ||
| end | ||
|
|
||
| broken && @test_broken kkt_errors!(KKTErrors(sol_dev), Scratch(sol_dev), sol_dev, milp_dev) isa KKTErrors |
Every instance of a batch now shares a single A, so BatchedGPUSparseMatrixCSR, the per-instance preconditioning, and BatchedDiagonal go away, and MILP is back to accepting plain matrices as A and At. Batching of the objective and the bounds stays. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
should_restart reduces the per-instance absolute KKT errors to one number before applying the three restart criteria; the reduction is now a batch_aggregation parameter (the mean by default) instead of being hard-coded, and Algorithm carries the restart parameters as a type parameter so the aggregation stays concretely typed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
|
||
| function iteration_allocations(state, milp, algo) | ||
| run_iterations!(state, milp, algo, 5) | ||
| return @allocated run_iterations!(state, milp, algo, 20) |
There was a problem hiding this comment.
I don't think we can actually guarantee this as is. For example on CUDA mapreduce may allocate if it decides to use multiple blocks: https://github.com/JuliaGPU/CUDA.jl/blob/959f3caa51f31846f8986554990af6d281e05928/CUDACore/src/mapreduce.jl#L266
It is why I had to implement the BatchQuadraticModels.batch_mapreduce!: see https://github.com/klamike/BatchQuadraticModels.jl/blob/main/src/batch_mapreduce.jl and CUDA overloads https://github.com/klamike/BatchQuadraticModels.jl/blob/main/ext/cuda/mapreduce.jl (namely the use of CUDA.reduce_block)
I don't think there is a way to do it while staying device agnostic though..
There was a problem hiding this comment.
We only run this test on the CPU, you are right that I believe all the GPU backends do rely on dynamic allocations for kernel launches. I think it's a good check though that at least the CPU case doesn't allocate to avoid any accidental allocations somewhere in the main loop.
Btw, do you know about GPUArrays.mapreducedim!? It allows you to reduce into an existing array. I think that could be used instead of your hand-rolled implementation and will probably be more optimized.
There was a problem hiding this comment.
It is not the kernel launch allocations I am talking about (these were actually fixed/reduced recently I believe, I remember seeing a PR)
it is that the mapreducedim! implementation (sometimes) uses temporaries instead of warp reduction. Note the first link I sent is to the CUDA implementation of GPUArrays.mapreducedim!
There was a problem hiding this comment.
Ah, sorry, I totally missed your first link, you're right! It might be nice if mapreducedim! allowed preallocating a scratch buffer that doesn't need to be reallocated every time, at least if the dimensions stayed the same, but I guess that's a more general issue and can always be addressed later
There was a problem hiding this comment.
I agree that keeping a no-allocation check on the CPU is a necessary safeguard, even if it fails on GPUs
|
|
||
| Decide whether the whole batch restarts, based on an aggregate of the per-column fixed-point residuals. | ||
|
|
||
| Since every instance of the batch restarts at the same time, the three usual criteria (sufficient decay, necessary decay without local progress, long inner loop) are applied to `params.batch_aggregation` (the mean by default) of the per-column absolute KKT errors instead of requiring each column to agree. |
There was a problem hiding this comment.
Do you have a sense of how invasive it would be to make termination/restart decisions on a per-instance basis? I guess we'd have to introduce some masking? (future work for sure; I understand you're following the nvidia paper on this.)
There was a problem hiding this comment.
We definitely talked about it, I don't think it would be too invasive. I believe what the paper does is permute the columns, so that the non-converged columns come first and then only run the iteration on those first n columns, which is basically the same as your masking idea
There was a problem hiding this comment.
Do we split the batch_aggregation reduction depending on whether we're considering restarts or termination? IIRC, the NVIDIA paper uses the mean for restarts and the max for termination, so perhaps we need two parameters with slightly more specific names than batch_aggregation?
|
|
||
| function is_feasible( | ||
| x, milp::MILP; | ||
| x::AbstractVector, milp::MILP; |
There was a problem hiding this comment.
a good way to track this would be to try to eliminate all usages of instance
| Extract the value of `val` for the `i`-th instance of the batch. | ||
| """ | ||
| instance_num(val::Number, ::Int) = val | ||
| instance_num(val::AbstractVector, i::Int) = val[i] |
There was a problem hiding this comment.
this will raise scalar indexing with gpu arrays right?
There was a problem hiding this comment.
Yeah, it would need an @allowscalar annotation at least. I think you typically will want to avoid calling instance with a GPU array anyways
There was a problem hiding this comment.
Does that mean the batched solver is not compatible with GPU? I believe I saw some calls to instance that happen in the loop.
There was a problem hiding this comment.
There shouldn't be any. I just grepped for it again and didn't find any calls that happened inside the loop
| rel_primal = primal / primal_scale | ||
| rel_dual = dual / dual_scale | ||
| rel_gap = gap / gap_scale | ||
| rel_primal = primal ./ primal_scale |
There was a problem hiding this comment.
perhaps we should reduce these to scalars instead? how does the progbar look with this?
There was a problem hiding this comment.
Yes, that's not a bad idea, especially with larger batches. Should we just take the average of the relative quantities? Though the maximum might be more interesting, maybe we show both?
There was a problem hiding this comment.
I think maximum is better since I believe the termination is currently based on maximum
There was a problem hiding this comment.
Maybe also a counter of how many instances are still violating the termination condition as well.
There was a problem hiding this comment.
Termination is, but restart uses the mean by default, so displaying both might make sense?
|
Thanks for the thorough review, very much appreciated! I left some replies in-line |
One value per column gets unreadable past a few instances, so the progress lines now show the maximum and the mean of the relative errors instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Reusing format_error keeps mean at the same column on every row, and matches how the final stats print their errors. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
3 participants
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.