feat: rework special concrete/traced rng handling (#1347)

* feat: rework special concrete/traced rng handling

* chore: bump reactant version

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.120"
+version = "0.2.121"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Overlay.jl

@@ -27,7 +27,7 @@ end
 end
 
 @reactant_overlay @noinline function TracedRandom.default_rng()
-    return TracedRNG(
+    return ReactantRNG(
         TracedUtils.promote_to(TracedRArray{UInt64,1}, TracedRandom.make_seed()), "DEFAULT"
     )
 end

src/Reactant.jl

@@ -160,8 +160,8 @@ use_overlayed_version(::TracedRArray) = true
 use_overlayed_version(::TracedRNumber) = true
 use_overlayed_version(::Number) = false
 use_overlayed_version(::MissingTracedValue) = true
-use_overlayed_version(::TracedRNG) = true
 use_overlayed_version(::AbstractArray{<:TracedRNumber}) = true
+use_overlayed_version(rng::ReactantRNG) = use_overlayed_version(rng.seed)
 
 function use_overlayed_version(x::AbstractArray)
     a = ancestor(x)

src/Tracing.jl

@@ -348,23 +348,6 @@ Base.@nospecializeinfer function traced_type_inner(
     end
 end
 
-Base.@nospecializeinfer function traced_type_inner(
-    @nospecialize(T::Type{<:ConcreteRNG}),
-    seen,
-    mode::TraceMode,
-    @nospecialize(track_numbers::Type),
-    @nospecialize(sharding),
-    @nospecialize(runtime)
-)
-    if mode == ConcreteToTraced
-        return TracedRNG
-    elseif mode == TracedToConcrete
-        return T
-    else
-        throw("Unsupported mode: $mode")
-    end
-end
-
 Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(T::Type{<:MissingTracedValue}), @nospecialize(args...)
 )
@@ -451,29 +434,6 @@ Base.@nospecializeinfer function traced_type_inner(
     end
 end
 
-Base.@nospecializeinfer function traced_type_inner(
-    @nospecialize(T::Type{<:TracedRNG}),
-    seen,
-    mode::TraceMode,
-    @nospecialize(track_numbers::Type),
-    @nospecialize(sharding),
-    @nospecialize(runtime)
-)
-    if mode == ConcreteToTraced
-        throw("TracedRNG cannot be traced")
-    elseif mode == TracedToConcrete
-        return ConcreteRNG{
-            traced_type_inner(
-                TracedRArray{UInt64,1}, seen, mode, track_numbers, sharding, runtime
-            ),
-        }
-    elseif mode == TracedTrack || mode == NoStopTracedTrack || mode == TracedSetPath
-        return T
-    else
-        throw("Unsupported mode: $mode")
-    end
-end
-
 Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(A::Type{AbstractArray}),
     seen,

src/Types.jl

@@ -87,12 +87,6 @@ const AnyTracedRVector{T} = AnyTracedRArray{T,1}
 const AnyTracedRMatrix{T} = AnyTracedRArray{T,2}
 const AnyTracedRVecOrMat{T} = Union{AnyTracedRVector{T},AnyTracedRMatrix{T}}
 
-## TracedRNG
-struct TracedRNG <: Random.AbstractRNG
-    seed::TracedRArray{UInt64,1}
-    algorithm::String
-end
-
 # Concrete Types
 ## ConcretePJRTNumber
 mutable struct ConcretePJRTNumber{T,D,S<:Sharding.ShardInfo} <: AbstractConcreteNumber{T}
@@ -442,12 +436,16 @@ function ConcreteIFRTArray{T,N}(x::AnyConcreteIFRTArray; kwargs...) where {T,N}
     )
 end
 
-## ConcreteRNG
-mutable struct ConcreteRNG{S<:AbstractConcreteArray} <: Random.AbstractRNG
+# RNGs
+struct ReactantRNG{S<:Union{<:AbstractConcreteArray{UInt64,1},TracedRArray{UInt64,1}}} <:
+       Random.AbstractRNG
     seed::S
-    const algorithm::String
+    algorithm::String
 end
 
+Base.@deprecate_binding ConcreteRNG ReactantRNG
+Base.@deprecate_binding TracedRNG ReactantRNG
+
 ## Aliases based on the set preferences
 if XLA.REACTANT_XLA_RUNTIME == "PJRT"
     const ConcreteRArray = ConcretePJRTArray

src/stdlibs/Random.jl

@@ -5,80 +5,45 @@ module TracedRandom
 # 2. https://github.com/JuliaRandom/Random123.jl/blob/master/src/common.jl
 
 using ..Reactant:
-    Reactant,
-    TracedRArray,
-    TracedRNumber,
-    ConcreteRNG,
-    TracedRNG,
-    AnyTracedRArray,
-    Reactant,
-    TracedUtils,
-    Ops,
-    AbstractConcreteArray,
-    AbstractConcreteNumber,
-    unwrapped_eltype
+    Reactant, TracedRArray, TracedRNumber, ReactantRNG, AnyTracedRArray, TracedUtils, Ops
 using Random: Random, AbstractRNG
 
 @noinline make_seed(rng::AbstractRNG=Random.RandomDevice()) =
     Random.rand!(rng, Vector{UInt64}(undef, 2))
 
-@noinline function Random.seed!(rng::TracedRNG, seed::Number)
+@noinline function Random.seed!(rng::ReactantRNG, 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))
-    return Random.seed!(
-        rng, TracedUtils.promote_to(TracedRArray{UInt64,1}, seed[1:length(rng.seed)])
-    )
+    return Random.seed!(rng, seed[1:length(rng.seed)])
 end
 
-@noinline function Random.seed!(rng::TracedRNG, seed::AbstractVector{<:Integer})
-    return Random.seed!(rng, UInt64.(seed))
-end
-
-@noinline function Random.seed!(rng::TracedRNG, seed::AbstractVector{UInt64})
-    return Random.seed!(rng, TracedUtils.promote_to(TracedRArray{UInt64,1}, seed))
-end
-
-@noinline function Random.seed!(rng::TracedRNG, seed::TracedRArray{UInt64,1})
-    copyto!(rng.seed, seed)
+@noinline function Random.seed!(rng::ReactantRNG, seed::AbstractVector)
+    rng.seed .= seed
     return rng
 end
 
-@noinline function Random.seed!(rng::ConcreteRNG, seed::Number)
-    seed isa AbstractConcreteNumber && (seed = unwrapped_eltype(seed)(seed))
-    seed = reinterpret(UInt64, Random.hash_seed(seed))
-    return Random.seed!(rng, Reactant.to_rarray(seed))
-end
-
-@noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractVector{<:Integer})
-    return Random.seed!(rng, seed)
-end
-
-@noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractVector{UInt64})
-    return Random.seed!(rng, Reactant.to_rarray(seed))
-end
+Base.copy(rng::ReactantRNG) = ReactantRNG(copy(rng.seed), rng.algorithm)
 
-@noinline function Random.seed!(rng::ConcreteRNG, seed::AbstractConcreteArray{UInt64,1})
-    Base.copyto!(rng.seed, seed)
-    return rng
+@noinline function ReactantRNG()
+    if Reactant.within_compile()
+        return ReactantRNG(TracedUtils.promote_to(TracedRArray{UInt64,1}, make_seed()))
+    else
+        return ReactantRNG(Reactant.to_rarray(make_seed()))
+    end
 end
+@noinline ReactantRNG(seed::AbstractVector) = ReactantRNG(seed, "DEFAULT")
 
-Base.copy(rng::ConcreteRNG) = ConcreteRNG(copy(rng.seed), rng.algorithm)
-Base.copy(rng::TracedRNG) = TracedRNG(copy(rng.seed), rng.algorithm)
+@noinline default_rng() = ReactantRNG()
 
-@noinline ConcreteRNG() = ConcreteRNG(Reactant.to_rarray(make_seed()))
-@noinline ConcreteRNG(seed::AbstractConcreteArray{UInt64,1}) = ConcreteRNG(seed, "DEFAULT")
-
-@noinline default_rng() = ConcreteRNG()
-
-@noinline rng_algorithm(rng::TracedRNG) = rng.algorithm
+@noinline rng_algorithm(rng::ReactantRNG) = rng.algorithm
 @noinline rng_algorithm(::AbstractRNG) = "DEFAULT"
 
 @noinline function internal_overload_rand!(
-    rng::TracedRNG, A::AnyTracedRArray{T,N}
+    rng::ReactantRNG{<:TracedRArray}, A::AnyTracedRArray{T,N}
 ) where {T,N}
     length(A) == 0 && return A
     res = Ops.rng_bit_generator(T, rng.seed, [size(A)...]; rng.algorithm)
@@ -88,7 +53,7 @@ Base.copy(rng::TracedRNG) = TracedRNG(copy(rng.seed), rng.algorithm)
 end
 
 @noinline function internal_overload_randn!(
-    rng::TracedRNG, A::AnyTracedRArray{T,N}
+    rng::ReactantRNG{<:TracedRArray}, A::AnyTracedRArray{T,N}
 ) where {T,N}
     length(A) == 0 && return A
     res = Ops.randn(T, rng.seed, [size(A)...]; rng.algorithm)
@@ -98,7 +63,7 @@ end
 end
 
 @noinline function internal_overload_randexp!(
-    rng::TracedRNG, A::AnyTracedRArray{T,N}
+    rng::ReactantRNG{<:TracedRArray}, A::AnyTracedRArray{T,N}
 ) where {T,N}
     length(A) == 0 && return A
     res = Ops.randexp(T, rng.seed, [size(A)...]; rng.algorithm)
@@ -114,25 +79,25 @@ for randfun in (:rand, :randn, :randexp)
 
     @eval begin
         @noinline function $(overload_randfun)(
-            rng::TracedRNG, ::Type{T}, dims::Dims
+            rng::ReactantRNG{<:TracedRArray}, ::Type{T}, dims::Dims
         ) where {T}
             return $(overload_randfun!)(
                 rng, TracedRArray{T,length(dims)}((), nothing, dims)
             )
         end
 
-        @noinline function $(overload_randfun)(rng::TracedRNG, dims::Dims)
+        @noinline function $(overload_randfun)(rng::ReactantRNG{<:TracedRArray}, dims::Dims)
             return $(overload_randfun)(rng, Float64, dims)
         end
 
         @noinline function $(overload_randfun)(
-            rng::TracedRNG, dim1::Integer, dims::Integer...
+            rng::ReactantRNG{<:TracedRArray}, dim1::Integer, dims::Integer...
         )
             return $(overload_randfun)(rng, Dims((dim1, dims...)))
         end
 
         @noinline function $(overload_randfun)(
-            rng::TracedRNG, ::Type{T}, dim1::Integer, dims::Integer...
+            rng::ReactantRNG{<:TracedRArray}, ::Type{T}, dim1::Integer, dims::Integer...
         ) where {T}
             return $(overload_randfun)(rng, T, Dims((dim1, dims...)))
         end
@@ -142,7 +107,9 @@ for randfun in (:rand, :randn, :randexp)
         end
 
         # scalars
-        @noinline function $(overload_randfun)(rng::TracedRNG, ::Type{T}=Float64) where {T}
+        @noinline function $(overload_randfun)(
+            rng::ReactantRNG{<:TracedRArray}, ::Type{T}=Float64
+        ) where {T}
             A = TracedUtils.promote_to(TracedRArray{T,0}, fill(T(0)))
             $(overload_randfun!)(rng, A)
             return TracedRNumber{T}((), A.mlir_data)
@@ -157,14 +124,14 @@ for randfun in (:rand, :randn, :randexp, :rand!, :randn!, :randexp!)
     internal_overload_randfun = Symbol(:internal_overload_, randfun)
     @eval begin
         @noinline function $(overload_randfun)(rng::AbstractRNG, args...)
-            rng = TracedRNG(
+            rng = ReactantRNG(
                 TracedUtils.promote_to(TracedRArray{UInt64,1}, make_seed(rng)),
                 rng_algorithm(rng),
             )
             return $(internal_overload_randfun)(rng, args...)
         end
 
-        @noinline function $(overload_randfun)(rng::TracedRNG, args...)
+        @noinline function $(overload_randfun)(rng::ReactantRNG, args...)
             return $(internal_overload_randfun)(rng, args...)
         end
     end

src/utils.jl

@@ -110,8 +110,7 @@ function should_rewrite_call(@nospecialize(ft))
             # Don't rewrite primitive ops, tracing utilities, or any MLIR-based functions
             if has_ancestor(mod, Reactant.Ops) ||
                 has_ancestor(mod, Reactant.TracedUtils) ||
-                has_ancestor(mod, Reactant.MLIR) ||
-                has_ancestor(mod, Reactant.TracedRandom)
+                has_ancestor(mod, Reactant.MLIR)
                 return false
             end
             if string(mod) == "CUDA"