feat: tracing Random.jl functionality correctly #363

avik-pal · 2024-12-11T04:30:37Z

TODOs

Tests
~~Overlay seems to act weird. Can't switch to non-reactant interpreter nicely~~ Using a workaround for now
Fix documentation build
Overlay Random123 to generate specific RNGs (Threefry and Philox)
Overlay all random number generators
- construct TracedRNG during tracing

avik-pal · 2024-12-11T07:59:55Z

julia> using Reactant, Random

julia> fn() = randn(Random.default_rng(), 2, 3)
fn (generic function with 1 method)

julia> @code_hlo optimize = false fn()
module {
  func.func @main() -> tensor<3x2xf64> {
    %c = stablehlo.constant dense<[9454987348304227925, 11257230962712577529]> : tensor<2xui64>
    %output_state, %output = stablehlo.rng_bit_generator %c, algorithm =  DEFAULT : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x3xui64>)
    %0 = stablehlo.convert %output : (tensor<2x3xui64>) -> tensor<2x3xf64>
    %cst = stablehlo.constant dense<1.8446744073709552E+19> : tensor<2x3xf64>
    %1 = stablehlo.divide %0, %cst : tensor<2x3xf64>
    %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<2x3xf64>
    %2 = stablehlo.multiply %1, %cst_0 : tensor<2x3xf64>
    %cst_1 = stablehlo.constant dense<1.000000e+00> : tensor<2x3xf64>
    %3 = stablehlo.subtract %2, %cst_1 : tensor<2x3xf64>
    %4 = chlo.erf_inv %3 : tensor<2x3xf64> -> tensor<2x3xf64>
    %cst_2 = stablehlo.constant dense<1.4142135623730951> : tensor<2x3xf64>
    %5 = stablehlo.multiply %4, %cst_2 : tensor<2x3xf64>
    %6 = stablehlo.transpose %5, dims = [1, 0] : (tensor<2x3xf64>) -> tensor<3x2xf64>
    return %6 : tensor<3x2xf64>
  }
}

julia> @code_hlo fn()
module {
  func.func @main() -> tensor<3x2xf64> {
    %cst = stablehlo.constant dense<1.4142135623730951> : tensor<2x3xf64>
    %cst_0 = stablehlo.constant dense<1.000000e+00> : tensor<2x3xf64>
    %cst_1 = stablehlo.constant dense<2.000000e+00> : tensor<2x3xf64>
    %cst_2 = stablehlo.constant dense<1.8446744073709552E+19> : tensor<2x3xf64>
    %c = stablehlo.constant dense<[17523564455668573441, 5342821220909967229]> : tensor<2xui64>
    %output_state, %output = stablehlo.rng_bit_generator %c, algorithm =  DEFAULT : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x3xui64>)
    %0 = stablehlo.convert %output : (tensor<2x3xui64>) -> tensor<2x3xf64>
    %1 = stablehlo.divide %0, %cst_2 : tensor<2x3xf64>
    %2 = stablehlo.multiply %1, %cst_1 : tensor<2x3xf64>
    %3 = stablehlo.subtract %2, %cst_0 : tensor<2x3xf64>
    %4 = chlo.erf_inv %3 : tensor<2x3xf64> -> tensor<2x3xf64>
    %5 = stablehlo.multiply %4, %cst : tensor<2x3xf64>
    %6 = stablehlo.transpose %5, dims = [1, 0] : (tensor<2x3xf64>) -> tensor<3x2xf64>
    return %6 : tensor<3x2xf64>
  }
}

avik-pal · 2024-12-11T14:00:51Z

This is kind of working now. Can I get an initial review?

src/Interpreter.jl

wsmoses

This looks good to me but should wait until the interpreter lands (which will simplify the override and can just use the macro)

test/ops.jl

avik-pal · 2024-12-15T04:05:57Z

The overlay mechanism fixed the previous unreachable instruction issue 🎉

src/Ops.jl

avik-pal · 2024-12-17T14:52:03Z

src/Overlay.jl

+            @warn "Directly writing to an array using Random.jl functions inside \
+                   ReactantInterpreter will generate a constant array in the IR. Use with \
+                   caution." maxlog = 1
+            return Random.$(randfun!)(rng, A)


@wsmoses my understanding was that this should call the non-overlayed version. But here I get

┌ Warning: Directly writing to an array using Random.jl functions inside ReactantInterpreter will generate a constant array in the IR. Use with caution. └ @ Reactant /mnt/software/lux/Reactant.jl/src/Overlay.jl:84 Unreachable reached at 0x7a64877b6b16 [1417783] signal 4 (2): Illegal instruction in expression starting at REPL[7]:1 fn2 at ./REPL[6]:4 [inlined] opaque closure at ./<missing>:0 unknown function (ip: 0x7a64877b6bff) fn2 at ./REPL[6]:2 [inlined] call_with_reactant at /mnt/software/lux/Reactant.jl/src/utils.jl:0 #8 at /mnt/software/lux/Reactant.jl/src/TracedUtils.jl:210 block! at /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201 unknown function (ip: 0x7a64877b6316) #make_mlir_fn#1 at /mnt/software/lux/Reactant.jl/src/TracedUtils.jl:197 make_mlir_fn at /mnt/software/lux/Reactant.jl/src/TracedUtils.jl:117 [inlined] #10 at /mnt/software/lux/Reactant.jl/src/Compiler.jl:295 [inlined] block! at /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201 #9 at /mnt/software/lux/Reactant.jl/src/Compiler.jl:294 [inlined] mmodule! at /mnt/software/lux/Reactant.jl/src/mlir/IR/Module.jl:92 unknown function (ip: 0x7a64877b5976) #compile_mlir!#8 at /mnt/software/lux/Reactant.jl/src/Compiler.jl:291 compile_mlir! at /mnt/software/lux/Reactant.jl/src/Compiler.jl:290 [inlined] #6 at /mnt/software/lux/Reactant.jl/src/Compiler.jl:285 [inlined] context! at /mnt/software/lux/Reactant.jl/src/mlir/IR/Context.jl:76 unknown function (ip: 0x7a64877b2a76) #compile_mlir#5 at /mnt/software/lux/Reactant.jl/src/Compiler.jl:283 compile_mlir at /mnt/software/lux/Reactant.jl/src/Compiler.jl:280 unknown function (ip: 0x7a64877b0a66) jl_apply at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/julia.h:2157 [inlined] do_call at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:126 eval_value at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:223 eval_stmt_value at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:174 [inlined] eval_body at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:663 jl_interpret_toplevel_thunk at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:821 jl_toplevel_eval_flex at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/toplevel.c:943 jl_toplevel_eval_flex at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/toplevel.c:886 eval_body at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:625 eval_body at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:539 jl_interpret_toplevel_thunk at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/interpreter.c:821 jl_toplevel_eval_flex at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/toplevel.c:943 jl_toplevel_eval_flex at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/toplevel.c:886 jl_toplevel_eval_flex at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/toplevel.c:886 ijl_toplevel_eval_in at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/toplevel.c:994 eval at ./boot.jl:430 [inlined] eval_user_input at /home/avikpal/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:245 repl_backend_loop at /home/avikpal/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:342 #start_repl_backend#59 at /home/avikpal/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:327 start_repl_backend at /home/avikpal/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:324 #run_repl#72 at /home/avikpal/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:483 run_repl at /home/avikpal/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:469 jfptr_run_repl_10705 at /mnt/.julia/compiled/v1.11/REPL/u0gqU_FGbh7.so (unknown line) #1150 at ./client.jl:446 jfptr_YY.1150_15174 at /mnt/.julia/compiled/v1.11/REPL/u0gqU_FGbh7.so (unknown line) jl_apply at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/julia.h:2157 [inlined] jl_f__call_latest at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/builtins.c:875 #invokelatest#2 at ./essentials.jl:1055 [inlined] invokelatest at ./essentials.jl:1052 [inlined] run_main_repl at ./client.jl:430 repl_main at ./client.jl:567 [inlined] _start at ./client.jl:541 jfptr__start_73406.1 at /home/avikpal/.julia/juliaup/julia-1.11.2+0.x64.linux.gnu/lib/julia/sys.so (unknown line) jl_apply at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/julia.h:2157 [inlined] true_main at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/jlapi.c:900 jl_repl_entrypoint at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/jlapi.c:1059 main at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/cli/loader_exe.c:58 unknown function (ip: 0x7a65341e2e07) __libc_start_main at /usr/lib/libc.so.6 (unknown line) unknown function (ip: 0x4010b8) Allocations: 48616895 (Pool: 48615481; Big: 1414); GC: 47 [1] 1417783 illegal hardware instruction (core dumped) julia --project=envs --threads=4 --check-bounds=yes

using this trick #369 (comment) seems to work correctly (though this has other weird edge-cases)

This is definitely a general issue with any kind of recursion of an overlayed function. Do we have a way to force using the NativeInterpreter?

ccing @gbaraldi and @aviatesk

avik-pal · 2024-12-17T14:52:40Z

src/stdlibs/Random.jl

+        function $(overload_randfun)(rng::AbstractRNG, args...)
+            # XXX: Ideally the following should just work but currently it gives an illegal
+            #      instruction error. Maybe an issue with Julia's AbsInt?
+            # seed_uint64 = rand(rng, UInt64, 2)


same problem here. calling this leads to the illegal instruction

github-actions · 2024-12-18T03:21:31Z

src/Overlay.jl

+            if T <: ReactantPrimitive
+                return TracedRandom.$(overload_randfun)(rng, T, dims)
+            end
+            error("Reactant doesn't support sampling of $(T) with the current interpreter.")


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

error("Reactant doesn't support sampling of $(T) with the current interpreter.")

return error(

"Reactant doesn't support sampling of $(T) with the current interpreter."

)

github-actions · 2024-12-18T03:21:31Z

src/Overlay.jl

+            if T <: ReactantPrimitive
+                return TracedRandom.$(overload_randfun)(rng, T, dim1, dims...)
+            end
+            error("Reactant doesn't support sampling of $(T) with the current interpreter.")


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

error("Reactant doesn't support sampling of $(T) with the current interpreter.")

return error(

"Reactant doesn't support sampling of $(T) with the current interpreter."

)

github-actions · 2024-12-18T03:21:31Z

src/Overlay.jl

+            if T <: ReactantPrimitive
+                return TracedRandom.$(overload_randfun)(rng, T)
+            end
+            error("Reactant doesn't support sampling of $(T) with the current interpreter.")


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

error("Reactant doesn't support sampling of $(T) with the current interpreter.")

return error(

"Reactant doesn't support sampling of $(T) with the current interpreter."

)

avik-pal · 2024-12-18T03:23:16Z

I am working around the AbsInt issues for now, but those would be nice to sort out, especially for handling the wrapped arrays.

Also, I am explicitly throwing errors in the cases where otherwise Julia would crash with an invalid instruction

This reverts commit 94e9576.

avik-pal force-pushed the ap/random_numbers branch from 99e0e76 to 7d17faf Compare December 11, 2024 13:54

avik-pal requested review from wsmoses, Pangoraw and mofeing December 11, 2024 14:01

avik-pal commented Dec 11, 2024

View reviewed changes

src/Interpreter.jl Outdated Show resolved Hide resolved

wsmoses reviewed Dec 12, 2024

View reviewed changes

mofeing reviewed Dec 13, 2024

View reviewed changes

test/ops.jl Show resolved Hide resolved

avik-pal force-pushed the ap/random_numbers branch 2 times, most recently from dd1f9e6 to e99089e Compare December 15, 2024 03:48

avik-pal linked an issue Dec 15, 2024 that may be closed by this pull request

Handling random numbers correctly #356

Closed

wsmoses reviewed Dec 15, 2024

View reviewed changes

src/Ops.jl Outdated Show resolved Hide resolved

avik-pal force-pushed the ap/random_numbers branch 3 times, most recently from 8af045f to 22f6810 Compare December 17, 2024 04:48

avik-pal added 13 commits December 17, 2024 17:40

refactor: move stdlib overloads to a different directory

5b40fdf

fix: Ops.rng_bit_generator

a584127

feat: initial prototype for random number generation

8699131

feat: add support for scalar sampling

9da77c7

feat: efficient sampling for non-native RNGs

6dcad12

fix: handling floating point sampling

bbf2a48

feat: use the override macro

e82c2c0

fix: use @noinline

4b7c504

feat: support randexp

7d73538

feat: override seeding inside interpreter

591320c

refactor: move things into a module

1b45638

refactor: rework how the overlays are implemented

51f91e4

docs: add internal api to the docs

5a94f14

avik-pal added 3 commits December 17, 2024 17:40

test: include floating point tests

d61d269

test: setup testing

36e56cb

feat: overlay all generators

6c21721

avik-pal force-pushed the ap/random_numbers branch from 9aec916 to 6c21721 Compare December 17, 2024 12:18

avik-pal added 2 commits December 17, 2024 18:43

test: ensure distributions are correct

d0064f2

test: overlay generation

9129ce2

avik-pal commented Dec 17, 2024

View reviewed changes

avik-pal added 3 commits December 17, 2024 21:08

fix: test whether we can call into the non-overlayed version

867324e

fix: try marking TracedRandom in whitelist

66b16c4

fix: workaround the AbsInt issues for now

e370200

avik-pal requested a review from wsmoses December 18, 2024 03:21

github-actions bot reviewed Dec 18, 2024

View reviewed changes

fix: throw errors for now instead of crashing

eaa1f30

avik-pal force-pushed the ap/random_numbers branch from ba82aa5 to eaa1f30 Compare December 18, 2024 03:23

avik-pal marked this pull request as ready for review December 18, 2024 03:24

This was referenced Dec 18, 2024

AbstractInterpreter and Recursive Calls into Overlay-ed Functions #393

Closed

Random Numbers & Reactant LuxDL/Lux.jl#1131

Closed

wsmoses merged commit 94e9576 into main Dec 18, 2024
35 of 54 checks passed

wsmoses deleted the ap/random_numbers branch December 18, 2024 16:39

wsmoses added a commit that referenced this pull request Dec 18, 2024

Revert "feat: tracing Random.jl functionality correctly (#363)"

8249408

This reverts commit 94e9576.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: tracing Random.jl functionality correctly #363

feat: tracing Random.jl functionality correctly #363

avik-pal commented Dec 11, 2024 •

edited

Loading

avik-pal commented Dec 11, 2024 •

edited

Loading

avik-pal commented Dec 11, 2024

wsmoses left a comment

avik-pal commented Dec 15, 2024

avik-pal Dec 17, 2024

avik-pal Dec 17, 2024 •

edited

Loading

avik-pal Dec 17, 2024

wsmoses Dec 17, 2024

avik-pal Dec 17, 2024

github-actions bot Dec 18, 2024

github-actions bot Dec 18, 2024

github-actions bot Dec 18, 2024

avik-pal commented Dec 18, 2024

feat: tracing Random.jl functionality correctly #363

feat: tracing Random.jl functionality correctly #363

Conversation

avik-pal commented Dec 11, 2024 • edited Loading

TODOs

avik-pal commented Dec 11, 2024 • edited Loading

avik-pal commented Dec 11, 2024

wsmoses left a comment

Choose a reason for hiding this comment

avik-pal commented Dec 15, 2024

avik-pal Dec 17, 2024

Choose a reason for hiding this comment

avik-pal Dec 17, 2024 • edited Loading

Choose a reason for hiding this comment

avik-pal Dec 17, 2024

Choose a reason for hiding this comment

wsmoses Dec 17, 2024

Choose a reason for hiding this comment

avik-pal Dec 17, 2024

Choose a reason for hiding this comment

github-actions bot Dec 18, 2024

Choose a reason for hiding this comment

github-actions bot Dec 18, 2024

Choose a reason for hiding this comment

github-actions bot Dec 18, 2024

Choose a reason for hiding this comment

avik-pal commented Dec 18, 2024

avik-pal commented Dec 11, 2024 •

edited

Loading

avik-pal commented Dec 11, 2024 •

edited

Loading

avik-pal Dec 17, 2024 •

edited

Loading