fori_loop_unroll - Branch - GalSim-developers/JAX-GalSim

perf: try using fori_loop but unrolling everything

#141Merged

Comparing

fori_loop_unroll

(

2dac00e

) with

main

(

3e45fb5

)

Untouched: 29

New: 2

Benchmarks

New

test_benchmark_spergel_calcfluxrad[compile]

tests/jax/test_benchmarks.py::test_benchmark_spergel_calcfluxrad[compile]

N/A

N/A15.4 s

test_benchmark_spergel_calcfluxrad[run]

tests/jax/test_benchmarks.py::test_benchmark_spergel_calcfluxrad[run]

N/A

N/A2.1 ms

Passed

test_benchmark_moffat_init[compile]

tests/jax/test_benchmarks.py::test_benchmark_moffat_init[compile]

+8%

500.6 ms464.9 ms

test_benchmark_interpimage_flux_frac[run]

tests/jax/test_benchmarks.py::test_benchmark_interpimage_flux_frac[run]

+4%

2.8 ms2.7 ms

test_benchmark_invert_ab_noraise[compile]

tests/jax/test_benchmarks.py::test_benchmark_invert_ab_noraise[compile]

+1%

1.6 s1.6 s

test_benchmark_rng_discard[compile]

tests/jax/test_benchmarks.py::test_benchmark_rng_discard[compile]

234.6 ms233.5 ms

test_benchmarks_lanczos_interp[kval-no_conserve_dc-run]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[kval-no_conserve_dc-run]

385.3 µs384.7 µs

test_benchmarks_lanczos_interp[kval-no_conserve_dc-compile]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[kval-no_conserve_dc-compile]

110.3 ms110.2 ms

test_benchmark_spergel_kvalue[run]

tests/jax/test_benchmarks.py::test_benchmark_spergel_kvalue[run]

137.3 ms137.2 ms

test_benchmark_spergel_conv[compile]

tests/jax/test_benchmarks.py::test_benchmark_spergel_conv[compile]

3.3 s3.3 s

test_benchmark_gaussian_init[compile]

tests/jax/test_benchmarks.py::test_benchmark_gaussian_init[compile]

15.8 ms15.8 ms

test_benchmarks_metacal[run]

tests/jax/test_benchmarks.py::test_benchmarks_metacal[run]

185 ms184.9 ms

test_benchmarks_interpolated_image[compile]

tests/jax/test_benchmarks.py::test_benchmarks_interpolated_image[compile]

1.8 s1.8 s

test_benchmark_spergel_init[compile]

tests/jax/test_benchmarks.py::test_benchmark_spergel_init[compile]

16 ms16 ms

test_benchmark_spergel_xvalue[run]

tests/jax/test_benchmarks.py::test_benchmark_spergel_xvalue[run]

24.5 s24.5 s

test_benchmark_spergel_xvalue[compile]

tests/jax/test_benchmarks.py::test_benchmark_spergel_xvalue[compile]

27.7 s27.7 s

test_benchmarks_lanczos_interp[xval-conserve_dc-compile]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[xval-conserve_dc-compile]

124.3 ms124.3 ms

test_benchmark_gaussian_init[run]

tests/jax/test_benchmarks.py::test_benchmark_gaussian_init[run]

85 µs85 µs

test_benchmark_spergel_conv[run]

tests/jax/test_benchmarks.py::test_benchmark_spergel_conv[run]

2.6 s2.6 s

test_benchmark_spergel_kvalue[compile]

tests/jax/test_benchmarks.py::test_benchmark_spergel_kvalue[compile]

727.1 ms727.3 ms

test_benchmarks_lanczos_interp[kval-conserve_dc-compile]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[kval-conserve_dc-compile]

110.4 ms110.4 ms

test_benchmark_moffat_init[run]

tests/jax/test_benchmarks.py::test_benchmark_moffat_init[run]

408.6 µs408.8 µs

test_benchmarks_lanczos_interp[xval-no_conserve_dc-run]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[xval-no_conserve_dc-run]

952.1 µs953.3 µs

test_benchmark_invert_ab_noraise[run]

tests/jax/test_benchmarks.py::test_benchmark_invert_ab_noraise[run]

720.3 µs721.2 µs

test_benchmark_rng_discard[run]

tests/jax/test_benchmarks.py::test_benchmark_rng_discard[run]

6.9 ms6.9 ms

test_benchmarks_lanczos_interp[xval-no_conserve_dc-compile]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[xval-no_conserve_dc-compile]

80.6 ms80.7 ms

test_benchmarks_lanczos_interp[kval-conserve_dc-run]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[kval-conserve_dc-run]

385.3 µs386.3 µs

test_benchmark_interpimage_flux_frac[compile]

tests/jax/test_benchmarks.py::test_benchmark_interpimage_flux_frac[compile]

-1%

234.4 ms235.7 ms

test_benchmarks_interpolated_image[run]

tests/jax/test_benchmarks.py::test_benchmarks_interpolated_image[run]

-1%

46.3 ms46.6 ms

test_benchmark_spergel_init[run]

tests/jax/test_benchmarks.py::test_benchmark_spergel_init[run]

-1%

83.8 µs85 µs

test_benchmarks_lanczos_interp[xval-conserve_dc-run]

tests/jax/test_benchmarks.py::test_benchmarks_lanczos_interp[xval-conserve_dc-run]

-9%

1.1 ms1.2 ms

Commits

Click on a commit to change the comparison range

Base

main

3e45fb5

+0.47%

perf: try using fori_loop but unrolling everything

0571c59

8 months ago

by beckermr

+0.17%

fix: cannot unroll discard since n is not static

5945ad1

8 months ago

by beckermr

-0.16%

test: reduce matrix

cca6a20

8 months ago

by beckermr

-0.46%

Update utils.py

e2fe74d

8 months ago

by beckermr

-53.42%

perf: try partial loop unrolling

0d67579

8 months ago

by beckermr

+53.84%

test: add spergel init benchmark

d15e76d

8 months ago

by beckermr

-0.45%

style: pre the commit

2dac00e

8 months ago

by beckermr

Home Terms Privacy Docs