Redo add_transform! for GPU backed Tensors 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

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
ext/TensorKitGPUArraysExt.jl	95.41% <100.00%> (+6.32%)	⬆️

🚀 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

[lkdvos]

Is this Base.searchsortedlast or Base.searchsortedfirst?

[lkdvos]

edit: I see now that this is used inside the kernel, which I assume is why that has to be like this

[kshyatt]

Should I add a comment about this to make it clearer?

[lkdvos]

This is slightly confusing to me, isn't this still computed in each kernel call below (line 371)? Keep in mind I might be completely misreading kernels here.

[kshyatt]

It's computed just once per kernel thread, though, which becomes more important for the generic case

[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.