implementation of ldiv?

[ctessum]

Hello, thanks for all your work on this!

I'm interested in solving a linear system of equations using a LU decomposotion and then a left division.

This is how it works on the CPU:

x = rand(Float32, 3, 3)
f = lu(x)
y = rand(Float32, 3)
# These are different implementations of left division
ldiv!(similar(y), f, y)
ldiv!(f, y)
f \ y

However, if I try it in Metal.jl, it doesn't work:

x = MtlArray(rand(Float32, 3, 3))
f = lu(x)
y = MtlArray(rand(Float32, 3))
ldiv!(f, y) # Fails
ldiv!(similar(y), f, y) # Fails
f \ y # Fails

All three options give the same error:

ERROR: MethodError: no method matching getrs!(::Char, ::MtlMatrix{Float32, Metal.PrivateStorage}, ::MtlVector{UInt32, Metal.PrivateStorage}, ::Vector{Float32})
The function `getrs!` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  getrs!(::AbstractChar, ::AbstractMatrix{Float32}, ::AbstractVector{Int64}, ::AbstractVecOrMat{Float32})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/lapack.jl:1043
  getrs!(::AbstractChar, ::AbstractMatrix{ComplexF32}, ::AbstractVector{Int64}, ::AbstractVecOrMat{ComplexF32})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/lapack.jl:1043
  getrs!(::AbstractChar, ::AbstractMatrix{ComplexF64}, ::AbstractVector{Int64}, ::AbstractVecOrMat{ComplexF64})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/lapack.jl:1043
  ...

Stacktrace:
 [1] ldiv!(A::LU{Float32, MtlMatrix{…}, MtlVector{…}}, B::Vector{Float32})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/lu.jl:498
 [2] ldiv(F::LU{Float32, MtlMatrix{…}, MtlVector{…}}, B::MtlVector{Float32, Metal.PrivateStorage})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/LinearAlgebra.jl:656
 [3] \(F::LU{Float32, MtlMatrix{…}, MtlVector{…}}, B::MtlVector{Float32, Metal.PrivateStorage})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/LinearAlgebra.jl:631
 [4] top-level scope

Since this error seems to be caused by the ipiv values in the Metal LU being UInt values instead of Int values, we can try it after converting the ipiv values to Ints:

f2 = LU(f.factors, Int.(f.ipiv), f.info)
ldiv!(f2, y) # Fails
ldiv!(similar(y), f2, y) # Fails
f2 \ y # Fails

This still fails, but it gives a different error:

ERROR: ArgumentError: cannot take the CPU address of a MtlMatrix{Float32, Metal.PrivateStorage}
Stacktrace:
 [1] unsafe_convert(::Type{Ptr{Float32}}, x::MtlMatrix{Float32, Metal.PrivateStorage})
   @ Metal ~/.julia/packages/Metal/uq8sb/src/array.jl:266
 [2] getrs!(trans::Char, A::MtlMatrix{…}, ipiv::MtlVector{…}, B::Vector{…})
   @ LinearAlgebra.LAPACK ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/lapack.jl:1056
 [3] ldiv!(A::LU{Float32, MtlMatrix{…}, MtlVector{…}}, B::Vector{Float32})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/lu.jl:498
 [4] ldiv(F::LU{Float32, MtlMatrix{…}, MtlVector{…}}, B::MtlVector{Float32, Metal.PrivateStorage})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/LinearAlgebra.jl:656
 [5] \(F::LU{Float32, MtlMatrix{…}, MtlVector{…}}, B::MtlVector{Float32, Metal.PrivateStorage})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/LinearAlgebra.jl:631
 [6] top-level scope

This error seems to suggest that in order for it to work, there needs to be a method of ldiv that is specialized for Metal.jl. Is this something that would be straightforward to do?

[christiangnrd]

Hi! If you’re interested you could look into finishing up #145.

[ctessum]

I'll take a look!