Static GPU compilation of Jacobian

SciML · Feb 15, 2024 · 05c6206 · 05c6206
1 parent f809d0a
commit 05c6206
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/lib/SimpleNonlinearSolve/.buildkite/pipeline.yml b/lib/SimpleNonlinearSolve/.buildkite/pipeline.yml
@@ -13,7 +13,7 @@ steps:
     agents:
       queue: "juliagpu"
       cuda: "*"
-    timeout_in_minutes: 30
+    timeout_in_minutes: 120
     # Don't run Buildkite if the commit message includes the text [skip tests]
     if: build.message !~ /\[skip tests\]/
 

diff --git a/lib/SimpleNonlinearSolve/Project.toml b/lib/SimpleNonlinearSolve/Project.toml
@@ -1,7 +1,7 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.4.2"
+version = "1.4.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/lib/SimpleNonlinearSolve/src/utils.jl b/lib/SimpleNonlinearSolve/src/utils.jl
@@ -40,17 +40,28 @@ function __pick_forwarddiff_chunk(x::StaticArray)
     end
 end
 
-function __get_jacobian_config(ad::AutoForwardDiff{CS}, f, x) where {CS}
+function __get_jacobian_config(ad::AutoForwardDiff{CS}, f::F, x) where {F, CS}
     ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}()
     tag = __standard_tag(ad.tag, x)
-    return ForwardDiff.JacobianConfig(f, x, ck, tag)
+    return __forwarddiff_jacobian_config(f, x, ck, tag)
 end
-function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!, y, x) where {CS}
+function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!::F, y, x) where {F, CS}
     ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}()
     tag = __standard_tag(ad.tag, x)
     return ForwardDiff.JacobianConfig(f!, y, x, ck, tag)
 end
 
+function __forwarddiff_jacobian_config(f::F, x, ck::ForwardDiff.Chunk, tag) where {F}
+    return ForwardDiff.JacobianConfig(f, x, ck, tag)
+end
+function __forwarddiff_jacobian_config(
+        f::F, x::SArray, ck::ForwardDiff.Chunk{N}, tag) where {F, N}
+    seeds = ForwardDiff.construct_seeds(ForwardDiff.Partials{N, eltype(x)})
+    duals = ForwardDiff.Dual{typeof(tag), eltype(x), N}.(x)
+    return ForwardDiff.JacobianConfig{typeof(tag), eltype(x), N, typeof(duals)}(seeds,
+        duals)
+end
+
 function __get_jacobian_config(ad::AutoPolyesterForwardDiff{CS}, args...) where {CS}
     x = last(args)
     return (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) :