[kshyatt]

For some of the PEPSKit GPU runs, I noticed KrylovKit was doing a lot of copying scalars back and forth and launching a lot of small kernels. The goal of this PR is to use a "all in one" fused approach as much as we can, which should cut down a lot on the number of copies and kernel launches we have to do.

In particular it adds an optional "dense" fused multiplication to two of the Gram-Schmidt procedures to copy all the orthonormal basis vectors into a matrix then operate on that in a broadcasted fashion. This should let us (e.g.) get all the Lanczos vectors in one call.

I added some JLArrays based tests too, so we don't need to turn on Buildkite for this. I can try to get some performance numbers if people are interested but it would be helpful to get some feedback on the overall approach first. I deliberately kept this relatively minimal -- there are probably other algos that could benefit from this approach but those could be looked at in a new PR.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.97%. Comparing base (d51fa25) to head (4e66a72).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #180      +/-   ##
==========================================
+ Coverage   88.62%   88.97%   +0.35%     
==========================================
  Files          36       36              
  Lines        3946     4027      +81     
==========================================
+ Hits         3497     3583      +86     
+ Misses        449      444       -5     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

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

[Jutho]

Haven't looked at the code yet. Is this an attempt at solving the inefficiency of doing Gram-Schmidt for GPU vectors? I thought this would require us to first introduce new inner methods that would not return the result as a scalar, which requires copying it to the CPU, and requires us to tap into cuBLAS functionality that is not even wrapped yet in CUDA.jl?

[kshyatt]

Is this an attempt at solving the inefficiency of doing Gram-Schmidt for GPU vectors?

Yes :).

I didn't follow #154 very closely and it seems it has stalled, so I guess some of this is duplicated from me starting over fresh. Looking over it, I think this covers a little bit more (MGS and Lanczos aren't hit by that PR I believe, please correct me if that's not the case).

I think the inner() issue @lkdvos called out there might apply here as well, e.g. if inner is not in fact inner(a, b) = a' * b this might be wrong. But maybe the DenseArray restriction prevents this?

Here also we do have to copy back to the CPU at a few points, I can work on trying to remove that restriction and add wrappers to cuBLAS if needed. The goal with this one was kind of testing the waters, and if we want to pursue it further, I'll go into those issues a bit more.

[kshyatt]

BTW at least for inner, CUBLAS doesn't provide a batched version. One could of course achieve something like that by using gemm.