add CUDA buffer support - #293
Merged
Merged
Conversation
kshyatt
reviewed
Aug 6, 2026
The prose said `BufferAllocator` while the example below it constructs a `CUDABufferAllocator`, which read as inconsistent.
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
JoeyT1994
pushed a commit
to JoeyT1994/TensorNetworkQuantumSimulator.jl
that referenced
this pull request
Aug 7, 2026
Describe one block of the double-layer contraction as an `ncon` network -- one
integer label per bond, read off each tensor's own `inds` -- instead of a
hand-written permutedims!/mul! chain, and compile that network once per message
into a list of pairwise `tensorcontract!` steps replayed per block. Temporaries
come from a `TensorOperations.BufferAllocator` arena that sizes itself, so
`message_scratch_length`, `scratch_buffer!` and the manual offset carving are
gone along with the fixed-shape kernel.
Three capability changes follow from the generic network:
* Any vertex degree. The old `length(ms) == 2` guard meant the kernel never
fired on a degree-4 vertex -- i.e. never on a square lattice interior, which
then took the ~6-factor `contract` path.
* A non-identity on-site operator is one more tensor in the network rather
than a fallback. The identity is still detected and dropped.
* `backend` is selectable via `Algorithm("blocked"; backend = ...)` and
otherwise resolved from the array type, so cuTENSOR is reachable.
Whether the closing layer is pre-permuted is a backend trait: the strided/Base
backends contract by permuting into temporaries and need it, cuTENSOR takes
arbitrary index modes and would pay a wasted factor for it.
Measured against the previous kernel (hexagonal(3,3) and grid(4,4), ComplexF64):
hex chi=64 blocked sweep 701 ms / 15.7 MiB -> 422 ms / 7.4 MiB
grid(4,4) chi=64 sweep 27.7 s / 30.7 GiB -> 18.3 s / 7.2 MiB
vertex_scalar hex chi=32 58.4 ms / 5.7 MiB -> 35.8 ms / 2.9 MiB
The arena tracks its high-water mark: 1.19 factors at chi=256, degree 3.
Adds `message_arena_stats` and `@debug` reporting for the arena, and takes
TensorOperations from QuantumKitHub/TensorOperations.jl#293 via `[sources]` for its
`BufferAllocator` device support.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
lkdvos
added a commit
that referenced
this pull request
Aug 7, 2026
Wire `ROCArray` and `JLArray` into the storage-generic `BufferAllocator` from #293, through `AMDBufferAllocator` and `JLBufferAllocator`. Unlike the CUDA extension, neither can serve temporaries via `unsafe_wrap`: `AMDGPU.unsafe_wrap` queries `hipPointerGetAttributes` on every call and would re-register host buffers, and `JLArrays` has no `unsafe_wrap` at all. Both instead share the buffer's refcounted `DataRef` at an offset, the way their own `reshape` does, which needs no API call and keeps the buffer alive by refcount. `JLArray` offsets are in elements rather than bytes, so element types whose size does not divide the alignment fall back on a regular allocation instead of landing on a truncated offset. The `JLArray` backing also makes the foreign-storage paths testable without GPU hardware, which they previously were not. Co-authored-by: Claude Opus 5 (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
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.
This PR extends the
BufferAllocatorsupport to also include CUDA-backed storage types.While CUDA itself hosts a memory pool as well, this can be used for a little more fine-grained control, for example when memory is really tight.
Will try and have a look to generalize to AMD too, but hope to do that in a separate PR.
Also for follow-up: I think the cuTENSOR implementations actually have a way to provide the workspace manually as well, so possibly that could take of a little more of the memory pressure.