Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Inference cache #405

Merged
merged 6 commits into from
Dec 23, 2024
Merged

Inference cache #405

merged 6 commits into from
Dec 23, 2024

Conversation

jumerckx
Copy link
Collaborator

@jumerckx jumerckx commented Dec 19, 2024 •
edited
Loading

This doesn't work yet! It only caches the top level codeinfo, some other absint method needs to be overwritten to do it properly.
Only tested on Julia 1.11
CI passes, and it seems like the cicache is being filled, I don't have a good test to show improvements though.

f(a, b) = a .* b
a = rand(10)
b = rand(10)

a_ra = Reactant.to_rarray(a)
b_ra = Reactant.to_rarray(b)

g(a, b) = @noinline f(a, b)

@code_hlo(g(a, b))

result = CC.get(CC.code_cache(Reactant.ReactantInterpreter()), Base.method_instance(f, Tuple{Vector{Float64}, Vector{Float64}}), nothing)

github-actions[bot]
github-actions bot reviewed Dec 19, 2024
View reviewed changes
src/utils.jl Outdated Show resolved Hide resolved
src/utils.jl Outdated Show resolved Hide resolved
@jumerckx jumerckx marked this pull request as ready for review December 19, 2024 22:43
@jumerckx @github-actions
Apply formatting suggestions
5475f10 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
wsmoses
wsmoses reviewed Dec 20, 2024
View reviewed changes
src/utils.jl Outdated Show resolved Hide resolved
wsmoses
wsmoses reviewed Dec 20, 2024
View reviewed changes
src/utils.jl Outdated
@static if VERSION < v"1.11-"
ir = Core.Compiler.run_passes(opt.src, opt, caller)
@static if VERSION < v"1.11"
ir, rt = CC.typeinf_ircode(interp, method, mi.specTypes, mi.sparam_vals, nothing)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah yeah so this is why I vendored the internals of that code. This causes it to re-run "get methodinstance", which is rather expensive. Can we continue to vendor code which doesn't re-run "get method instance" for 1.10?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay

@mofeing
Copy link
Collaborator

mofeing commented Dec 20, 2024

does this fix #353?

Jules Merckx added 2 commits December 22, 2024 23:13
remove debug logging
dfc272d
vendor in type inference code for v1.10
115da17 
To avoid having to build a MethodInstance twice (performance hazard)
@jumerckx
Copy link
Collaborator Author

does this fix #353?

You mean the segfaults or something else? I doubt this will do anything to the segfault but you never know. Is there a MWE or is still only observed in CI?

@wsmoses wsmoses merged commit 2759c3c into main Dec 23, 2024
36 of 38 checks passed
@wsmoses wsmoses deleted the jm/cicache branch December 23, 2024 00:35
@wsmoses
Copy link
Member

wsmoses commented Dec 23, 2024

does this fix #353?

You mean the segfaults or something else? I doubt this will do anything to the segfault but you never know. Is there a MWE or is still only observed in CI?

Yeah I think this is unrelated (and just compile time speedups)

jumerckx added a commit to jumerckx/Reactant.jl that referenced this pull request Jan 1, 2025
@jumerckx
Squashed commit of the following:
de86c26 
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments

@wsmoses wsmoses wsmoses approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@jumerckx @mofeing @wsmoses