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

0.3.0: Performance Improvement #2

Merged
merged 12 commits into from
Jun 12, 2023
Merged

0.3.0: Performance Improvement #2

merged 12 commits into from
Jun 12, 2023

Commits on Jun 12, 2023

  1. refactor(jit): make all jitted functions "inline=True" 

    This should be more beneficial, as followed by the discussions in Jax repository, see
jax-ml/jax#6584 jax-ml/jax#6681 jax-ml/jax#9298 jax-ml/jax#9342
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    d4b9201 View commit details
    Browse the repository at this point in the history

  2. perf(pipeline): try to render 4/2/1 rows per batch using vmap to redu… 

    …ce fori_loop iterations
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    2af8030 View commit details
    Browse the repository at this point in the history

  3. feat(_meta_utils): simple way to add multiple trace annotations toget… 

    …her for functions

add `@ad_tracing_name` to most functions to assist profiling
also bump to Python 3.10

BREAKING CHANGE: Now requires Python 3.10
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    696bb1a View commit details
    Browse the repository at this point in the history

  4. perf(pipeline): big refactor to not updating per rows, but renders al… 

    …l rows then concat and merge
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    1ace375 View commit details
    Browse the repository at this point in the history

  5. perf(pipeline): using scan + unroll (equiv map + unroll) 

    This is very similar to map + vmap (minibatch processing) as the inner
loop is too complex
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    03b4da9 View commit details
    Browse the repository at this point in the history

  6. build(pyproject): try to relax jax{lib}'s verion constraint to ">=0.3… 

    ….25,<5.0.0"
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    69a2735 View commit details
    Browse the repository at this point in the history

  7. test(pre-gen-brax): example inputs for profiling

    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    688cafd View commit details
    Browse the repository at this point in the history

  8. perf: try to eliminate all lax.cond 

    under `vmap`, `lax.cond` are lowered to `select_n` in HLO which leads to execution in both branches,
thus fails to 1) save computation when possible; 2) prevent unexpected values to be
produced/unexpected branches to be executed (defensive), thus let the non-dummy branch to be
executed anyway and only rule-out garbage value at the final stage all together to try to improve
performance. See google/brax#8409 for more details about unconditional executation of cond under
vmap
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    e868968 View commit details
    Browse the repository at this point in the history

  9. fix(pipeline): gl_FrontFacing: fix its determination in pipeline 

    `True` if NOT back-facing
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    e51b618 View commit details
    Browse the repository at this point in the history

  10. perf: added extra stage in pipeline, aiming to interpolate and shade … 

    …only one fragment per pixel
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    7ec527b View commit details
    Browse the repository at this point in the history

  11. docs(changelog): expose option loop_unroll; dependency version change 

    Bump minimum Python version from 3.9 to 3.10;
lower minimum jax & jaxlib to 0.3.25.
    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    c04fdde View commit details
    Browse the repository at this point in the history

  12. build(pyproject): bump to 0.3.0

    @JoeyTeng
    JoeyTeng committed Jun 12, 2023
    Configuration menu
    Copy the full SHA
    3f15f30 View commit details
    Browse the repository at this point in the history