Overload cholesky method in order to prevent scalar indexing for Diagonal

[simsurace]

Makes this work without scalar indexing:

using CUDA, LinearAlgebra
CUDA.allowscalar(false)

n = 1024
d = CUDA.rand(n)
D = Diagonal(d)
cholesky(D)

Performance can probably be improved.

[simsurace]

Unfortunately I have no experience with oneAPI.jl. Will look into the failing tests later, if people think that this is a useful PR.

[maleadt]

Thanks for the PR!

Unfortunately I have no experience with oneAPI.jl. Will look into the failing tests later, if people think that this is a useful PR.

The failure is a generic GPU array failure where Base (specifically BLAS) functionality is called where a specialization is needed:

  ArgumentError: cannot take the host address of a oneArray{Float32, 2}
  Stacktrace:
    [1] unsafe_convert(#unused#::Type{Ptr{Float32}}, x::oneArray{Float32, 2})
      @ oneAPI ~/.cache/julia-buildkite-plugin/depots/c9f52312-b528-44e4-9501-6d408762012b/packages/oneAPI/KbEvp/src/array.jl:176
    [2] potrf!(uplo::Char, A::oneArray{Float32, 2})
      @ LinearAlgebra.LAPACK /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/lapack.jl:3088
    [3] _chol!
      @ /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:172 [inlined]
    [4] cholesky!(A::Hermitian{Float32, oneArray{Float32, 2}}, ::Val{false}; check::Bool)
      @ LinearAlgebra /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:252
    [5] cholesky!(A::oneArray{Float32, 2}, ::Val{false}; check::Bool)
      @ LinearAlgebra /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:285
    [6] #cholesky#134
      @ /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:378 [inlined]
    [7] cholesky (repeats 2 times)
      @ /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/cholesky.jl:378 [inlined]
    [8] macro expansion
      @ /var/lib/buildkite-agent/builds/gpuci3/julialang/gpuarrays-dot-jl/test/testsuite/linalg.jl:107 [inlined]

No GPUArrays in that stack trace, so maybe this PR is missing an overload (or relying on CUDA.jl-specific functionality, which isn't allowed for generic GPUArrays.jl implementations)?

[simsurace]

Ah, looks like cholesky is generally broken for oneArrays.
EDIT: will add tests for cholesky without the Diagonal wrapper.

[simsurace]

I think the reason for the failure is that the generic LinearAlgebra.LAPACK.potrf! method that calls into BLAS does not work for oneArrays, a specific overload analogous to the one in CUDA/lib/cusolver/dense.jl is missing. Shall I open an issue in oneAPI.jl? I presume the overload should be in there, not here.

[maleadt]

Those kind of overloads (replacing LinearAlgebra.BLAS and LinearAlgebra.LAPACK methods with GPU-specific ones) is something I'm trying to get rid of, actually, because CUBLAS really isn't a drop-in replacement for BLAS so it's a bad place for putting overrides.

That said, unless GPUArrays.jl can provide a generic potrf this is a bad test, so maybe it's better to move this implementation to CUDA.jl where the required functionality is available.

[simsurace]

So just to clarify: the generic cholesky overload introduced in this PR works on diagonal oneArrays and CuArrays. The problematic method was the one being called on dense oneArrays in the previous tests. I fixed the tests to not rely on these methods anymore. With this modification, everything passes.

[simsurace]

This modification avoids calling a dense cholesky method, which is not available for oneArrays