feat: tracing Random.jl functionality correctly

Project.toml

@@ -14,6 +14,7 @@ Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
@@ -25,8 +26,8 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 
-[sources.ReactantCore]
-path = "lib/ReactantCore"
+[sources]
+ReactantCore = {path = "lib/ReactantCore"}
 
 [extensions]
 ReactantAbstractFFTsExt = "AbstractFFTs"
@@ -48,6 +49,7 @@ LinearAlgebra = "1.10"
 NNlib = "0.9.26"
 OrderedCollections = "1"
 Preferences = "1.4"
+Random = "1.10"
 ReactantCore = "0.1.2"
 Reactant_jll = "0.0.26"
 Scratch = "1.2"

src/Interpreter.jl

@@ -98,6 +98,18 @@ function set_reactant_abi(
         end
     end
 
+    # ensures we are not generating a constant array in the trace
+    # https://github.com/EnzymeAD/Reactant.jl/issues/356
+    if (f === Random.default_rng || f === default_rng) && length(argtypes) == 1
+        arginfo2 = ArgInfo(
+            fargs isa Nothing ? nothing : Any[:($(default_rng_inside_interpreter))],
+            Any[Core.Const(default_rng_inside_interpreter)],
+        )
+        return abstract_call_known(
+            interp, default_rng_inside_interpreter, arginfo2, si, sv, max_methods
+        )
+    end
+
     return Base.@invoke abstract_call_known(
         interp::AbstractInterpreter,
         f::Any,

src/Ops.jl

@@ -10,6 +10,7 @@ using ..Reactant:
     ConcreteRNumber,
     TracedRArray,
     TracedRNumber,
+    ReactantPrimitive,
     mlir_type,
     mlir_stacktrace
 
@@ -981,18 +982,46 @@ end
 
 # random ops
 function rng_bit_generator(
+    ::Type{T},
     seed::TracedRArray{UInt64,1},
     shape;
     algorithm::String="DEFAULT",
     location=mlir_stacktrace("rng_bit_generator", @__FILE__, @__LINE__),
-)
-    output = MLIR.IR.TensorType(TracedRArray{UInt64,1}, shape)
+) where {T<:Integer}
+    @assert algorithm in ("DEFAULT", "PHILOX", "THREE_FRY")
+    if algorithm == "PHILOX"
+        @assert length(seed) ∈ (2, 3)
+    elseif algorithm == "THREE_FRY"
+        @assert length(seed) == 2
+    end
+
+    output = MLIR.IR.TensorType(shape, MLIR.IR.Type(T))
+    output_state = MLIR.IR.TensorType(size(seed), MLIR.IR.Type(UInt64))
     rng_algorithm = MLIR.API.stablehloRngAlgorithmAttrGet(MLIR.IR.context(), algorithm)
-    op = stablehlo.rng_bit_generator(seed.mlir_data; output, rng_algorithm, location)
+    op = stablehlo.rng_bit_generator(
+        seed.mlir_data; output, output_state, rng_algorithm, location
+    )
     return (;
-        output_state=TracedRArray{UInt64,1}((), MLIR.IR.result(op, 1), MLIR.IR.size(seed)),
-        output=TracedRArray{T,length(shape)}((), MLIR.IR.result(op, 2), shape),
+        output_state=TracedRArray{UInt64,1}((), MLIR.IR.result(op, 1), size(seed)),
+        output=TracedRArray{T,length(shape)}((), MLIR.IR.result(op, 2), Tuple(shape)),
+    )
+end
+
+function rng_bit_generator(
+    ::Type{T},
+    seed::TracedRArray{UInt64,1},
+    shape;
+    algorithm::String="DEFAULT",
+    location=mlir_stacktrace("rng_bit_generator", @__FILE__, @__LINE__),
+) where {T<:Union{Float16,Float32,Float64}}
+    nbits = sizeof(T) * 8
+    uT = nbits == 16 ? UInt16 : (nbits == 32 ? UInt32 : UInt64)
+    (; output_state, output) = rng_bit_generator(uT, seed, shape; algorithm, location)
+    output = divide(
+        convert(TracedRArray{T,ndims(output)}, output),
+        constant(fill(T(typemax(uT)), Tuple(shape)); location),
     )
+    return (; output_state, output)
 end
 
 # functional ops

src/Reactant.jl

@@ -3,6 +3,8 @@ module Reactant
 using ReactantCore: ReactantCore, @trace, MissingTracedValue
 
 using LinearAlgebra: LinearAlgebra
+using Random: Random
+
 using Adapt: Adapt, WrappedArray
 using GPUArraysCore: GPUArraysCore, @allowscalar, allowscalar # keep this import to allow users to do `Reactant.allowscalar(false)`
 
@@ -97,7 +99,9 @@ include("TracedRArray.jl")
 
 include("Ops.jl")
 
-include("linear_algebra.jl")
+# StdLib Overloads
+include("stdlibs/LinearAlgebra.jl")
+include("stdlibs/Random.jl")
 
 const TracedType = Union{TracedRArray,TracedRNumber,MissingTracedValue}
 

src/stdlibs/Random.jl

@@ -0,0 +1,109 @@
+# Implementation based on the following:
+# 1. https://github.com/JuliaGPU/CUDA.jl/blob/master/src/random.jl
+# 2. https://github.com/JuliaRandom/Random123.jl/blob/master/src/common.jl#L125
+
+mutable struct TracedRNG <: Random.AbstractRNG
+    seed::Union{ConcreteRArray{UInt64,1},TracedRArray{UInt64,1}}
+    const algorithm::String
+end
+
+function Random.seed!(rng::TracedRNG, seed::Number)
+    seed = reinterpret(UInt64, Random.hash_seed(seed))
+    # TODO: Using `seed!` inside tracing should generate a TracedRArray
+    return Random.seed!(rng, ConcreteRArray(seed[1:length(rng.seed)]))
+end
+
+function Random.seed!(
+    rng::TracedRNG, seed::Union{ConcreteRArray{UInt64,1},TracedRArray{UInt64,1}}
+)
+    rng.seed = seed
+    return rng
+end
+
+make_seed() = rand(Random.RandomDevice(), UInt64, 2)
+
+TracedRNG() = TracedRNG(ConcreteRArray(make_seed()))
+TracedRNG(seed::ConcreteRArray{UInt64,1}) = TracedRNG(seed, "DEFAULT")
+
+default_rng() = TracedRNG()
+function default_rng_inside_interpreter()
+    return TracedRNG(promote_to(TracedRArray{UInt64,1}, make_seed()), "DEFAULT")
+end
+
+# XXX: Currently we get an illegal instruction if we don't call Random.default_rng()
+
+function Random.rand!(rng::TracedRNG, A::AnyTracedRArray{T,N}) where {T,N}
+    length(A) == 0 && return A
+    res = Ops.rng_bit_generator(T, rng.seed, [size(A)...]; rng.algorithm)
+    rng.seed = res.output_state
+    set_mlir_data!(A, res.output.mlir_data)
+    return A
+end
+
+function Random.randn!(rng::TracedRNG, A::AnyTracedRArray{T,N}) where {T,N}
+    length(A) == 0 && return A
+    Random.rand!(rng, A)
+    scaled_uniform = Ops.subtract(
+        Ops.multiply(A, Ops.constant(fill(T(2), size(A)))),
+        Ops.constant(fill(T(1), size(A))),
+    )
+    probit = Ops.erf_inv(scaled_uniform)
+    rand_normal = Ops.multiply(probit, Ops.constant(fill(sqrt(T(2)), size(A))))
+    set_mlir_data!(A, rand_normal.mlir_data)
+    return A
+end
+
+for randfun in (:rand, :randn)
+    randfun! = Symbol(randfun, :!)
+    @eval begin
+        function Random.$(randfun)(rng::TracedRNG, ::Type{T}, dims::Dims) where {T}
+            return Random.$(randfun!)(rng, TracedRArray{T,length(dims)}((), nothing, dims))
+        end
+
+        function Random.$(randfun)(rng::TracedRNG, dims::Dims)
+            return Random.$(randfun)(rng, Float64, dims)
+        end
+
+        function Random.$(randfun)(rng::TracedRNG, dim1::Integer, dims::Integer...)
+            return Random.$(randfun)(rng, Dims((dim1, dims...)))
+        end
+
+        function Random.$(randfun)(
+            rng::TracedRNG, ::Type{T}, dim1::Integer, dims::Integer...
+        ) where {T}
+            return Random.$(randfun)(rng, T, Dims((dim1, dims...)))
+        end
+
+        Random.$(randfun!)(A::AnyTracedRArray) = Random.$(randfun!)(default_rng(), A)
+
+        # scalars
+        function Random.$(randfun)(rng::TracedRNG, ::Type{T}=Float64) where {T}
+            A = promote_to(TracedRArray{T,0}, fill(T(0)))
+            Random.$(randfun!)(rng, A)
+            return A[]
+        end
+
+        # Non-Traced RNGs if used will lead to disastrous performance. We attempt to fix
+        # that but with a warning
+        function Random.$(randfun!)(rng::Random.AbstractRNG, A::AnyTracedRArray)
+            @warn "`rng` is not a `TracedRNG`. We will use this to seed the `TracedRNG` \
+                   instead of generating samples from this RNG type." maxlog = 1
+            seed = promote_to(TracedRArray{UInt64,1}, rand(rng, UInt64, 2))
+            trng = TracedRNG(seed, "DEFAULT")
+            return Random.$(randfun!)(trng, A)
+        end
+    end
+end
+
+# resolve ambiguities
+function Random.randn(rng::TracedRNG, T::Random.BitFloatType)
+    A = promote_to(TracedRArray{T,0}, fill(T(0)))
+    Random.randn!(rng, A)
+    return A[]
+end
+
+# TODO: At some later point we might want to implement the sampler API as well since it
+#       makes all RNG implementation work by default. From the post-optimize IR we need to
+#       confirm that the dynamic_update_slice calls are optimized away into a single
+#       `stablehlo.rng_bit_generator` call -- confirm that this should be the case based on
+#       how the seeding should work?

test/ops.jl

@@ -538,8 +538,34 @@ end
 end
 
 @testset "rng_bit_generator" begin
-    # seed = ConcreteRArray([0, 0])
-    # @jit Ops.rng_bit_generator(seed, [2])
+    genInt32(seed) = Ops.rng_bit_generator(Int32, seed, [2, 4])
+    genInt64(seed) = Ops.rng_bit_generator(Int64, seed, [2, 4])
+    genUInt64(seed) = Ops.rng_bit_generator(UInt64, seed, [2, 4])
+
+    @testset for (alg, sz) in
+                 [("DEFAULT", 2), ("PHILOX", 2), ("PHILOX", 3), ("THREE_FRY", 2)]
+        seed = ConcreteRArray(zeros(UInt64, sz))
+
+        res = @jit genInt32(seed)
+        @test res.output_state !== seed
+        @test size(res.output_state) == (sz,)
+        @test res.output isa ConcreteRArray{Int32,2}
+        @test size(res.output) == (2, 4)
+
+        seed = res.output_state
+        res = @jit genInt64(seed)
+        @test res.output_state !== seed
+        @test size(res.output_state) == (sz,)
+        @test res.output isa ConcreteRArray{Int64,2}
+        @test size(res.output) == (2, 4)
+
+        seed = res.output_state
+        res = @jit genUInt64(seed)
+        @test res.output_state !== seed
+        @test size(res.output_state) == (sz,)
+        @test res.output isa ConcreteRArray{UInt64,2}
+        @test size(res.output) == (2, 4)
+    end
 end
 
 @testset "round_nearest_afz" begin