Generic _axpy!

SciML · Oct 24, 2023 · 02cf007 · 02cf007
1 parent 30f7873
commit 02cf007
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 4 deletions.
diff --git a/src/broyden.jl b/src/broyden.jl
@@ -78,7 +78,7 @@ function perform_step!(cache::GeneralBroydenCache{true})
 
     mul!(_vec(du), J⁻¹, -_vec(fu))
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(fu2, u, p)
 
     cache.internalnorm(fu2) < cache.abstol && (cache.force_stop = true)

diff --git a/src/klement.jl b/src/klement.jl
@@ -115,7 +115,7 @@ function perform_step!(cache::GeneralKlementCache{true})
 
     # Line Search
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(cache.fu2, u, p)
 
     cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true)

diff --git a/src/lbroyden.jl b/src/lbroyden.jl
@@ -93,7 +93,7 @@ function perform_step!(cache::LimitedMemoryBroydenCache{true})
     T = eltype(u)
 
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(cache.fu2, u, p)
 
     cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true)

diff --git a/src/raphson.jl b/src/raphson.jl
@@ -97,7 +97,7 @@ function perform_step!(cache::NewtonRaphsonCache{true})
 
     # Line Search
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(-α, du, u)
+    _axpy!(-α, du, u)
     f(cache.fu1, u, p)
 
     cache.internalnorm(fu1) < cache.abstol && (cache.force_stop = true)

diff --git a/src/utils.jl b/src/utils.jl
@@ -260,6 +260,11 @@ _try_factorize_and_check_singular!(::Nothing, x) = _issingular(x), false
 _reshape(x, args...) = reshape(x, args...)
 _reshape(x::Number, args...) = x
 
+@generated function _axpy!(α, x, y)
+    hasmethod(axpy!, Tuple{typeof(α), typeof(x), typeof(y)}) && return :(axpy!(α, x, y))
+    return :(y .= α * x)
+end
+
 # Needs Square Matrix
 """
     needs_square_A(alg)