Add rules for BigFloat integer constructors

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))

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