Conversation
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
Benchmark:
|
| metric (steady-state) | DefaultAllocator |
BufferAllocator |
change |
|---|---|---|---|
| wall time | 65.5 s | 50.4 s | −23% |
| GC time | 24.2 s (37%) | 10.3 s (20.5%) | −57% |
| allocated | 110.85 GiB | 39.85 GiB | −64% (−71 GiB) |
| GC pauses | ~311 | ~137 | −56% |
| full collections | 22 | 9 | −59% |
| ground-state energy | −138.940086126678 | −138.940086126678 | Δ = 0 |
Of the ~111 GiB a sweep churns through, ~71 GiB is recyclable effective-operator contraction scratch — with BufferAllocator it now comes from the reused buffer instead of the GC heap. The residual ~40 GiB (Krylov basis vectors, the gauge-step SVD, environment updates, result tensors) is not managed by the operator allocator and is therefore identical for both — that's the floor.
Raw output
model=Heisenberg S=1 L=100 χ=128 sweeps=5 krylovdim=20 reps=3
julia=1.12.6 julia_threads=1 gc_threads=1 BLAS_threads=16
DefaultAllocator | rep 1/3 | 66.79 s | GC 25.48 s ( 38.1%) | 110.85 GiB | 320 pauses | 23 full
DefaultAllocator | rep 2/3 | 65.62 s | GC 24.24 s ( 36.9%) | 110.85 GiB | 310 pauses | 22 full
DefaultAllocator | rep 3/3 | 65.41 s | GC 24.14 s ( 36.9%) | 110.85 GiB | 313 pauses | 22 full
BufferAllocator | rep 1/3 | 50.91 s | GC 10.34 s ( 20.3%) | 39.85 GiB | 135 pauses | 9 full
BufferAllocator | rep 2/3 | 50.37 s | GC 10.33 s ( 20.5%) | 39.85 GiB | 137 pauses | 9 full
BufferAllocator | rep 3/3 | 50.38 s | GC 10.32 s ( 20.5%) | 39.85 GiB | 137 pauses | 9 full
energy check: |Δ| = 0.00e+00
| # ------- | ||
| function (H::JordanMPO_AC_Hamiltonian)(x::MPSTensor) | ||
| backend, allocator = H.backend, H.allocator | ||
| y = ismissing(H.A) ? zerovector(x) : H.A(x) |
There was a problem hiding this comment.
As far as I can tell, this continuing block's action doesn't go through an @plansor with backends or allocators passed (this particular line dispatches to line 77 of mpo_derivatives.jl). I think this will occur in every algorithm which calculates a Galerkin error, as there through the projection operator * x you end up here. This isn't an issue within the actual eigsolve calls which matter in e.g. DMRG, so I'm not sure how severe this is
There was a problem hiding this comment.
I guess this should be automatically resolved once the leading_boundary implementations also get backend and allocator support.
There was a problem hiding this comment.
Oh I think you are right, only the prepared ones get the allocator/backend support. I'll make sure to update this, although I would really like to simplify this mess of code at some point...
leburgel
left a comment
There was a problem hiding this comment.
Looks great to me. Aside from the comments I am quite curious: do you have any idea of how this affects the compilation time?
| # ------- | ||
| function (H::JordanMPO_AC_Hamiltonian)(x::MPSTensor) | ||
| backend, allocator = H.backend, H.allocator | ||
| y = ismissing(H.A) ? zerovector(x) : H.A(x) |
There was a problem hiding this comment.
I guess this should be automatically resolved once the leading_boundary implementations also get backend and allocator support.
lkdvos
left a comment
There was a problem hiding this comment.
I'm not sure this affects compilation times that much, in the sense that of course we have to compile stuff multiple times for different backends and allocators, but since that was previously not possible this isn't really a regression.
I'm still playing around a bit with some compilation time stuff too though, for example I think the higher-level algorithm implementations don't actually need to be specialized which might recover some of this.
| # ------- | ||
| function (H::JordanMPO_AC_Hamiltonian)(x::MPSTensor) | ||
| backend, allocator = H.backend, H.allocator | ||
| y = ismissing(H.A) ? zerovector(x) : H.A(x) |
There was a problem hiding this comment.
Oh I think you are right, only the prepared ones get the allocator/backend support. I'll make sure to update this, although I would really like to simplify this mess of code at some point...
`C_/AC_/AC2_projection` build their effective operator with `prepare = false` and apply it once, but had no way to pass a contraction backend or a scratch allocator down to it. Add both as keyword arguments. Note that `AC2_projection`'s trailing `kwargs...` are forwarded to `AC2(above, site; kind)`, not to the Hamiltonian, so the two new arguments have to be named explicitly rather than swept up by the splat. Also add `_derivative_kwargs`, which drops the allocator from the keyword arguments when it is `nothing`. Algorithms that may sweep concurrently use this to fall back on the per-call default of the derivative constructors (a buffer private to that local update) instead of sharing their own stateful one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The unprepared `MPODerivativeOperator` applied itself through bare `@plansor`
calls, so it had nowhere to put a backend or an allocator: passing either to
`AC_hamiltonian(...; prepare = false)` was silently a no-op for every generic MPO.
Only `JordanMPO_*` and the prepared `PrecomputedDerivative` carried them.
Add both as fields and use them in all nine action methods. The constructors take
them as optional trailing positional arguments defaulting to `DefaultBackend()`
and `DefaultAllocator()`, so this operator still holds no scratch space of its own
unless explicitly given some, and existing call sites keep their behaviour.
`@plansor` brackets its generated code with `allocator_checkpoint!`/
`allocator_reset!` whenever an allocator is supplied, so a shared buffer is scoped
per application and its offset is released on the way out.
The type aliases gain trailing parameters as a result, which means
`prepared_operator_type` has to match them as a subtype (`::Type{<:MPO_..}`)
rather than as an exact pattern; the latter would no longer match a concrete type.
Also let the generic `MPO_AC_/AC2_Hamiltonian` blocks nested in the `JordanMPO_*`
operators (the "continuing" `A`/`AA` fields) inherit the outer backend and
allocator, so the Jordan matvec no longer drops to the default allocator for its
dense sub-block.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the fields to `IDMRG`/`IDMRG2` and thread them through the local updates of all three drivers built on them: `find_groundstate`, the statmech `leading_boundary`, and `approximate!`. Both sweeps are serial, so the allocator is shared across every local update of a solve. That makes it unsafe to reuse one algorithm across concurrent solves, which the docstrings now warn about. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the fields to `TDVP`/`TDVP2` and thread them through the local integration steps. The finite sweeps are serial and share the allocator throughout. The infinite `TDVP` sweep already branched on whether the scheduler is serial: the serial branch shares `alg.allocator`, while the concurrent one passes only `alg.backend` and leaves each local update with its own private buffer, since a `BufferAllocator` is not thread-safe. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the fields and thread them into the per-site local update. The unit-cell sweep is driven by the global scheduler, so `alg.allocator` is only handed to the local updates when that scheduler is serial; otherwise each update falls back on a private buffer via `_derivative_kwargs`. The inner AC/C eigensolves of the `parallel = true` path always run concurrently with each other, so those take the backend alone. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the fields and thread them into the per-site local update, on the same scheduler-dependent terms as VUMPS. The `approximate` variant already dispatches its local update on `::SerialScheduler` versus the generic case, so the split falls out of the existing methods: the serial one shares the allocator, the concurrent one passes only the backend. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the fields and give `GrassmannMPS.fg` matching keyword arguments, forwarded from both `find_groundstate` and the statmech `leading_boundary`. The `FiniteMPS` gradient is a plain serial `map` and shares the allocator across sites. The infinite and multiline gradients sweep under the global scheduler, so they gate on it exactly as VUMPS does; the scheduler is read once and passed to both the gate and the sweep so the two cannot disagree. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the fields to `OptimalExpand`, `SketchedExpand` and `VUMPSSvdCut`, and thread them into the effective operators each builds. `SketchedExpand` only becomes able to use them now that `MPODerivativeOperator` carries the fields, since it applies a bare `MPO_AC_Hamiltonian` rather than a prepared operator. `RandExpand` and `SvdCut` are left alone: neither builds a derivative operator at all. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`DMRG3S` builds a bare `MPO_AC_Hamiltonian` to form its perturbation, so it can only pick up a backend and allocator now that `MPODerivativeOperator` has the fields. It is a gauge wrapper rather than a standalone algorithm, so it takes them from the enclosing `DMRG` instead of carrying its own. The 7-argument `gauge!` entry point gains the two keyword arguments. A plain, non-expanding gauge step touches neither `H` nor `envs` and so has no contraction to route them to: the fallback accepts and drops them, which lets `local_update!` pass them unconditionally. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the fields to `ChepigaAnsatz`/`ChepigaAnsatz2` and `DynamicalDMRG`, keyword arguments to `exact_diagonalization`, and forward the existing `DMRG`/`DMRG2` fields through the finite `approximate!`. The `ChepigaAnsatz` constructors keep their old positional forms working via methods that fill in the defaults. `contract_mpo_expval`'s generic fallback in expval.jl is deliberately left at the defaults: reaching it would mean threading both arguments through `expectation_value`, and it is a single contraction per site rather than a repeated matvec. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`TensorOperations.BufferAllocator` is a stateful bump buffer, so reusing one across the local updates of an algorithm avoids re-allocating scratch space every iteration, but it cannot be shared by concurrent tasks. An `AllocatorPool` reconciles the two by combining a factory for creating allocators with a pool of idle ones: every unit of concurrent work checks one out for the duration of its work and hands it back afterwards. The pool retains at most `maxsize` idle allocators and remembers the largest buffer it has seen, so an allocator created for a task that joins late starts out pre-sized instead of having to re-learn how much scratch space the contractions need. `with_allocators` marks the outermost solve and releases the buffers of every pool reachable from the algorithm when it returns. The marker is a `ScopedValue` rather than a counter on the pool: scoped values are inherited by child tasks (and thus by threaded sweeps) but are private to a task tree, so concurrent solves cannot observe each other's state. `acquire!` pins the factory result to the pool's allocator type. Without it the allocator infers as `Any` -- the default factory `BufferAllocator` is a `UnionAll`, so calling it is a dynamic call -- and every contraction the allocator is handed to becomes a dynamic call as well. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Replace the `allocator` field of the algorithm structs with an `allocatorpool`, and check an allocator out of it per unit of concurrent work: per local update in a threaded sweep, or once around the sweep when that sweep is serial. Previously a stateful allocator could only be handed to a local update on a provably serial path. `_derivative_kwargs` dropped it on every other path, so the derivative constructors fell back on a fresh `BufferAllocator` per call -- which is worse than no buffer at all, since it starts at size 0 and every contraction in that call falls back on Julia's allocator anyway. Threaded sweeps therefore got no buffer reuse whatsoever, which is exactly the configuration where it matters most. That mechanism, the five `scheduler isa SerialScheduler` gates that drove it, and the docstring warnings that an algorithm cannot be reused across concurrent solves are all gone. The `parallel = true` branches of `_localupdate_vumps_step!` and `_localupdate_vomps_step!` are removed. They only existed to keep the AC and C solves off a shared allocator, were unreachable (every caller hardcoded `parallel = false`), and nested `@spawn` inside an OhMyThreads region. The low-level derivative constructors now default to `DefaultAllocator()` for the same reason their operators already did: a throwaway `BufferAllocator` per call buys nothing and is itself garbage. `excitations(::FiniteExcited)` is recursive, so it opens the solve scope once and recurses through a plain helper; a closure at every level would defeat inference. Along the way, `exact_diagonalization` now holds its allocator for as long as the effective Hamiltonian refers to it, and `time_evolve` no longer reads `t` before it is assigned when logging at verbosity >= 2. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The contraction backend and the scratch-space allocator were the one pair of MPSKit settings whose defaults were spelled inline at every construction site rather than routed through `Defaults`. Give them a home: * `backend()` returns `TensorOperations.DefaultBackend`, matching the shape of `alg_svd()`/`alg_orth()`/`alg_gauge()`. * `buffering[]` / `set_buffering!` sit next to `set_scheduler!`, which is the right neighbour: both are execution-resource knobs that do not change the answer, unlike the numerical parameters carried on the algorithms. The switch is a `Bool` rather than an allocator object on purpose. A `Ref` or `ScopedValue` holding a pool would be abstractly typed, so the allocator drawn from it would infer as `Any` and turn every downstream contraction into a dynamic call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
An algorithm carried its own `allocatorpool`, constructed before any state was
in scope. That ordering was the root problem: at the moment the default was
chosen MPSKit could not know where the tensors would live, so a GPU-backed
state was silently handed a host bump buffer, and there was no signal on which
to base any automatic choice. Create the pool at first contact with the state
instead, and both problems disappear rather than needing to be patched.
* `allocatortype(::Type{M})` maps a storage type to an allocator type: host
storage gets `BufferAllocator`, anything else `DefaultAllocator`, which
allocates through the storage type itself and is therefore correct on any
device. This trait is the extension point, keyed on the axis that actually
determines the answer.
* `allocatorpool(::Type{A})` keeps one lazily created pool per *allocator*
type, so e.g. `Vector{Float64}` and `Vector{ComplexF64}` share a byte buffer.
The `::AllocatorPool{A}` assertion on the registry lookup is load-bearing:
without it the pool infers as `Any` and every contraction handed the
allocator becomes a dynamic call.
* `with_allocator(f, x)` takes anything with a `storagetype`, normally the
state being operated on.
* `tforeach_allocated`/`tmap_allocated!` check out one allocator per *task*
rather than per iteration, with the scheduler deciding the granularity. This
replaces three hand-written and mutually inconsistent patterns, including the
explicit `scheduler isa SerialScheduler` branch in `_timestep_infinite` and
its duplicated serial/parallel bodies.
* Buffers are no longer released when a solve returns. Pools outlive any single
call, so scope-based release stopped being meaningful; `release_allocators!()`
hands the memory back on demand. This removes `_IN_SOLVE`, `with_allocators`,
its ~30 wrappers and the `release_allocators!(::Algorithm)` `nfields` walk,
which also had a latent hole: it only reached pools stored directly in
`Algorithm` fields, never one behind a `Vector` or `NamedTuple`.
* `__init__` empties the registry so a downstream `@compile_workload` cannot
serialize a live buffer, or a `maxsize` sized for the precompiling process.
`backend` stays a field: it is stateless, has no lifetime, and
`TensorOperations` already resolves `DefaultBackend` from the argument types,
so it needs no device handling of its own.
Removing a struct field and rewriting its uses cannot be split into
individually building commits, which is why this touches 22 files at once. Only
the leaf-level `backend::AbstractBackend = DefaultBackend()` defaults in
`derivatives/` and `utility.jl` are left alone; those are always called with an
explicit backend from the algorithm layer.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
One `allocatortype` method each, which is all the trait needs. Without them a CUDA- or ROCm-backed state falls back on `DefaultAllocator`: correct, but it leaves scratch space to the garbage collector. `CUDAAllocator`/`AMDAllocator` call `unsafe_free!` on temporaries instead, returning them to the vendor's caching pool, so this makes GPU better off than before rather than merely unbroken. Pooling itself buys nothing on device: both allocators are stateless singletons, CUDA.jl already runs a caching allocator, and its streams are task-local. The pool composes with them regardless, since the reset/size/presize traits fall through to their no-op branches. The CUDA memory kind is taken from the state's own storage rather than hardcoded: the zero-argument `CUDAAllocator()` defaults `Mout` to `UnifiedMemory`, which would give a contraction's output a different storage type - and hence a different `TensorMap` type - than a `DeviceMemory`-backed state. `TensorOperations` moves to 5.6.2, the first release with `AMDAllocator`. `ScopedValues` is dropped; nothing uses it now that the solve scope is gone. The CUDA path is verified against a real device. AMD is not: no ROCm hardware was available, so that method rests on the Buildkite `rocm` queue. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`test/misc/allocatorpool.jl` keeps the `AllocatorPool` unit tests unchanged and replaces the solve-scope tests, which no longer have a subject, with the trait, the lazy registry, `set_buffering!`, and inference assertions on the path that actually matters - `acquire!` is what hands the allocator to the contractions, so that is what must infer concretely. The `tforeach_allocated` matrix runs eight scheduler configurations and asserts the checkout is per task. Getting the bound right took two tries: `maxsize` caps how many idle allocators are *retained*, not how many can be live, and the task count is the chunk count rather than `nthreads()` - a `chunksize`-based scheduler spawns one task per chunk and can exceed the thread count. The bound is now computed the same way the primitive chunks, so it is exact rather than timing-dependent. `test/gpu/` gained the algorithm coverage it never had: it exercised construction, state and operator algebra and `expectation_value`, but never ran an algorithm, so nothing checked that the allocator handed to a local update matches the device the state lives on. The new tests run DMRG and TDVP against a CPU reference and confirm the state does not migrate off the device. The shared body lives under `setup/` because `runtests.jl` filters that prefix out of discovery; anywhere under `test/gpu/` it would be collected and run on its own with `ArrType` undefined. `infer_rt` wraps `Base.infer_return_type`, which is 1.11+, over `Core.Compiler.return_type`, which MPSKit already uses elsewhere and which the 1.10 LTS still has. Buildkite runs the GPU suites on the LTS. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Rewrite the "Allocators and buffer reuse" section: the pool is no longer a field to configure, `set_buffering!` is the knob, `release_allocators!()` the lever, and `allocatortype` the extension point. Adds a note on GPU-backed states, including why `backend` needs no equivalent treatment - `DefaultBackend` is a placeholder that resolves from the argument types, so it already selects cuTENSOR for `CuArray`-backed states. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Runs `test/gpu/cuda/` including the new algorithm tests. Useful for exercising cluster GPU models; the partition and `--gres` are placeholders. Uses the juliaup 1.12.6 install on GPFS directly, since `module load julia` gives 1.11.2 and the depot is precompiled under 1.12.6. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Update here: I'm still working on what I want to do with the design here, and for now I'm thinking that this might be overengineered right now, so I'm going to try and slim this down slightly. |
The pooling machinery existed for one reason: `BufferAllocator` holds a
mutable bump-pointer offset and is not thread-safe, so concurrent work
needed one buffer per task - hence a registry of pools, a checkout
protocol, and helpers that re-chunked a scheduler to learn its task
granularity.
Threaded work now gets a `ManualAllocator` instead. It bypasses the
garbage collector just as a bump buffer does, but holds no state at all,
so every task can share one. That collapses the choice of allocator to a
single dispatch on (storage type, is-it-shared) that each algorithm makes
once at the start of a run:
default_allocator(::Type{<:HostStorage}, ::SerialScheduler) = BufferAllocator()
default_allocator(::Type{<:HostStorage}, ::Scheduler) = ManualAllocator()
default_allocator(::Type, ::Scheduler) = DefaultAllocator()
Gone with the pool: `AllocatorPool`, the pool registry, `acquire!`/
`release!`, `with_allocator`, `tforeach_allocated`/`tmap_allocated!`,
`nochunking` and the `OhMyThreads.Schedulers` internals it needed,
`release_allocators!` and `allocatortype`. Also gone are the CUDA and
AMDGPU extensions: device storage falls through to `DefaultAllocator`,
which builds intermediates from the storage type itself.
Two things keep this inferable. `Defaults.buffering` becomes a
compile-time preference, so the selector is pure dispatch rather than a
branch returning a union. And since `Defaults.scheduler[]` is abstractly
typed, the allocator is always built inside a function that *takes* the
scheduler, where Julia has specialized on its concrete type; the two
sites that read the global in the same body (`_timestep_infinite` and
`GrassmannMPS.fg`) gained that boundary.
Note that dropping the pool costs little: a fresh `BufferAllocator` sizes
itself from the first `tensoralloc` and is warm after one or two
contractions, so cross-call reuse was worth a handful of allocations per
solve, not per update.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Alright, update again here, I think this is now more or less settling on a design I'm happy with. I kept the backends selectable, being a hook for making changes, but actually decided against making the allocator a configurable thing, mostly because it is just too easy to accidentally do something wrong with that. I think it is fair to always take the bumper where possible, and this now also handles the GPU case, so the main thing left for me to do is clean up the tests and rebase. Also note that this now means that parameter sweeps in parallel would work perfectly fine, and would probably have way better performance because the GC needs to run less often. |
This PR threads the
backendandallocators through the DMRG implementations.As a result, we can use a single buffer that captures all temporary intermediate tensors in the effective hamiltonian applications, which is reused across sites, hopefully alleviating quite a bit of pressure from the garbage collector.
The one thing about this that is somewhat unfortunate is that this turns the algorithms into a non-threadsafe construction, and reusing the same algorithm in a parallel parameter sweep would yield wrong results.
In principle I can try to circumvent this, but I'm wondering how much that is worth it in the end?
Side note: The goal is of course to also do this for the rest of the package, but it is slightly easier to focus on a single algorithm first.