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?
There was a problem hiding this comment.
I think the GPU arrays don't support stack 😭
|
|
||
| # 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
There was a problem hiding this comment.
It's a limitation of the drivers, you can work around it by padding the 3D Array with zeros as needed
|
I mostly agree with Lukas here, I think I would prefer
|
This is already implemented in a separate branch I have over at TensorKit, I can move it over here of course! |
e20efbd to
cb6a8aa
Compare
|
Latest commit gets rid of the
|
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.