gpu - explore 3D arrays in cp.linalg.solve and cp.eigh

[lastephey]

Stephen found that both cp.linalg.solve and cp.eigh accept 3D arrays.

Examine performance of 2D vs 3D arrays at scales relevant to desi/gpu_specter to determine if switching to 3D arrays would be a performance win.

[lastephey]

You can see these cupy functions in action here:

https://github.com/sbailey/gpu_specter/blob/hackathon/gpu_extract.py#L715

We are using cupy to do matrix manipulation-- we know from profiling that cp.eigh (especially) and cp.linalg.solve are time consuming. It makes sense to see if there is something we can do to speed up how we use these functions. Specifically the point of this study is to learn whether these functions make sense to use in 2D or 3D.

Doing a scan over matrix size (100 - 10,000) in both 2D and 3D on the gpu will help us understand how best to use these functions.

[dmargala]

According to the cupy docs, cupy.linalg.eigh seems limited to 2D arrays.

https://docs-cupy.chainer.org/en/stable/reference/generated/cupy.linalg.eigh.html

It looks like cupy.linalg.solve can handle higher dimensional arrays though.

https://docs-cupy.chainer.org/en/stable/reference/generated/cupy.linalg.solve.html

[lastephey]

Thanks for looking at this. Starting with cupy.linalg.solve seems good for now.

@sbailey did we just misremember about cupy eigh?

[sbailey]

nb/Explore-3D-Array-Performance.ipynb added in PR #27 concluded:

tl;dr: there does not seem to be a performance boost using 3d arrays with cupy.linalg.solve

a possible followup is to check if CPU BLAS has a natively 3D eigh interface, and if so ask cuBLAS / cupy to do the same, but overall unfortunately this doesn't appear to be a fruitful avenue to pursue, especially given how major a bookkeeping restructuring it would require. Closing this ticket.