Specialise GradedSpace functions based on storage type - #511
Conversation
Codecov Report❌ Patch coverage is
... and 8 files with indirect coverage changes 🚀 New features to boost your workflow:
|
lkdvos
left a comment
lkdvos
left a comment
There was a problem hiding this comment.
Thanks for starting to have a look at this!
I am trying to go over this in slightly more detail, but let me put some global remarks/thoughts I have here:
The first thing relates to the NTuple case for large N. I absolutely agree that the compile time of that is untenable, and should be addressed, and while I would probably do something similar to you, there is something that looks a bit off to first constructing a vector, only to then immediately convert it to a tuple. I would definitely prefer to keep the small N values non-allocating and fully specialized, if at all possible, since otherwise this is a lot of extra code to maintain and we might also consider just keeping everything SectorDict and optimizing that.
In some sense, if we are already allocating a vector, we might as well just store the result as a vector and see if the extra pointer indirections actually matter, which I'm not expecting them to do.
So basically, I think it might be useful to just have a GradedSpace{I, Vector{Int}} implementation that has the exact same semantics as the tuple version, but avoids the compile-time issues.
@assume_effects :foldable function sectorstoragetype(::Type{I}) where {I <: Sector}
if Base.IteratorSize(values(I)) isa Union{HasLength, HasShape}
N = length(values(I))
return N <= 10 ? NTuple{N, Int} : Vector{Int}
else
return SectorDict{I, Int}
end
end
Base.getindex(::SpaceTable, I::Type{<:Sector}) = GradedSpace{I, sectorstoragetype(I)}A different thing that could be relevant is that in principle we could also reduce the compile time issues for the smaller N values by "blocking" the compiled types, e.g. NTuple{N, Int} for the next value in the set N = 1,2,4,8,16,32 which basically trades some compilation for storage efficiency (see e.g. the packages SmallCollections.jl or similar. I do however think that this probably hints at the Vector{Int} approach being more appropriate anyways.
A final comment is that I think one of the inefficiencies of the implementations e.g. for truncated factorizations is probably that we are always constructing these as SectorDicts, which is a bit wasteful in the case of NTuple storage. There might be additional optimizations in that realm as well.
…into bd/gradedspace-storage
|
Short summary of the recent commits:
Some comments of the review I didn't immediately address, and why:
There was a suggestion to run some DMRG code with these changes which I haven't done yet. I believe since that will be dominated by LAPACK, the only thing I'd have to look for is no notable regression. So if things remain in the same ballpark (or just compile at all, since that was an issue), then I think things are fine. |
|
My apologies for the slow review. I find the current code quite difficult to properly review, as I am not entirely convinced of the structure. I do like the modifications to However, I am less convinced by the I still have to read through truncation.jl to see if there is anything better I could come up with, which is what I will do next. The reason for bringing this up and being difficult is that I want to ensure that the code remains logically structured and therefore easier to maintain. |
|
In particular, is all the code of |
`FullVectorDict`/`sectormap` existed only to give the truncation code a dense
"sector => index" map for tuple-backed spaces. Their call sites all go through
`pairs(::SectorVector)`, which materialises a `SectorDict` anyway, so the dense
map bought little for a lot of `Sector`-specific machinery in `auxiliary/`.
Every truncation index map is a plain `SectorDict` again; instead
`pairs(::SectorVector)` is built directly from the already-sorted structure
rather than by repeated sorted insertion. The storage-specialised
`truncate_space` methods are kept, since they never used the dense map.
Also from the review:
- `⊖` for dict storage reuses `_sortedmerge`, whose `_keepunmatched` trait is
generalised into per-side `_unmatched1`/`_unmatched2` hooks
- `_ntuple_storage_threshold` -> `_NTUPLE_STORAGE_THRESHOLD`, lowered to 8
- `ZNSpace{N}` deprecated in favour of `Vect[ZNIrrep{N}]`, which the alias can
no longer track once `N` exceeds the threshold
- drop the over-strong `I == Trivial` assert in `truncate_space`, and name the
index variables `ind`/`inds` instead of shadowing the sector type `I`
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`valtype(::SectorVector)` claims a `SubArray`, but `view` of a GPU array is itself a `CuArray`/`ROCArray`, so pinning the element type to `valtype` broke every GPU factorization. Take the element type from the views instead. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- add `TupleGradedSpace{I,N}` / `DictGradedSpace{I}` aliases for the two
variants `sectorstoragetype` selects between, and use them for dispatch
- `SubtractDims` doubles as the unmatched handler for `⊖`, and is used for the
tuple variant as well, so both paths share one callable
- `_sortedmerge` takes the unmatched handlers as arguments; `mergewith` selects
them inline
- `pairs(::SectorVector)` is lazy, like `blocks(::AbstractTensorMap)`: every
consumer only iterates, and lookups go through the vector itself
- document `sectorstoragetype`, whose docstring interpolates the threshold
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
I think this is ready for another round of review. I've incorporated the code style suggestions and variable names, removed the Small thing I noted, the type alias |
| newdims = MutableNTuple(ntuple(_ -> 0, StaticLength(N))) | ||
| for (c, ind) in pairs(inds) | ||
| d = dim(V, c) | ||
| n_write = findindex(vals, c) | ||
| @inbounds newdims[n_write] = _blocklength(d, ind) | ||
| end |
There was a problem hiding this comment.
Do we need this construction via MutableNTuple? How about
| newdims = MutableNTuple(ntuple(_ -> 0, StaticLength(N))) | |
| for (c, ind) in pairs(inds) | |
| d = dim(V, c) | |
| n_write = findindex(vals, c) | |
| @inbounds newdims[n_write] = _blocklength(d, ind) | |
| end | |
| newdims = ntuple(N) do n | |
| c = vals[n] | |
| d = V.dims[n] | |
| ind = inds[c] | |
| return _blocklength(d, ind) | |
| end |
or thus as onliner
| newdims = MutableNTuple(ntuple(_ -> 0, StaticLength(N))) | |
| for (c, ind) in pairs(inds) | |
| d = dim(V, c) | |
| n_write = findindex(vals, c) | |
| @inbounds newdims[n_write] = _blocklength(d, ind) | |
| end | |
| newdims = ntuple(n->_blocklength(V.dims[n], inds[vals[n]]), N) |
So we are then not iterating the inds dictionary, but since we anyway know that N is pretty small in this case, I do wonder how severe the key lookups are.
(We are also not needing to do findindex(vals, c) anymore, which might sometimes be a simple calculation, but sometimes also an inefficient lookup).
There was a problem hiding this comment.
Good call, although it requires a little more care because not all inds[vals[n]] will work, but I think I can actually work around this because we always have either V.dims[n] == 0 or the key is present
Jutho
left a comment
Jutho
left a comment
There was a problem hiding this comment.
I left some smaller final comments, but otherwise approved.
|
I have not actually approved since I hadn't checked wether automerge was on, and wanted to give you time to look at my suggestions. However, feel free to merge after having looked at them. |
Co-authored-by: Jutho <Jutho@users.noreply.github.com>
…inds Address Jutho's final review round on #511: - replace `StaticLength(N)` with plain `N` (constant propagation makes the old TupleTools artefact unnecessary now that N is already a type parameter), and drop the now-dead `TupleTools: StaticLength` imports. - rewrite `truncate_space(::TupleGradedSpace, inds)` as a pure `ntuple` closure instead of mutating a `MutableNTuple`, guarding zero-dimension sectors before indexing into `inds` (which only holds keys for sectors with nonzero dimension). - `truncate_space(::DictGradedSpace, inds)` no longer collects and sorts: `inds` (whether a `SectorDict` or a `SectorVector`, depending on the truncation strategy) already iterates in sorted order by sector. - drop the now-unused `MutableNTuple`/`findindex` imports in the factorizations submodule. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
`pairs(::SectorVector)` is intentionally lazy (a `Base.Generator`), so comparing it directly to a `Dict` with `==` silently evaluates to `false` rather than erroring, since no `==` is defined between those unrelated iterator types. Collect it before comparing. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
4 participants
This is somewhat connected to what I was doing in QuantumKitHub/TensorKitSectors.jl#106, but beneficial for all sector types. Since I've firsthand experienced how my code transitioned from being unusable to performing well by going from
NTuplestorage toSectorDictstorage, I thought it was about time to look at these storage paths.I had two options going into this. The one I'm still looking into is seeing whether there's a cutoff (or range) where
NTuplestorage severely starts underperforming. The other approach I'm taking in this PR is to specialiseGradedSpacefunctions and constructors based on their storage type. The overarching problems I tried to fix were the following:NTupleconstructor could take unboundedly long to compile for sector types with many sectorsNTupleversions of these functions were not making use of the fact that the sector types match, so there's efficient ways to accessing the sectors, knowing always how many there are as well.SectorDictversions of these functions were doing more per-pair dictionary work than the algorithm actually needs (extra hash lookups, double work, etc)Summary of changes I made:
GradedSpace{I, NTuple{N,Int}}constructor: the old constructor built up the dims tuple viaTupleTools.setindex, which fell back to Base'sntuple(f, Val(N)). This requires compiling this for every distinctN, which I found to scale terribly withN. I first tried building into a vector and then annotating the splat into a tuple, but it turns out that it has a cost that scales withN, which dominated for large enoughN. So now I directly convert the vector through a Base iterator-to-tuple constructor which Julia specialised to make faster depending on N. TheNTuplefuseandtruncate_spacemake use of this as well.dimspecialisation: in general I tried avoiding constructingsectors(V)where possible, and just directly checking theNTupledirectly (throughvalues(I)) or the pairs inSectorDict.⊕,⊖,infimum,supremumspecialisations:NTuplestorage: the two spaces here are always of the same sector type, so their tuples are aligned. I could just do the appropriatemapwithout looking up sectors. Againsectors(V)is the plague.SectorDictstorage: I made use of how the keys are sorted here to merge them in an appropriate way depending on the what the function actually wanted to achieve. These structurally looked the same, so I refactored them into_sortedmerge. This outperforms having to work directly with aSectorDict.fuseSectorDictpath: previously a bunch ofgets andsetindex!s were done in the double for-loop on theSectorDict, which accumulated inefficiently due to lookup cost for this type of dictionary. PlainDicts don't have this, so I just do the accumulation in this and then sort at the end. The complexity hasn't changed since there's still two for-loops, but there's a speedup.truncate_spaceSectorDictpath: same structure asfuseforSectorDicts, but now with vectors because you don't have to look up anything along the way.SectorDict's_searchsortedfirst: I looked into when this was implemented, and goes back to 2019 back when product sectors didn't even exist. So I guess back thenNwas always fairly small, and the linear search was more efficient. However, it seems now that's not particularly the case, so I took the liberty of having it default to Base's method.Benchmarks
I tested Julia 1.10.10 (LTS) and 1.12.6 (stable) since I think those are the two versions most people are on. For the
NTuplestorage sector types, I testedN = 2 / 8 / 64 / 256 / 1296withZ2Irrep / ZNIrrep{8} / Z4Irrep⊠^3 / Z4Irrep⊠^4 / ZNIrrep{6}⊠^6. I also testedN = 15625withZNIrrep{5}⊠^6where possible, which is important to mention. For theSectorDictI testedU1Irrepwith charges-6:6,-50:50, and-200:200(13/101/401 sectors).And here the many many numbers. I spared my sanity by having a robot friend write this in markdown.
Constructor compile time (the
Valeffect)Details
N=15625before does not complete (at least within 5 minutes on my laptop) on either version. After: 973 ms (1.10.10), 1.00 s (1.12.6).NTuple storage (after above compilation time)
Details
And now just the
N=15625case separately, also just after only as before doesn't finish:On the
⊖/1.12.6 number: the first call at this N takes ~290-353s on 1.12.6 specifically (reproduced 3×), but the second call in the same session takes ~1.2s, matching 1.10.10's steady-state ~1.1s for the same op at the same N. Only compiling the whole⊖method together on 1.12.6 is this slow. I didn't look deeper into this, also since its use-case is extremely limited for this large N.SectorDict U1Irrep
Details
(Negative) conclusions/remarks from the benchmarks:
Nthe constructor is still somewhat slower (0.45x-0.99x), but the alternative, per the first table, is to make largeNimpossible to compile in reasonable time.truncate_spaceis still slightly slower atN=2(0.29x/0.86x) but wins fromN=8up.NTuplefusescales the way it does withNI have no clue, but it's at least better for reasonable ranges.⊖on 1.12.6 is a one-time compile bottleneck, but afterwards does fine (see note above).SectorDictwins scale with sector count.All in all, these are improvements, notably the
SectorDict. So it makes you wonder if there's in fact some cutoffNabove which you want to saySizeUnknownto theSectorValues' length 🤔