Restore GPU Broyden inverse initialization

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/NonlinearSolveBase/src/utils.jl

@@ -235,6 +235,9 @@ linsolve_workspace(A::Diagonal) = nothing, A
 linsolve_workspace(A::SMatrix) = nothing, A
 function linsolve_workspace(A::AbstractMatrix)
     LinearAlgebra.checksquare(A)
+    # Backends without fast scalar indexing may not implement factorization solves with
+    # a matrix right-hand side. Leave those arrays on their native `pinv` path.
+    ArrayInterface.fast_scalar_indexing(A) || return nothing, A
     # the linear solve promotes e.g. integer eltypes; the buffers must hold the promoted
     # values. The `size` form of `similar` gives a dense buffer for structured input
     # (e.g. Tridiagonal), which the generically-dense inverse and identity RHS need.
@@ -283,8 +286,14 @@ end
 function linsolve_identity!!(workspace, A::AbstractMatrix)
     LinearAlgebra.checksquare(A)
     # a `nothing` workspace (the preinverted-init path in NonlinearSolveQuasiNewton)
-    # builds a transient linear-solve cache; it allocates, but runs at most once per solve
-    ws = workspace === nothing ? first(linsolve_workspace(A)) : workspace
+    # builds a transient linear-solve cache; it allocates, but runs at most once per solve.
+    # Sparse dispatch gets a chance to construct its dense workspace here, while backend
+    # arrays that cannot use the cached matrix-RHS solve stay on their native `pinv`.
+    ws = workspace
+    if ws === nothing
+        ws, workspace_A = linsolve_workspace(A)
+        ws === nothing && return pinv(workspace_A)
+    end
     # the previous solve may have consumed the RHS buffer, so refill it with I. The
     # lincache call copies A into its internal buffer before overwriting the solution
     # buffer, so A is never mutated and passing the previously returned inverse as A is

lib/NonlinearSolveBase/test/linsolve_workspace.jl

@@ -1,8 +1,18 @@
 using NonlinearSolveBase: NonlinearSolveBase, Utils
+using ArrayInterface
 using LinearSolve  # linsolve_identity!! routes matrices through the LinearSolve default
 using LinearAlgebra
 using Test
 
+struct NoFastScalarMatrix{T} <: AbstractMatrix{T}
+    data::Matrix{T}
+end
+
+Base.size(A::NoFastScalarMatrix) = size(A.data)
+Base.getindex(A::NoFastScalarMatrix, i::Int, j::Int) = A.data[i, j]
+Base.copy(A::NoFastScalarMatrix) = copy(A.data)
+ArrayInterface.fast_scalar_indexing(::Type{<:NoFastScalarMatrix}) = false
+
 @testset "nonsingular input matches inv" begin
     n = 20
     A_general = rand(n, n) + n * I
@@ -46,6 +56,13 @@ end
     @test Utils.linsolve_identity!!(workspace, workspace.rhs) ≈ inv(A_general)
 end
 
+@testset "arrays without fast scalar indexing use pinv" begin
+    A = NoFastScalarMatrix(rand(5, 5))
+    workspace, A_ret = Utils.linsolve_workspace(A)
+    @test workspace === nothing && A_ret === A
+    @test Utils.linsolve_identity!!(workspace, A) ≈ pinv(A.data)
+end
+
 @testset "singular input takes the pivoted-QR rescue" begin
     # The result is the LinearSolve default algorithm's least-squares generalized
     # inverse from its singular-LU → pivoted-QR rescue, NOT the SVD `pinv` (an

test/gpu/cuda_tests__item1.jl

@@ -1,6 +1,7 @@
 using NonlinearSolve
 
 using CUDA, NonlinearSolve, LinearSolve, StableRNGs, ADTypes
+using LinearAlgebra: norm
 
 if CUDA.functional()
     CUDA.allowscalar(false)
@@ -12,6 +13,7 @@ if CUDA.functional()
     linear_f(du, u, p) = (du .= A * u .+ b)
 
     prob = NonlinearProblem(linear_f, u0)
+    iip_prob = prob
 
     # ForwardDiff uses scalar indexing which doesn't work on GPU
     # Use AutoFiniteDiff for GPU-compatible Jacobian computation
@@ -39,10 +41,21 @@ if CUDA.functional()
     linear_f(u, p) = A * u .+ b
 
     prob = NonlinearProblem{false}(linear_f, u0)
+    oop_prob = prob
 
     @testset "[OOP] GPU Solvers" begin
         @testset "$(nameof(typeof(alg)))" for alg in SOLVERS
             @test_nowarn sol = solve(prob, alg; abstol = 1.0f-5, reltol = 1.0f-5)
         end
     end
+
+    @testset "Broyden inverse initialization" begin
+        for prob in (iip_prob, oop_prob)
+            sol = solve(
+                prob, Broyden(; linesearch = LiFukushimaLineSearch());
+                abstol = 1.0f-5, reltol = 1.0f-5
+            )
+            @test norm(A * sol.u + b) < 1.0f-4
+        end
+    end
 end