Merge pull request #564 from jClugstor/alias_API

Use Aliasing API for alias_A and alias_b

Project.toml

@@ -105,7 +105,7 @@ RecursiveArrayTools = "3.8"
 RecursiveFactorization = "0.2.14"
 Reexport = "1"
 SafeTestsets = "0.1"
-SciMLBase = "2.26.3"
+SciMLBase = "2.70"
 SciMLOperators = "0.3.7"
 Setfield = "1"
 SparseArrays = "1.10"

benchmarks/applelu.jl

@@ -39,9 +39,9 @@ for i in 1:length(ns)
     for j in 1:length(algs)
         bt = @belapsed solve(prob, $(algs[j])).u setup=(prob = LinearProblem(copy(A),
             copy(b);
-            u0 = copy(u0),
-            alias_A = true,
-            alias_b = true))
+            u0 = copy(u0), 
+            alias = LinearAliasSpecifier(alias_A = true, alias_b = true)
+            ))
         push!(res[j], luflop(n) / bt / 1e9)
     end
 end

benchmarks/cudalu.jl

@@ -34,8 +34,7 @@ for i in 1:length(ns)
         bt = @belapsed solve(prob, $(algs[j])).u setup=(prob = LinearProblem(copy(A),
             copy(b);
             u0 = copy(u0),
-            alias_A = true,
-            alias_b = true))
+            alias = LinearAliasSpecifier(alias_A = true, alias_b = true)))
         push!(res[j], luflop(n) / bt / 1e9)
     end
 end

benchmarks/lu.jl

@@ -41,8 +41,7 @@ for i in 1:length(ns)
         bt = @belapsed solve(prob, $(algs[j])).u setup=(prob = LinearProblem(copy(A),
             copy(b);
             u0 = copy(u0),
-            alias_A = true,
-            alias_b = true))
+            alias = LinearAliasSpecifier(alias_A = true, alias_b = true)))
         push!(res[j], luflop(n) / bt / 1e9)
     end
 end

benchmarks/metallu.jl

@@ -34,8 +34,7 @@ for i in 1:length(ns)
         bt = @belapsed solve(prob, $(algs[j])).u setup=(prob = LinearProblem(copy(A),
             copy(b);
             u0 = copy(u0),
-            alias_A = true,
-            alias_b = true))
+            alias = LinearAliasSpecifier(alias_A = true, alias_b = true)))
         GC.gc()
         push!(res[j], luflop(n) / bt / 1e9)
     end

benchmarks/sparselu.jl

@@ -69,8 +69,7 @@ function run_and_plot(; dims = [1, 2, 3], kmax = 12)
                     copy($A),
                     copy($b);
                     u0 = copy($u0),
-                    alias_A = true,
-                    alias_b = true))
+                    alias = LinearAliasSpecifier(alias_A = true, alias_b = true)))
                 push!(res[dim][j], bt)
             end
         end

docs/src/basics/common_solver_opts.md

@@ -6,15 +6,11 @@ in order to give composability. These are also the options taken at `init` time.
 The following are the options these algorithms take, along with their defaults.
 
 ## General Controls
-
-  - `alias_A::Bool`: Whether to alias the matrix `A` or use a copy by default. When `true`,
-    algorithms like LU-factorization can be faster by reusing the memory via `lu!`,
-    but care must be taken as the original input will be modified. Default is `true` if the
-    algorithm is known not to modify `A`, otherwise is `false`.
-  - `alias_b::Bool`: Whether to alias the matrix `b` or use a copy by default. When `true`,
-    algorithms can write and change `b` upon usage. Care must be taken as the
-    original input will be modified. Default is `true` if the algorithm is known not to
-    modify `b`, otherwise `false`.
+  - `alias::LinearAliasSpecifier`: Holds the fields `alias_A` and `alias_b` which specify
+    whether to alias the matrices `A` and `b` respectively. When these fields are `true`,
+    `A` and `b` can be written to and changed by the solver algorithm. When fields are `nothing`
+    the default behavior is used, which is to default to `true` when the algorithm is known 
+    not to modify the matrices, and false otherwise.
   - `verbose`: Whether to print extra information. Defaults to `false`.
   - `assumptions`: Sets the assumptions of the operator in order to effect the default
     choice algorithm. See the [Operator Assumptions page for more details](@ref assumptions).

ext/LinearSolveHYPREExt.jl

@@ -6,7 +6,7 @@ using HYPRE: HYPRE, HYPREMatrix, HYPRESolver, HYPREVector
 using LinearSolve: HYPREAlgorithm, LinearCache, LinearProblem, LinearSolve,
                    OperatorAssumptions, default_tol, init_cacheval, __issquare,
                    __conditioning, LinearSolveAdjoint
-using SciMLBase: LinearProblem, SciMLBase
+using SciMLBase: LinearProblem, LinearAliasSpecifier, SciMLBase
 using UnPack: @unpack
 using Setfield: @set!
 
@@ -55,7 +55,7 @@ end
 
 function SciMLBase.init(prob::LinearProblem, alg::HYPREAlgorithm,
         args...;
-        alias_A = false, alias_b = false,
+        alias = LinearAliasSpecifier(),
         # TODO: Implement eltype for HYPREMatrix in HYPRE.jl? Looks useful
         #       even if it is not AbstractArray.
         abstol = default_tol(prob.A isa HYPREMatrix ? HYPRE_Complex :
@@ -72,6 +72,45 @@ function SciMLBase.init(prob::LinearProblem, alg::HYPREAlgorithm,
         kwargs...)
     @unpack A, b, u0, p = prob
 
+    if haskey(kwargs, :alias_A) || haskey(kwargs, :alias_b)
+        aliases = LinearAliasSpecifier()
+
+        if haskey(kwargs, :alias_A)
+            message = "`alias_A` keyword argument is deprecated, to set `alias_A`,
+            please use an ODEAliasSpecifier, e.g. `solve(prob, alias = LinearAliasSpecifier(alias_A = true))"
+            Base.depwarn(message, :init)
+            Base.depwarn(message, :solve)
+            aliases = LinearAliasSpecifier(alias_A = values(kwargs).alias_A)
+        end
+
+        if haskey(kwargs, :alias_b)
+            message = "`alias_b` keyword argument is deprecated, to set `alias_b`,
+            please use an ODEAliasSpecifier, e.g. `solve(prob, alias = LinearAliasSpecifier(alias_b = true))"
+            Base.depwarn(message, :init)
+            Base.depwarn(message, :solve)
+            aliases = LinearAliasSpecifier(
+                alias_A = aliases.alias_A, alias_b = values(kwargs).alias_b)
+        end
+    else
+        if alias isa Bool
+            aliases = LinearAliasSpecifier(alias = alias)
+        else
+            aliases = alias
+        end
+    end
+
+    if isnothing(aliases.alias_A)
+        alias_A = false
+    else
+        alias_A = aliases.alias_A
+    end
+
+    if isnothing(aliases.alias_b)
+        alias_b = false
+    else
+        alias_b = aliases.alias_b
+    end
+
     A = A isa HYPREMatrix ? A : HYPREMatrix(A)
     b = b isa HYPREVector ? b : HYPREVector(b)
     u0 = u0 isa HYPREVector ? u0 : (u0 === nothing ? nothing : HYPREVector(u0))

src/LinearSolve.jl

@@ -12,7 +12,7 @@ using LinearAlgebra
 using SparseArrays
 using SparseArrays: AbstractSparseMatrixCSC, nonzeros, rowvals, getcolptr
 using LazyArrays: @~, BroadcastArray
-using SciMLBase: AbstractLinearAlgorithm
+using SciMLBase: AbstractLinearAlgorithm, LinearAliasSpecifier
 using SciMLOperators
 using SciMLOperators: AbstractSciMLOperator, IdentityOperator
 using Setfield

src/common.jl

@@ -139,8 +139,7 @@ __init_u0_from_Ab(::SMatrix{S1, S2}, b) where {S1, S2} = zeros(SVector{S2, eltyp
 
 function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
         args...;
-        alias_A = default_alias_A(alg, prob.A, prob.b),
-        alias_b = default_alias_b(alg, prob.A, prob.b),
+        alias = LinearAliasSpecifier(),
         abstol = default_tol(real(eltype(prob.b))),
         reltol = default_tol(real(eltype(prob.b))),
         maxiters::Int = length(prob.b),
@@ -152,6 +151,45 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
         kwargs...)
     (; A, b, u0, p) = prob
 
+    if haskey(kwargs,:alias_A) || haskey(kwargs,:alias_b)
+        aliases = LinearAliasSpecifier()
+
+        if haskey(kwargs, :alias_A)
+            message = "`alias_A` keyword argument is deprecated, to set `alias_A`,
+            please use an ODEAliasSpecifier, e.g. `solve(prob, alias = LinearAliasSpecifier(alias_A = true))"
+            Base.depwarn(message, :init)
+            Base.depwarn(message, :solve)
+            aliases = LinearAliasSpecifier(alias_A = values(kwargs).alias_A)
+        end
+
+        if haskey(kwargs, :alias_b)
+            message = "`alias_b` keyword argument is deprecated, to set `alias_b`,
+            please use an ODEAliasSpecifier, e.g. `solve(prob, alias = LinearAliasSpecifier(alias_b = true))"
+            Base.depwarn(message, :init)
+            Base.depwarn(message, :solve)
+            aliases = LinearAliasSpecifier(alias_A = aliases.alias_A, alias_b = values(kwargs).alias_b)
+        end 
+    else
+        if alias isa Bool
+            aliases = LinearAliasSpecifier(alias = alias)
+        else
+            aliases = alias
+        end
+    end
+
+    if isnothing(aliases.alias_A)
+        alias_A = default_alias_A(alg, prob.A, prob.b)
+    else
+        alias_A = aliases.alias_A
+    end
+
+    if isnothing(aliases.alias_b)
+        alias_b = default_alias_b(alg, prob.A, prob.b)
+    else
+        alias_b = aliases.alias_b
+    end
+
+
     A = if alias_A || A isa SMatrix
         A
     elseif A isa Array

test/gpu/cuda.jl

@@ -84,9 +84,9 @@ prob2 = LinearProblem(transpose(A), b)
 
 @testset "Adjoint/Transpose Type: $(alg)" for alg in (NormalCholeskyFactorization(),
     CholeskyFactorization(), LUFactorization(), QRFactorization(), nothing)
-    sol = solve(prob1, alg; alias_A = false)
+    sol = solve(prob1, alg; alias = LinearAliasSpecifier(alias = LinearAliasSpecifier(alias_A = false)))
     @test norm(A' * sol.u .- b) < 1e-5
 
-    sol = solve(prob2, alg; alias_A = false)
+    sol = solve(prob2, alg; alias = LinearAliasSpecifier(alias_A = false))
     @test norm(transpose(A) * sol.u .- b) < 1e-5
 end

test/resolve.jl

@@ -27,7 +27,7 @@ for alg in vcat(InteractiveUtils.subtypes(AbstractDenseFactorization),
         alg in [LDLtFactorization] && (A = SymTridiagonal(A))
         b = [1.0, 2.0]
         prob = LinearProblem(A, b)
-        linsolve = init(prob, alg(), alias_A = false, alias_b = false)
+        linsolve = init(prob, alg(), alias = LinearAliasSpecifier(alias_A = false, alias_b = false))
         @test solve!(linsolve).u ≈ [-2.0, 1.5]
         @test !linsolve.isfresh
         @test solve!(linsolve).u ≈ [-2.0, 1.5]
@@ -48,7 +48,7 @@ end
 A = Diagonal([1.0, 4.0])
 b = [1.0, 2.0]
 prob = LinearProblem(A, b)
-linsolve = init(prob, DiagonalFactorization(), alias_A = false, alias_b = false)
+linsolve = init(prob, DiagonalFactorization(), alias = LinearAliasSpecifier(alias_A = false, alias_b = false))
 @test solve!(linsolve).u ≈ [1.0, 0.5]
 @test solve!(linsolve).u ≈ [1.0, 0.5]
 A = Diagonal([1.0, 4.0])
@@ -59,7 +59,7 @@ A = Symmetric([1.0 2.0
                2.0 1.0])
 b = [1.0, 2.0]
 prob = LinearProblem(A, b)
-linsolve = init(prob, BunchKaufmanFactorization(), alias_A = false, alias_b = false)
+linsolve = init(prob, BunchKaufmanFactorization(), alias = LinearAliasSpecifier(alias_A = false, alias_b = false))
 @test solve!(linsolve).u ≈ [1.0, 0.0]
 @test solve!(linsolve).u ≈ [1.0, 0.0]
 A = Symmetric([1.0 2.0