Skip to content

use backends and allocators in DMRG - #467

Draft
lkdvos wants to merge 22 commits into
mainfrom
reuse-buf
Draft

use backends and allocators in DMRG#467
lkdvos wants to merge 22 commits into
mainfrom
reuse-buf

Conversation

@lkdvos

@lkdvos lkdvos commented Jul 20, 2026

Copy link
Copy Markdown
Member

This PR threads the backend and allocators 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.

@lkdvos
lkdvos marked this pull request as draft July 20, 2026 19:09
@lkdvos
lkdvos requested review from borisdevos and leburgel July 20, 2026 19:09
@codecov

codecov Bot commented Jul 20, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 92.64706% with 5 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
.../algorithms/derivatives/hamiltonian_derivatives.jl 90.90% 5 Missing ⚠️
Files with missing lines Coverage Δ
src/algorithms/derivatives/mpo_derivatives.jl 72.53% <100.00%> (ø)
...c/algorithms/derivatives/projection_derivatives.jl 100.00% <100.00%> (ø)
src/algorithms/groundstate/dmrg.jl 95.23% <100.00%> (ø)
.../algorithms/derivatives/hamiltonian_derivatives.jl 93.36% <90.90%> (+0.06%) ⬆️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@lkdvos
lkdvos marked this pull request as ready for review July 20, 2026 22:42
@lkdvos

lkdvos commented Jul 21, 2026

Copy link
Copy Markdown
Member Author

Benchmark: DefaultAllocator vs BufferAllocator

End-to-end DMRG2 ground-state run measuring the effect of reusing a BufferAllocator across the sweep, driven by benchmark/allocator_gc.sbatch / benchmark/allocator_gc_run.jl.

Setup: spin-1 Heisenberg chain (Trivial), L=100, χ=128, 5 fixed sweeps, krylovdim=20, non-adaptive eigensolver (identical work for both allocators). Single Julia thread (so GC is single-threaded and its wall time is a clean number), MKL with 16 threads. julia 1.12.6.

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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess this should be automatically resolved once the leading_boundary implementations also get backend and allocator support.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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...

Comment thread src/algorithms/derivatives/hamiltonian_derivatives.jl

@leburgel leburgel left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great to me. Aside from the comments I am quite curious: do you have any idea of how this affects the compilation time?

Comment thread src/algorithms/groundstate/dmrg.jl Outdated
# -------
function (H::JordanMPO_AC_Hamiltonian)(x::MPSTensor)
backend, allocator = H.backend, H.allocator
y = ismissing(H.A) ? zerovector(x) : H.A(x)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess this should be automatically resolved once the leading_boundary implementations also get backend and allocator support.

@lkdvos lkdvos left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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...

lkdvos and others added 13 commits July 27, 2026 15:22
`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>
lkdvos and others added 6 commits July 29, 2026 15:10
`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>
lkdvos and others added 2 commits July 30, 2026 16:17
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>
@lkdvos
lkdvos marked this pull request as draft July 30, 2026 20:46
@lkdvos

lkdvos commented Jul 30, 2026

Copy link
Copy Markdown
Member Author

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>
@lkdvos

lkdvos commented Jul 31, 2026

Copy link
Copy Markdown
Member Author

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.
I'll try to do some more before/after tests to actually see if this matters.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants