TracedUnitRange (#1004)

* copy-paste definitions from range.jl

* add tracing methods

* fixes

* test indexing with Julia integer

* fixes and test for tracedrnumber getindex

* formatting

* only convert tracedranges if the type eltype actually changes

* Define `unitrange_last` with `ifelse`

* Define rounding method for integers

* fix creation of traced values

* disambiguate getindex

* RNumber should've been ReactantPrimitive

* `Adapt.parent_type`

* fix

* same fix for TraceStepRangeLen

* more getindex

* fix for 1.10

* Formatting

---------

Co-authored-by: Mosè Giordano <[email protected]>
Co-authored-by: William S. Moses <[email protected]>

Author: 3 people

src/TracedRNumber.jl

@@ -437,11 +437,14 @@ function Base.float(x::TracedRNumber{T}) where {T}
     return TracedUtils.promote_to(TracedRNumber{float(T)}, x)
 end
 
-using Reactant: ReactantFloat
+using Reactant: ReactantFloat, ReactantInt
 
 Base.round(A::TracedRNumber{<:ReactantFloat}) = Ops.round_nearest_even(A)
+Base.round(A::TracedRNumber{<:ReactantInt}) = A
 Base.floor(A::TracedRNumber{<:ReactantFloat}) = Ops.floor(A)
+Base.floor(A::TracedRNumber{<:ReactantInt}) = A
 Base.ceil(A::TracedRNumber{<:ReactantFloat}) = Ops.ceil(A)
+Base.ceil(A::TracedRNumber{<:ReactantInt}) = A
 
 function Base.unsafe_trunc(
     T::Type{<:Reactant.ReactantInt}, x::TracedRNumber{<:Reactant.ReactantFloat}
@@ -499,6 +502,81 @@ function Base.getindex(
     return Base.unsafe_getindex(r, i)
 end
 
+function unitrange_last(start::Integer, stop::Integer)
+    return ifelse(stop >= start, stop, convert(typeof(stop), start - oneunit(start - stop)))
+end
+function unitrange_last(start, stop)
+    return ifelse(
+        stop >= start,
+        convert(typeof(stop), start + floor(stop - start)),
+        convert(typeof(stop), start - oneunit(start - stop)),
+    )
+end
+
+struct TracedUnitRange{T} <: AbstractUnitRange{T}
+    start::T
+    stop::T
+    function TracedUnitRange{T}(start::T, stop::T) where {T}
+        return new(start, unitrange_last(start, stop))
+    end
+end
+function Adapt.parent_type(::Type{TracedUnitRange{T}}) where {T}
+    return TracedUnitRange{T}
+end
+function TracedUnitRange{T}(start, stop) where {T}
+    return TracedUnitRange{T}(convert(T, start), convert(T, stop))
+end
+TracedUnitRange(start::T, stop::T) where {T} = TracedUnitRange{T}(start, stop)
+function TracedUnitRange(start, stop)
+    startstop_promoted = promote(start, stop)
+    not_sametype((start, stop), startstop_promoted)
+    return TracedUnitRange(startstop_promoted...)
+end
+function Base._in_unit_range(
+    v::TracedUnitRange, val, i::Union{Integer,TracedRNumber{<:Integer}}
+)
+    return (i > 0) & (val <= v.stop) & (val >= v.start)
+end
+
+function _traced_unitrange_getindex(v::TracedUnitRange{T}, i) where {T}
+    val = convert(T, v.start + (i - oneunit(i)))
+    # TODO: we should have error messages at some point.
+    # @boundscheck Base._in_unit_range(v, val, i) || throw_boundserror(v, i)
+    return val
+end
+
+function Base._getindex(v::TracedUnitRange, i::TracedRNumber{<:Integer})
+    return _traced_unitrange_getindex(v, i)
+end
+Base.getindex(v::TracedUnitRange, i::Integer) = _traced_unitrange_getindex(v, i)
+Base.getindex(r::TracedUnitRange, i::TracedRNumber) = Base._getindex(r, i)
+function Base.getindex(r::Base.UnitRange, i::I) where {I<:TracedRNumber{<:Integer}}
+    val = convert(I, r.start + (i - oneunit(i)))
+    # TODO: we should have error messages at some point.
+    # @boundscheck Base._in_unit_range(v, val, i) || throw_boundserror(v, i)
+    return val
+end
+
+function Base.promote_rule(
+    a::Type{TracedUnitRange{T1}}, b::Type{TracedUnitRange{T2}}
+) where {T1,T2}
+    return el_same(promote_type(T1, T2), a, b)
+end
+TracedUnitRange{T}(r::TracedUnitRange{T}) where {T<:Real} = r
+TracedUnitRange{T}(r::TracedUnitRange) where {T<:Real} = TracedUnitRange{T}(r.start, r.stop)
+
+function Base.promote_rule(
+    a::Type{TracedUnitRange{T1}}, ::Type{UR}
+) where {T1,UR<:AbstractUnitRange}
+    return promote_rule(a, TracedUnitRange{eltype(UR)})
+end
+function TracedUnitRange{T}(r::AbstractUnitRange) where {T<:Real}
+    return TracedUnitRange{T}(first(r), last(r))
+end
+TracedUnitRange(r::AbstractUnitRange) = TracedUnitRange(first(r), last(r))
+
+AbstractUnitRange{T}(r::TracedUnitRange) where {T} = TracedUnitRange{T}(r)
+
 struct TracedStepRangeLen{T,R,S,L} <: AbstractRange{T}
     ref::R
     step::S

src/TracedUtils.jl

@@ -312,23 +312,29 @@ function make_mlir_fn(
 
     seen_results = OrderedIdDict()
 
-    traced_result = Reactant.make_tracer(
-        seen_results,
-        result,
-        (:result,),
-        concretein ? Reactant.NoStopTracedTrack : Reactant.TracedSetPath;
-        runtime,
-    )
-
-    # marks buffers to be donated
-    for i in 1:N
-        Reactant.make_tracer(
+    MLIR.IR.activate!(fnbody)
+    traced_result = try
+        traced_result = Reactant.make_tracer(
             seen_results,
-            traced_args[i],
-            concretein ? (:resargs, i) : (),
-            Reactant.NoStopTracedTrack;
+            result,
+            (:result,),
+            concretein ? Reactant.NoStopTracedTrack : Reactant.TracedSetPath;
             runtime,
         )
+
+        # marks buffers to be donated
+        for i in 1:N
+            Reactant.make_tracer(
+                seen_results,
+                traced_args[i],
+                concretein ? (:resargs, i) : (),
+                Reactant.NoStopTracedTrack;
+                runtime,
+            )
+        end
+        traced_result
+    finally
+        MLIR.IR.deactivate!(fnbody)
     end
 
     linear_results = Reactant.TracedType[]

src/Tracing.jl

@@ -1868,33 +1868,68 @@ end
 end
 
 function Reactant.traced_type_inner(
-    @nospecialize(RT::Type{<:StepRangeLen}),
+    @nospecialize(RT::Type{<:UnitRange{<:ReactantPrimitive}}),
     seen,
     mode::Reactant.TraceMode,
     track_numbers::Type,
     sharding,
     runtime,
 )
-    if !(Number <: track_numbers)
-        modified_track_numbers = Number
+    (T,) = RT.parameters
+    newT = Reactant.traced_type_inner(T, seen, mode, track_numbers, sharding, runtime)
+    if T == newT
+        return RT
     else
-        modified_track_numbers = track_numbers
+        return TracedRNumberOverrides.TracedUnitRange{newT}
     end
+end
+
+function Reactant.make_tracer(
+    seen,
+    @nospecialize(prev::UnitRange),
+    @nospecialize(path),
+    mode;
+    @nospecialize(sharding = Sharding.NoSharding()),
+    kwargs...,
+)
+    Reactant.Sharding.is_sharded(sharding) && error("Cannot specify sharding for UnitRange")
+    if mode == Reactant.TracedToTypes
+        push!(path, Core.Typeof(prev))
+        make_tracer(seen, prev.start, path, mode; kwargs...)
+        make_tracer(seen, prev.stop, path, mode; kwargs...)
+        return nothing
+    end
+    newstart = Reactant.make_tracer(
+        seen, prev.start, Reactant.append_path(path, :start), mode; kwargs...
+    )
+    newstop = Reactant.make_tracer(
+        seen, prev.stop, Reactant.append_path(path, :stop), mode; kwargs...
+    )
+    if typeof(newstart) == typeof(prev.start) && typeof(newstop) == typeof(prev.stop)
+        return prev
+    else
+        return TracedRNumberOverrides.TracedUnitRange(newstart, newstop)
+    end
+end
+
+function Reactant.traced_type_inner(
+    @nospecialize(RT::Type{<:StepRangeLen}),
+    seen,
+    mode::Reactant.TraceMode,
+    track_numbers::Type,
+    sharding,
+    runtime,
+)
     T, R, S, L = RT.parameters
-    return TracedRNumberOverrides.TracedStepRangeLen{
-        Reactant.traced_type_inner(
-            T, seen, mode, modified_track_numbers, sharding, runtime
-        ),
-        Reactant.traced_type_inner(
-            R, seen, mode, modified_track_numbers, sharding, runtime
-        ),
-        Reactant.traced_type_inner(
-            S, seen, mode, modified_track_numbers, sharding, runtime
-        ),
-        Reactant.traced_type_inner(
-            L, seen, mode, modified_track_numbers, sharding, runtime
-        ),
-    }
+    newT = Reactant.traced_type_inner(T, seen, mode, track_numbers, sharding, runtime)
+    newR = Reactant.traced_type_inner(R, seen, mode, track_numbers, sharding, runtime)
+    newS = Reactant.traced_type_inner(S, seen, mode, track_numbers, sharding, runtime)
+    newL = Reactant.traced_type_inner(L, seen, mode, track_numbers, sharding, runtime)
+    if T == newT && R == newR && S == newS && L == newL
+        return RT
+    else
+        return TracedRNumberOverrides.TracedStepRangeLen{newT,newR,newS,newL}
+    end
 end
 
 function Reactant.make_tracer(
@@ -1909,44 +1944,30 @@ function Reactant.make_tracer(
         error("Cannot specify sharding for StepRangeLen")
     if mode == Reactant.TracedToTypes
         push!(path, Core.Typeof(prev))
-        make_tracer(seen, prev.ref, path, mode; kwargs...)
-        make_tracer(seen, prev.step, path, mode; kwargs...)
-        make_tracer(seen, prev.len, path, mode; kwargs...)
-        make_tracer(seen, prev.offset, path, mode; kwargs...)
+        make_tracer(seen, prev.ref, path, mode; sharding, kwargs...)
+        make_tracer(seen, prev.step, path, mode; sharding, kwargs...)
+        make_tracer(seen, prev.len, path, mode; sharding, kwargs...)
+        make_tracer(seen, prev.offset, path, mode; sharding, kwargs...)
         return nothing
     end
-    return TracedRNumberOverrides.TracedStepRangeLen(
-        Reactant.make_tracer(
-            seen,
-            prev.ref,
-            Reactant.append_path(path, :ref),
-            mode;
-            kwargs...,
-            track_numbers=Number,
-        ),
-        Reactant.make_tracer(
-            seen,
-            prev.step,
-            Reactant.append_path(path, :step),
-            mode;
-            kwargs...,
-            track_numbers=Number,
-        ),
-        Reactant.make_tracer(
-            seen,
-            prev.len,
-            Reactant.append_path(path, :len),
-            mode;
-            kwargs...,
-            track_numbers=Number,
-        ),
-        Reactant.make_tracer(
-            seen,
-            prev.offset,
-            Reactant.append_path(path, :offset),
-            mode;
-            kwargs...,
-            track_numbers=Number,
-        ),
+    newref = Reactant.make_tracer(
+        seen, prev.ref, Reactant.append_path(path, :ref), mode; sharding, kwargs...
+    )
+    newstep = Reactant.make_tracer(
+        seen, prev.step, Reactant.append_path(path, :step), mode; sharding, kwargs...
+    )
+    newlen = Reactant.make_tracer(
+        seen, prev.len, Reactant.append_path(path, :len), mode; sharding, kwargs...
+    )
+    newoffset = Reactant.make_tracer(
+        seen, prev.offset, Reactant.append_path(path, :offset), mode; sharding, kwargs...
     )
+    if typeof(newref) == typeof(prev.ref) &&
+        typeof(newstep) == typeof(prev.step) &&
+        typeof(newlen) == typeof(prev.len) &&
+        typeof(newoffset) == typeof(prev.offset)
+        return prev
+    else
+        return TracedRNumberOverrides.TracedStepRangeLen(newref, newstep, newlen, newoffset)
+    end
 end

test/basic.jl

@@ -993,20 +993,35 @@ end
 
 @testset "Fractional index" begin
     times = 0:0.01:4.5
+    @test times isa Base.StepRangeLen
     res = @jit fractional_idx(times, ConcreteRNumber(2.143))
     @test res[1] == 0.29999999999997334
     @test res[2] == 215
     @test res[3] == 216
 end
 
 @testset "Traced fractional index" begin
-    times = Reactant.to_rarray(0:0.01:4.5)
+    times = Reactant.to_rarray(0:0.01:4.5; track_numbers=Number)
+    @test times isa Reactant.TracedRNumberOverrides.TracedStepRangeLen
     res = @jit fractional_idx(times, ConcreteRNumber(2.143))
     @test res[1] == 0.29999999999997334
     @test res[2] == 215
     @test res[3] == 216
 end
 
+function unitrange_test(r, i)
+    return r[i]
+end
+@testset "Unitrange" begin
+    x = 2:10
+    @test (@jit unitrange_test(x, 3)) == 4
+    @test (@jit unitrange_test(x, Reactant.ConcreteRNumber(4))) == 5
+
+    x = Reactant.to_rarray(2:10; track_numbers=Number)
+    @test (@jit unitrange_test(x, 3)) == 4
+    @test (@jit unitrange_test(x, Reactant.ConcreteRNumber(4))) == 5
+end
+
 mulpi(x) = π * x
 
 @testset "Irrational promotion" begin