Improve type stability

Author: mtfishman

src/algorithms.jl

@@ -54,45 +54,48 @@ function _show_alg(io::IO, alg::Algorithm)
 end
 
 @doc """
-    MatrixAlgebraKit.select_algorithm(f, A; kwargs...)
+    MatrixAlgebraKit.select_algorithm(f, A, alg=nothing; kwargs...)
 
-Given some keyword arguments and an input `A`, decide on an algorithm to use for
-implementing the function `f` on inputs of type `A`.
+Decide on an algorithm to use for implementing the function `f` on inputs of type `A`.
 
-In general, if an algorithm is specified explicitly through the `alg` keyword argument
-(either as an algorithm type, an algorithm name as a Symbol, or as an algorithm object),
-that algorithm will be used instead of selecting it automatically. However, that
-behavior may be modified for factorization functions and/or matrix types.
+If `alg` is `nothing` (the default value), an algorithm will be selected automatically
+with [`MatrixAlgebra.default_algorithm`](@ref) and the keyword arguments will be passed
+to the algorithm constructor.
 
-When the `alg` keyword argument is not provided, a default algorithm specified by
-[`default_algorithm`](@ref) will be used.
+If `alg` is a `NamedTuple`, an algorithm will be selected automatically
+with [`default_algorithm`](@ref) and `alg` will be passed to the algorithm
+as keyword arguments. In that case, keyword arguments can't be passed
+to `MatrixAlgebraKit.select_algorithm`
+
+If `alg` is an `AbstractAlgorithm`, it will be returned as-is. In that case, keyword arguments
+can't be passed to `MatrixAlgebraKit.select_algorithm`.
 """
 function select_algorithm end
 
-function select_algorithm(f, A; alg=nothing, kwargs...)
+Base.@constprop :aggressive function select_algorithm(f::F, A, alg::Alg=nothing; kwargs...) where {F,Alg}
     return _select_algorithm(f, A, alg; kwargs...)
 end
 
-function _select_algorithm(f, A, alg::Nothing; kwargs...)
+function _select_algorithm(f::F, A, alg::Nothing; kwargs...) where {F}
     return default_algorithm(f, A; kwargs...)
 end
-function _select_algorithm(f, A, alg::AbstractAlgorithm; kwargs...)
-    isempty(kwargs) ||
-        throw(ArgumentError("Additional keyword arguments are not allowed when an algorithm is specified."))
-    return alg
-end
-function _select_algorithm(f, A, alg::Symbol; kwargs...)
-    return _select_algorithm(f, A, Algorithm{alg}(; kwargs...))
+Base.@constprop :aggressive function _select_algorithm(f::F, A, alg::Symbol; kwargs...) where {F}
+    return Algorithm{alg}(; kwargs...)
 end
-function _select_algorithm(f, A, alg::Type; kwargs...)
-    return _select_algorithm(f, A, alg(; kwargs...))
+Base.@constprop :aggressive function _select_algorithm(f::F, A, ::Type{Alg}; kwargs...) where {F,Alg}
+    return Alg(; kwargs...)
 end
-function _select_algorithm(f, A::AbstractMatrix, alg::NamedTuple; kwargs...)
+function _select_algorithm(f::F, A, alg::NamedTuple; kwargs...) where {F}
     isempty(kwargs) ||
         throw(ArgumentError("Additional keyword arguments are not allowed when algorithm parameters are specified."))
-    return select_algorithm(f, A; alg...)
+    return default_algorithm(f, A; alg...)
 end
-function _select_algorithm(f, A, alg; kwargs...)
+function _select_algorithm(f::F, A, alg::AbstractAlgorithm; kwargs...) where {F}
+    isempty(kwargs) ||
+        throw(ArgumentError("Additional keyword arguments are not allowed when an algorithm is specified."))
+    return alg
+end
+function _select_algorithm(f::F, A, alg; kwargs...) where {F}
     return throw(ArgumentError("Unknown alg $alg"))
 end
 
@@ -171,13 +174,15 @@ macro functiondef(f)
 
     return esc(quote
                    # out of place to inplace
-                   $f(A; kwargs...) = $f!(copy_input($f, A); kwargs...)
+                   Base.@constprop :aggressive $f(A; kwargs...) = $f!(copy_input($f, A); kwargs...)
                    $f(A, alg::AbstractAlgorithm) = $f!(copy_input($f, A), alg)
 
                    # fill in arguments
-                   $f!(A; kwargs...) = $f!(A, select_algorithm($f!, A; kwargs...))
-                   function $f!(A, out; kwargs...)
-                       return $f!(A, out, select_algorithm($f!, A; kwargs...))
+                   Base.@constprop :aggressive function $f!(A; alg=nothing, kwargs...)
+                       return $f!(A, select_algorithm($f!, A, alg; kwargs...))
+                   end
+                   Base.@constprop :aggressive function $f!(A, out; alg=nothing, kwargs...)
+                       return $f!(A, out, select_algorithm($f!, A, alg; kwargs...))
                    end
                    function $f!(A, alg::AbstractAlgorithm)
                        return $f!(A, initialize_output($f!, A, alg), alg)

src/implementations/orthnull.jl

@@ -89,22 +89,22 @@ function left_orth!(A::AbstractMatrix, VC; trunc=nothing,
         throw(ArgumentError("truncation not supported for left_orth with kind=$kind"))
     end
     if kind == :qr
-        alg_qr′ = _select_algorithm(qr_compact!, A, alg_qr)
+        alg_qr′ = select_algorithm(qr_compact!, A, alg_qr)
         return qr_compact!(A, VC, alg_qr′)
     elseif kind == :polar
         size(A, 1) >= size(A, 2) ||
             throw(ArgumentError("`left_orth!` with `kind = :polar` only possible for `(m, n)` matrix with `m >= n`"))
-        alg_polar′ = _select_algorithm(left_polar!, A, alg_polar)
+        alg_polar′ = select_algorithm(left_polar!, A, alg_polar)
         return left_polar!(A, VC, alg_polar′)
     elseif kind == :svd && isnothing(trunc)
-        alg_svd′ = _select_algorithm(svd_compact!, A, alg_svd)
+        alg_svd′ = select_algorithm(svd_compact!, A, alg_svd)
         V, C = VC
         S = Diagonal(initialize_output(svd_vals!, A, alg_svd′))
         U, S, Vᴴ = svd_compact!(A, (V, S, C), alg_svd′)
         return U, lmul!(S, Vᴴ)
     elseif kind == :svd
-        alg_svd′ = _select_algorithm(svd_compact!, A, alg_svd)
-        alg_svd_trunc = select_algorithm(svd_trunc!, A; trunc, alg=alg_svd′)
+        alg_svd′ = select_algorithm(svd_compact!, A, alg_svd)
+        alg_svd_trunc = select_algorithm(svd_trunc!, A, alg_svd′; trunc)
         V, C = VC
         S = Diagonal(initialize_output(svd_vals!, A, alg_svd_trunc.alg))
         U, S, Vᴴ = svd_trunc!(A, (V, S, C), alg_svd_trunc)
@@ -122,22 +122,22 @@ function right_orth!(A::AbstractMatrix, CVᴴ; trunc=nothing,
         throw(ArgumentError("truncation not supported for right_orth with kind=$kind"))
     end
     if kind == :lq
-        alg_lq′ = _select_algorithm(lq_compact!, A, alg_lq)
+        alg_lq′ = select_algorithm(lq_compact!, A, alg_lq)
         return lq_compact!(A, CVᴴ, alg_lq′)
     elseif kind == :polar
         size(A, 2) >= size(A, 1) ||
             throw(ArgumentError("`right_orth!` with `kind = :polar` only possible for `(m, n)` matrix with `m <= n`"))
-        alg_polar′ = _select_algorithm(right_polar!, A, alg_polar)
+        alg_polar′ = select_algorithm(right_polar!, A, alg_polar)
         return right_polar!(A, CVᴴ, alg_polar′)
     elseif kind == :svd && isnothing(trunc)
-        alg_svd′ = _select_algorithm(svd_compact!, A, alg_svd)
+        alg_svd′ = select_algorithm(svd_compact!, A, alg_svd)
         C, Vᴴ = CVᴴ
         S = Diagonal(initialize_output(svd_vals!, A, alg_svd′))
         U, S, Vᴴ = svd_compact!(A, (C, S, Vᴴ), alg_svd′)
         return rmul!(U, S), Vᴴ
     elseif kind == :svd
-        alg_svd′ = _select_algorithm(svd_compact!, A, alg_svd)
-        alg_svd_trunc = select_algorithm(svd_trunc!, A; trunc, alg=alg_svd′)
+        alg_svd′ = select_algorithm(svd_compact!, A, alg_svd)
+        alg_svd_trunc = select_algorithm(svd_trunc!, A, alg_svd′; trunc)
         C, Vᴴ = CVᴴ
         S = Diagonal(initialize_output(svd_vals!, A, alg_svd_trunc.alg))
         U, S, Vᴴ = svd_trunc!(A, (C, S, Vᴴ), alg_svd_trunc)
@@ -167,15 +167,15 @@ function left_null!(A::AbstractMatrix, N; trunc=nothing,
         throw(ArgumentError("truncation not supported for left_null with kind=$kind"))
     end
     if kind == :qr
-        alg_qr′ = _select_algorithm(qr_null!, A, alg_qr)
+        alg_qr′ = select_algorithm(qr_null!, A, alg_qr)
         return qr_null!(A, N, alg_qr′)
     elseif kind == :svd && isnothing(trunc)
-        alg_svd′ = _select_algorithm(svd_full!, A, alg_svd)
+        alg_svd′ = select_algorithm(svd_full!, A, alg_svd)
         U, _, _ = svd_full!(A, alg_svd′)
         (m, n) = size(A)
         return copy!(N, view(U, 1:m, (n + 1):m))
     elseif kind == :svd
-        alg_svd′ = _select_algorithm(svd_full!, A, alg_svd)
+        alg_svd′ = select_algorithm(svd_full!, A, alg_svd)
         U, S, _ = svd_full!(A, alg_svd′)
         trunc′ = trunc isa TruncationStrategy ? trunc :
                  trunc isa NamedTuple ? null_truncation_strategy(; trunc...) :
@@ -194,15 +194,15 @@ function right_null!(A::AbstractMatrix, Nᴴ; trunc=nothing,
         throw(ArgumentError("truncation not supported for right_null with kind=$kind"))
     end
     if kind == :lq
-        alg_lq′ = _select_algorithm(lq_null!, A, alg_lq)
+        alg_lq′ = select_algorithm(lq_null!, A, alg_lq)
         return lq_null!(A, Nᴴ, alg_lq′)
     elseif kind == :svd && isnothing(trunc)
-        alg_svd′ = _select_algorithm(svd_full!, A, alg_svd)
+        alg_svd′ = select_algorithm(svd_full!, A, alg_svd)
         _, _, Vᴴ = svd_full!(A, alg_svd′)
         (m, n) = size(A)
         return copy!(Nᴴ, view(Vᴴ, (m + 1):n, 1:n))
     elseif kind == :svd
-        alg_svd′ = _select_algorithm(svd_full!, A, alg_svd)
+        alg_svd′ = select_algorithm(svd_full!, A, alg_svd)
         _, S, Vᴴ = svd_full!(A, alg_svd′)
         trunc′ = trunc isa TruncationStrategy ? trunc :
                  trunc isa NamedTuple ? null_truncation_strategy(; trunc...) :

src/implementations/truncation.jl

@@ -32,17 +32,16 @@ Trivial truncation strategy that keeps all values, mostly for testing purposes.
 """
 struct NoTruncation <: TruncationStrategy end
 
-function select_truncation(trunc::TruncationStrategy)
-    return trunc
-end
-function select_truncation(trunc::NamedTuple)
-    return TruncationStrategy(; trunc...)
-end
-function select_truncation(trunc::Nothing)
-    return NoTruncation()
-end
 function select_truncation(trunc)
-    return throw(ArgumentError("Unknown truncation strategy: $trunc"))
+    if isnothing(trunc)
+        return NoTruncation()
+    elseif trunc isa NamedTuple
+        return TruncationStrategy(; trunc...)
+    elseif trunc isa TruncationStrategy
+        return trunc
+    else
+        return throw(ArgumentError("Unknown truncation strategy: $trunc"))
+    end
 end
 
 # TODO: how do we deal with sorting/filters that treat zeros differently

src/interface/eig.jl

@@ -99,11 +99,11 @@ for f in (:eig_full, :eig_vals)
     end
 end
 
-function select_algorithm(::typeof(eig_trunc), A; kwargs...)
-    return select_algorithm(eig_trunc!, A; kwargs...)
+function select_algorithm(::typeof(eig_trunc), A, alg; kwargs...)
+    return select_algorithm(eig_trunc!, A, alg; kwargs...)
 end
-function select_algorithm(::typeof(eig_trunc!), A; trunc=nothing, kwargs...)
-    alg_eig = select_algorithm(eig_full!, A; kwargs...)
+function select_algorithm(::typeof(eig_trunc!), A, alg; trunc=nothing, kwargs...)
+    alg_eig = select_algorithm(eig_full!, A, alg; kwargs...)
     return TruncatedAlgorithm(alg_eig, select_truncation(trunc))
 end
 

src/interface/eigh.jl

@@ -98,11 +98,11 @@ for f in (:eigh_full, :eigh_vals)
     end
 end
 
-function select_algorithm(::typeof(eigh_trunc), A; kwargs...)
-    return select_algorithm(eigh_trunc!, A; kwargs...)
+function select_algorithm(::typeof(eigh_trunc), A, alg; kwargs...)
+    return select_algorithm(eigh_trunc!, A, alg; kwargs...)
 end
-function select_algorithm(::typeof(eigh_trunc!), A; trunc=nothing, kwargs...)
-    alg_eigh = select_algorithm(eigh_full!, A; kwargs...)
+function select_algorithm(::typeof(eigh_trunc!), A, alg; trunc=nothing, kwargs...)
+    alg_eigh = select_algorithm(eigh_full!, A, alg; kwargs...)
     return TruncatedAlgorithm(alg_eigh, select_truncation(trunc))
 end
 

src/interface/svd.jl

@@ -102,11 +102,11 @@ for f in (:svd_full, :svd_compact, :svd_vals)
     end
 end
 
-function select_algorithm(::typeof(svd_trunc), A; kwargs...)
-    return select_algorithm(svd_trunc!, A; kwargs...)
+function select_algorithm(::typeof(svd_trunc), A, alg; kwargs...)
+    return select_algorithm(svd_trunc!, A, alg; kwargs...)
 end
-function select_algorithm(::typeof(svd_trunc!), A; trunc=nothing, kwargs...)
-    alg_svd = select_algorithm(svd_compact!, A; kwargs...)
+function select_algorithm(::typeof(svd_trunc!), A, alg; trunc=nothing, kwargs...)
+    alg_svd = select_algorithm(svd_compact!, A, alg; kwargs...)
     return TruncatedAlgorithm(alg_svd, select_truncation(trunc))
 end
 

test/eig.jl

@@ -8,28 +8,16 @@ using MatrixAlgebraKit: diagview
 @testset "eig_full! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
     rng = StableRNG(123)
     m = 54
-    for alg in
-        (LAPACK_Simple(), LAPACK_Expert(), LAPACK_Simple, LAPACK_Expert, :LAPACK_Simple,
-         :LAPACK_Expert)
+    for alg in (LAPACK_Simple(), LAPACK_Expert())
         A = randn(rng, T, m, m)
         Tc = complex(T)
 
-        alg′ = if (alg isa Type) || (alg isa Symbol)
-            # These cases aren't inferable right now.
-            MatrixAlgebraKit.select_algorithm(eig_full!, A; alg)
-        else
-            @constinferred MatrixAlgebraKit.select_algorithm(eig_full!, A; alg)
-        end
-
-        D, V = if (alg isa Type) || (alg isa Symbol)
-            # These cases aren't inferable right now.
-            eig_full(A; alg)
-        else
-            @constinferred eig_full(A; alg)
-        end
+        D, V = @constinferred eig_full(A; alg)
         @test eltype(D) == eltype(V) == Tc
         @test A * V ≈ V * D
 
+        alg′ = @constinferred MatrixAlgebraKit.select_algorithm(eig_full!, A, alg)
+
         Ac = similar(A)
         D2, V2 = @constinferred eig_full!(copy!(Ac, A), (D, V), alg′)
         @test D2 === D
@@ -40,6 +28,21 @@ using MatrixAlgebraKit: diagview
         @test eltype(Dc) == Tc
         @test D ≈ Diagonal(Dc)
     end
+
+    # Test other alg inputs.
+    A = randn(rng, T, m, m)
+    Tc = complex(T)
+    D, V = @constinferred eig_full(A; alg=:LAPACK_Simple)
+    @test eltype(D) == eltype(V) == Tc
+    @test A * V ≈ V * D
+
+    A = randn(rng, T, m, m)
+    Tc = complex(T)
+    ## Inference is broken for this case for some reason.
+    ## D, V = @constinferred eig_full(A; alg=LAPACK_Simple)
+    D, V = eig_full(A; alg=LAPACK_Simple)
+    @test eltype(D) == eltype(V) == Tc
+    @test A * V ≈ V * D
 end
 
 @testset "eig_trunc! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)