refactor: rework TracedRNG to be similar to other types (#448)

Author: avik-pal

docs/src/api/internal.md

@@ -8,5 +8,4 @@ These functions are not part of the public API and are subject to change at any
 
 ```@docs
 Reactant.REDUB_ARGUMENTS_NAME
-Reactant.within_reactant_interpreter
 ```

src/Overlay.jl

@@ -3,15 +3,6 @@
 #       correctly. Once that (https://github.com/timholy/Revise.jl/issues/646) is resolved
 #       we should move all the reactant_overrides to relevant files.
 
-# Helper Function to determine if we are inside the ReactantInterpreter
-"""
-    within_reactant_interpreter()
-
-Returns `true` if we are currently inside the ReactantInterpreter.
-"""
-@noinline within_reactant_interpreter() = false
-@reactant_overlay @noinline within_reactant_interpreter() = true
-
 # Compiling within a compile should return simply the original function
 @reactant_overlay function Compiler.compile(
     f, args; client=nothing, optimize=true, sync=false
@@ -37,6 +28,12 @@ end
     return call_with_reactant(TracedRandom.default_rng)
 end
 
+@reactant_overlay @noinline function TracedRandom.default_rng()
+    return TracedRNG(
+        TracedUtils.promote_to(TracedRArray{UInt64,1}, TracedRandom.make_seed()), "DEFAULT"
+    )
+end
+
 ## Only problematic edge case here is the direct `<randfun!>(rng, A::AbstractArray)` call
 ## We can't directly overlay that call without breaking the semantics of inplace update
 for randfun in (:rand, :randn, :randexp)

src/Reactant.jl

@@ -175,8 +175,13 @@ include("TracedRArray.jl")
 
 include("ConcreteRArray.jl")
 
+mutable struct ConcreteRNG <: Random.AbstractRNG
+    seed::ConcreteRArray{UInt64,1}
+    const algorithm::String
+end
+
 mutable struct TracedRNG <: Random.AbstractRNG
-    seed::Union{ConcreteRArray{UInt64,1},TracedRArray{UInt64,1}}
+    seed::TracedRArray{UInt64,1}
     const algorithm::String
 end
 

src/Tracing.jl

@@ -227,6 +227,16 @@ function traced_type(
     end
 end
 
+function traced_type(::Type{<:ConcreteRNG}, seen, ::Val{mode}, track_numbers) where {mode}
+    if mode == ConcreteToTraced
+        return TracedRNG
+    elseif mode == TracedToConcrete
+        return ConcreteRNG
+    else
+        throw("Unsupported mode: $mode")
+    end
+end
+
 function traced_type(
     ::Type{T}, seen::ST, ::Val{mode}, track_numbers
 ) where {ST,T<:TracedType,mode}
@@ -246,6 +256,18 @@ function traced_type(
     end
 end
 
+function traced_type(::Type{T}, seen, ::Val{mode}, track_numbers) where {T<:TracedRNG,mode}
+    if mode == ConcreteToTraced
+        throw("TracedRNG cannot be traced")
+    elseif mode == TracedToConcrete
+        return ConcreteRNG
+    elseif mode == TracedTrack || mode == TracedSetPath
+        return T
+    else
+        throw("Unsupported mode: $mode")
+    end
+end
+
 function traced_type(::Type{T}, seen, mode, track_numbers) where {T<:XLAArray}
     throw("XLA $T array cannot be traced")
 end

src/stdlibs/Random.jl

@@ -8,12 +8,15 @@ using ..Reactant:
     Reactant,
     TracedRArray,
     TracedRNumber,
+    ConcreteRNG,
     TracedRNG,
     AnyTracedRArray,
     Reactant,
     TracedUtils,
     Ops,
-    ConcreteRArray
+    ConcreteRArray,
+    ConcreteRNumber,
+    unwrapped_eltype
 using Random: Random, AbstractRNG
 
 @noinline function make_seed(rng::AbstractRNG=Random.RandomDevice())
@@ -25,44 +28,55 @@ using Random: Random, AbstractRNG
     return seed
 end
 
-function Random.seed!(rng::TracedRNG, seed::Number)
+@noinline function Random.seed!(rng::TracedRNG, seed::Number)
     if seed isa TracedRNumber
         error("Passing in `TracedRNumber` as a seed is not supported. Please pass in a \
                `TracedRArray` of the appropriate size instead.")
     end
 
     seed = reinterpret(UInt64, Random.hash_seed(seed))
-    seed = if Reactant.within_reactant_interpreter()
-        TracedUtils.promote_to(TracedRArray{UInt64,1}, seed[1:length(rng.seed)])
-    else
-        ConcreteRArray(seed[1:length(rng.seed)])
-    end
-    return Random.seed!(rng, seed)
+    return Random.seed!(
+        rng, TracedUtils.promote_to(TracedRArray{UInt64,1}, seed[1:length(rng.seed)])
+    )
 end
 
-function Random.seed!(rng::TracedRNG, seed::AbstractArray{<:Integer,1})
+@noinline function Random.seed!(rng::TracedRNG, seed::AbstractVector{<:Integer})
     return Random.seed!(rng, UInt64.(seed))
 end
 
-function Random.seed!(rng::TracedRNG, seed::AbstractArray{UInt64,1})
+@noinline function Random.seed!(rng::TracedRNG, seed::AbstractVector{UInt64})
     return Random.seed!(rng, TracedUtils.promote_to(TracedRArray{UInt64,1}, seed))
 end
 
-function Random.seed!(
-    rng::TracedRNG, seed::Union{ConcreteRArray{UInt64,1},TracedRArray{UInt64,1}}
-)
+@noinline function Random.seed!(rng::TracedRNG, seed::TracedRArray{UInt64,1})
     rng.seed = seed
     return rng
 end
 
-@noinline TracedRNG() = TracedRNG(ConcreteRArray(make_seed()))
-@noinline TracedRNG(seed::ConcreteRArray{UInt64,1}) = TracedRNG(seed, "DEFAULT")
+@noinline function Random.seed!(rng::ConcreteRNG, seed::Number)
+    seed isa ConcreteRNumber && (seed = unwrapped_eltype(seed)(seed))
+    seed = reinterpret(UInt64, Random.hash_seed(seed))
+    return Random.seed!(rng, ConcreteRArray(seed))
+end
+
+@noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractVector{<:Integer})
+    return Random.seed!(rng, seed)
+end
 
-@noinline function default_rng()
-    Reactant.within_reactant_interpreter() || return TracedRNG()
-    return TracedRNG(TracedUtils.promote_to(TracedRArray{UInt64,1}, make_seed()), "DEFAULT")
+@noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractVector{UInt64})
+    return Random.seed!(rng, ConcreteRArray(seed))
 end
 
+@noinline function Random.seed!(rng::ConcreteRNG, seed::ConcreteRArray{UInt64,1})
+    rng.seed = seed
+    return rng
+end
+
+@noinline ConcreteRNG() = ConcreteRNG(ConcreteRArray(make_seed()))
+@noinline ConcreteRNG(seed::ConcreteRArray{UInt64,1}) = ConcreteRNG(seed, "DEFAULT")
+
+@noinline default_rng() = ConcreteRNG()
+
 @noinline rng_algorithm(rng::TracedRNG) = rng.algorithm
 @noinline rng_algorithm(::AbstractRNG) = "DEFAULT"