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

Modularize Bazel build #421

Merged
merged 27 commits into from
Dec 30, 2024
Merged

Modularize Bazel build #421

merged 27 commits into from
Dec 30, 2024

Conversation

mofeing
Copy link
Collaborator

@mofeing mofeing commented Dec 23, 2024 •
edited
Loading

Working on #370, i realized it was maybe time to refactor the build and make it more modular, like other Bazel-based projects have.

This PR moves the dependencies to different files in a modular-manner. This organization also works if we move to higher versions of Bazel, where WORKSPACE is deprecated and Bazelmod must be used.

  • BUILD should most probably be empty, but still required so that Bazel recognizes that folder as a package (like __init__.py in Python).
  • workspace.bzl defines a repo function that imports the repository. This repo function is then called from WORKSPACE.
  • *.BUILD in the case of local repositories or repos which are not Bazel-build, this file should wrap it to make Bazel-compatible. check out Refactor IFRT bindings to CxxWrap #370 how i did it for Julia and libcxxwrap_julia.

Btw, I've also started modularizing the sources of ReactantExtra into different files. The reason is that currently we only have 1 big file, "API.cpp", that if we divide into different files, compile time gets reduced a lot for the files involving ReactantExtra. It's nice for development. For example, when moving the IFRT part to another file (still wip), compile time (caching all XLA and MLIR deps) reduced from 60s to 30s in my laptop.

@mofeing mofeing requested a review from wsmoses December 23, 2024 10:07
wsmoses
wsmoses reviewed Dec 23, 2024
View reviewed changes
deps/ReactantExtra/third_party/hedron_compile_commands/workspace.bzl Outdated
"""Loads Hedron's Compile Commands Extractor for Bazel."""
# https://github.com/hedronvision/bazel-compile-commands-extractor

load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
Copy link
Member

Choose a reason for hiding this comment

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

Ideally we should make this export from enzyme-jax itself

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

do you mean importing this rule from enzyme-jax itself? not 100% how can it be done right now as Enzyme-JAX is (i.e. without the modular deps), but i guess it doesn't matter much if we leave it like it is and remove it in the future once we "modularize" the Enzyme-JAX build

Copy link
Member

Choose a reason for hiding this comment

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

yeah fair nuff, we should probably fix that up then use the corresponding imports here

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, lemme do the same kind of PR in Enzyme-JAX

@mofeing mofeing mentioned this pull request Dec 24, 2024
Merged
@mofeing
Copy link
Collaborator Author

mofeing commented Dec 24, 2024

The PR for the modularization of the Bazel deps in Enzyme-JAX is open now. Once it's merged, I'll refactor this PR on top of it.

@glou-nes glou-nes mentioned this pull request Dec 28, 2024
Draft
mofeing added 22 commits December 28, 2024 23:50
@mofeing
organize platforms and toolchains
5906ecb
@mofeing
hardcode libcxxwrap_julia path
12845ef
@mofeing
format code
0273a35
@mofeing
remove outdated hardcoded symbolic links
750703b
@mofeing
add third party bazel wrapper to libcxxwrap_julia
038045d
@mofeing
readd platforms
54c167e
@mofeing
some small fixes
ee5908c
@mofeing
first step on moving externals to modular organization
eb05d33
@mofeing
refactor libcxxwrap_julia on top of cc_import
4e9e3eb
@mofeing
use modular workspaces
3631928
@mofeing
add libcxxwrap_julia as dependency
aaf730c
@mofeing
hardcode julia dep
4b1a16a
@mofeing
export reactant_* functions
0763788
@mofeing
downgrade libcxxwrap_julia to v0.13.3
29ea440
@mofeing
fix major version when linking to libcxxwrap_julia
4a305ac
@mofeing
remove legacy export
0901a97
@mofeing
move API.cpp to new src/ folder to start modularizing code
985edd4
@mofeing
export register_julia_module from libcxxwrap_julia
c810660
@mofeing
fix symbol visibility
2d799ef
@mofeing
clean code
c8c6a86 
remove libcxxwrap and julia deps
@mofeing
format code
eadaa06
@mofeing
import hedron compile commands from Enzyme-JAX
d98e9e5
@mofeing mofeing requested a review from wsmoses December 28, 2024 22:52
@mofeing
Copy link
Collaborator Author

mofeing commented Dec 28, 2024

@wsmoses this PR should be ready. do you want to move the commit vars to a workspace.bzl, just like in Enzyme-JAX?

@wsmoses
Copy link
Member

wsmoses commented Dec 28, 2024

Yeah just because we will be constantly doing version bumps doing this all in one place will be best

@mofeing
Copy link
Collaborator Author

mofeing commented Dec 28, 2024

done

@wsmoses
Copy link
Member

wsmoses commented Dec 29, 2024

sorry @mofeing I'm going to ask that we merge #427 first since that contains jll changes that have been already built in the last jll in yggdrassil but haven't been merged (and we should merge as is and then rebase on top of it, otherwise we won't have a way to rebuild this present jll if need be)

wsmoses
wsmoses approved these changes Dec 29, 2024
View reviewed changes
Copy link
Member

@wsmoses wsmoses left a comment

Choose a reason for hiding this comment

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

LGTM but yeah we need to land #427 first [which hopefully happens tonight] per reason above

@mofeing
Copy link
Collaborator Author

mofeing commented Dec 29, 2024

okay, no problem. i will wait for it to land.

just ping me whenever is ready.

@mofeing mofeing mentioned this pull request Dec 29, 2024
Open
@mofeing mofeing merged commit 9375f57 into main Dec 30, 2024
21 of 39 checks passed
@mofeing mofeing deleted the modularize-bazel branch December 30, 2024 21:56
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 <[email protected]>
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 <[email protected]>
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 <[email protected]>
Date:   Mon Dec 30 16:33:52 2024 -0500

    chore: bump jll (EnzymeAD#437)

commit 25abfe4
Author: Avik Pal <[email protected]>
Date:   Mon Dec 30 12:03:20 2024 -0500

    fix: try building with cudnn 9.4 (EnzymeAD#436)

commit 241fd14
Author: Avik Pal <[email protected]>
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 <[email protected]>

commit 8e4c095
Author: Avik Pal <[email protected]>
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 <[email protected]>
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 <[email protected]>
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 <[email protected]>
Date:   Sun Dec 29 00:15:51 2024 +0000

    Format code

commit eeaf86c
Author: glounes <[email protected]>
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 <[email protected]>
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 <[email protected]>
    Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

commit 057e6b8
Author: Avik Pal <[email protected]>
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 <[email protected]>
Date:   Tue Dec 24 10:43:18 2024 +0100

    Bump Reactant_jll to v0.0.32

commit a02fd5b
Author: William Moses <[email protected]>
Date:   Mon Dec 23 21:58:40 2024 -0500

    Update WORKSPACE

commit 6e1710d
Author: William Moses <[email protected]>
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 <[email protected]>
Date:   Mon Dec 23 18:50:34 2024 -0500

    Fix error on global (EnzymeAD#422)

commit 695cc80
Author: William Moses <[email protected]>
Date:   Mon Dec 23 13:59:31 2024 -0500

    Update Project.toml

commit 38916f5
Author: Avik Pal <[email protected]>
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 <[email protected]>
Date:   Sun Dec 22 23:49:03 2024 -0500

    bump enzymexla commit

commit 5b89b56
Author: William Moses <[email protected]>
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 <[email protected]>
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 <[email protected]>
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 <[email protected]>

commit f9c43ad
Author: William S. Moses <[email protected]>
Date:   Sun Dec 22 19:33:25 2024 -0500

    Bump enzymexla
wsmoses added a commit that referenced this pull request Jan 1, 2025
@wsmoses
Revert "Modularize Bazel build (#421)"
475d0b4 
This reverts commit 9375f57.
wsmoses added a commit that referenced this pull request Jan 2, 2025
@wsmoses
Revert modularize bazel build to unbreak yggy (#453)
a49bd13 
* Revert "Update workspace.bzl"

This reverts commit 7534508.

* Revert "Modularize Bazel build (#421)"

This reverts commit 9375f57.

* Revert bazel redo
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@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.
2 participants
@mofeing @wsmoses