fix: handle traced array returns inside objects #417

avik-pal · 2024-12-21T17:48:14Z

fixes #416

wsmoses · 2024-12-21T19:00:22Z

src/Compiler.jl

+                    setfield!(
+                        $unflatcode,
+                        $(Meta.quot(path[end])),
+                        ConcreteRArray{eltype($final_val),ndims($final_val)}(


this feels like it can lead to an issue of the concreterarray aliases with another one

the loop here is over the concrete_res_name which is returned from the XLA Executable. Won't that ensure it is not aliased?

Yeah sorry it’s not aliasing of the insides but the aliasing of objects we’re passing into compile.

for example something like

x = to_rarray([3])
T = (x,)
Function ip(m)
m[1] .= 2
M
End
res = @jit ip(m)
m[1] .== 6
pointer(m[1]) == pointer(res[1])

So in this case the actual issue is we want to make sure it has the same aliasing properties

So constructing a new concreterarray will not Alia’s anything. If this part of the code deals with new allocations then yeah it’s fine (and correct), otherwise we should make sure things have the same internal aliasing structure

Added this as a testcase

wsmoses · 2024-12-23T18:33:33Z

src/Compiler.jl

+                        if $final_val isa TracedRArray
+                            traced_setfield!(
+                                $unflatcode,
+                                $(Meta.quot(path[end])),


shouldn't this set the inner data ptr? it looks liked traced_setfield is just setfield?

ah wait I understand we are somehow getting a traced value here.

we should presumably cache the tracedrarray in an iddict in case the same the same tracedrarray/number is returned twice and use the same concreterarray [and not create a new one]

I suppse an extension to your case would be something like

tmp = x .+ 1 stret.st = (tmp, tmp)

and obviously we should have the same concreterarray for both terms of the pair

src/Compiler.jl

commit 6556944 Author: Avik Pal <avikpal@mit.edu> Date: Tue Dec 31 08:53:39 2024 -0500 feat: support Base.stack (EnzymeAD#433) * refactor: use scatter for generating diagm * refactor: directly generate the region for simple_scatter_op * feat: generalize diagm * feat: support Base.stack * fix: incorrect rebase * test: stack tests commit 9375f57 Author: Sergio Sánchez Ramírez <15837247+mofeing@users.noreply.github.com> Date: Mon Dec 30 22:56:14 2024 +0100 Modularize Bazel build (EnzymeAD#421) * organize platforms and toolchains * hardcode libcxxwrap_julia path * format code * remove outdated hardcoded symbolic links * add third party bazel wrapper to libcxxwrap_julia * readd platforms * some small fixes * first step on moving externals to modular organization * refactor libcxxwrap_julia on top of `cc_import` * use modular workspaces * add `libcxxwrap_julia` as dependency * hardcode julia dep * export `reactant_*` functions * downgrade libcxxwrap_julia to v0.13.3 * fix major version when linking to libcxxwrap_julia * remove legacy export * move `API.cpp` to new `src/` folder to start modularizing code * export `register_julia_module` from libcxxwrap_julia * fix symbol visibility * clean code remove libcxxwrap and julia deps * format code * import hedron compile commands from Enzyme-JAX * move deps commits to `workspace.bzl` commit 3244204 Author: Avik Pal <avikpal@mit.edu> Date: Mon Dec 30 16:33:52 2024 -0500 chore: bump jll (EnzymeAD#437) commit 25abfe4 Author: Avik Pal <avikpal@mit.edu> Date: Mon Dec 30 12:03:20 2024 -0500 fix: try building with cudnn 9.4 (EnzymeAD#436) commit 241fd14 Author: Avik Pal <avikpal@mit.edu> Date: Mon Dec 30 04:54:50 2024 -0500 feat: indexing using traced values (EnzymeAD#434) * feat: indexing using traced values * feat: implement repeat inner * feat: support scalar linear indexing + tests * fix: regression in cartesian index support * Update src/TracedRArray.jl commit 7d2b898 Author: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Date: Sun Dec 29 20:28:23 2024 -0500 Regenerate MLIR Bindings (EnzymeAD#435) Co-authored-by: mofeing <15837247+mofeing@users.noreply.github.com> commit 8e4c095 Author: Avik Pal <avikpal@mit.edu> Date: Sun Dec 29 13:48:27 2024 -0500 feat: add support for the remaining wrapper types (EnzymeAD#369) * feat: add materialize_traced_array for all other wrappers * refactor: use scatter for generating diagm * refactor: directly generate the region for simple_scatter_op * feat: generalize diagm * feat: efficient non-contiguous setindex * fix: non-contiguous indexing is now supported * feat: implement set_mlir_data for the remaining types * refactor: use `Ops.gather_getindex` to implement diag * fix: noinline ops * fix: incorrect rebase * fix: dispatches * fix: diagm for repeated indices and initial tests * fix: higher dimensional indexing + tests * fix: matrix multiplication of wrapper types * fix: de-specialize 3 arg mul! commit d4e7c76 Author: William Moses <gh@wsmoses.com> Date: Sat Dec 28 23:05:51 2024 -0500 CUDA kernels take 3 (EnzymeAD#427) * CUDA take 3 * conditional run cuda * Update test/integration/cuda.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * bump enzymexla * fix * fix gpu reg * Update BUILD * Update BUILD * Update Project.toml * Update ReactantCUDAExt.jl * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * fix reactant method blocker * Update ReactantCUDAExt.jl * only do compile * use names in cache * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * cleanup further gc issues * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * fix --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> commit b0a58bd Author: Avik Pal <avikpal@mit.edu> Date: Sat Dec 28 21:53:29 2024 -0500 fix: make eltype of Traced/Concrete Arrays to be respective RNumbers (EnzymeAD#426) * feat: overlay eltype conversion * fix: overload the main methods * fix: make eltype of Traced/Concrete Arrays to be respective RNumbers * fix: handle more cases * fix: tracing of wrapped types * fix: arrayinterface overload * fix: python call commit f079a9d Author: mofeing <15837247+mofeing@users.noreply.github.com> Date: Sun Dec 29 00:15:51 2024 +0000 Format code commit eeaf86c Author: glounes <130663280+glou-nes@users.noreply.github.com> Date: Sat Dec 28 04:24:38 2024 +0100 `stablehlo.sort` Ops (EnzymeAD#374) * `stablehlo.sort` Ops * review * use `return_dialect` * feedback * fix test GPU commit 925544f Author: William Moses <gh@wsmoses.com> Date: Tue Dec 24 15:34:52 2024 -0500 Cuv2 (EnzymeAD#423) * Kernel-supporting jll * fix rulescc * adapt to hedron dep * init target * fixup * additional fixups * fixup * fix * registry utils * callname * reg * fix * fix bld * cleanup * no pip * fix * force rules python to older version before bug * fixup jll * with proto * fix * fix * Update WORKSPACE * more deps for apple * bump * fix * workspace bump * workspace * Update Compiler.jl * Update ReactantCUDAExt.jl * Update ReactantCUDAExt.jl * Update ReactantCUDAExt.jl * Update ReactantCUDAExt.jl * Update ReactantCUDAExt.jl * Update Project.toml * Update ReactantCUDAExt.jl * Update Project.toml * Update Project.toml * fix * Update ReactantCUDAExt.jl * Update ReactantCUDAExt.jl * Update ReactantCUDAExt.jl * Update cuda.jl * Update cuda.jl * Update cuda.jl * Cuda kernel v2 * Update Project.toml * Update API.cpp * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: William Moses <wsmoses@cyclops.juliacomputing.io> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> commit 057e6b8 Author: Avik Pal <avikpal@mit.edu> Date: Tue Dec 24 21:46:23 2024 +0530 fix: handle traced array returns inside objects (EnzymeAD#417) * fix: handle traced array returns inside objects * test: add EnzymeAD#416 as a test * fix: propagate track_numbers correctly * fix: aliasing and add a test * test: use updated API for the tests * feat: cache new arrays * fix: traced_getfield commit 0b6dafc Author: Sergio Sánchez Ramírez <15837247+mofeing@users.noreply.github.com> Date: Tue Dec 24 10:43:18 2024 +0100 Bump Reactant_jll to v0.0.32 commit a02fd5b Author: William Moses <gh@wsmoses.com> Date: Mon Dec 23 21:58:40 2024 -0500 Update WORKSPACE commit 6e1710d Author: William Moses <gh@wsmoses.com> Date: Mon Dec 23 19:48:12 2024 -0500 disable absint of absint (EnzymeAD#424) * disable absint of absint * no typeinf ext * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> commit 228732f Author: William Moses <gh@wsmoses.com> Date: Mon Dec 23 18:50:34 2024 -0500 Fix error on global (EnzymeAD#422) commit 695cc80 Author: William Moses <gh@wsmoses.com> Date: Mon Dec 23 13:59:31 2024 -0500 Update Project.toml commit 38916f5 Author: Avik Pal <avikpal@mit.edu> Date: Mon Dec 23 21:46:22 2024 +0530 feat: add zero and fill! for ConcreteRArray (EnzymeAD#420) * feat: add zero and fill! for ConcreteRArray * test: add tests commit 6571d54 Author: William S. Moses <gh@wsmoses.com> Date: Sun Dec 22 23:49:03 2024 -0500 bump enzymexla commit commit 5b89b56 Author: William Moses <gh@wsmoses.com> Date: Sun Dec 22 21:23:33 2024 -0500 Fix ReactantPythonCallExt.jl (EnzymeAD#419) * Fix ReactantPythonCallExt.jl * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> commit 4cc000c Author: William Moses <gh@wsmoses.com> Date: Sun Dec 22 19:55:58 2024 -0500 Improve reactant error messages (EnzymeAD#418) * Improve reactant error messages * More exported symbols * Update XLA.jl * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> commit 2759c3c Author: jumerckx <31353884+jumerckx@users.noreply.github.com> Date: Mon Dec 23 01:35:17 2024 +0100 Inference cache (EnzymeAD#405) * add inference cache * start from `typeinf_ircode` * julia 1.10 * Apply formatting suggestions Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * remove debug logging * vendor in type inference code for v1.10 To avoid having to build a MethodInstance twice (performance hazard) --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Jules Merckx <jumerckx@mac.local> commit f9c43ad Author: William S. Moses <gh@wsmoses.com> Date: Sun Dec 22 19:33:25 2024 -0500 Bump enzymexla

avik-pal force-pushed the ap/enzyme_patch branch 2 times, most recently from 3633105 to 80bfb44 Compare December 21, 2024 18:28

wsmoses reviewed Dec 21, 2024

View reviewed changes

avik-pal force-pushed the ap/enzyme_patch branch 2 times, most recently from c0b8f27 to b8b5e44 Compare December 22, 2024 09:02

avik-pal marked this pull request as ready for review December 22, 2024 12:02

avik-pal requested a review from wsmoses December 22, 2024 16:01

avik-pal force-pushed the ap/enzyme_patch branch from 9f6f50c to 5395887 Compare December 23, 2024 08:55

avik-pal marked this pull request as draft December 23, 2024 12:07

avik-pal force-pushed the ap/enzyme_patch branch from 5395887 to dab1255 Compare December 23, 2024 18:10

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

avik-pal force-pushed the ap/enzyme_patch branch 2 times, most recently from 7b2bf0d to f94034d Compare December 23, 2024 18:30

wsmoses reviewed Dec 23, 2024

View reviewed changes

avik-pal force-pushed the ap/enzyme_patch branch from f94034d to afbf199 Compare December 23, 2024 18:46

avik-pal and others added 6 commits December 24, 2024 10:41

fix: handle traced array returns inside objects

1fd966e

test: add #416 as a test

e557f75

fix: propagate track_numbers correctly

c003fa9

fix: aliasing and add a test

c154237

test: use updated API for the tests

327d907

feat: cache new arrays

ba67359

avik-pal force-pushed the ap/enzyme_patch branch from 00ae662 to ba67359 Compare December 24, 2024 05:11

github-actions bot reviewed Dec 24, 2024

View reviewed changes

src/Compiler.jl Outdated Show resolved Hide resolved

avik-pal force-pushed the ap/enzyme_patch branch 3 times, most recently from 3256abe to b89d355 Compare December 24, 2024 06:47

avik-pal mentioned this pull request Dec 24, 2024

fix: update to new reactant changes LuxDL/Lux.jl#1140

Merged

4 tasks

wsmoses approved these changes Dec 24, 2024

View reviewed changes

fix: traced_getfield

884f378

avik-pal force-pushed the ap/enzyme_patch branch from b89d355 to 884f378 Compare December 24, 2024 15:34

avik-pal merged commit 057e6b8 into main Dec 24, 2024
36 of 39 checks passed

avik-pal deleted the ap/enzyme_patch branch December 24, 2024 16:16

fix: handle traced array returns inside objects #417

fix: handle traced array returns inside objects #417

Uh oh!

Conversation

avik-pal commented Dec 21, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wsmoses Dec 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

wsmoses Dec 22, 2024 •

edited

Loading