Redo add_transform! for GPU backed tensors with general symmetries to be much, much faster

[kshyatt]

Should close #508

The short version of what I've done here is rework add_transform! into two kernels with associated helper structures, one for the Abelian case (which is quite simple) and one for the more generic case. Instead of parallelizing over the blocks, which are of pretty different sizes, I parallelize over the output data in data_dst. This means a lot more annoying bookkeeping but also a much more uniform workload across the GPU and better usage of the large number of available threads. There are still probably some performance optimizations to be done but I thought this was a) already hard enough to understand and b) pretty compelling!

Here are the results for the sample script I linked in the issue above, on an AMD MI210:

GPU permute time (ms)

case	trees	elems	bytes	main	ksh/add_transform	speedup
SU2 iPEPO tensor	42	329	2.6 KiB	1.915	0.033	58×
trivial symmetry	1	9604	75.0 KiB	0.022	0.020	1.1×
SU2 jmax=1	42	672	5.3 KiB	1.964	0.033	60×
SU2 jmax=2	323	5168	40.4 KiB	13.818	0.034	406×
SU2 jmax=3	1364	21824	170.5 KiB	53.537	0.046	1164×
U1	70	2226	17.4 KiB	1.064	0.020	53×

Cost per fusion tree (GPU, µs/tree)

case	trees	main	ksh/add_transform
SU2 iPEPO tensor	42	45.6	0.79
SU2 jmax=1	42	46.8	0.79
SU2 jmax=2	323	42.8	0.105
SU2 jmax=3	1364	39.2	0.034
U1	70	15.2	0.29

So this should hopefully finally let people doing complicated stuff (e.g. anything involving SU(2)) really benefit from the GPU. All the tests passed for me locally.

I still have some lingering questions about where stuff should live. I added some new caches of GPU objects as well to avoid sending things back to the GPU that don't need to go, but maybe those should live in the main package?

[kshyatt]

I'm also running the GPU branch of PEPSKit.jl against this to see how much/if it helps

[kshyatt]

Well, the short story for now is it doesn't help much because nearly all the time in our PEPSKit runs is spent on the CPU (GPU utilization is quite low). I'll look into that separately but I still think this is a nice bit of progress :)

[codecov]

Codecov Report

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

Files with missing lines	Patch %	Lines
ext/TensorKitGPUArraysExt.jl	98.33%	1 Missing ⚠️

Files with missing lines	Coverage Δ
ext/TensorKitGPUArraysExt.jl	95.48% <98.33%> (+1.79%)	⬆️

... and 5 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[kshyatt]

OK I've got one more thing to push here to make it play nicely with PEPSKit (cache-related) but most of this is ready to go I think

[kshyatt]

OK!!! I modified stuff a bit for the DEVICE_TRANSFORMER_CACHE because we need the GPUArrays caching allocator to really see the benefits of this on PEPSKit (otherwise we choke to death on allocations). The @uncached here is to protect the cached GPU-side transformer infos from getting reaped while we still expect them to be live.

[lkdvos]

have only looked at the abelian part for now, looks really cool! I didn't really realize that indeed we can just spawn a thread per element. I am definitely wondering how this affects the case where the tensors are a bit larger and the number of blocks not so severe, for example what would effectively be the case of more MPS-like contractions, for which it might be reasonable to try and generalize/run the benchmarks that are in this repository on GPU as well to get a sense about overall performance implications?

Given that every element is read/written exactly once, it could also be cool to just express these numbers in terms of actual read/write bandwidth percentages, for which 🤖 might be able to help? Not that I'm expecting or requiring any huge results, but I'd be very interested to know how much is left on the table

[kshyatt]

I am definitely wondering how this affects the case where the tensors are a bit larger and the number of blocks not so severe, for example what would effectively be the case of more MPS-like contractions, for which it might be reasonable to try and generalize/run the benchmarks that are in this repository on GPU as well to get a sense about overall performance implications?

Yeah, it's a really good question. I focused on the cases I looked at in the linked issue, but if we have some others I can run I will certainly do that to get a clearer picture.

[kshyatt]

I didn't really realize that indeed we can just spawn a thread per element.

I did do a bit of testing with these kernels on CPU arrays and there's no real benefit there, this approach is really GPU specific I guess.

[kshyatt]

OK, we need a patch version bump at TO. Then, are there any objections to merging this?

[lkdvos]

Overall definitely a cool set of improvements! I'm still trying to fully wrap my head around this, left some small comments but also maybe a request: could we split this into 2 PRs and treat the Abelian case first to make review simpler?

The main motivation here is that I think mostly the caching could use some work. It seems like we might want to include a storagetype parameter in the hash as well as in the transformer data, which should unify the implementation of that across the board?

[kshyatt]

I'll split out the Abelian case into a new PR but keep this open to let others keep using it w PEPSKit

[lkdvos]

@Jutho do we have more strong opinions on this? I think this is about as good as we can currently get it, so good to go for me and we can always modify later.