Typed rounding (#619)

* Typed rounding

* fixup

* reorder

* don't throw err

* fixup

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix parsing

* fixtest

* testfix

* Update src/utils.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

src/TracedRNumber.jl

@@ -289,6 +289,57 @@ Base.round(A::TracedRNumber{<:ReactantFloat}) = Ops.round_nearest_even(A)
 Base.floor(A::TracedRNumber{<:ReactantFloat}) = Ops.floor(A)
 Base.ceil(A::TracedRNumber{<:ReactantFloat}) = Ops.ceil(A)
 
+function Base.unsafe_trunc(
+    T::Type{<:Reactant.ReactantInt}, x::TracedRNumber{<:Reactant.ReactantFloat}
+)
+    return Ops.convert(TracedRNumber{T}, x)
+end
+
+for Ti in (Int8, Int16, Int32, Int64, Int128, UInt8, UInt16, UInt32, UInt64, UInt128)
+    for Tf in (Float16, Float32, Float64)
+        if Ti <: Unsigned || sizeof(Ti) < sizeof(Tf)
+            # Here `Tf(typemin(Ti))-1` is exact, so we can compare the lower-bound
+            # directly. `Tf(typemax(Ti))+1` is either always exactly representable, or
+            # rounded to `Inf` (e.g. when `Ti==UInt128 && Tf==Float32`).
+            @eval begin
+                function Base.trunc(::Type{$Ti}, x::TracedRNumber{$Tf})
+                    # TODO throw error within traced
+                    # if $(Tf(typemin(Ti))-one(Tf)) < x < $(Tf(typemax(Ti))+one(Tf))
+                    return Base.unsafe_trunc($Ti, x)
+                    # else
+                    #     throw(Base.InexactError(:trunc, $Ti, x))
+                    # end
+                end
+            end
+        else
+            # Here `eps(Tf(typemin(Ti))) > 1`, so the only value which can be truncated to
+            # `Tf(typemin(Ti)` is itself. Similarly, `Tf(typemax(Ti))` is inexact and will
+            # be rounded up. This assumes that `Tf(typemin(Ti)) > -Inf`, which is true for
+            # these types, but not for `Float16` or larger integer types.
+            @eval begin
+                function Base.trunc(::Type{$Ti}, x::TracedRNumber{$Tf})
+                    # TODO throw error within traced
+                    # if $(Tf(typemin(Ti))) <= x < $(Tf(typemax(Ti)))
+                    return Base.unsafe_trunc($Ti, x)
+                    # else
+                    #     throw(Base.InexactError(:trunc, $Ti, x))
+                    # end
+                end
+            end
+        end
+    end
+end
+
+function Base.round(::Type{T}, x::TracedRNumber{<:AbstractFloat}) where {T<:Integer}
+    return trunc(T, Base.round(x))
+end
+function Base.floor(::Type{T}, x::TracedRNumber{<:AbstractFloat}) where {T<:Integer}
+    return trunc(T, Base.floor(x))
+end
+function Base.ceil(::Type{T}, x::TracedRNumber{<:AbstractFloat}) where {T<:Integer}
+    return trunc(T, Base.ceil(x))
+end
+
 # Concatenation. Numbers in Julia are handled in a much less generic fashion than arrays
 Base.vcat(x::TracedRNumber...) = Base.typed_vcat(Base.promote_eltypeof(x...), x...)
 function Base.typed_vcat(::Type{T}, x::TracedRNumber...) where {T}

src/utils.jl

@@ -1,5 +1,5 @@
 
-function apply(f, args...; kwargs...)
+function apply(f::F, args...; kwargs...) where {F}
     return f(args...; kwargs...)
 end
 

test/basic.jl

@@ -588,8 +588,13 @@ end
 end
 
 @testset "$op" for op in [:round, :ceil, :floor]
+    intop = Symbol("int_$op")
     for x in (rand(Float32, (3, 3)), rand(Float64))
-        @eval @test @jit($op.(ConcreteRNumber.($x))) == $op.($x)
+        @eval begin
+            @test @jit($op.(ConcreteRNumber.($x))) == $op.($x)
+            $intop(x) = $op(Int, x)
+            @test @jit($intop.(ConcreteRNumber.($x))) == $intop.($x)
+        end
     end
 end