diff --git a/src/internal_rules/bigfloat.jl b/src/internal_rules/bigfloat.jl index b534b80f6e..d013b8a1f0 100644 --- a/src/internal_rules/bigfloat.jl +++ b/src/internal_rules/bigfloat.jl @@ -67,6 +67,88 @@ function EnzymeRules.reverse( return () end +# `BigFloat(x::Integer)` lowers to `mpfr_set_si` / `mpfr_set_ui` (for `Clong` / `Culong`) +# or `mpfr_set_z` (via `BigInt`), for which there is no derivative, so any active use +# of an integer-constructed BigFloat otherwise fails. Note `Clong === Int32` on Windows, +# so `BigFloat(::Int64)` takes the `BigInt` path there. +_bigfloat_rounding_val(r::Const{Base.MPFR.MPFRRoundingMode}) = r.val + +# Callable struct constructing `BigFloat(x, rs...; kwargs...)`; each call allocates a +# fresh BigFloat, so batched shadows built with `ntuple` don't alias one another. +struct BigFloatIntCtor{T <: Integer, R <: Tuple, K} + x::T + rs::R + kwargs::K +end +@inline (f::BigFloatIntCtor)(_...) = BigFloat(f.x, f.rs...; f.kwargs...) + +function EnzymeRules.forward( + config::EnzymeRules.FwdConfig, + Ty::Const{Type{BigFloat}}, + RT::Type{<:Union{DuplicatedNoNeed, Duplicated, BatchDuplicated, BatchDuplicatedNoNeed}}, + x::Const{Ti}, + rs::Const{Base.MPFR.MPFRRoundingMode}...; + kwargs..., + ) where {Ti <: Integer} + rvals = map(_bigfloat_rounding_val, rs) + primal = BigFloatIntCtor(x.val, rvals, kwargs) + shadow = BigFloatIntCtor(zero(Ti), rvals, kwargs) + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + return remove_innerty(RT)(primal(), shadow()) + else + return BatchDuplicated(primal(), ntuple(shadow, Val(EnzymeRules.width(config)))) + end + elseif EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + return shadow() + else + return ntuple(shadow, Val(EnzymeRules.width(config))) + end + elseif EnzymeRules.needs_primal(config) + return primal() + else + return nothing + end +end + +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfig, + Ty::Const{Type{BigFloat}}, + RT::Type{<:Union{DuplicatedNoNeed, Duplicated, BatchDuplicated, BatchDuplicatedNoNeed}}, + x::Const{Ti}, + rs::Const{Base.MPFR.MPFRRoundingMode}...; + kwargs..., + ) where {Ti <: Integer} + rvals = map(_bigfloat_rounding_val, rs) + primal = if EnzymeRules.needs_primal(config) + BigFloatIntCtor(x.val, rvals, kwargs)() + else + nothing + end + shadow = if RT <: Const + nothing + elseif EnzymeRules.width(config) == 1 + BigFloatIntCtor(zero(Ti), rvals, kwargs)() + else + ntuple(BigFloatIntCtor(zero(Ti), rvals, kwargs), Val(EnzymeRules.width(config))) + end + return EnzymeRules.AugmentedReturn(primal, shadow, nothing) +end + +function EnzymeRules.reverse( + config::EnzymeRules.RevConfig, + Ty::Const{Type{BigFloat}}, + RT::Type{<:Union{DuplicatedNoNeed, Duplicated, BatchDuplicated, BatchDuplicatedNoNeed}}, + tape, + x::Const{<:Integer}, + rs::Const{Base.MPFR.MPFRRoundingMode}...; + kwargs..., + ) + # the integer argument carries no derivative + return ntuple(Returns(nothing), Val(1 + length(rs))) +end + EnzymeRules.@easy_rule(+(a::BigFloat, b::Number), (1,1)) EnzymeRules.@easy_rule(+(a::Number, b::BigFloat), (1,1)) EnzymeRules.@easy_rule(+(a::BigFloat, b::BigFloat), (1,1)) diff --git a/test/rules/internal_rules/bigfloat.jl b/test/rules/internal_rules/bigfloat.jl index c287cbe100..d4f996b29b 100644 --- a/test/rules/internal_rules/bigfloat.jl +++ b/test/rules/internal_rules/bigfloat.jl @@ -31,3 +31,43 @@ using Test @test autodiff(Enzyme.Forward, cos, Duplicated, Duplicated(a, da))[:1] ≈ -sin(a) * da @test autodiff(Enzyme.Forward, tan, Duplicated, Duplicated(a, da))[:1] ≈ autodiff(Enzyme.Forward, tan, Duplicated, Duplicated(af64, daf64))[1] end + +@testset "BigFloat integer constructors" begin + a = BigFloat(1.234) + da = BigFloat(-0.23) + + # BigFloat(::Clong) / BigFloat(::Culong) construct a constant, so the shadow is zero + @test autodiff(Enzyme.Forward, BigFloat, Duplicated, Const(2))[:1] ≈ 0 + @test autodiff(Enzyme.Forward, BigFloat, Duplicated, Const(UInt(2)))[:1] ≈ 0 + + # ... but the primal value must still be right, and propagate through arithmetic + mul2(x) = x * BigFloat(2) + @test autodiff(Enzyme.Forward, mul2, Duplicated, Duplicated(a, da))[:1] ≈ 2 * da + umul2(x) = x * BigFloat(UInt(2)) + @test autodiff(Enzyme.Forward, umul2, Duplicated, Duplicated(a, da))[:1] ≈ 2 * da + + # on Windows `Clong === Int32`, so `BigFloat(::Int64)` takes the generic + # `Integer -> BigInt` path; `Int128` and `BigInt` exercise it on every platform + @test autodiff(Enzyme.Forward, BigFloat, Duplicated, Const(Int128(2)))[:1] ≈ 0 + @test autodiff(Enzyme.Forward, BigFloat, Duplicated, Const(big(2)))[:1] ≈ 0 + i128mul2(x) = x * BigFloat(Int128(2)) + @test autodiff(Enzyme.Forward, i128mul2, Duplicated, Duplicated(a, da))[:1] ≈ 2 * da + dres = autodiff(Enzyme.Forward, i128mul2, BatchDuplicated, BatchDuplicated(a, (da, 2da)))[1] + @test dres[1] ≈ 2 * da + @test dres[2] ≈ 4 * da + + # the constructor also has to keep working under a non-default precision + setprecision(BigFloat, 512) do + @test precision(BigFloat(3)) == 512 + mul3(x) = x * BigFloat(3) + @test autodiff(Enzyme.Forward, mul3, Duplicated, Duplicated(a, da))[:1] ≈ 3 * da + end + + # `eps(::Type{BigFloat})` is `nextfloat(BigFloat(1)) - BigFloat(1)`, so it goes + # through the integer constructor; without a rule this fails to compile. + # NOTE: using the result of `eps(BigFloat)` in an *active* position (e.g. + # `x * eps(BigFloat)`) still segfaults -- Enzyme miscompiles the BigFloat + # construction inside `nextfloat`/`_duplicate` when it has to differentiate it. + epsconst(x) = x * x + zero(eps(BigFloat)) + @test autodiff(Enzyme.Forward, epsconst, Duplicated, Duplicated(a, da))[:1] ≈ 2 * a * da +end