Specialise GradedSpace functions based on storage type - #511
Conversation
| return i | ||
| end | ||
| _searchsortedfirst(v::Vector, k) = searchsortedfirst(v, k) | ||
| # function _searchsortedfirst(v::Vector, k) |
There was a problem hiding this comment.
This is definitely a non-trivial change here though that might actually require a bit of benchmarking in various cases, as the question is basically "do we use linear or binary search".
I seem to recall @Jutho saying that he measured this and found that the linear search is overall faster, although I can definitely see how this has to depend on the length of the vector.
I do however agree with the point that it could be great to just get rid of a bunch of code in general, and if we find something that actually maintains a competitive dictionary type it would be wonderful if we could outsource this in its entirety. I can understand that for standard MPS it is probably more relevant to focus on the small length cases, while for tensors with more legs or symmetries with more sectors it probably ends up flipping, and at some point there might actually be a case for just switching to either Dict or Dictionary, deleting this code here. Unfortunately it is really hard to measure the total effect this will have since there is soo many different usecases and regimes 😢
There was a problem hiding this comment.
I admit I wasn't very clear in my benchmarks, and I don't have tests literally showing the benefit to switching back to binary search, but it was noticeable. I'll see to showing this more transparently. I'm predicting a crossover at fairly small number of keys though where linear search starts better, but then binary takes over, but I believe even with the smallest case I tested of 13, binary already outperformed.
At some point I was indeed looking at starting from Dict and then converting to SectorDict, which ended up working in e.g. fuse. For things like this, I did decide based on global improvement, so no cutoff where something else started working better. Indeed, at some point I was losing my mind about all these cutoffs, so I just stuck to what always worked best, not necessarily the best case depending on the scenario.
Codecov Report❌ Patch coverage is
... and 3 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.
| n_read = findindex(vals, dualV ? dual(c) : c) # dual-adjusted index for reading V.dims | ||
| n_write = findindex(vals, c) # output is never dual, so c is fine as-is | ||
| newdims[n_write] = _blocklength(V.dims[n_read], ind) # dim(c) = dim(dual(c)) |
There was a problem hiding this comment.
| n_read = findindex(vals, dualV ? dual(c) : c) # dual-adjusted index for reading V.dims | |
| n_write = findindex(vals, c) # output is never dual, so c is fine as-is | |
| newdims[n_write] = _blocklength(V.dims[n_read], ind) # dim(c) = dim(dual(c)) | |
| d = dim(V, c) | |
| n_write = findindex(vals, c) # output is never dual, so c is fine as-is | |
| newdims[n_write] = _blocklength(d, ind) # dim(c) = dim(dual(c)) |
I don't think there's a benefit for not using the generic codepath for n_read here.
I also am slightly confused by the comments: I know that we are mostly using this function for the factorizations in which case the space is non-dual, but then I don't know why the input could be dual? I also don't really understnd why dim(c) = dim(dual(c)) is relevant here.
| function ⊕(V₁::GradedSpace{I, <:SectorDict}, V₂::GradedSpace{I, <:SectorDict}) where {I <: Sector} | ||
| dual1 = isdual(V₁) | ||
| dual1 == isdual(V₂) || throw(SpaceMismatch("Direct sum of a vector space and a dual space does not exist")) | ||
| ks, vs = _sortedmerge(V₁.dims.keys, V₁.dims.values, V₂.dims.keys, V₂.dims.values, +, identity, identity) |
There was a problem hiding this comment.
Can we route this through Base.mergewith(+, V₁.dims, V₂.dims), and then implement a specialized version of that for the sorted vector dict? I do like the idea of keeping this a bit closer to the AbstractDict functionality, such that if we ever end up switching out the dictionary type this is not too much work.
|
|
||
| function fuse(V₁::GradedSpace{I, <:SectorDict}, V₂::GradedSpace{I, <:SectorDict}) where {I <: Sector} | ||
| dual1, dual2 = isdual(V₁), isdual(V₂) | ||
| acc = Dict{I, Int}() # SectorDict `get` within the double for loop accumulates O(N^2) ` findindex` calls -> sort afterwards |
There was a problem hiding this comment.
Is the main purpose of this specialization bypassing the bad scaling of repeated insertion into a SectorDict being slow, or is the point to avoid calling dim(V, c), or is it the combination?
For the latter and readability, it might actually make sense to introduce a new function, similar to Base.pairs that produces something that iterates like "zip(sectors(V), dims(V))" without having to rehash a bunch of things. (Which looks like it might be useful in a bunch of these functions).
For the former, do you think we could get away with a specialized SectorDict(generator) constructor that does precisely this, or alternatively try out the approach where we just accumulate everything into vectors, and then "sort+merge"?
| function fuse(V₁::GradedSpace{I, NTuple{N, Int}}, V₂::GradedSpace{I, NTuple{N, Int}}) where {I <: Sector, N} | ||
| vals = values(I) | ||
| dual1, dual2 = isdual(V₁), isdual(V₂) | ||
| newdims = zeros(Int, N) |
There was a problem hiding this comment.
I do think that probably here we want to restrict to tuple as well, and use Base.setindex instead. (worst case scenario, there is also this: https://github.com/Jutho/TupleTools.jl/blob/a57ef2a1189604bba10ac08b2e2bae02d29d1a3e/src/TupleTools.jl#L37-L43
| return i | ||
| end | ||
| _searchsortedfirst(v::Vector, k) = searchsortedfirst(v, k) | ||
| # function _searchsortedfirst(v::Vector, k) |
There was a problem hiding this comment.
This is definitely a non-trivial change here though that might actually require a bit of benchmarking in various cases, as the question is basically "do we use linear or binary search".
I seem to recall @Jutho saying that he measured this and found that the linear search is overall faster, although I can definitely see how this has to depend on the length of the vector.
I do however agree with the point that it could be great to just get rid of a bunch of code in general, and if we find something that actually maintains a competitive dictionary type it would be wonderful if we could outsource this in its entirety. I can understand that for standard MPS it is probably more relevant to focus on the small length cases, while for tensors with more legs or symmetries with more sectors it probably ends up flipping, and at some point there might actually be a case for just switching to either Dict or Dictionary, deleting this code here. Unfortunately it is really hard to measure the total effect this will have since there is soo many different usecases and regimes 😢
3 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 🤔