diff --git a/README.md b/README.md index 0a41032..cc15b4d 100644 --- a/README.md +++ b/README.md @@ -45,15 +45,28 @@ Results are framed by **astronomy instrument** (HST, Euclid, JWST, …) rather t **Likelihood breakdown** — latest per-step decompositions: -| Cell | Instrument | Inversion path | Step-sum total | PyAutoLens version | -|------|------------|----------------|----------------|--------------------| -| `cluster/image_plane` | — | dense (mapping) | 4.99 s | v2026.7.6.649 | -| `cluster/source_plane` | — | dense (mapping) | 3.1 ms | v2026.7.6.649 | -| `imaging/delaunay` | hst | dense (mapping) | 11.29 s | v2026.5.29.4 | -| `imaging/delaunay` | hst | sparse (w-tilde) | 7.96 s | v2026.5.29.4 | -| `imaging/mge` | hst | dense (mapping) | 178.2 ms | v2026.5.29.4 | -| `imaging/pixelization` | hst | dense (mapping) | 7.79 s | v2026.5.29.4 | -| `imaging/pixelization` | hst | sparse (w-tilde) | 8.94 s | v2026.5.29.4 | +| Cell | Instrument | Platform | Inversion path | Step-sum total | PyAutoLens version | +|------|------------|----------|----------------|----------------|--------------------| +| `cluster/image_plane` | — | local_cpu_fp64 | dense (mapping) | 4.99 s | v2026.7.6.649 | +| `cluster/source_plane` | — | local_cpu_fp64 | dense (mapping) | 3.1 ms | v2026.7.6.649 | +| `datacube/delaunay` | alma_high | hpc_a100_fp64 | dense (mapping) | — | v2026.7.6.649 | +| `datacube/delaunay` | alma_high | hpc_a100_mp | dense (mapping) | — | v2026.7.6.649 | +| `datacube/inversion` | alma_high | hpc_a100_fp64 | dense (mapping) | — | v2026.7.6.649 | +| `datacube/inversion` | alma_high | hpc_a100_mp | dense (mapping) | — | v2026.7.6.649 | +| `imaging/delaunay` | hst | local_cpu_fp64 | dense (mapping) | 10.07 s | v2026.7.6.649 | +| `imaging/delaunay` | hst | local_cpu_fp64 | sparse (w-tilde) | 8.81 s | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_fp64 | dense (mapping) | 96.6 ms | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_fp64 | sparse (w-tilde) | 98.0 ms | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_mp | dense (mapping) | 96.8 ms | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_mp | sparse (w-tilde) | 95.5 ms | v2026.7.6.649 | +| `imaging/mge` | hst | local_cpu_fp64 | dense (mapping) | 179.5 ms | v2026.7.6.649 | +| `imaging/mge` | hst | hpc_a100_fp64 | dense (mapping) | 7.8 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | local_cpu_fp64 | dense (mapping) | 8.65 s | v2026.7.6.649 | +| `imaging/pixelization` | hst | local_cpu_fp64 | sparse (w-tilde) | 10.17 s | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_fp64 | dense (mapping) | 57.6 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_fp64 | sparse (w-tilde) | 57.9 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_mp | dense (mapping) | 56.4 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_mp | sparse (w-tilde) | 55.5 ms | v2026.7.6.649 | The tables above are auto-generated by `scripts/build_readme.py` from the artifacts under [`results/`](./results/README.md) — never edit them by hand; run `python scripts/build_readme.py` after a profiling run and commit the result (CI checks idempotence via `--check`). Narrative context — per-cell "where to optimize next" recommendations and the mp-vs-fp64 verdicts — lives in [`likelihood_runtime/OPTIMIZATION_NOTES.md`](./likelihood_runtime/OPTIMIZATION_NOTES.md). @@ -108,7 +121,7 @@ This repo is being built in phases (bootstrap history now archived in `PyAutoMin | 4 | Top-level + per-section README dashboard with instrument framing | ✓ shipped | | 5 | GitHub Actions for lint + profile re-runs + README refresh | ✓ shipped (`lint.yml` per-PR; `profile.yml` manual/on-release) | | 6 | Design lock-in + results/dashboard groundwork ([#52](https://github.com/PyAutoLabs/autolens_profiling/issues/52)) | in progress | -| 7 | **PreOptimizationTimes** baseline campaign (vram-first, then runtime + breakdown) | queued | +| 7 | **PreOptimizationTimes** baseline campaign (vram-first, then runtime + breakdown) | ✓ shipped (runtime [#56](https://github.com/PyAutoLabs/autolens_profiling/issues/56); breakdown + dashboard [#59](https://github.com/PyAutoLabs/autolens_profiling/issues/59); laptop-GPU legs extend in a later re-run) | ### Future enhancements (Phase 4 follow-ups) diff --git a/_profile_cli.py b/_profile_cli.py index a117d83..f106716 100644 --- a/_profile_cli.py +++ b/_profile_cli.py @@ -146,6 +146,13 @@ def device_info_dict() -> dict: info = { "backend": jax.default_backend(), "device": str(jax.devices()[0]), + # Environment provenance: a stray XLA_FLAGS (e.g. disabling + # constant_folding) or thread pinning silently rescales every timing + # in a result by integer factors — record them so drift between runs + # is attributable (found the hard way: autolens_profiling#59). + "xla_flags": os.environ.get("XLA_FLAGS") or None, + "omp_num_threads": os.environ.get("OMP_NUM_THREADS") or None, + "cpu_count": os.cpu_count(), } if info["backend"] == "gpu": try: diff --git a/hpc/batch_gpu/submit_breakdown_datacube_delaunay_a100_alma_high_fp64 b/hpc/batch_gpu/submit_breakdown_datacube_delaunay_a100_alma_high_fp64 new file mode 100644 index 0000000..07af59c --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_datacube_delaunay_a100_alma_high_fp64 @@ -0,0 +1,48 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: datacube/delaunay x alma_high x fp64. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_datacube_delaunay_alma_high_fp64 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/datacube/delaunay" +echo "Instrument: alma_high" +echo "Precision: fp64" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/datacube/delaunay.py \ + --config-name hpc_a100_fp64 \ + --instrument alma_high + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_datacube_delaunay_a100_alma_high_mp b/hpc/batch_gpu/submit_breakdown_datacube_delaunay_a100_alma_high_mp new file mode 100644 index 0000000..d637583 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_datacube_delaunay_a100_alma_high_mp @@ -0,0 +1,49 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: datacube/delaunay x alma_high x mp. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_datacube_delaunay_alma_high_mp +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/datacube/delaunay" +echo "Instrument: alma_high" +echo "Precision: mp" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/datacube/delaunay.py \ + --config-name hpc_a100_mp \ + --instrument alma_high \ + --use-mixed-precision + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_datacube_inversion_setup_decompose_a100_alma_high_fp64 b/hpc/batch_gpu/submit_breakdown_datacube_inversion_setup_decompose_a100_alma_high_fp64 new file mode 100644 index 0000000..5f93272 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_datacube_inversion_setup_decompose_a100_alma_high_fp64 @@ -0,0 +1,48 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: datacube/inversion_setup_decompose x alma_high x fp64. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_datacube_inversion_setup_decompose_alma_high_fp64 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/datacube/inversion_setup_decompose" +echo "Instrument: alma_high" +echo "Precision: fp64" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/datacube/inversion_setup_decompose.py \ + --config-name hpc_a100_fp64 \ + --instrument alma_high + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_datacube_inversion_setup_decompose_a100_alma_high_mp b/hpc/batch_gpu/submit_breakdown_datacube_inversion_setup_decompose_a100_alma_high_mp new file mode 100644 index 0000000..30f4f43 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_datacube_inversion_setup_decompose_a100_alma_high_mp @@ -0,0 +1,49 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: datacube/inversion_setup_decompose x alma_high x mp. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_datacube_inversion_setup_decompose_alma_high_mp +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/datacube/inversion_setup_decompose" +echo "Instrument: alma_high" +echo "Precision: mp" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/datacube/inversion_setup_decompose.py \ + --config-name hpc_a100_mp \ + --instrument alma_high \ + --use-mixed-precision + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64 b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64 new file mode 100644 index 0000000..d55f05d --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64 @@ -0,0 +1,47 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/delaunay x hst x fp64. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_delaunay_hst_fp64 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/delaunay" +echo "Instrument: hst" +echo "Precision: fp64" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/delaunay.py \ + --config-name hpc_a100_fp64 + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64_sparse b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64_sparse new file mode 100644 index 0000000..a4af499 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64_sparse @@ -0,0 +1,48 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/delaunay x hst x fp64. +# SPARSE inversion path (w-tilde operator). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_delaunay_hst_fp64_sparse +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/delaunay" +echo "Instrument: hst" +echo "Precision: fp64" +echo "Inversion path: sparse (w-tilde)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/delaunay.py \ + --config-name hpc_a100_fp64 \ + --sparse + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_mp b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_mp new file mode 100644 index 0000000..9244260 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_mp @@ -0,0 +1,48 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/delaunay x hst x mp. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_delaunay_hst_mp +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/delaunay" +echo "Instrument: hst" +echo "Precision: mp" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/delaunay.py \ + --config-name hpc_a100_mp \ + --use-mixed-precision + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_mp_sparse b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_mp_sparse new file mode 100644 index 0000000..65d80e3 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_mp_sparse @@ -0,0 +1,49 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/delaunay x hst x mp. +# SPARSE inversion path (w-tilde operator). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_delaunay_hst_mp_sparse +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/delaunay" +echo "Instrument: hst" +echo "Precision: mp" +echo "Inversion path: sparse (w-tilde)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/delaunay.py \ + --config-name hpc_a100_mp \ + --sparse \ + --use-mixed-precision + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_mge_a100_hst_fp64 b/hpc/batch_gpu/submit_breakdown_imaging_mge_a100_hst_fp64 new file mode 100644 index 0000000..6ef9f97 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_mge_a100_hst_fp64 @@ -0,0 +1,47 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/mge x hst x fp64. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_mge_hst_fp64 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/mge" +echo "Instrument: hst" +echo "Precision: fp64" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/mge.py \ + --config-name hpc_a100_fp64 + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_fp64 b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_fp64 new file mode 100644 index 0000000..209d5b2 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_fp64 @@ -0,0 +1,47 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/pixelization x hst x fp64. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_pixelization_hst_fp64 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/pixelization" +echo "Instrument: hst" +echo "Precision: fp64" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/pixelization.py \ + --config-name hpc_a100_fp64 + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_fp64_sparse b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_fp64_sparse new file mode 100644 index 0000000..a103e54 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_fp64_sparse @@ -0,0 +1,48 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/pixelization x hst x fp64. +# SPARSE inversion path (w-tilde operator). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_pixelization_hst_fp64_sparse +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/pixelization" +echo "Instrument: hst" +echo "Precision: fp64" +echo "Inversion path: sparse (w-tilde)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/pixelization.py \ + --config-name hpc_a100_fp64 \ + --sparse + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_mp b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_mp new file mode 100644 index 0000000..4fe3c03 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_mp @@ -0,0 +1,48 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/pixelization x hst x mp. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_pixelization_hst_mp +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/pixelization" +echo "Instrument: hst" +echo "Precision: mp" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/pixelization.py \ + --config-name hpc_a100_mp \ + --use-mixed-precision + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_mp_sparse b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_mp_sparse new file mode 100644 index 0000000..a027508 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_pixelization_a100_hst_mp_sparse @@ -0,0 +1,49 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/pixelization x hst x mp. +# SPARSE inversion path (w-tilde operator). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_imaging_pixelization_hst_mp_sparse +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/pixelization" +echo "Instrument: hst" +echo "Precision: mp" +echo "Inversion path: sparse (w-tilde)" + +cd $AP_ROOT + +python3 likelihood_breakdown/imaging/pixelization.py \ + --config-name hpc_a100_mp \ + --sparse \ + --use-mixed-precision + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_interferometer_delaunay_a100_alma_high_fp64 b/hpc/batch_gpu/submit_breakdown_interferometer_delaunay_a100_alma_high_fp64 new file mode 100644 index 0000000..d8e33b3 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_interferometer_delaunay_a100_alma_high_fp64 @@ -0,0 +1,48 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: interferometer/delaunay x alma_high x fp64. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_interferometer_delaunay_alma_high_fp64 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/interferometer/delaunay" +echo "Instrument: alma_high" +echo "Precision: fp64" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/interferometer/delaunay.py \ + --config-name hpc_a100_fp64 \ + --instrument alma_high + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_interferometer_delaunay_a100_alma_high_mp b/hpc/batch_gpu/submit_breakdown_interferometer_delaunay_a100_alma_high_mp new file mode 100644 index 0000000..9430924 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_interferometer_delaunay_a100_alma_high_mp @@ -0,0 +1,49 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: interferometer/delaunay x alma_high x mp. +# DENSE inversion path (mapping matrices). +# Per-step JIT decomposition (PreOptimizationTimes campaign, autolens_profiling#59). +# Compile-heavy: every pipeline step is JIT-compiled in isolation, so wall +# time is dominated by XLA compiles, not the steady-state timing loop. + +#SBATCH -J breakdown_interferometer_delaunay_alma_high_mp +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=4:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err +#SBATCH --mail-type=END,FAIL +#SBATCH --mail-user=james.w.nightingale@durham.ac.uk + +export AP_ROOT=/mnt/ral/jnightin/autolens_profiling +source $AP_ROOT/activate.sh + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/interferometer/delaunay" +echo "Instrument: alma_high" +echo "Precision: mp" +echo "Inversion path: dense (mapping)" + +cd $AP_ROOT + +python3 likelihood_breakdown/interferometer/delaunay.py \ + --config-name hpc_a100_mp \ + --instrument alma_high \ + --use-mixed-precision + +echo "Finished." +date diff --git a/likelihood_breakdown/README.md b/likelihood_breakdown/README.md index c094f5b..9d51f7c 100644 --- a/likelihood_breakdown/README.md +++ b/likelihood_breakdown/README.md @@ -11,15 +11,28 @@ For *how long the likelihood actually takes* on production hardware — i.e. a s ## Latest results -| Cell | Instrument | Inversion path | Step-sum total | PyAutoLens version | -|------|------------|----------------|----------------|--------------------| -| `cluster/image_plane` | — | dense (mapping) | 4.99 s | v2026.7.6.649 | -| `cluster/source_plane` | — | dense (mapping) | 3.1 ms | v2026.7.6.649 | -| `imaging/delaunay` | hst | dense (mapping) | 11.29 s | v2026.5.29.4 | -| `imaging/delaunay` | hst | sparse (w-tilde) | 7.96 s | v2026.5.29.4 | -| `imaging/mge` | hst | dense (mapping) | 178.2 ms | v2026.5.29.4 | -| `imaging/pixelization` | hst | dense (mapping) | 7.79 s | v2026.5.29.4 | -| `imaging/pixelization` | hst | sparse (w-tilde) | 8.94 s | v2026.5.29.4 | +| Cell | Instrument | Platform | Inversion path | Step-sum total | PyAutoLens version | +|------|------------|----------|----------------|----------------|--------------------| +| `cluster/image_plane` | — | local_cpu_fp64 | dense (mapping) | 4.99 s | v2026.7.6.649 | +| `cluster/source_plane` | — | local_cpu_fp64 | dense (mapping) | 3.1 ms | v2026.7.6.649 | +| `datacube/delaunay` | alma_high | hpc_a100_fp64 | dense (mapping) | — | v2026.7.6.649 | +| `datacube/delaunay` | alma_high | hpc_a100_mp | dense (mapping) | — | v2026.7.6.649 | +| `datacube/inversion` | alma_high | hpc_a100_fp64 | dense (mapping) | — | v2026.7.6.649 | +| `datacube/inversion` | alma_high | hpc_a100_mp | dense (mapping) | — | v2026.7.6.649 | +| `imaging/delaunay` | hst | local_cpu_fp64 | dense (mapping) | 10.07 s | v2026.7.6.649 | +| `imaging/delaunay` | hst | local_cpu_fp64 | sparse (w-tilde) | 8.81 s | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_fp64 | dense (mapping) | 96.6 ms | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_fp64 | sparse (w-tilde) | 98.0 ms | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_mp | dense (mapping) | 96.8 ms | v2026.7.6.649 | +| `imaging/delaunay` | hst | hpc_a100_mp | sparse (w-tilde) | 95.5 ms | v2026.7.6.649 | +| `imaging/mge` | hst | local_cpu_fp64 | dense (mapping) | 179.5 ms | v2026.7.6.649 | +| `imaging/mge` | hst | hpc_a100_fp64 | dense (mapping) | 7.8 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | local_cpu_fp64 | dense (mapping) | 8.65 s | v2026.7.6.649 | +| `imaging/pixelization` | hst | local_cpu_fp64 | sparse (w-tilde) | 10.17 s | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_fp64 | dense (mapping) | 57.6 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_fp64 | sparse (w-tilde) | 57.9 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_mp | dense (mapping) | 56.4 ms | v2026.7.6.649 | +| `imaging/pixelization` | hst | hpc_a100_mp | sparse (w-tilde) | 55.5 ms | v2026.7.6.649 | Auto-generated by `scripts/build_readme.py` from the versioned artifacts under `results/breakdown/`. @@ -37,11 +50,27 @@ For each pipeline step (e.g. *ray-trace grids* → *blurred mapping matrix* → 3. Asserts the JIT output matches the eager FitImaging / FitInterferometer reference at `rtol=1e-4`, so the per-step decomposition is provably equivalent to the production path. 4. Emits a JSON with `{steps: {name: per_call_s}, total_step_by_step: ...}` and a single horizontal bar chart sorted by step cost. -## Single-config rationale +## Platform policy: aim for GPU, fall back to CPU -The per-step JIT compile is expensive — for a 1000-vertex Delaunay cell it's tens of minutes of compile time. Running the same 10-step decomposition six times (CPU/GPU/A100 × fp64/mp) burns roughly an hour of compile per cell, for marginal new signal: cross-hardware comparisons live in the runtime package. Default is **CPU fp64**. +**A breakdown should always aim to be a full GPU breakdown.** Production +likelihoods run on GPU, and a step's bottleneck shape genuinely changes across +backends (XLA fusion behaviour, sparse-precision-matrix layout, callback +boundaries), so a CPU-only decomposition can point optimization at the wrong +step. The canonical decomposition for a cell is therefore **A100 fp64** (plus +mixed precision where the source supports it), dispatched via the +`hpc/batch_gpu/submit_breakdown_*` scripts. -Opt in to GPU with `--gpu` when you suspect a step's bottleneck shape changes on different backends (XLA fusion behaviour, sparse-precision-matrix layout, etc.). This is the only multi-config knob the breakdown scripts expose. +CPU fp64 remains the fallback when no GPU is available (laptop-only sessions, +HPC downtime) — it is cheap to run locally and still catches the +step-dominance picture for compile-bound and callback-bound steps, but treat it +as provisional until the GPU decomposition confirms it. + +The per-step JIT compile is expensive — for a 1000-vertex Delaunay cell it's +tens of minutes of compile time per config — so the grid stays deliberately +small: one representative instrument per dataset class, GPU (+mp) as the +target, CPU as the fallback. Cross-hardware *runtime* comparisons stay in the +runtime package; this package answers where the time goes, on the hardware +that matters. ## XLA fusion caveat @@ -109,10 +138,10 @@ From the autolens_profiling root: # Default (CPU fp64, sma/hst as appropriate to the cell) python likelihood_breakdown/imaging/mge.py -# GPU backend -python likelihood_breakdown/imaging/mge.py --gpu +# GPU backend (local GPU; the A100 tier goes via hpc/batch_gpu/submit_breakdown_*) +JAX_PLATFORM_NAME=cuda JAX_PLATFORMS=cuda,cpu python likelihood_breakdown/imaging/mge.py -# Mixed precision (rare in breakdown — only when investigating fp32-induced step changes) +# Mixed precision (paired with the fp64 run on GPU tiers; supported by all cells except imaging/mge) python likelihood_breakdown/imaging/mge.py --use-mixed-precision ``` diff --git a/likelihood_breakdown/imaging/delaunay.py b/likelihood_breakdown/imaging/delaunay.py index 8e4d5bb..2a6a840 100644 --- a/likelihood_breakdown/imaging/delaunay.py +++ b/likelihood_breakdown/imaging/delaunay.py @@ -665,6 +665,86 @@ def blurred_mm_from_params(params_tree): print(f" blurred_mapping_matrix (JIT) shape: {bmm_jit.shape}") +# --------------------------------------------------------------------------- +# Optional: four-way split of the inversion-setup block (--split-setup) +# --------------------------------------------------------------------------- +# Nested prefix-JITs of the same staged computation: params -> step-5 output, +# -> step-6, -> step-7, -> step-8. Successive differences attribute the +# combined block's cost to border relocation / triangulation+interpolation / +# mapping matrix / PSF convolution. The differences inherit the fusion caveat +# (XLA may move work across prefix boundaries, so small negatives are noise), +# and every prefix pays the ray-trace preamble (~0.3 ms, measured separately +# in steps 1-2) which lands in the first difference. + +_setup_split: dict | None = None + +if "--split-setup" in sys.argv: + print("\n--- Inversion setup four-way split (--split-setup) ---") + + def _setup_prefix_fn(upto): + def fn(pt): + t = al.Tracer(galaxies=list(pt.galaxies)) + traced_source = t.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp)[-1] + traced_mesh = t.traced_grid_2d_list_from( + grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp + )[-1] + relocated = border_relocator.relocated_grid_from(grid=traced_source, xp=jnp) + relocated_mesh = border_relocator.relocated_mesh_grid_from( + grid=traced_source, mesh_grid=traced_mesh, xp=jnp + ) + if upto == 5: + return relocated.array, relocated_mesh.array + interp = al.InterpolatorDelaunay( + mesh=pixelization_obj.mesh, + mesh_grid=relocated_mesh, + data_grid=relocated, + xp=jnp, + ) + m = al.Mapper( + interpolator=interp, + image_plane_mesh_grid=image_plane_mesh_grid, + xp=jnp, + ) + if upto == 6: + return ( + m.pix_indexes_for_sub_slim_index, + m.pix_weights_for_sub_slim_index, + ) + mm = m.mapping_matrix + if upto == 7: + return mm + return dataset.psf.convolved_mapping_matrix_from( + mapping_matrix=mm, mask=dataset.mask, xp=jnp + ) + + return fn + + _prefix_labels = { + 5: "Border relocation", + 6: "Triangulation + interpolation", + 7: "Mapping matrix", + 8: "Blurred mapping matrix (PSF)", + } + _prefix_per_call = {} + for _upto in (5, 6, 7, 8): + jit_profile(_setup_prefix_fn(_upto), f"setup_prefix_{_upto}", params_tree) + _prefix_per_call[_upto] = timer.records[-1][1] / 10 + + _setup_split = {} + _prev = 0.0 + for _upto in (5, 6, 7, 8): + _setup_split[_prefix_labels[_upto]] = _prefix_per_call[_upto] - _prev + _prev = _prefix_per_call[_upto] + + print( + " prefix per-call: " + + ", ".join(f"5..{u}={_prefix_per_call[u] * 1000:.2f} ms" for u in (5, 6, 7, 8)) + ) + for _label, _dt in _setup_split.items(): + print(f" {_label}: {_dt * 1000:8.2f} ms") + _combined = dict(likelihood_steps)["Inversion setup (steps 5-8 combined)"] + print(f" (combined single-JIT reference: {_combined * 1000:.2f} ms)") + bmm_jnp = bmm_ref # Use the reference matrices for linear algebra steps print(f" blurred_mapping_matrix shape: {blurred_mapping_matrix.shape}") @@ -997,6 +1077,9 @@ def compute_log_evidence( "total_step_by_step": step_total, } +if _setup_split is not None: + breakdown_summary["setup_split"] = {k: float(v) for k, v in _setup_split.items()} + dict_path, chart_path = resolve_output_paths( _cli, default_dir=_workspace_root / "results" / "breakdown" / "imaging", diff --git a/likelihood_breakdown/interferometer/delaunay.py b/likelihood_breakdown/interferometer/delaunay.py index 5e1b8ba..17c7727 100644 --- a/likelihood_breakdown/interferometer/delaunay.py +++ b/likelihood_breakdown/interferometer/delaunay.py @@ -82,7 +82,7 @@ _cli = parse_profile_cli() -instrument = "sma" # <-- change this to profile a different instrument +instrument = _cli.instrument or "sma" # --instrument overrides (default sma) hilbert_pixels = 1500 # 1500-tier production fiducial (matches imaging/datacube) regularization_coefficient = 1.0 diff --git a/results/breakdown/datacube/delaunay_hpc_a100_fp64.json b/results/breakdown/datacube/delaunay_hpc_a100_fp64.json new file mode 100644 index 0000000..cdb9a0b --- /dev/null +++ b/results/breakdown/datacube/delaunay_hpc_a100_fp64.json @@ -0,0 +1,90 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 33375 MiB, 81920 MiB" + }, + "instrument": "alma_high", + "model": "delaunay", + "n_channels": 34, + "dense_breakdown": false, + "configuration": { + "pixel_scale_arcsec": 0.025, + "mask_radius_arcsec": 3.5, + "real_space_shape": [ + 800, + 800 + ], + "visibilities_per_channel": 5000000, + "hilbert_pixels": 1500, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "regularization_coefficient": 1.0, + "transformer_chunk_size": 1000000 + }, + "cube_log_evidence_eager": -2048327890.5681412, + "log_evidence_per_channel_eager": [ + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653, + -60244937.95788653 + ], + "steps_cube_cost": { + "Ray-trace data grid (shared)": 0.00013197309999668505, + "Ray-trace mesh grid (shared)": 0.0001361580998491263, + "Inversion setup, incl. NUFFT (per channel \u00d7 34)": 3.0903537686037454, + "Data vector D, sparse (per channel \u00d7 34)": 0.02265050759924634, + "Curvature matrix F, sparse (per channel \u00d7 34)": 3.1518002481978327, + "Regularization matrix H (shared)": 0.0, + "Reconstruction NNLS (per channel \u00d7 34)": 0.6901009749988589, + "fast_chi_squared + log evidence (per channel \u00d7 34)": 0.077919228004248 + }, + "per_channel_costs": { + "inversion_setup": 0.09089275790011016, + "data_vector": 0.0006661913999778335, + "curvature_matrix": 0.09270000729993626, + "reconstruction": 0.020297087499966438, + "log_evidence": 0.002291742000124941 + }, + "inversion_setup_decomposition": { + "mapper_plus_L_invariant": 0.09131777300008252, + "data_vector_matmul_variant": 0.0 + }, + "total_step_by_step_cube": 7.033092858603777, + "shared_lwl_savings_estimate": 3.0591002408978967, + "shared_invariant_savings_estimate": 6.07258674990062 +} \ No newline at end of file diff --git a/results/breakdown/datacube/delaunay_hpc_a100_fp64.png b/results/breakdown/datacube/delaunay_hpc_a100_fp64.png new file mode 100644 index 0000000..13dd2e2 Binary files /dev/null and b/results/breakdown/datacube/delaunay_hpc_a100_fp64.png differ diff --git a/results/breakdown/datacube/delaunay_hpc_a100_mp.json b/results/breakdown/datacube/delaunay_hpc_a100_mp.json new file mode 100644 index 0000000..5759ed1 --- /dev/null +++ b/results/breakdown/datacube/delaunay_hpc_a100_mp.json @@ -0,0 +1,90 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 33375 MiB, 81920 MiB" + }, + "instrument": "alma_high", + "model": "delaunay", + "n_channels": 34, + "dense_breakdown": false, + "configuration": { + "pixel_scale_arcsec": 0.025, + "mask_radius_arcsec": 3.5, + "real_space_shape": [ + 800, + 800 + ], + "visibilities_per_channel": 5000000, + "hilbert_pixels": 1500, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "regularization_coefficient": 1.0, + "transformer_chunk_size": 1000000 + }, + "cube_log_evidence_eager": -2048327890.5799994, + "log_evidence_per_channel_eager": [ + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353, + -60244937.9582353 + ], + "steps_cube_cost": { + "Ray-trace data grid (shared)": 0.00013103119999868795, + "Ray-trace mesh grid (shared)": 0.00017612669998925412, + "Inversion setup, incl. NUFFT (per channel \u00d7 34)": 3.1698899924002033, + "Data vector D, sparse (per channel \u00d7 34)": 0.015085351203379106, + "Curvature matrix F, sparse (per channel \u00d7 34)": 3.1523015747978205, + "Regularization matrix H (shared)": 0.0, + "Reconstruction NNLS (per channel \u00d7 34)": 0.6942325224015804, + "fast_chi_squared + log evidence (per channel \u00d7 34)": 0.07882488259565434 + }, + "per_channel_costs": { + "inversion_setup": 0.09323205860000598, + "data_vector": 0.00044368680009938545, + "curvature_matrix": 0.0927147521999359, + "reconstruction": 0.02041860360004648, + "log_evidence": 0.0023183788998721866 + }, + "inversion_setup_decomposition": { + "mapper_plus_L_invariant": 0.09144079910001893, + "data_vector_matmul_variant": 0.0017912594999870551 + }, + "total_step_by_step_cube": 7.110641481298626, + "shared_lwl_savings_estimate": 3.059586822597885, + "shared_invariant_savings_estimate": 6.077133192898509 +} \ No newline at end of file diff --git a/results/breakdown/datacube/delaunay_hpc_a100_mp.png b/results/breakdown/datacube/delaunay_hpc_a100_mp.png new file mode 100644 index 0000000..10eb18b Binary files /dev/null and b/results/breakdown/datacube/delaunay_hpc_a100_mp.png differ diff --git a/results/breakdown/datacube/inversion_hpc_a100_fp64.json b/results/breakdown/datacube/inversion_hpc_a100_fp64.json new file mode 100644 index 0000000..46cf154 --- /dev/null +++ b/results/breakdown/datacube/inversion_hpc_a100_fp64.json @@ -0,0 +1,36 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 33325 MiB, 81920 MiB" + }, + "instrument": "alma_high", + "model": "delaunay", + "purpose": "decompose the datacube inversion-setup step into channel-invariant vs channel-variant sub-steps (sparse w-tilde route)", + "configuration": { + "visibilities": 5000000, + "delaunay_vertices": 1500, + "use_mixed_precision": false + }, + "note": "cutpoints are a dependency TREE not a linear chain \u2014 XLA DCE prunes per leaf, so C(curvature) and D(data_vector) are parallel children of B(mapper+L); deltas subtract the common parent B, never C from D", + "raw_cutpoints_per_call_s": { + "A_trace": 0.0002395286999671953, + "B_mapper": 0.09454641280008218, + "C_curvature": 0.18481450480012426, + "D_data_vector": 0.09209770440011197 + }, + "sub_steps_per_call_s": { + "ray_trace": 0.0002395286999671953, + "mapper_plus_mapping_matrix_L": 0.09430688410011498, + "curvature_F_LtWtL": 0.09026809200004209, + "data_vector_D": -0.0024487083999702036 + }, + "per_channel_inversion_total_per_call_s": 0.18236579640015405, + "channel_invariant_per_call_s": 0.18481450480012426, + "channel_variant_per_call_s": -0.0024487083999702036, + "shareable_fraction_of_inversion_total": 1.013427454316034 +} \ No newline at end of file diff --git a/results/breakdown/datacube/inversion_hpc_a100_mp.json b/results/breakdown/datacube/inversion_hpc_a100_mp.json new file mode 100644 index 0000000..b75dd6f --- /dev/null +++ b/results/breakdown/datacube/inversion_hpc_a100_mp.json @@ -0,0 +1,36 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 33325 MiB, 81920 MiB" + }, + "instrument": "alma_high", + "model": "delaunay", + "purpose": "decompose the datacube inversion-setup step into channel-invariant vs channel-variant sub-steps (sparse w-tilde route)", + "configuration": { + "visibilities": 5000000, + "delaunay_vertices": 1500, + "use_mixed_precision": true + }, + "note": "cutpoints are a dependency TREE not a linear chain \u2014 XLA DCE prunes per leaf, so C(curvature) and D(data_vector) are parallel children of B(mapper+L); deltas subtract the common parent B, never C from D", + "raw_cutpoints_per_call_s": { + "A_trace": 0.00021824800005560975, + "B_mapper": 0.09195884380005737, + "C_curvature": 0.1846953169999324, + "D_data_vector": 0.09318047489996388 + }, + "sub_steps_per_call_s": { + "ray_trace": 0.00021824800005560975, + "mapper_plus_mapping_matrix_L": 0.09174059580000175, + "curvature_F_LtWtL": 0.09273647319987505, + "data_vector_D": 0.0012216310999065144 + }, + "per_channel_inversion_total_per_call_s": 0.18591694809983894, + "channel_invariant_per_call_s": 0.1846953169999324, + "channel_variant_per_call_s": 0.0012216310999065144, + "shareable_fraction_of_inversion_total": 0.9934291568768088 +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649.json b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649.json new file mode 100644 index 0000000..9de6019 --- /dev/null +++ b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649.json @@ -0,0 +1,34 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "cpu", + "device": "TFRT_CPU_0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": "1", + "cpu_count": 8 + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace data grid": 0.002923889999874518, + "Ray-trace mesh grid": 0.000754400000005262, + "Lens light images (pre-PSF)": 4.809999991266523e-05, + "Blurred image (PSF convolution)": 0.0024762500001088484, + "Profile-subtracted image": 0.00022619999981543515, + "Inversion setup (steps 5-8 combined)": 3.913500949999798, + "Data vector (D)": 0.05481417000009969, + "Curvature matrix (F)": 4.233989429999928, + "Regularization matrix (H)": 0.0470290000012028, + "Regularized reconstruction": 1.683930590000091, + "Mapped recon + log evidence": 0.13277568000012252 + }, + "total_step_by_step": 10.07246866000096 +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649.png b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649.png new file mode 100644 index 0000000..e304b4a Binary files /dev/null and b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649.png differ diff --git a/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649_sparse.json b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649_sparse.json new file mode 100644 index 0000000..dfe3ada --- /dev/null +++ b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649_sparse.json @@ -0,0 +1,34 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "cpu", + "device": "TFRT_CPU_0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": "1", + "cpu_count": 8 + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "inversion_path": "sparse" + }, + "steps": { + "Ray-trace data grid": 0.011347219999879598, + "Ray-trace mesh grid": 0.0006453899997723056, + "Lens light images (pre-PSF)": 3.168000002915505e-05, + "Blurred image (PSF convolution)": 0.0015794700000697048, + "Profile-subtracted image": 0.00021048999988124706, + "Inversion setup (steps 5-8 combined)": 3.084574600000269, + "Data vector (D)": 0.027584440000282484, + "Curvature matrix (F)": 4.196036459999959, + "Regularization matrix (H)": 0.04668609999862383, + "Regularized reconstruction": 1.3451809400001367, + "Mapped recon + log evidence": 0.09821528000029503 + }, + "total_step_by_step": 8.8120920699992 +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649_sparse.png b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649_sparse.png new file mode 100644 index 0000000..b2b78a0 Binary files /dev/null and b/results/breakdown/imaging/delaunay_breakdown_hst_v2026.7.6.649_sparse.png differ diff --git a/results/breakdown/imaging/delaunay_hpc_a100_fp64.json b/results/breakdown/imaging/delaunay_hpc_a100_fp64.json new file mode 100644 index 0000000..10aa0c4 --- /dev/null +++ b/results/breakdown/imaging/delaunay_hpc_a100_fp64.json @@ -0,0 +1,41 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4645 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace data grid": 0.00017328279991488671, + "Ray-trace mesh grid": 0.00010508720006328077, + "Lens light images (pre-PSF)": 5.106359985802555e-05, + "Blurred image (PSF convolution)": 0.00072238800003106, + "Profile-subtracted image": 0.00012498710002546433, + "Inversion setup (steps 5-8 combined)": 0.04113771680003993, + "Data vector (D)": 0.0003055739000046742, + "Curvature matrix (F)": 0.004818168800011335, + "Regularization matrix (H)": 0.01439599100012856, + "Regularized reconstruction": 0.03254721299999801, + "Mapped recon + log evidence": 0.002242607500011218 + }, + "total_step_by_step": 0.09662407970008643, + "setup_split": { + "Border relocation": 0.0008993409001050168, + "Triangulation + interpolation": 0.02662460360006662, + "Mapping matrix": 0.0063388323997060035, + "Blurred mapping matrix (PSF)": 0.006100841000079529 + } +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_hpc_a100_fp64.png b/results/breakdown/imaging/delaunay_hpc_a100_fp64.png new file mode 100644 index 0000000..9e85302 Binary files /dev/null and b/results/breakdown/imaging/delaunay_hpc_a100_fp64.png differ diff --git a/results/breakdown/imaging/delaunay_hpc_a100_fp64_sparse.json b/results/breakdown/imaging/delaunay_hpc_a100_fp64_sparse.json new file mode 100644 index 0000000..90a6c44 --- /dev/null +++ b/results/breakdown/imaging/delaunay_hpc_a100_fp64_sparse.json @@ -0,0 +1,35 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4701 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "inversion_path": "sparse" + }, + "steps": { + "Ray-trace data grid": 0.00011405620007280959, + "Ray-trace mesh grid": 0.00015693100003772998, + "Lens light images (pre-PSF)": 5.8916599937219874e-05, + "Blurred image (PSF convolution)": 0.0007329633999688667, + "Profile-subtracted image": 8.804649996818625e-05, + "Inversion setup (steps 5-8 combined)": 0.04142457910002122, + "Data vector (D)": 0.00031941799998094214, + "Curvature matrix (F)": 0.0047273813000174416, + "Regularization matrix (H)": 0.01542312200035667, + "Regularized reconstruction": 0.03274251569991975, + "Mapped recon + log evidence": 0.002231738899990887 + }, + "total_step_by_step": 0.0980196687002717 +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_hpc_a100_fp64_sparse.png b/results/breakdown/imaging/delaunay_hpc_a100_fp64_sparse.png new file mode 100644 index 0000000..1634452 Binary files /dev/null and b/results/breakdown/imaging/delaunay_hpc_a100_fp64_sparse.png differ diff --git a/results/breakdown/imaging/delaunay_hpc_a100_mp.json b/results/breakdown/imaging/delaunay_hpc_a100_mp.json new file mode 100644 index 0000000..8b4fdd9 --- /dev/null +++ b/results/breakdown/imaging/delaunay_hpc_a100_mp.json @@ -0,0 +1,35 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4653 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace data grid": 9.903330001179712e-05, + "Ray-trace mesh grid": 0.00015739800001028924, + "Lens light images (pre-PSF)": 8.48646000122244e-05, + "Blurred image (PSF convolution)": 0.001022397699944122, + "Profile-subtracted image": 0.0001505960999566014, + "Inversion setup (steps 5-8 combined)": 0.03988879589996941, + "Data vector (D)": 0.00032698599998184366, + "Curvature matrix (F)": 0.004787345900058426, + "Regularization matrix (H)": 0.01549665200036543, + "Regularized reconstruction": 0.03255944599995928, + "Mapped recon + log evidence": 0.002271496799949091 + }, + "total_step_by_step": 0.09684501230021852 +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_hpc_a100_mp.png b/results/breakdown/imaging/delaunay_hpc_a100_mp.png new file mode 100644 index 0000000..1b01a40 Binary files /dev/null and b/results/breakdown/imaging/delaunay_hpc_a100_mp.png differ diff --git a/results/breakdown/imaging/delaunay_hpc_a100_mp_sparse.json b/results/breakdown/imaging/delaunay_hpc_a100_mp_sparse.json new file mode 100644 index 0000000..e4e5e7d --- /dev/null +++ b/results/breakdown/imaging/delaunay_hpc_a100_mp_sparse.json @@ -0,0 +1,35 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4709 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "delaunay_vertices": 1500, + "edge_zeroed_pixels": 0, + "inversion_path": "sparse" + }, + "steps": { + "Ray-trace data grid": 0.00015994799996406074, + "Ray-trace mesh grid": 0.00011522330005391268, + "Lens light images (pre-PSF)": 7.366250001723529e-05, + "Blurred image (PSF convolution)": 0.0006731299000421132, + "Profile-subtracted image": 0.00011623330001384602, + "Inversion setup (steps 5-8 combined)": 0.03962325540005622, + "Data vector (D)": 0.0003264400000261958, + "Curvature matrix (F)": 0.0048954911000691935, + "Regularization matrix (H)": 0.0144903880000129, + "Regularized reconstruction": 0.03275093279999055, + "Mapped recon + log evidence": 0.0022801256000093417 + }, + "total_step_by_step": 0.09550482990025558 +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_hpc_a100_mp_sparse.png b/results/breakdown/imaging/delaunay_hpc_a100_mp_sparse.png new file mode 100644 index 0000000..470b1b6 Binary files /dev/null and b/results/breakdown/imaging/delaunay_hpc_a100_mp_sparse.png differ diff --git a/results/breakdown/imaging/mge_breakdown_hst_v2026.7.6.649.json b/results/breakdown/imaging/mge_breakdown_hst_v2026.7.6.649.json new file mode 100644 index 0000000..e20d9d7 --- /dev/null +++ b/results/breakdown/imaging/mge_breakdown_hst_v2026.7.6.649.json @@ -0,0 +1,30 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "cpu", + "device": "TFRT_CPU_0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": "1", + "cpu_count": 8 + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "linear_gaussians": 0, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace grids": 0.002696519999881275, + "Mapping matrix": 0.06340354000021761, + "Blurred mapping matrix": 0.10661125000005996, + "Data vector (D)": 0.0003620499999669846, + "Curvature matrix (F)": 0.0047935300000972346, + "Reconstruction (NNLS)": 0.0006198899998707929, + "Mapped reconstructed image": 0.0006326000002445653, + "Chi-squared & log likelihood": 0.00035997999984829223 + }, + "total_step_by_step": 0.1794793600001867 +} \ No newline at end of file diff --git a/results/breakdown/imaging/mge_breakdown_hst_v2026.7.6.649.png b/results/breakdown/imaging/mge_breakdown_hst_v2026.7.6.649.png new file mode 100644 index 0000000..5387001 Binary files /dev/null and b/results/breakdown/imaging/mge_breakdown_hst_v2026.7.6.649.png differ diff --git a/results/breakdown/imaging/mge_hpc_a100_fp64.json b/results/breakdown/imaging/mge_hpc_a100_fp64.json new file mode 100644 index 0000000..50d5c0b --- /dev/null +++ b/results/breakdown/imaging/mge_hpc_a100_fp64.json @@ -0,0 +1,31 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 555 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "linear_gaussians": 0, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace grids": 0.00013788409996777772, + "Mapping matrix": 0.0011563078000108363, + "Blurred mapping matrix": 0.003171063999980106, + "Data vector (D)": 8.137940003507537e-05, + "Curvature matrix (F)": 0.00010563240002738893, + "Reconstruction (NNLS)": 0.0028374931000143986, + "Mapped reconstructed image": 0.0001686049000454659, + "Chi-squared & log likelihood": 0.00010258929996780353 + }, + "total_step_by_step": 0.007760955000048853 +} \ No newline at end of file diff --git a/results/breakdown/imaging/mge_hpc_a100_fp64.png b/results/breakdown/imaging/mge_hpc_a100_fp64.png new file mode 100644 index 0000000..f5e9887 Binary files /dev/null and b/results/breakdown/imaging/mge_hpc_a100_fp64.png differ diff --git a/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649.json b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649.json new file mode 100644 index 0000000..318cdd4 --- /dev/null +++ b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649.json @@ -0,0 +1,37 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "cpu", + "device": "TFRT_CPU_0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": "1", + "cpu_count": 8 + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "mesh_shape": [ + 39, + 39 + ], + "source_pixels": 1521, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace grids": 0.002827370000159135, + "Lens light images (pre-PSF)": 6.536000000778586e-05, + "Blurred image (PSF convolution)": 0.001444900000205962, + "Profile-subtracted image": 0.00016214999996009283, + "Overlay grid (source pixel centres)": 0.0003320800002256874, + "Inversion setup (steps 4-8 combined)": 2.5466561199998976, + "Data vector (D)": 0.022497959999964224, + "Curvature matrix (F)": 4.172744299999977, + "Regularization matrix (H)": 0.0029215799997473367, + "Regularized reconstruction": 1.705186699999831, + "Mapped recon + log evidence": 0.1941322399998171 + }, + "total_step_by_step": 8.648970759999793 +} \ No newline at end of file diff --git a/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649.png b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649.png new file mode 100644 index 0000000..f01d56d Binary files /dev/null and b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649.png differ diff --git a/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649_sparse.json b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649_sparse.json new file mode 100644 index 0000000..0fa38b1 --- /dev/null +++ b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649_sparse.json @@ -0,0 +1,37 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "cpu", + "device": "TFRT_CPU_0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": "1", + "cpu_count": 8 + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "mesh_shape": [ + 39, + 39 + ], + "source_pixels": 1521, + "inversion_path": "sparse" + }, + "steps": { + "Ray-trace grids": 0.002189679999719374, + "Lens light images (pre-PSF)": 2.727999999478925e-05, + "Blurred image (PSF convolution)": 0.0013315799998963485, + "Profile-subtracted image": 0.00018315000015718398, + "Overlay grid (source pixel centres)": 0.00028071000015188473, + "Inversion setup (steps 4-8 combined)": 2.6214837200001058, + "Data vector (D)": 0.02361151000004611, + "Curvature matrix (F)": 4.646836339999936, + "Regularization matrix (H)": 0.006793530000140891, + "Regularized reconstruction": 2.6747720700001083, + "Mapped recon + log evidence": 0.19714425000020128 + }, + "total_step_by_step": 10.174653820000458 +} \ No newline at end of file diff --git a/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649_sparse.png b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649_sparse.png new file mode 100644 index 0000000..c3711ab Binary files /dev/null and b/results/breakdown/imaging/pixelization_breakdown_hst_v2026.7.6.649_sparse.png differ diff --git a/results/breakdown/imaging/pixelization_hpc_a100_fp64.json b/results/breakdown/imaging/pixelization_hpc_a100_fp64.json new file mode 100644 index 0000000..9037029 --- /dev/null +++ b/results/breakdown/imaging/pixelization_hpc_a100_fp64.json @@ -0,0 +1,38 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4645 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "mesh_shape": [ + 39, + 39 + ], + "source_pixels": 1521, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace grids": 0.00015067610002006404, + "Lens light images (pre-PSF)": 5.268260001685121e-05, + "Blurred image (PSF convolution)": 0.0007426443999975163, + "Profile-subtracted image": 0.0001132882999627327, + "Overlay grid (source pixel centres)": 0.00010292930001014611, + "Inversion setup (steps 4-8 combined)": 0.011578287100019225, + "Data vector (D)": 0.00031442210001841886, + "Curvature matrix (F)": 0.004852892399958364, + "Regularization matrix (H)": 0.00013811220005663927, + "Regularized reconstruction": 0.0372411424999882, + "Mapped recon + log evidence": 0.0023399961999530207 + }, + "total_step_by_step": 0.05762707320000118 +} \ No newline at end of file diff --git a/results/breakdown/imaging/pixelization_hpc_a100_fp64.png b/results/breakdown/imaging/pixelization_hpc_a100_fp64.png new file mode 100644 index 0000000..d9a6d55 Binary files /dev/null and b/results/breakdown/imaging/pixelization_hpc_a100_fp64.png differ diff --git a/results/breakdown/imaging/pixelization_hpc_a100_fp64_sparse.json b/results/breakdown/imaging/pixelization_hpc_a100_fp64_sparse.json new file mode 100644 index 0000000..4472c71 --- /dev/null +++ b/results/breakdown/imaging/pixelization_hpc_a100_fp64_sparse.json @@ -0,0 +1,38 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4701 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "mesh_shape": [ + 39, + 39 + ], + "source_pixels": 1521, + "inversion_path": "sparse" + }, + "steps": { + "Ray-trace grids": 0.0001741967999805638, + "Lens light images (pre-PSF)": 6.866470002933056e-05, + "Blurred image (PSF convolution)": 0.0007241815000270435, + "Profile-subtracted image": 0.00012297709999984364, + "Overlay grid (source pixel centres)": 0.00012518029998318526, + "Inversion setup (steps 4-8 combined)": 0.01180446270000175, + "Data vector (D)": 0.0003395748000002641, + "Curvature matrix (F)": 0.004875196200009669, + "Regularization matrix (H)": 0.0001594619000570674, + "Regularized reconstruction": 0.037105540199991086, + "Mapped recon + log evidence": 0.0023849688000154856 + }, + "total_step_by_step": 0.057884405000095285 +} \ No newline at end of file diff --git a/results/breakdown/imaging/pixelization_hpc_a100_fp64_sparse.png b/results/breakdown/imaging/pixelization_hpc_a100_fp64_sparse.png new file mode 100644 index 0000000..90f20be Binary files /dev/null and b/results/breakdown/imaging/pixelization_hpc_a100_fp64_sparse.png differ diff --git a/results/breakdown/imaging/pixelization_hpc_a100_mp.json b/results/breakdown/imaging/pixelization_hpc_a100_mp.json new file mode 100644 index 0000000..9ed6eca --- /dev/null +++ b/results/breakdown/imaging/pixelization_hpc_a100_mp.json @@ -0,0 +1,38 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4653 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "mesh_shape": [ + 39, + 39 + ], + "source_pixels": 1521, + "inversion_path": "dense" + }, + "steps": { + "Ray-trace grids": 0.00014402019996850868, + "Lens light images (pre-PSF)": 7.778340004733763e-05, + "Blurred image (PSF convolution)": 0.0010373194999374392, + "Profile-subtracted image": 9.399930004292401e-05, + "Overlay grid (source pixel centres)": 8.846850005284068e-05, + "Inversion setup (steps 4-8 combined)": 0.009818835200076138, + "Data vector (D)": 0.0003231530000448402, + "Curvature matrix (F)": 0.004977759700068418, + "Regularization matrix (H)": 0.00010096730002260302, + "Regularized reconstruction": 0.037331846799952474, + "Mapped recon + log evidence": 0.0024208797999563103 + }, + "total_step_by_step": 0.056415032700169836 +} \ No newline at end of file diff --git a/results/breakdown/imaging/pixelization_hpc_a100_mp.png b/results/breakdown/imaging/pixelization_hpc_a100_mp.png new file mode 100644 index 0000000..8298192 Binary files /dev/null and b/results/breakdown/imaging/pixelization_hpc_a100_mp.png differ diff --git a/results/breakdown/imaging/pixelization_hpc_a100_mp_sparse.json b/results/breakdown/imaging/pixelization_hpc_a100_mp_sparse.json new file mode 100644 index 0000000..89bf4e5 --- /dev/null +++ b/results/breakdown/imaging/pixelization_hpc_a100_mp_sparse.json @@ -0,0 +1,38 @@ +{ + "autolens_version": "2026.7.6.649", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 4709 MiB, 81920 MiB" + }, + "instrument": "hst", + "configuration": { + "pixel_scale_arcsec": 0.05, + "mask_radius_arcsec": 3.5, + "image_pixels_masked": 15361, + "over_sampled_pixels": 17980, + "mesh_shape": [ + 39, + 39 + ], + "source_pixels": 1521, + "inversion_path": "sparse" + }, + "steps": { + "Ray-trace grids": 0.0001519421000011789, + "Lens light images (pre-PSF)": 4.994759992769104e-05, + "Blurred image (PSF convolution)": 0.0006980666999879759, + "Profile-subtracted image": 0.00012471030004235217, + "Overlay grid (source pixel centres)": 0.0001309291999859852, + "Inversion setup (steps 4-8 combined)": 0.009859460000006948, + "Data vector (D)": 0.00030537110005752767, + "Curvature matrix (F)": 0.004806445700069162, + "Regularization matrix (H)": 0.000127485200027877, + "Regularized reconstruction": 0.036963547199957245, + "Mapped recon + log evidence": 0.0023011826000583825 + }, + "total_step_by_step": 0.05551908770012232 +} \ No newline at end of file diff --git a/results/breakdown/imaging/pixelization_hpc_a100_mp_sparse.png b/results/breakdown/imaging/pixelization_hpc_a100_mp_sparse.png new file mode 100644 index 0000000..72f7f63 Binary files /dev/null and b/results/breakdown/imaging/pixelization_hpc_a100_mp_sparse.png differ diff --git a/results/notes/preopt_breakdown_baseline.md b/results/notes/preopt_breakdown_baseline.md new file mode 100644 index 0000000..bcda98b --- /dev/null +++ b/results/notes/preopt_breakdown_baseline.md @@ -0,0 +1,98 @@ +# PreOptimizationTimes — likelihood_breakdown baseline (polish phase 4) + +Status: laptop-CPU fallback tier COMPLETE + A100 imaging tier COMPLETE +(2026-07-10); A100 alma_high cells (interferometer + datacube) in flight +after a nufftax venv install on RAL. Task: autolens_profiling#59. + +## A100 canonical tier — imaging (jobs 330062–330070) + +| Cell (hst) | Path | A100 fp64 | A100 mp | CPU fp64 | GPU speedup | Top step (A100) | +|------------|------|----------:|--------:|---------:|------------:|-----------------| +| `imaging/mge` | dense | 7.8 ms | — | 179.5 ms | 23× | Blurred mapping matrix (41%) | +| `imaging/pixelization` | dense | 57.6 ms | 56.4 ms | 8.65 s | 150× | Regularized reconstruction (65%) | +| `imaging/pixelization` | sparse | 57.9 ms | 55.5 ms | 10.17 s | 176× | Regularized reconstruction (64%) | +| `imaging/delaunay` | dense | 97.6 ms | 96.8 ms | 10.07 s | 103× | Inversion setup 5–8 (43%) | +| `imaging/delaunay` | sparse | 98.0 ms | 95.5 ms | 8.81 s | 90× | Inversion setup 5–8 (42%) | + +### Four-way split of the delaunay inversion-setup block (A100, job 330079) + +Measured via the opt-in `--split-setup` prefix-JIT decomposition (prefix-sum +39.96 ms vs 41.14 ms combined — faithful, no fusion redistribution): + +| Piece | A100 | share of full likelihood | +|---|---:|---:| +| Border relocation | 0.90 ms | ~1% | +| **Triangulation + interpolation** | **26.62 ms** | **~27% — top target** | +| Mapping matrix build | 6.34 ms | ~7% | +| Blurred mapping matrix (PSF) | 6.10 ms | ~6% | + +The qhull host callback is a few ms of the 26.6 ms at most (1500 points + +transfer round-trip); the work item is the JAX-side visibility-walk point +location + barycentric interpolation. On laptop CPU the same split is +convolution-dominated (PSF-convolved mapping matrix 3.37 s of 3.88 s) — +another instance of CPU decompositions pointing at the wrong GPU target. + +**The headline validates the GPU-first platform policy: the bottleneck moves.** +On CPU every mesh cell is dominated by the Curvature matrix (F) at 42–48%; +on the A100, F collapses into the noise and pixelization is dominated by the +**Regularized reconstruction (NNLS solve, ~65%)** while delaunay is dominated +by **inversion setup**. Optimization driven from the CPU decomposition would +have targeted the wrong step. Mixed precision is flat at this problem scale +(≤2% vs fp64 on every cell) — the mp win reported by the runtime package at +larger visibility counts does not appear in hst-scale imaging breakdowns. + +## The baseline (laptop CPU fp64, policy env, quiet machine) + +| Cell | Instrument | Path | Step-sum total | vs May rows | Top step | +|------|-----------|------|---------------:|------------:|----------| +| `imaging/mge` | hst | dense | 179.5 ms | 1.01× | Blurred mapping matrix (59%) | +| `imaging/pixelization` | hst | dense | 8.65 s | 1.11× | Curvature matrix F (48%) | +| `imaging/pixelization` | hst | sparse | 10.17 s | 1.14× | Curvature matrix F (46%) | +| `imaging/delaunay` | hst | dense | 10.07 s | 0.89× | Curvature matrix F (42%) | +| `imaging/delaunay` | hst | sparse | 8.81 s | 1.11× | Curvature matrix F (48%) | + +Verdict vs the May (v2026.5.29.4) rows: **no library-level drift** — +0.89–1.14× is measurement scatter. The Curvature matrix (F) step is the +dominant optimization target on every mesh cell; MGE is convolution-bound. + +## Cells deferred to the A100 tier (GPU-only on this laptop) + +`interferometer/delaunay @ alma_high`, `datacube/delaunay @ alma_high`, +`datacube/inversion_setup_decompose @ alma_high`: the NUFFT precision-operator +construction exceeded a 2 h timeout twice under ambient load (consistent with +the phase-3 infeasibility map for alma-tier visibility counts on CPU). Per the +GPU-first platform policy these land via the A100 submits when RAL returns. + +## Measurement methodology (read before comparing runs) + +The canonical environment for breakdown numbers is: + +- `XLA_FLAGS="--xla_disable_hlo_passes=constant_folding"` **explicitly + exported in the shell** — this is autoconf library policy + (`autoconf/jax_wrapper.py` sets it at import and recommends the terminal + export). Do not rely on the in-process set: whether XLA honours it depends + on import-vs-backend-init order. The A100 submits export it explicitly. +- A **quiet machine**. This is not a nicety: the first pass of this campaign + ran with ~12/15 GB in use from parallel sessions and came back **2.5–5.3× + slow across every step uniformly** — initially misread as a library + regression. Uniform per-step inflation (including trivially cheap steps) is + the contention signature; targeted regressions concentrate in specific steps. +- Result JSONs record `xla_flags` / `omp_num_threads` / `cpu_count` + (`device` dict) so env drift between runs is attributable. Note the + recorded `xla_flags` is the env at JSON-write time — autoconf re-sets it + in-process — so shell-level provenance is canonical. + +Measured flag effect (back-to-back `likelihood_runtime/imaging/mge.py` pair, +equal load): single-JIT 0.141 s → 0.218 s (**1.54×** with the flag); +**vmap per-call insensitive** (0.110 s both ways). Constant folding only has +something to fold when inputs are concrete constants — batch tracers defeat +it, which is why the runtime package's vmap numbers are the honest +cross-config comparator regardless of the flag. + +## Cross-package sanity (fusion caveat) + +Step-sums vs the phase-3 runtime full-pipeline single-JIT for the same cells: +pixelization 8.65 s vs 13.7 s and delaunay 10.07 s vs 16.7 s — the runtime +numbers were taken on the 2026-07-08/09 overnight matrix and carry their own +load conditions, so treat cross-package ratios qualitatively until the A100 +tier provides both packages on the same quiet hardware. diff --git a/scripts/build_readme.py b/scripts/build_readme.py index 3c816e4..60d9883 100644 --- a/scripts/build_readme.py +++ b/scripts/build_readme.py @@ -82,6 +82,18 @@ r"\.json$" ) +# Config-tagged artifact filename (sweep-style `--config-name` outputs, e.g. +# the A100 tier):