Specialise GradedSpace functions based on storage type

[borisdevos]

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 NTuple storage to SectorDict storage, 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 NTuple storage severely starts underperforming. The other approach I'm taking in this PR is to specialise GradedSpace functions and constructors based on their storage type. The overarching problems I tried to fix were the following:

• the NTuple constructor could take unboundedly long to compile for sector types with many sectors
• the NTuple versions 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.
• the SectorDict versions 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 via TupleTools.setindex, which fell back to Base's ntuple(f, Val(N)). This requires compiling this for every distinct N, which I found to scale terribly with N. 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 with N, which dominated for large enough N. So now I directly convert the vector through a Base iterator-to-tuple constructor which Julia specialised to make faster depending on N. The NTuple fuse and truncate_space make use of this as well.
• dim specialisation: in general I tried avoiding constructing sectors(V) where possible, and just directly checking the NTuple directly (through values(I)) or the pairs in SectorDict.
• ⊕, ⊖, infimum, supremum specialisations:
  • NTuple storage: the two spaces here are always of the same sector type, so their tuples are aligned. I could just do the appropriate map without looking up sectors. Again sectors(V) is the plague.
  • SectorDict storage: 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 a SectorDict.
• fuse SectorDict path: previously a bunch of gets and setindex!s were done in the double for-loop on the SectorDict, which accumulated inefficiently due to lookup cost for this type of dictionary. Plain Dicts 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_space SectorDict path: same structure as fuse for SectorDicts, but now with vectors because you don't have to look up anything along the way.
• Restored binary search in 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 then N was 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 NTuple storage sector types, I tested N = 2 / 8 / 64 / 256 / 1296 with Z2Irrep / ZNIrrep{8} / Z4Irrep⊠^3 / Z4Irrep⊠^4 / ZNIrrep{6}⊠^6. I also tested N = 15625 with ZNIrrep{5}⊠^6 where possible, which is important to mention. For the SectorDict I tested U1Irrep with 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 Val effect)

Details

Julia	before	after	speedup
1.10.10	104 ms / 144 ms / 697 ms / 1.55 s / 14.3 s	158 ms / 157 ms / 274 ms / 342 ms / 135 ms	0.66x / 0.91x / 2.5x / 4.5x / 106x
1.12.6	170 ms / 170 ms / 1.01 s / 1.64 s / 11.2 s	102 ms / 152 ms / 247 ms / 239 ms / 136 ms	1.67x / 1.12x / 4.1x / 6.9x / 83x

N=15625 before 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

operation	Julia	before	after	speedup
constructor	1.10.10	468 ns / 2.67 µs / 465 µs / 3.83 ms / 99.2 ms	1.03 µs / 5.57 µs / 41.7 µs / 375 µs / 5.30 ms	0.45x / 0.48x / 11.1x / 10.2x / 18.7x
constructor	1.12.6	466 ns / 2.47 µs / 177 µs / 2.78 ms / 72.4 ms	538 ns / 2.50 µs / 32.2 µs / 184 µs / 2.59 ms	0.87x / 0.99x / 5.5x / 15.1x / 28.0x
dim	1.10.10	19.0 ns / 45.3 ns / 21.3 µs / 239 µs / 4.43 ms	10.8 ns / 11.8 ns / 22.8 µs / 306 µs / 5.22 ms	1.8x / 3.8x / 0.93x / 0.78x / 0.85x
dim	1.12.6	19.5 ns / 47.9 ns / 12.7 µs / 260 µs / 4.76 ms	5.7 ns / 5.1 ns / 7.41 µs / 148 µs / 2.19 ms	3.4x / 9.4x / 1.7x / 1.8x / 2.2x
⊕	1.10.10	520 ns / 1.04 µs / 259 µs / 4.45 ms / 102 ms	10.0 ns / 10.8 ns / 4.56 µs / 16.7 µs / 112 µs	52.0x / 96.4x / 56.8x / 267x / 909x
⊕	1.12.6	305 ns / 843 ns / 220 µs / 3.31 ms / 94.9 ms	5.1 ns / 5.6 ns / 3.35 µs / 16.2 µs / 107 µs	59.7x / 151x / 65.7x / 205x / 886x
⊖	1.10.10	138 ns / 384 ns / 260 µs / 5.24 ms / 106 ms	202 ns / 328 ns / 48.1 µs / 651 µs / 11.1 ms	0.68x / 1.17x / 5.40x / 8.05x / 9.56x
⊖	1.12.6	76.3 ns / 222 ns / 205 µs / 3.37 ms / 84.9 ms	56.2 ns / 131 ns / 21.1 µs / 296 µs / 5.18 ms	1.4x / 1.7x / 9.7x / 11.4x / 16.4x
infimum	1.10.10	542 ns / 1.54 µs / 287 µs / 4.76 ms / 102 ms	10.8 ns / 11.3 ns / 4.57 µs / 17.3 µs / 125 µs	50.2x / 136x / 62.9x / 275x / 820x
infimum	1.12.6	326 ns / 907 ns / 218 µs / 3.41 ms / 87.6 ms	5.1 ns / 5.7 ns / 3.35 µs / 12.6 µs / 106 µs	63.9x / 159x / 65.2x / 270x / 825x
supremum	1.10.10	311 ns / 688 ns / 265 µs / 4.48 ms / 99.1 ms	10.8 ns / 11.3 ns / 4.59 µs / 17.2 µs / 114 µs	28.7x / 60.8x / 57.8x / 261x / 866x
supremum	1.12.6	219 ns / 556 ns / 216 µs / 3.30 ms / 83.1 ms	5.6 ns / 7.3 ns / 3.39 µs / 12.4 µs / 105 µs	39.0x / 76.2x / 63.7x / 267x / 790x
fuse	1.10.10	268 ns / 1.86 µs / 2.86 ms / 280 ms / 25.9 s	111 ns / 281 ns / 2.56 ms / 135 ms / 12.4 s	2.4x / 6.6x / 1.1x / 2.1x / 2.1x
fuse	1.12.6	173 ns / 1.69 µs / 2.45 ms / 286 ms / 13.2 s	16.2 ns / 65.7 ns / 1.16 ms / 135 ms / 12.6 s	10.7x / 25.7x / 2.1x / 2.1x / 1.05x
truncate_space	1.10.10	41.9 ns / 300 ns / 247 µs / 4.52 ms / 230 ms	145 ns / 217 ns / 32.9 µs / 525 µs / 9.68 ms	0.29x / 1.4x / 7.5x / 8.6x / 23.8x
truncate_space	1.12.6	40.1 ns / 130 ns / 199 µs / 3.20 ms / 91.3 ms	46.8 ns / 80.4 ns / 18.4 µs / 251 µs / 4.50 ms	0.86x / 1.6x / 10.8x / 12.7x / 20.3x

And now just the N=15625 case separately, also just after only as before doesn't finish:

operation	1.10.10	1.12.6
constructor (warm)	533 ms	526 ms
dim	554 ms	580 ms
⊕	1.08 ms	1.28 ms
⊖	1.14 s	290.7 s cold (see note below)
infimum	937 µs	1.32 ms
supremum	869 µs	1.01 ms
truncate_space	972 ms	1.17 s
fuse	not run (O(N^2))

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

operation	Julia	N=13 before→after	N=101 before→after	N=401 before→after
constructor	1.10.10	7.38 µs → 7.02 µs (1.05x)	56.0 µs → 45.2 µs (1.24x)	344 µs → 203 µs (1.69x)
constructor	1.12.6	5.29 µs → 4.27 µs (1.24x)	37.3 µs → 27.3 µs (1.37x)	199 µs → 112 µs (1.78x)
dim	1.10.10	279 ns → 21.6 ns (12.9x)	12.7 µs → 45.1 ns (282x)	170 µs → 87.5 ns (1941x)
dim	1.12.6	129 ns → 12.3 ns (10.5x)	6.65 µs → 74.1 ns (89.8x)	75.3 µs → 222 ns (340x)
⊕	1.10.10	5.67 µs → 872 ns (6.50x)	91.0 µs → 4.98 µs (18.3x)	967 µs → 21.8 µs (44.4x)
⊕	1.12.6	2.43 µs → 201 ns (12.1x)	46.8 µs → 1.01 µs (46.1x)	435 µs → 3.71 µs (117x)
⊖	1.10.10	2.83 µs → 1.52 µs (1.86x)	66.4 µs → 11.3 µs (5.88x)	774 µs → 54.7 µs (14.1x)
⊖	1.12.6	1.53 µs → 566 ns (2.71x)	29.9 µs → 4.77 µs (6.27x)	375 µs → 21.8 µs (17.2x)
infimum	1.10.10	6.25 µs → 866 ns (7.22x)	80.1 µs → 4.51 µs (17.7x)	816 µs → 22.3 µs (36.6x)
infimum	1.12.6	4.99 µs → 206 ns (24.2x)	38.4 µs → 963 ns (39.9x)	381 µs → 4.49 µs (84.8x)
supremum	1.10.10	4.86 µs → 856 ns (5.68x)	71.5 µs → 4.47 µs (16.0x)	767 µs → 22.2 µs (34.6x)
supremum	1.12.6	3.06 µs → 198 ns (15.5x)	32.4 µs → 997 ns (32.5x)	344 µs → 3.62 µs (94.9x)
fuse	1.10.10	20.9 µs → 13.2 µs (1.58x)	6.50 ms → 694 µs (9.36x)	335 ms → 11.1 ms (30.3x)
fuse	1.12.6	16.1 µs → 5.87 µs (2.74x)	3.65 ms → 333 µs (11.0x)	171 ms → 5.62 ms (30.5x)
truncate_space	1.10.10	2.44 µs → 2.31 µs (1.06x)	36.4 µs → 17.8 µs (2.05x)	441 µs → 108 µs (4.07x)
truncate_space	1.12.6	1.07 µs → 890 ns (1.21x)	18.3 µs → 7.68 µs (2.39x)	206 µs → 35.6 µs (5.80x)

(Negative) conclusions/remarks from the benchmarks:

• For very small N the constructor is still somewhat slower (0.45x-0.99x), but the alternative, per the first table, is to make large N impossible to compile in reasonable time.
• truncate_space is still slightly slower at N=2 (0.29x/0.86x) but wins from N=8 up.
• Why NTuple fuse scales the way it does with N I 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).
• SectorDict wins scale with sector count.

All in all, these are improvements, notably the SectorDict. So it makes you wonder if there's in fact some cutoff N above which you want to say SizeUnknown to the SectorValues' length 🤔

[kshyatt]

Can this be gotten rid of?

[codecov]

Codecov Report

❌ Patch coverage is 99.35065% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/spaces/gradedspace.jl	98.94%	1 Missing ⚠️

Files with missing lines	Coverage Δ
src/TensorKit.jl	17.24% <ø> (ø)
src/auxiliary/dicts.jl	70.55% <100.00%> (+6.91%)	⬆️
src/factorizations/factorizations.jl	84.61% <ø> (ø)
src/factorizations/truncation.jl	87.98% <100.00%> (+1.42%)	⬆️
src/spaces/gradedspace.jl	94.53% <98.94%> (+2.00%)	⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.