Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add LinearOperator constructor for CUDA #321

Merged
merged 3 commits into from
Apr 20, 2024

Conversation

nHackel
Copy link
Contributor

@nHackel nHackel commented Apr 8, 2024

This PR adds a method for the LinearOperator function which has a correct default storage type for CUDA's CuArray.

This should (partially, it does not fix the upstream error in NLPModels.jl) fix #307, in particular this error is resolved:

julia> using CUDA, LinearOperators

julia> N = 8
8

julia> A = CUDA.rand(N, N);

julia> x = CUDA.rand(N);

julia> (A + opEye(Float32, N, S = typeof(x)));
ERROR: LinearOperatorException("storage types cannot be promoted to a concrete type")
Stacktrace:
 [1] +(op1::LinearOperator{Float32, Int64, LinearOperators.var"#5#8"{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}, LinearOperators.var"#6#9"{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}, LinearOperators.var"#7#10"{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}, Vector{Float32}}, op2::LinearOperator{Float32, Int64, LinearOperators.var"#135#136"{Int64}, LinearOperators.var"#135#136"{Int64}, LinearOperators.var"#135#136"{Int64}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}})
   @ LinearOperators ~/.julia/packages/LinearOperators/lK3j5/src/operations.jl:190
 [2] +(M::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, op::LinearOperator{Float32, Int64, LinearOperators.var"#135#136"{Int64}, LinearOperators.var"#135#136"{Int64}, LinearOperators.var"#135#136"{Int64}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}})
   @ LinearOperators ~/.julia/packages/LinearOperators/lK3j5/src/operations.jl:196
 [3] top-level scope
   @ REPL[14]:1
 [4] top-level scope
   @ ~/.julia/packages/CUDA/htRwP/src/initialization.jl:206

The method is defined in a package extension for CUDA. I've also tried to make it available for older Julia Version with Requires, similar to how it was done for the ChainRulesExt.

@nHackel
Copy link
Contributor Author

nHackel commented Apr 11, 2024

Another option could be to try to write an extension for AbstractGPUArrays, however this will require some workarounds to strip the Array type of its parameters.

One way might be with the workaround shown here or maybe with a typeof(similar(...)). If such a workaround but generically for different GPUs sounds fine, I can close this PR and try it in a new one

Copy link
Contributor

github-actions bot commented Apr 11, 2024

Package name latest stable
CaNNOLeS.jl
DCISolver.jl
FletcherPenaltySolver.jl
JSOSolvers.jl
Krylov.jl
NLPModels.jl
NLPModelsModifiers.jl
PROPACK.jl
Percival.jl
QuadraticModels.jl
SolverTools.jl

@geoffroyleconte
Copy link
Member

Hi @nHackel , thank you for this PR!

I was wondering if something like

A = CUDA.rand(10,10)
tA = typeof(A)
S = tA.name.wrapper{eltype(A), 1, tA.parameters[3:end]...}
test_vec = S(undef, 10)

would work for any matrix type (including other GPUArrays, but I can only test with CUDA) to avoid having CUDA in the extensions and solve your issue?
However it is not very pretty and adds a bit of allocations so I'm not sure it is the best solution:

julia> @allocated tA.parameters[3:end]
128

I think that

S = tA.name.wrapper{eltype(A), 1, tA.parameters[3]}

might work at least for CUDA but probably not for CPU arrays.

Let me know if you don't see an obvious answer, we can probably merge this PR and open an issue to improve genericity in the future.

@nHackel
Copy link
Contributor Author

nHackel commented Apr 11, 2024

Hello @geoffroyleconte, thanks for looking into the PR!

I thought about the workarounds a bit more and I fear one issue we would run into if we try to make it very generic (either just for GPU or even all arrays) are all the various matrix types where the correct storage/result type is not directly a 1-dimensional version of it.

For example a sparse matrix would result in a dense vector in the end. While one could solve that with a new case for AbstractSparseMatrix, the same thing could happen for any number of array types we don't know about.

Copy link
Member

@geoffroyleconte geoffroyleconte left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes that's a good point, I mainly used dense vectors even with sparse arrays by there could applications where sparse vectors are needed.

@dpo does this looks good to you?

@nHackel
Copy link
Contributor Author

nHackel commented Apr 11, 2024

I thought about adding the method storage_type(M::CuMatrix{T}) where {T} = CuVector{T} to the PR just now. Then I had the idea that we could then do the following and remove the constructor I defined here so far:

function LinearOperator(
  M::AbstractMatrix{T};
  symmetric = false,
  hermitian = false,
  S = storage_type(M) # previously Vector{T},

for "all" arrays here.

That change would be a bit more invasive, since it affects the fallback for all matrix types. Let me know what you think, I can either just add the storage_type for CuArray or the invasive change

@geoffroyleconte
Copy link
Member

geoffroyleconte commented Apr 11, 2024

Yes that's a great idea! Adding new Array types would then require few lines of code.

Copy link
Member

@dpo dpo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you!

@dpo dpo merged commit be5e5bc into JuliaSmoothOptimizers:main Apr 20, 2024
34 of 35 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

ERROR: LinearOperatorException("storage types cannot be promoted to a concrete type")
3 participants