Batched SVD support for ROCSOLVER and CUSOLVER - #275
Conversation
|
TODOs here:
Since I'll be out for 3 weeks everyone should feel free to just push to this. |
e174ddb to
6394bc8
Compare
lkdvos
left a comment
lkdvos
left a comment
There was a problem hiding this comment.
Overall looks great, I think the main thing I am wondering about is about the interface decision around wether we implement this as a set of BatchedAlg versions, or rather as a set of batched_f(...) functions. I definitely like using dispatch for switching between the strided and non-strided inputs, but I am wondering if there might be a benefit to really having a batched_svd_compact etc function.
This is also partially since that allows us to have a CPU version for this as well, so we can just offload all of this from TensorKit to here, (and possibly play with multithreading?).
| ) | ||
| summary_str = testargs_summary(T, sz) | ||
| return @testset "svd_compact! $summary_str batch_size $batch_size" begin | ||
| As = [instantiate_matrix(T, sz) for bi in 1:batch_size] |
There was a problem hiding this comment.
should we try and vary the sizes within the batch?
| function _collect_batch(As::AbstractVector{<:AbstractMatrix}) | ||
| B = similar(first(As), (size(first(As))..., length(As))) | ||
| for (i, A) in enumerate(As) | ||
| copyto!(view(B, :, :, i), A) | ||
| end | ||
| return B | ||
| end | ||
| device_batch(As::AbstractVector{<:Array}) = _collect_batch(As) | ||
| device_batch(As::AbstractVector{<:CuArray}) = _collect_batch(As) | ||
| device_batch(As::AbstractVector{<:ROCArray}) = _collect_batch(As) |
There was a problem hiding this comment.
Is this the same as stack(As; dims = 3), and does that work on both devices and host arrays?
|
|
||
| # batched varieties | ||
| function check_input(::typeof(svd_full!), A::AbstractVector{<:AbstractMatrix}, USVᴴ, ::AbstractAlgorithm) | ||
| @assert all(==(size(first(A))), size.(A)) |
There was a problem hiding this comment.
Why do we require the sizes to be the same? I'm assuming this is some restriction of some of the drivers, but that does kind of seems to contradict the requirements we need in TensorKit
| function gesvdj_batched! end | ||
| function gesvdx_batched! end | ||
|
|
||
|
|
There was a problem hiding this comment.
Isn't this already defined on line 240-241?
|
I mostly agree with Lukas here, I think I would prefer
|
3 participants
Basically what it says on the tin. For CTMRG and other algorithms, we're getting absolutely slaughtered on GPU performance for TensorMaps with sectors because we have to spin up huge numbers of very small SVDs. I'm wrapping the batched SVDs each library provides to try to address this. Extremely open to comments but I wanted to get this rolling so I can unblock others.