Merge #93

93: Prep alleviating latency issues r=charleskawczynski a=charleskawczynski

This PR takes a couple steps to tackle some latency issues:
 - [Avoid UnionAll::DataType in struct](52c9ff0), which was shown to improve inference [here](CliMA/Thermodynamics.jl#71), and a smaller MWE [here](https://gist.github.com/charleskawczynski/61efc7b769e14eab7231e686773b9d0e)
 - [Allow SparseContainers to work with Arrays](c5c53ee)

Co-authored-by: Charles Kawczynski <[email protected]>

Author: bors[bot]

src/solvers/newtons_method.jl

@@ -327,7 +327,7 @@ end
     KrylovMethod(;
         jacobian_free_jvp = nothing,
         forcing_term = ConstantForcing(0),
-        type = Krylov.GmresSolver,
+        type = Val(Krylov.GmresSolver),
         args = (20,),
         kwargs = (;),
         solve_kwargs = (;),
@@ -347,7 +347,8 @@ where `x_prototype` is `similar` to `x` (and also to `Δx` and `f`).
 
 This is primarily a wrapper for a `Krylov.KrylovSolver` from `Krylov.jl`. In
 `allocate_cache`, the solver is constructed with
-`solver = type(l, l, args..., Krylov.ktypeof(x_prototype); kwargs...)`, where
+`solver = type(l, l, args..., Krylov.ktypeof(x_prototype); kwargs...)` (note
+that `type` must be passed through in a `Val` struct), where
 `l = length(x_prototype)` and `Krylov.ktypeof(x_prototype)` is a subtype of
 `DenseVector` that can be used to store `x_prototype`. By default, the solver
 is a `Krylov.GmresSolver` with a Krylov subspace size of 20 (the default Krylov
@@ -387,17 +388,17 @@ each iteration of the Krylov method. If a debugger is specified, it is run
 before the call to `Kyrlov.solve!`.
 """
 Base.@kwdef struct KrylovMethod{
+    T <: Val,
     J <: Union{Nothing, JacobianFreeJVP},
     F <: ForcingTerm,
-    T <: Type,
     A <: Tuple,
     K <: NamedTuple,
     S <: NamedTuple,
     D <: Union{Nothing, KrylovMethodDebugger},
 }
+    type::T = Val(Krylov.GmresSolver)
     jacobian_free_jvp::J = nothing
     forcing_term::F = ConstantForcing(0)
-    type::T = Krylov.GmresSolver
     args::A = (20,)
     kwargs::K = (;)
     solve_kwargs::S = (;)
@@ -406,9 +407,12 @@ Base.@kwdef struct KrylovMethod{
     debugger::D = nothing
 end
 
+solver_type(::KrylovMethod{Val{T}}) where {T} = T
+
 function allocate_cache(alg::KrylovMethod, x_prototype)
-    (; jacobian_free_jvp, forcing_term, type, args, kwargs, debugger) = alg
-    @assert alg.type isa Type{<:Krylov.KrylovSolver}
+    (; jacobian_free_jvp, forcing_term, args, kwargs, debugger) = alg
+    type = solver_type(alg)
+    @assert type isa Type{<:Krylov.KrylovSolver}
     l = length(x_prototype)
     return (;
         jacobian_free_jvp_cache = isnothing(jacobian_free_jvp) ? nothing :
@@ -421,8 +425,9 @@ function allocate_cache(alg::KrylovMethod, x_prototype)
 end
 
 function run!(alg::KrylovMethod, cache, Δx, x, f!, f, n, j = nothing)
-    (; jacobian_free_jvp, forcing_term, type, solve_kwargs) = alg
+    (; jacobian_free_jvp, forcing_term, solve_kwargs) = alg
     (; disable_preconditioner, verbose, debugger) = alg
+    type = solver_type(alg)
     (; jacobian_free_jvp_cache, forcing_term_cache, solver, debugger_cache) =
         cache
     jΔx!(jΔx, Δx) = isnothing(jacobian_free_jvp) ? mul!(jΔx, j, Δx) :

src/sparse_containers.jl

@@ -25,14 +25,15 @@ struct SparseContainer{SIM, T}
     function SparseContainer(
             compressed_data::T,
             sparse_index_map::Tuple
-        ) where {N, ET, T <: NTuple{N, ET}}
+        ) where {T}
         @assert all(map(x-> eltype(compressed_data) .== typeof(x), compressed_data))
         return new{sparse_index_map, T}(compressed_data)
     end
 end
 
 Base.parent(sc::SparseContainer) = sc.data
 sc_eltype(::Type{NTuple{N, T}}) where {N, T} = T
+sc_eltype(::Type{T}) where {ET, T <: AbstractArray{ET}} = ET
 sc_eltype(::SparseContainer{SIM, T}) where {SIM, T} = sc_eltype(T)
 @inline function Base.getindex(sc::SparseContainer, i::Int)
     return _getindex_sparse(sc, Val(i))::sc_eltype(sc)

test/sparse_containers.jl

@@ -19,4 +19,24 @@ using Test
 
     @test_throws ErrorException("No index 2 found in sparse index map (1, 3, 5, 7)") v[2]
     @test_throws ErrorException("No index 8 found in sparse index map (1, 3, 5, 7)") v[8]
+    @inferred v[7]
+
+    a1 = ones(3) .* 1
+    a2 = ones(3) .* 2
+    a3 = ones(3) .* 3
+    a4 = ones(3) .* 4
+    v = SparseContainer([a1,a2,a3,a4], (1,3,5,7))
+    @test v[1] == ones(3) .* 1
+    @test v[3] == ones(3) .* 2
+    @test v[5] == ones(3) .* 3
+    @test v[7] == ones(3) .* 4
+
+    @test parent(v)[1] == ones(3) .* 1
+    @test parent(v)[2] == ones(3) .* 2
+    @test parent(v)[3] == ones(3) .* 3
+    @test parent(v)[4] == ones(3) .* 4
+
+    @test_throws ErrorException("No index 2 found in sparse index map (1, 3, 5, 7)") v[2]
+    @test_throws ErrorException("No index 8 found in sparse index map (1, 3, 5, 7)") v[8]
+    @inferred v[7]
 end