diff --git a/src/dual.jl b/src/dual.jl index 6d13dec3..48d5d0e3 100644 --- a/src/dual.jl +++ b/src/dual.jl @@ -625,7 +625,7 @@ end @generated function calc_fma_xyz(x::Dual{T,<:Any,N}, y::Dual{T,<:Any,N}, z::Dual{T,<:Any,N}) where {T,N} - ex = Expr(:tuple, [:(fma(value(x), partials(y)[$i], fma(value(y), partials(x)[$i], partials(z)[$i]))) for i in 1:N]...) + ex = Expr(:tuple, [:(_fma_partial(value(x), partials(y, $i), _fma_partial(value(y), partials(x, $i), partials(z, $i)))) for i in 1:N]...) return quote $(Expr(:meta, :inline)) v = fma(value(x), value(y), value(z)) @@ -642,7 +642,7 @@ end @generated function calc_fma_xz(x::Dual{T,<:Any,N}, y::Real, z::Dual{T,<:Any,N}) where {T,N} - ex = Expr(:tuple, [:(fma(partials(x)[$i], y, partials(z)[$i])) for i in 1:N]...) + ex = Expr(:tuple, [:(_fma_partial(y, partials(x, $i), partials(z, $i))) for i in 1:N]...) return quote $(Expr(:meta, :inline)) v = fma(value(x), y, value(z)) @@ -667,7 +667,7 @@ end @generated function calc_muladd_xyz(x::Dual{T,<:Any,N}, y::Dual{T,<:Any,N}, z::Dual{T,<:Any,N}) where {T,N} - ex = Expr(:tuple, [:(muladd(value(x), partials(y)[$i], muladd(value(y), partials(x)[$i], partials(z)[$i]))) for i in 1:N]...) + ex = Expr(:tuple, [:(_muladd_partial(value(x), partials(y, $i), _muladd_partial(value(y), partials(x, $i), partials(z, $i)))) for i in 1:N]...) return quote $(Expr(:meta, :inline)) v = muladd(value(x), value(y), value(z)) @@ -684,7 +684,7 @@ end @generated function calc_muladd_xz(x::Dual{T,<:Any,N}, y::Real, z::Dual{T,<:Any,N}) where {T,N} - ex = Expr(:tuple, [:(muladd(partials(x)[$i], y, partials(z)[$i])) for i in 1:N]...) + ex = Expr(:tuple, [:(_muladd_partial(y, partials(x, $i), partials(z, $i))) for i in 1:N]...) return quote $(Expr(:meta, :inline)) v = muladd(value(x), y, value(z)) diff --git a/src/partials.jl b/src/partials.jl index 49f8364b..66a29b26 100644 --- a/src/partials.jl +++ b/src/partials.jl @@ -117,9 +117,19 @@ if NANSAFE_MODE_ENABLED y = partial / x return iszero(partial) ? zero(y) : y end + @inline function _fma_partial(x::Real, partial::Real, acc::Real) + y = fma(x, partial, acc) + return iszero(partial) ? oftype(y, acc) : y + end + @inline function _muladd_partial(x::Real, partial::Real, acc::Real) + y = muladd(x, partial, acc) + return iszero(partial) ? oftype(y, acc) : y + end else @inline _mul_partial(partial::Real, x::Real) = partial * x @inline _div_partial(partial::Real, x::Real) = partial / x + @inline _fma_partial(x::Real, partial::Real, acc::Real) = fma(x, partial, acc) + @inline _muladd_partial(x::Real, partial::Real, acc::Real) = muladd(x, partial, acc) end # edge cases where N == 0 # diff --git a/test/DualTest.jl b/test/DualTest.jl index 67d9c9f7..e7429be5 100644 --- a/test/DualTest.jl +++ b/test/DualTest.jl @@ -761,4 +761,18 @@ end end end +@testset "fma and muladd: consistency with `x * y + z`" begin # issue #860 + for f in (fma, muladd) + for (x, y, z) in ( + (Dual{TestTag}(Inf, 0.0), Dual{TestTag}(2.0, 0.0), Dual{TestTag}(0.0, 1.0)), + (Dual{TestTag}(1.0, 0.0), Inf, Dual{TestTag}(0.0, 1.0)), + (Inf, Dual{TestTag}(1.0, 0.0), Dual{TestTag}(0.0, 1.0)), + ) + d = f(x, y, z) + @test value(d) === Inf + @test isequal(d, x * y + z) + end + end +end + end # module