Fix JET analysis (#502) (#503)

* Fix JET analysis (#502)

* Add regression tests

* Run tests also on prerelease version of Julia

* Update version number

* Add missing newline character

Co-authored-by: Mosè Giordano <765740+giordano@users.noreply.github.com>

---------

Co-authored-by: Mosè Giordano <765740+giordano@users.noreply.github.com>

Author: devmotion

.github/workflows/ci.yml

@@ -17,6 +17,7 @@ jobs:
           - 'min'
           - 'lts'
           - '1'
+          - 'pre'
         os:
           - ubuntu-latest
           - macOS-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "SpecialFunctions"
 uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
-version = "2.6.0"
+version = "2.6.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -19,6 +19,7 @@ SpecialFunctionsChainRulesCoreExt = "ChainRulesCore"
 ChainRulesCore = "0.9.44, 0.10, 1"
 ChainRulesTestUtils = "0.6.8, 0.7, 1"
 IrrationalConstants = "0.1, 0.2"
+JET = "0.9, 0.10"
 LogExpFunctions = "0.3.2"
 OpenLibm_jll = "0.7, 0.8"
 OpenSpecFun_jll = "0.5"
@@ -27,8 +28,9 @@ julia = "1.10"
 [extras]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ChainRulesCore", "ChainRulesTestUtils", "Random", "Test"]
+test = ["ChainRulesCore", "ChainRulesTestUtils", "JET", "Random", "Test"]

src/logabsgamma/e_lgamma_r.jl

@@ -92,32 +92,37 @@ function _logabsgamma(x::Float64)
     lx = ux % UInt32
 
     #= purge off +-inf, NaN, +-0, tiny and negative arguments =#
-    signgam = 1
-    isneg = hx < Int32(0)
     ix = hx & 0x7fffffff
-    ix ≥ 0x7ff00000 && return x * x, signgam
-    ix | lx == 0x00000000 && return Inf, signgam
+    ix ≥ 0x7ff00000 && return x * x, 1
+    ix | lx == 0x00000000 && return Inf, 1
+    isneg = hx < Int32(0)
     if ix < 0x3b900000 #= |x|<2**-70, return -log(|x|) =#
         if isneg
-            signgam = -1
-            return -log(-x), signgam
+            return -log(-x), -1
         else
-            return -log(x), signgam
+            return -log(x), 1
         end
     end
+
+    # remaining cases
     if isneg
         # ix ≥ 0x43300000 && return Inf, signgam #= |x|>=2**52, must be -integer =#
         t = sinpi(x)
-        iszero(t) && return Inf, signgam #= -integer =#
-        nadj = logπ - log(abs(t * x))
-        if t < 0.0; signgam = -1; end
-        x = -x
+        iszero(t) && return Inf, 1 #= -integer =#
+        r = logπ - log(abs(t * x)) - _logabsgamma_pos(-x, ix)
+        signgam = t < 0.0 ? -1 : 1
+    else
+        r = _logabsgamma_pos(x, ix)
+        signgam = 1
     end
+    return r, signgam
+end
+function _logabsgamma_pos(x::Float64, ix::UInt32)
     if ix < 0x40000000     #= x < 2.0 =#
         i = round(x, RoundToZero)
         f = x - i
         if f == 0.0 #= purge off 1; 2 handled by x < 8.0 branch =#
-            return 0.0, signgam
+            return 0.0
         elseif i == 1.0
             r = 0.0
             c = 1.0
@@ -169,8 +174,5 @@ function _logabsgamma(x::Float64)
     else #= 2^58 ≤ x ≤ Inf =#
         r = muladd(x, log(x), -x)
     end
-    if isneg
-        r = nadj - r
-    end
-    return r, signgam
+    return r
 end

src/logabsgamma/e_lgammaf_r.jl

@@ -20,33 +20,37 @@ function _logabsgamma(x::Float32)
     hx = reinterpret(Int32, x)
 
     #= purge off +-inf, NaN, +-0, tiny and negative arguments =#
-    signgam = 1
-    isneg = hx < Int32(0)
     ix = hx & 0x7fffffff
-    ix ≥ 0x7f800000 && return x * x, signgam
-    ix == 0x00000000 && return Inf32, signgam
+    ix ≥ 0x7f800000 && return x * x, 1
+    ix == 0x00000000 && return Inf32, 1
+    isneg = hx < Int32(0)
     if ix < 0x35000000 #= |x|<2**-21, return -log(|x|) =#
         if isneg
-            signgam = -1
-            return -log(-x), signgam
+            return -log(-x), -1
         else
-            return -log(x), signgam
+            return -log(x), 1
         end
     end
+
+    # remaining cases
     if isneg
         # ix ≥ 0x4b000000 && return Inf32, signgam #= |x|>=2**23, must be -integer =#
         t = sinpi(x)
-        t == 0.0f0 && return Inf32, signgam #= -integer =#
-        nadj = logπ - log(abs(t * x))
-        if t < 0.0f0; signgam = -1; end
-        x = -x
+        t == 0.0f0 && return Inf32, 1 #= -integer =#
+        r = logπ - log(abs(t * x)) - _logabsgamma_pos(-x, ix)
+        signgam = copysign(1, t)
+    else
+        r = _logabsgamma_pos(x, ix)
+        signgam = 1
     end
-
+    return r, signgam
+end
+function _logabsgamma_pos(x::Float32, ix::UInt32)
     if ix < 0x40000000 #= x < 2.0 =#
         i = round(x, RoundToZero)
         f = x - i
         if f == 0.0f0 #= purge off 1; 2 handled by x < 8.0 branch =#
-            return 0.0f0, signgam
+            return 0.0f0
         elseif i == 1.0f0
             r = 0.0f0
             c = 1.0f0
@@ -99,8 +103,5 @@ function _logabsgamma(x::Float32)
         #= 2^58 ≤ x ≤ Inf =#
         r = muladd(x, log(x), -x)
     end
-    if isneg
-        r = nadj - r
-    end
-    return r, signgam
+    return r
 end

test/logabsgamma.jl

@@ -156,3 +156,11 @@ x = 8.000001f0
 # (i.e. prevfloat(8.000001f0) == 8.0f0)
 # We still check appropriate behavior at 8.0f0
 @test ulp(8.0f0) < 0.4006594736129046
+
+@testset "JET" begin
+    # issue #502
+    JET.@test_call logabsgamma(1.0)
+    JET.@test_opt logabsgamma(1.0)
+    JET.@test_call logabsgamma(1f0)
+    JET.@test_opt logabsgamma(1f0)
+end

test/runtests.jl

@@ -7,6 +7,8 @@ using Random
 using Test
 using Base.MathConstants: γ
 
+using JET: JET
+
 using SpecialFunctions: AmosException, f64
 
 # useful test functions for relative error, which differ from isapprox