Merge #180

180: Use ClimaComms to determine array type r=charleskawczynski a=charleskawczynski

This PR changes from using `Krylov.ktypeof` to `ClimaComms.array_type` to determine the array type passed to Krylov.jl

Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.7.4"
+version = "0.7.5"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -18,7 +18,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 CUDA = "3, 4"
-ClimaComms = "0.3, 0.4, 0.5"
+ClimaComms = "0.4, 0.5"
 DataStructures = "0.18"
 DiffEqBase = "6"
 DiffEqCallbacks = "2"

docs/Manifest.toml

@@ -190,9 +190,9 @@ version = "0.4.2"
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DiffEqBase", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "Rotations", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"]
-git-tree-sha1 = "ec80e55b9a3b9e8fa99e75108923ff9b863480f3"
+git-tree-sha1 = "e8cfd4aeb2099d55331677f37f2447cffa4bf0f4"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.10.33"
+version = "0.10.34"
 
 [[deps.ClimaCorePlots]]
 deps = ["ClimaCore", "RecipesBase", "StaticArrays", "TriplotBase"]
@@ -633,9 +633,9 @@ version = "1.12.2+2"
 
 [[deps.HTTP]]
 deps = ["Base64", "CodecZlib", "ConcurrentUtilities", "Dates", "Logging", "LoggingExtras", "MbedTLS", "NetworkOptions", "OpenSSL", "Random", "SimpleBufferStream", "Sockets", "URIs", "UUIDs"]
-git-tree-sha1 = "41f7dfb2b20e7e8bf64f6b6fae98f4d2df027b06"
+git-tree-sha1 = "ba9eca9f8bdb787c6f3cf52cb4a404c0e349a0d1"
 uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
-version = "1.9.4"
+version = "1.9.5"
 
 [[deps.HarfBuzz_jll]]
 deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "Graphite2_jll", "JLLWrappers", "Libdl", "Libffi_jll", "Pkg"]

perf/Manifest.toml

@@ -168,9 +168,9 @@ version = "0.4.2"
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DiffEqBase", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "Rotations", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"]
-git-tree-sha1 = "ec80e55b9a3b9e8fa99e75108923ff9b863480f3"
+git-tree-sha1 = "e8cfd4aeb2099d55331677f37f2447cffa4bf0f4"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.10.33"
+version = "0.10.34"
 
 [[deps.ClimaTimeSteppers]]
 deps = ["CUDA", "ClimaComms", "DataStructures", "DiffEqBase", "DiffEqCallbacks", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "SciMLBase", "StaticArrays"]

src/ClimaTimeSteppers.jl

@@ -50,6 +50,7 @@ using LinearAlgebra
 using LinearOperators
 using StaticArrays
 using CUDA
+import ClimaComms
 
 export AbstractAlgorithmName, AbstractAlgorithmConstraint, Unconstrained, SSP
 

src/nl_solvers/newtons_method.jl

@@ -415,6 +415,12 @@ end
 
 solver_type(::KrylovMethod{Val{T}}) where {T} = T
 
+# TODO: should we make ClimaCore depend on Krylov.jl and overload
+# Krylov.ktypeof for `FieldVector`s?
+ct_ktypeof(x) = #= x::ClimaCore.Fields.FieldVector =#
+    ClimaComms.array_type(x){eltype(parent(x))}
+ct_ktypeof(x::Vector) = Krylov.ktypeof(x)
+
 function allocate_cache(alg::KrylovMethod, x_prototype)
     (; jacobian_free_jvp, forcing_term, args, kwargs, debugger) = alg
     type = solver_type(alg)
@@ -423,7 +429,7 @@ function allocate_cache(alg::KrylovMethod, x_prototype)
         jacobian_free_jvp_cache = isnothing(jacobian_free_jvp) ? nothing :
                                   allocate_cache(jacobian_free_jvp, x_prototype),
         forcing_term_cache = allocate_cache(forcing_term, x_prototype),
-        solver = type(l, l, args..., Krylov.ktypeof(x_prototype); kwargs...),
+        solver = type(l, l, args..., ct_ktypeof(x_prototype); kwargs...),
         debugger_cache = isnothing(debugger) ? nothing : allocate_cache(debugger, x_prototype),
     )
 end

test/Project.toml

@@ -25,7 +25,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [compat]
 Aqua = "0.6"
 CUDA = "3, 4"
-ClimaComms = "0.3, 0.4, 0.5"
+ClimaComms = "0.4, 0.5"
 ClimaCore = "0.10"
 DataStructures = "0.18"
 DiffEqBase = "6"