# CuArrays with KrylovKit

**URL:** <https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116>\
**Category:** Numerics\
**Created:** [November 12, 2023, 4:02pm UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116 "2023-11-12T16:02:49Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [November 12, 2023, 4:02pm UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/1 "2023-11-12T16:02:49Z")

</div>

What is the recommended way of diagonalizing linear operators on CuArrays using KrylovKit? I am targeting the first 100 eigenvalues of a large matrix - Operating directly on CuArrays leads to a very large allocation on the GPU. I do this the following way:

```julia
    function gpu_wrapper!(y_d::CuArray{Tcalc}, x_d::CuArray{Tcalc})
        CUDA.@sync begin
            CUDA.@cuda threads=num_threads blocks=num_blocks two_body_operator_action_gpu!(y_d, x_d, creation_ops_d, creation_op_phase_helpers_d, annhilation_ops_d, annhilation_phase_helpers_d, coulomb_matrix_elements_d, lll_basis_d)
        end
        return
    end

```

I wrap this function using LinearMaps and pass it to Krylovkit’s eigensolve with an initial guess provided as a CuArray. (I set the krylodim to 110) - The memory usage on the GPU went beyond about 36GB before I shut the process down. However, if I do something like below:

```julia
    function gpu_wrapper!(y::Vector{Tcalc}, x::Vector{Tcalc})
        x_d = CuArray{Tcalc}(undef, size(lll_basis,1))
        y_d = CuArray{Tcalc}(undef, size(lll_basis,1))
        copy!(x_d, x)
        CUDA.@sync begin
            CUDA.@cuda threads=num_threads blocks=num_blocks two_body_operator_action_gpu!(y_d, x_d, creation_ops_d, creation_op_phase_helpers_d, annhilation_ops_d, annhilation_phase_helpers_d, coulomb_matrix_elements_d, lll_basis_d)
        end
        copy!(y, y_d)
        return
    end

```

The memory usage on the GPU is fairly contained ~ its about 2.5-3.5 GB.

The only problem is, as you may expect, this implementation is slower than the earlier one - for the exact operator I am considering, the time spent copying the array to the device is about 1/4 the total time.

---

<div class="post-metadata">

**Author:** ![juthohaegeman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juthohaegeman/32/8620_2.png) [@juthohaegeman](https://discourse.julialang.org/u/juthohaegeman)\
**Post date:** [November 12, 2023, 5:52pm UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/2 "2023-11-12T17:52:06Z")

</div>

What is the size/length of the vectors involved? And what is `Tcalc`: single or double precision, real or complex?

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [November 12, 2023, 6:03pm UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/3 "2023-11-12T18:03:14Z")

</div>

The dimension of the vectors is 7142568 ~ 10^7. Tcalc is Float32.

---

<div class="post-metadata">

**Author:** ![juthohaegeman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juthohaegeman/32/8620_2.png) [@juthohaegeman](https://discourse.julialang.org/u/juthohaegeman)\
**Post date:** [November 12, 2023, 8:19pm UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/4 "2023-11-12T20:19:37Z")

</div>

I think that this memory usage can then be expected: the memory size of 110 vectors of length 7142568 with `Float32` numbers is 3142729920 bytes, or thus about 3 GB.

Note that KrylovKit stores the full Krylov subspace, even with the Lanczos algorithm (actually, Lanczos is merely the name used for the orthogonalisation process in building the Krylov subspace, the actual eigenvalue algorithm being used is rather called Krylov-Schur). There is some basic support for building a Krylov subspace where the older vectors are discarded, but this is not really being used in the `eigsolve` routine.

---

<div class="post-metadata">

**Author:** ![juthohaegeman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juthohaegeman/32/8620_2.png) [@juthohaegeman](https://discourse.julialang.org/u/juthohaegeman)\
**Post date:** [November 12, 2023, 8:21pm UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/5 "2023-11-12T20:21:47Z")

</div>

That was too quick, I now see that you say 36 GB instead of 3 GB. Is that the total amount of allocations being performed, or really the memory being used at a given point in the proces?

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [November 12, 2023, 10:34pm UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/6 "2023-11-12T22:34:16Z")

</div>

That’s the total (I read it off from Nvidia-smi).

> I think that this memory usage can then be expected: the memory size of 110 vectors of length 7142568 with Float32 numbers is 3142729920 bytes, or thus about 3 GB.  
> This was my estimate as well.

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [November 14, 2023, 2:18am UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/7 "2023-11-14T02:18:51Z")

</div>

I checked the memory usage of KrylovKit’s eigsolve alone for a vector of similar dimension(7040196), with 30 Krylov subspace vectors (the default) and solving for one eigenvalue using CUDA.@time - It was about 8 GB. With 5 Krylov subspace vectors, it was about 20 GB.

Also, the memory usage by Julia remains at 20 GB even after eigsolve has been called. I run the program under a let block as:

```julia
let
   
    num_particles = 14
    Lztotal = 0//1
    Q = (num_particles//2-2)//2 + (num_particles-1)
    num_eigvals = 1
    
    ham = generate_hamiltonian(num_particles, Lztotal, Q)
    x0 = randn(Tcalc, size(ham, 1))
    x0 ./= sqrt(x0'*x0)
    x0_d = CuArray(x0)
# KrylovKit.eigsolve(ham, x0_d, num_eigvals, :SR; tol=1e-8, issymmetric=true, krylovdim=max(30, num_eigvals + 10), maxiter=1000)
    CUDA.@profile CUDA.@time KrylovKit.eigsolve(ham, x0_d, num_eigvals, :SR; tol=1e-8, issymmetric=true, krylovdim=5, maxiter=1000)
end

```

This is the output from CUDA.@time for 5 krylov subspace vectors:  
233.890395 seconds (46.51 k CPU allocations: 2.661 MiB) (743 GPU allocations: 19.486 GiB, 0.00% memmgmt time)

Out[51]:

Maybe there’s something in my code? I am pasting it here just for reference:

```julia
function two_body_operator_action_gpu!(y, x, creation_ops, creation_op_phase_helpers, annhilation_ops, annhilation_phase_helpers, two_body_op_matrix_elements, lll_basis)
    
    index = (blockIdx().x - 1) *blockDim().x + threadIdx().x
    stride = gridDim().x * blockDim().x
    
    i = index
    while (i <= length(y))
# CUDA.@cuprintln(i)
        @inbounds y[i] = zero(Tcalc)
        @inbounds fock_state = lll_basis[i]
        op_iter = 1
        while op_iter <= length(annhilation_ops)
            @inbounds annhilated_state = ~annhilation_ops[op_iter] & fock_state
            @inbounds if isequal(annhilation_ops[op_iter], annhilation_ops[op_iter] & fock_state) && isequal(creation_ops[op_iter], creation_ops[op_iter] & (~annhilated_state))
                @inbounds y[i] += x[searchsortedfirst(lll_basis, annhilated_state | creation_ops[op_iter])] * two_body_op_matrix_elements[op_iter] * (-1)^(count_ones(annhilation_phase_helpers[op_iter] & fock_state) + count_ones(creation_op_phase_helpers[op_iter] & annhilated_state))
            end
            
            op_iter += 1
        end
        i += stride
    end
    return    
end

```

lll\_basis is of the same dimension (and type UInt32) as the vectors passed to eigsolve. The other vectors are orders of magnitude smaller.

---

<div class="post-metadata">

**Author:** ![juthohaegeman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/juthohaegeman/32/8620_2.png) [@juthohaegeman](https://discourse.julialang.org/u/juthohaegeman)\
**Post date:** [November 14, 2023, 9:36am UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/8 "2023-11-14T09:36:26Z")

</div>

I am still confused by what your numbers mean. For example, the total allocations reported by `@time` (definitely `Base.@time`, I assume also for `CUDA.@time`) is the total memory that was requested over the whole duration of that block.

One particular aspect of KrylovKit is that (currently) it allocates new vectors for the Krylov subspace over the different iterations, rather then recycling the old ones. This is also clear from the fact that function that implements the lineair operator is to be given as a function with a single argument (the input) and that returns the new output. So actually this is the function that allocates the new vector; the further orthonormalisation of this vector with respect to the current basis of the Krylov subspace happens in-place and should not cause new allocations.

But this means that the total number of allocations scales roughly as “memory of a single vector \* krylovdim \* numiter” (not exactly, as not the whole Krylov subspace is rebuilt for each iteration).

However, the vectors from a previous iteration that are discarded should be cleaned up by the GC, and the memory required at a given point in time should not exceed “memory of a single vector \* krylovdim”. However, it depends on how aggressive the GC cleans up the old vectors how much the instantaneous memory usage of Julia will be.

I am not sure if this answers your question?

PS: Having the possibility to recylce the vectors (which thus requires the linear operator to be specified as a two-argument function where the first argument is the output vector to which the result needs to be written) is on the to do list. I hope to get to it somewhere in the near future.

---

<div class="post-metadata">

**Author:** ![Gattu\_Mytraya](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gattu_mytraya/32/45429_2.png) [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)\
**Post date:** [November 14, 2023, 10:21am UTC](https://discourse.julialang.org/t/cuarrays-with-krylovkit/106116/9 "2023-11-14T10:21:51Z")

</div>

That does answer my question.