Merge remote-tracking branch 'upstream/main' into fix/1236-matrixfree-woperator

Author: AJ0070

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearSolve"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
-version = "5.13.0"
+version = "5.13.1"
 authors = ["SciML"]
 
 [deps]
@@ -71,6 +71,7 @@ PartitionedSolvers = "11b65f7f-80ac-401b-9ef2-3db765482d62"
 ParU_jll = "9e0b026c-e8ce-559c-a2c4-6a3d5c955bc9"
 PETSc = "ace2c81b-2b5f-4b1e-a30d-d662738edfe0"
 PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
 SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac"
@@ -115,6 +116,7 @@ LinearSolveParUExt = "ParU_jll"
 LinearSolvePETScExt = ["PETSc", "SparseMatricesCSR"]
 LinearSolvePETScMPIExt = ["PETSc", "PartitionedArrays", "SparseMatricesCSR"]
 LinearSolvePureUMFPACKExt = "PureUMFPACK"
+LinearSolveReactantExt = "Reactant"
 LinearSolveRecursiveFactorizationExt = ["RecursiveFactorization", "TriangularSolve"]
 LinearSolveSparspakExt = "Sparspak"
 LinearSolveSpecializingFactorizationsExt = "SpecializingFactorizations"
@@ -186,6 +188,7 @@ PrecompileTools = "1.2"
 Preferences = "1.4"
 PureKLU = "1.4"
 PureUMFPACK = "1"
+Reactant = "0.2.230"
 Random = "1.10"
 RecursiveArrayTools = "4"
 RecursiveFactorization = "0.2.27"

ext/LinearSolveReactantExt.jl

@@ -0,0 +1,46 @@
+module LinearSolveReactantExt
+
+using LinearSolve: LinearSolve, OperatorAssumptions
+using Reactant: Reactant
+using SciMLBase: SciMLBase, LinearProblem, LinearSolution, ReturnCode
+
+Reactant._parent_type(::Type{T}) where {T <: LinearSolution} = T
+
+function reactant_solve(prob::LinearProblem, args...; kwargs...)
+    output_size = prob.u0 === nothing ? (size(prob.A, 2), Base.tail(size(prob.b))...) :
+        size(prob.u0)
+    output_eltype = prob.u0 === nothing ? eltype(prob.b) : eltype(prob.u0)
+    inputs = prob.u0 === nothing ? (prob.A, prob.b) : (prob.A, prob.b, prob.u0)
+    callback = let prob = prob, args = args, kwargs = kwargs
+        function (output, A, b, u0...)
+            runtime_prob = isempty(u0) ? SciMLBase.remake(prob; A, b) :
+                SciMLBase.remake(prob; A, b, u0 = only(u0))
+            sol = SciMLBase.solve(runtime_prob, args...; kwargs...)
+            copyto!(output, sol.u)
+            return nothing
+        end
+    end
+    u = Reactant.Ops.julia_callback(
+        callback, ((output_eltype, output_size),), inputs...
+    )
+    alg = if isempty(args) || first(args) === nothing
+        assump = get(
+            kwargs, :assump,
+            OperatorAssumptions(size(prob.A, 1) == size(prob.A, 2))
+        )
+        LinearSolve.defaultalg(prob.A, prob.b, assump)
+    else
+        first(args)
+    end
+    return SciMLBase.build_linear_solution(
+        alg, u, nothing, nothing; retcode = ReturnCode.Success
+    )
+end
+
+Reactant.@reactant_overlay function SciMLBase.solve(
+        prob::LinearProblem, args...; kwargs...
+    )
+    return reactant_solve(prob, args...; kwargs...)
+end
+
+end

lib/LinearSolveAutotune/Project.toml

@@ -1,7 +1,7 @@
 name = "LinearSolveAutotune"
 uuid = "67398393-80e8-4254-b7e4-1b9a36a3c5b6"
 authors = ["SciML"]
-version = "1.14.0"
+version = "1.14.1"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

lib/LinearSolvePyAMG/Project.toml

@@ -1,7 +1,7 @@
 name = "LinearSolvePyAMG"
 uuid = "7a56c47d-7ab1-4e99-b0e3-2952e463d64a"
 authors = ["SciML"]
-version = "1.3.2"
+version = "1.3.3"
 
 [deps]
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"

src/factorization.jl

@@ -1054,7 +1054,13 @@ function do_factorization(alg::QRFactorization, A, b, u)
         # `alg.pivot` so the return type is determined by the static
         # `QRFactorization{P}` parameter (otherwise this branch returns
         # `Union{QRCompactWY, QRPivoted}` depending on `alg.inplace`).
-        if A isa GPUArraysCore.AnyGPUArray || is_cusparse(A) || issparsematrixcsc(A)
+        if A isa GPUArraysCore.AnyGPUArray && !is_cusparse(A) && is_underdetermined(A)
+            # A GPU `qr` factors a wide matrix happily, but solving with the result
+            # builds `UpperTriangular(R)` on an `R` that is not square and throws.
+            # Going through `Aᵀ` turns it back into a triangular solve and gives the
+            # minimum-norm solution, which is what dense `\` returns on the CPU.
+            fact = MinNormQR(qr(copy(transpose(A))))
+        elseif A isa GPUArraysCore.AnyGPUArray || is_cusparse(A) || issparsematrixcsc(A)
             fact = qr(A)
         elseif alg.inplace
             if A isa Symmetric
@@ -1076,7 +1082,13 @@ function init_cacheval(
         maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
         assumptions::OperatorAssumptions
     )
-    return ArrayInterface.qr_instance(convert(AbstractMatrix, A), alg.pivot)
+    A_ = convert(AbstractMatrix, A)
+    # Matches the wide GPU branch of `do_factorization`: the slot has to be typed for
+    # what will be stored in it, not for a plain `QR`.
+    if A_ isa GPUArraysCore.AnyGPUArray && !is_cusparse(A_) && is_underdetermined(A_)
+        return MinNormQR(qr(copy(transpose(A_))))
+    end
+    return ArrayInterface.qr_instance(A_, alg.pivot)
 end
 
 function init_cacheval(
@@ -1192,6 +1204,12 @@ function init_cacheval(
         alg::CholeskyFactorization, A::GPUArraysCore.AnyGPUArray, b, u, Pl,
         Pr, maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions
     )
+    # `cholesky` needs a square matrix, and the default solver initializes this slot for
+    # every `A` before it knows which algorithm it will use. Instance an empty
+    # factorization for a non-square `A`, rather than returning `nothing`, so the return
+    # type does not depend on the runtime value of `assumptions.issq`.
+    # See https://github.com/SciML/NonlinearSolve.jl/issues/746
+    assumptions.issq || return cholesky(similar(A, 0, 0); check = false)
     return cholesky(A; check = false)
 end
 

src/sparsearrays.jl

@@ -1133,6 +1133,13 @@ function LinearSolve.init_cacheval(
         nothing
     elseif LinearSolve.is_cusparse_csr(A) && !LinearSolve.cudss_loaded(A)
         nothing
+    elseif A isa LinearSolve.GPUArraysCore.AnyGPUArray && !assumptions.issq
+        # `cholesky_instance` needs a square matrix, and the default solver reaches this
+        # slot for every `A` before it knows which algorithm it will use. A 0x0 is square,
+        # so instance that instead of returning `nothing`. Sparse CPU input is left alone:
+        # `cholesky_instance` handles a non-square sparse `A` and `solve!` stores a real
+        # factorization into the slot afterwards.
+        ArrayInterface.cholesky_instance(convert(AbstractMatrix, similar(A, 0, 0)))
     else
         ArrayInterface.cholesky_instance(convert(AbstractMatrix, A))
     end

test/GPU/cuda.jl

@@ -219,3 +219,39 @@ end
     @test LinearSolve.defaultalg(Wc, bc, OperatorAssumptions(true)).alg ===
         LinearSolve.DefaultAlgorithmChoice.LHLFactorization
 end
+
+# Two separate failures kept a non-square GPU `A` from solving at all.
+#
+# `_init_default_cacheval` builds a cacheval for every algorithm slot before it knows
+# which one it will use, and two of those slots called `cholesky` on `A` itself, so
+# `init` threw `DimensionMismatch` before any algorithm ran. Then for a wide `A`, the
+# QR solve itself threw: a GPU `qr` factors a wide matrix, but solving with the result
+# builds `UpperTriangular(R)` on a non-square `R`.
+# See https://github.com/SciML/NonlinearSolve.jl/issues/746 and #857.
+@testset "Non-square GPU matrices" begin
+    tall = CUDACore.adapt(CuArray, Float32[1 2; 3 4; 5 6; 7 8])
+    btall = CUDACore.adapt(CuArray, Float32[1, 2, 3, 4])
+    wide = CUDACore.adapt(CuArray, Float32[1 2 3 4; 5 6 7 8])
+    bwide = CUDACore.adapt(CuArray, Float32[1, 2])
+
+    @testset "$name" for (name, A, b) in (("tall", tall, btall), ("wide", wide, bwide))
+        ref = Array(A) \ Array(b)
+
+        # `init` is where every slot gets built, and where this used to throw.
+        cache = init(LinearProblem(A, b))
+        sol = solve!(cache)
+        @test SciMLBase.successful_retcode(sol)
+        @test Array(sol.u) ≈ ref rtol = 1.0e-4
+
+        @test LinearSolve.defaultalg(A, b, OperatorAssumptions(false)).alg ===
+            LinearSolve.DefaultAlgorithmChoice.QRFactorization
+        @test Array(solve(LinearProblem(A, b), QRFactorization()).u) ≈ ref rtol = 1.0e-4
+    end
+
+    # The wide solve has to be the minimum-norm one, which is what dense `\` gives on
+    # the CPU. Agreeing on the residual alone would not distinguish it from any other
+    # point on the solution manifold.
+    xwide = Array(solve(LinearProblem(wide, bwide)).u)
+    @test norm(xwide) ≈ norm(Array(wide) \ Array(bwide)) rtol = 1.0e-4
+    @test Array(wide) * xwide ≈ Array(bwide) rtol = 1.0e-4
+end

test/Reactant/Project.toml

@@ -0,0 +1,13 @@
+[deps]
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+LinearSolve = {path = "../.."}
+
+[compat]
+LinearSolve = "5"
+Reactant = "0.2.280"
+Test = "1.10"
+julia = "1.10"

test/Reactant/reactant.jl

@@ -0,0 +1,48 @@
+using LinearSolve
+using Reactant
+using Test
+
+function solve_once(A, b)
+    return solve(LinearProblem(A, b)).u
+end
+
+A = Float32[4 1; 2 3]
+b = Float32[1, 2]
+expected = A \ b
+A_reactant = Reactant.to_rarray(A)
+b_reactant = Reactant.to_rarray(b)
+
+sol = @jit solve(LinearProblem(A_reactant, b_reactant))
+@test Array(sol.u) ≈ expected
+@test Array(@jit solve_once(A_reactant, b_reactant)) ≈ expected
+@testset "preserves the LinearSolve operation" begin
+    @test occursin(
+        "reactant_julia_callback", repr(@code_hlo solve_once(A_reactant, b_reactant))
+    )
+
+    qr_sol = @jit solve(LinearProblem(A_reactant, b_reactant), QRFactorization())
+    @test qr_sol.alg isa QRFactorization
+    @test Array(qr_sol.u) ≈ expected
+end
+
+@testset "default LU preserves QR safety fallback" begin
+    A_singular = Float32[1 1; 1 1]
+    b_singular = Float32[1, 2]
+    expected_singular = Float32[0.75, 0.75]
+
+    lu_sol = solve(LinearProblem(A_singular, b_singular), LUFactorization())
+    @test lu_sol.retcode == ReturnCode.Failure
+    @test solve(LinearProblem(A_singular, b_singular)).u ≈ expected_singular
+
+    A_singular_reactant = Reactant.to_rarray(A_singular)
+    b_singular_reactant = Reactant.to_rarray(b_singular)
+    lu_sol_reactant = @jit solve(
+        LinearProblem(A_singular_reactant, b_singular_reactant), LUFactorization()
+    )
+    @test Array(lu_sol_reactant.u) == lu_sol.u
+
+    sol_singular = @jit solve(
+        LinearProblem(A_singular_reactant, b_singular_reactant)
+    )
+    @test Array(sol_singular.u) ≈ expected_singular
+end

test/qa/Project.toml

@@ -26,6 +26,7 @@ Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b"
 RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -67,6 +68,7 @@ Metal = "1.4"
 Mooncake = "0.5.15"
 PureUMFPACK = "1"
 RecursiveFactorization = "0.2.26"
+Reactant = "0.2.230"
 SafeTestsets = "0.1, 1"
 SciMLTesting = "2.8"
 SparseArrays = "1.10"