diff --git a/hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly new file mode 100644 index 000000000..bac50c6a8 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly @@ -0,0 +1,108 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/delaunay_nn x hst x fp64 — PyAutoArray#536 +# A/B leg: **feature**. +# +# Verification run for the ConstantSplit split-stencil assembly compaction +# (PyAutoArray issue #536 / branch feature/delaunay-nn-constant-split-assembly: +# the JAX path of pixel_splitted_regularization_matrix_from scatters only the +# first SPLIT_REG_COMPACT_WIDTH = 12 columns of every (4P, K) stencil row +# instead of the full padded (4P, K, K) outer product, supplements the +# SPLIT_REG_WIDE_ROW_BUDGET = 256 widest rows at full width, and NaNs the +# matrix out if more rows exceed the compact width than the budget holds). +# Both legs of the A/B run in the same session on the same node class with the +# same everything except the PyAutoArray checkout prepended to PYTHONPATH: +# +# control -> 180c8a40 (the branch's merge base = PyAutoArray main) +# feature -> feature/delaunay-nn-constant-split-assembly (04152bb7) +# +# The predecessor A/B on this cell (#532/#533, merge base d7c96762) is recorded +# in results/notes/delaunay_nn_launch_latency.md; its legs are kept as +# submit_..._launch_latency{,_control} and are comparable to these rows only +# insofar as the merge base moved by that one merge. +# +# The shared /mnt/ral/jnightin/PyAuto/PyAutoArray install is NOT touched — long +# science runs (subhalo validation, Euclid DR1) share it. The private checkouts +# live under /mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/ +# and are prepended to the PYTHONPATH that $AP_ROOT/activate.sh exports, so only +# PyAutoArray is overridden and Nerves/Fit/Galaxy/Lens stay on the shared tree. +# +# WITNESS. The target row is "Regularization matrix (H, ConstantSplit +# assembly)" in the --split-setup table (JSON regularization_split / +# regularization_split_vmap): 10.0 ms per call at vmap 16 on the #533 shipped +# leg, and the witness is < 3 ms per call. The params->H prefix +# (regularization_matrix_prefix_s) carries it: 16.4 ms per call at vmap 16 on +# that leg, witness < 11 ms. Judge on the params->H prefix and the assembly +# row; the H *step* row is a prefix difference and shifts with attribution +# (the caveat in results/notes/delaunay_walk_early_exit.md). +# +# --split-setup four-way split of the inversion-setup block (+ the +# split-point Sibson cut of the params->H interval) +# --vmap-batch 16 batched re-timing at the Nautilus reference n_batch + +#SBATCH -J bd_delnn_asm +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=12:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err + +export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)" +echo "AP_ROOT=$AP_ROOT" +source $AP_ROOT/activate.sh + +# --- PyAutoArray A/B override (private checkout, shared install untouched) --- +AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/PyAutoArray_feature +export PYTHONPATH=$AA_ROOT:$PYTHONPATH + +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 + +# PYAUTO_SIBSON_QUERY_CHUNK is deliberately left UNSET on both A/B legs, as in +# the #533 run whose rows these are compared against: the library default +# (4096, chosen by that run's sweep) is what production uses. The script +# records the effective SIBSON_QUERY_CHUNK under "sibson" in the result JSON. + +nvidia-smi + +# MIG backstop. A no-op since the fleet's MIG mode was confirmed Disabled +# (2026-09-05); kept because a silent removal re-opens the trap. Only bounces +# when the job was dispatched with --requeue. +source "$AP_ROOT/hpc/batch_gpu/_gpu_preflight.sh" + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/delaunay_nn" +echo "Instrument: hst" +echo "Precision: fp64" +echo "A/B leg: assembly (feature/delaunay-nn-constant-split-assembly, 04152bb7)" +echo "Flags: --split-setup --vmap-batch 16 --config-name hpc_a100_fp64_assembly" +echo "Sibson chunk: PYAUTO_SIBSON_QUERY_CHUNK=${PYAUTO_SIBSON_QUERY_CHUNK:-unset (library default)}" + +echo "=== source revisions ===" +printf '%-14s %s\n' "PyAutoArray(AB)" "$(git -C $AA_ROOT rev-parse HEAD 2>&1)" +for r in PyAutoNerves PyAutoFit PyAutoGalaxy PyAutoLens; do + printf '%-14s %s\n' "$r" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)" +done +echo "AP_ROOT rev: $(git -C $AP_ROOT rev-parse HEAD 2>&1)" + +cd $AP_ROOT + +echo "=== import provenance ===" +python3 -c "import autoarray, autolens; print('autoarray.__file__ =', autoarray.__file__); print('autoarray.__version__ =', autoarray.__version__); print('autolens.__file__ =', autolens.__file__)" + +python3 -u scripts/imaging/likelihood_breakdown/delaunay_nn.py \ + --config-name hpc_a100_fp64_assembly \ + --split-setup \ + --vmap-batch 16 + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly_control b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly_control new file mode 100644 index 000000000..c6261d8d9 --- /dev/null +++ b/hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly_control @@ -0,0 +1,108 @@ +#!/bin/bash -l +# +# A100 likelihood-breakdown profile: imaging/delaunay_nn x hst x fp64 — PyAutoArray#536 +# A/B leg: **control**. +# +# Verification run for the ConstantSplit split-stencil assembly compaction +# (PyAutoArray issue #536 / branch feature/delaunay-nn-constant-split-assembly: +# the JAX path of pixel_splitted_regularization_matrix_from scatters only the +# first SPLIT_REG_COMPACT_WIDTH = 12 columns of every (4P, K) stencil row +# instead of the full padded (4P, K, K) outer product, supplements the +# SPLIT_REG_WIDE_ROW_BUDGET = 256 widest rows at full width, and NaNs the +# matrix out if more rows exceed the compact width than the budget holds). +# Both legs of the A/B run in the same session on the same node class with the +# same everything except the PyAutoArray checkout prepended to PYTHONPATH: +# +# control -> 180c8a40 (the branch's merge base = PyAutoArray main) +# feature -> feature/delaunay-nn-constant-split-assembly (04152bb7) +# +# The predecessor A/B on this cell (#532/#533, merge base d7c96762) is recorded +# in results/notes/delaunay_nn_launch_latency.md; its legs are kept as +# submit_..._launch_latency{,_control} and are comparable to these rows only +# insofar as the merge base moved by that one merge. +# +# The shared /mnt/ral/jnightin/PyAuto/PyAutoArray install is NOT touched — long +# science runs (subhalo validation, Euclid DR1) share it. The private checkouts +# live under /mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/ +# and are prepended to the PYTHONPATH that $AP_ROOT/activate.sh exports, so only +# PyAutoArray is overridden and Nerves/Fit/Galaxy/Lens stay on the shared tree. +# +# WITNESS. The target row is "Regularization matrix (H, ConstantSplit +# assembly)" in the --split-setup table (JSON regularization_split / +# regularization_split_vmap): 10.0 ms per call at vmap 16 on the #533 shipped +# leg, and the witness is < 3 ms per call. The params->H prefix +# (regularization_matrix_prefix_s) carries it: 16.4 ms per call at vmap 16 on +# that leg, witness < 11 ms. Judge on the params->H prefix and the assembly +# row; the H *step* row is a prefix difference and shifts with attribution +# (the caveat in results/notes/delaunay_walk_early_exit.md). +# +# --split-setup four-way split of the inversion-setup block (+ the +# split-point Sibson cut of the params->H interval) +# --vmap-batch 16 batched re-timing at the Nautilus reference n_batch + +#SBATCH -J bd_delnn_asm_ctl +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=12:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err + +export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)" +echo "AP_ROOT=$AP_ROOT" +source $AP_ROOT/activate.sh + +# --- PyAutoArray A/B override (private checkout, shared install untouched) --- +AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/PyAutoArray_control +export PYTHONPATH=$AA_ROOT:$PYTHONPATH + +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 + +# PYAUTO_SIBSON_QUERY_CHUNK is deliberately left UNSET on both A/B legs, as in +# the #533 run whose rows these are compared against: the library default +# (4096, chosen by that run's sweep) is what production uses. The script +# records the effective SIBSON_QUERY_CHUNK under "sibson" in the result JSON. + +nvidia-smi + +# MIG backstop. A no-op since the fleet's MIG mode was confirmed Disabled +# (2026-09-05); kept because a silent removal re-opens the trap. Only bounces +# when the job was dispatched with --requeue. +source "$AP_ROOT/hpc/batch_gpu/_gpu_preflight.sh" + +echo "==========================================" +date +echo "Cell: likelihood_breakdown/imaging/delaunay_nn" +echo "Instrument: hst" +echo "Precision: fp64" +echo "A/B leg: control (merge base 180c8a40)" +echo "Flags: --split-setup --vmap-batch 16 --config-name hpc_a100_fp64_assembly_control" +echo "Sibson chunk: PYAUTO_SIBSON_QUERY_CHUNK=${PYAUTO_SIBSON_QUERY_CHUNK:-unset (library default)}" + +echo "=== source revisions ===" +printf '%-14s %s\n' "PyAutoArray(AB)" "$(git -C $AA_ROOT rev-parse HEAD 2>&1)" +for r in PyAutoNerves PyAutoFit PyAutoGalaxy PyAutoLens; do + printf '%-14s %s\n' "$r" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)" +done +echo "AP_ROOT rev: $(git -C $AP_ROOT rev-parse HEAD 2>&1)" + +cd $AP_ROOT + +echo "=== import provenance ===" +python3 -c "import autoarray, autolens; print('autoarray.__file__ =', autoarray.__file__); print('autoarray.__version__ =', autoarray.__version__); print('autolens.__file__ =', autolens.__file__)" + +python3 -u scripts/imaging/likelihood_breakdown/delaunay_nn.py \ + --config-name hpc_a100_fp64_assembly_control \ + --split-setup \ + --vmap-batch 16 + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_delaunay_nn_assembly_bench_a100 b/hpc/batch_gpu/submit_delaunay_nn_assembly_bench_a100 new file mode 100644 index 000000000..a4ab9eedf --- /dev/null +++ b/hpc/batch_gpu/submit_delaunay_nn_assembly_bench_a100 @@ -0,0 +1,65 @@ +#!/bin/bash -l +# +# A100 direct bench of the DelaunayNN ConstantSplit assembly — PyAutoArray#536. +# +# The companion to the breakdown/runtime A/B legs: instead of timing the +# assembly as a difference of two compiled prefixes, this times the assembly +# function itself on the real (4P, K) HST split stencil tables committed at +# results/delaunay_nn/tables_hst_1500.npz. `V0` is the pre-#536 full padded +# (4P, K, K) scatter and `V6_library` is the shipped library path with its +# default compaction, so the V0-vs-V6_library row is the direct statement of +# what the change bought, free of any prefix-attribution caveat. +# +# Runs on the **feature** private checkout (the library default constants are +# what V6_library exercises); V0 is reproduced from the same code by forcing +# compact_width = K, so both rows come from one checkout and no control +# checkout is needed here. +# +# Cheap: no dataset, no imaging fit, minutes not hours. + +#SBATCH -J delnn_asm_bench +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=1:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err + +export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)" +echo "AP_ROOT=$AP_ROOT" +source $AP_ROOT/activate.sh + +AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/PyAutoArray_feature +export PYTHONPATH=$AA_ROOT:$PYTHONPATH + +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 + +# MIG backstop; no-op since 2026-09-05, only bounces under --requeue. +source "$AP_ROOT/hpc/batch_gpu/_gpu_preflight.sh" + +echo "==========================================" +date +printf 'PyAutoArray(AB) %s\n' "$(git -C $AA_ROOT rev-parse HEAD 2>&1)" +echo "AP_ROOT rev: $(git -C $AP_ROOT rev-parse HEAD 2>&1)" + +cd $AP_ROOT +python3 -c "import autoarray; print('autoarray.__file__ =', autoarray.__file__)" + +python3 -u scripts/misc/delaunay_nn/assembly_bench.py \ + --batch 16 \ + --repeats 10 \ + --label "A100 fp64, PyAutoArray#536 A/B session" \ + --out results/delaunay_nn/assembly_bench_a100_fp64.json + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_runtime_imaging_delaunay_nn_a100_hst_fp64_assembly b/hpc/batch_gpu/submit_runtime_imaging_delaunay_nn_a100_hst_fp64_assembly new file mode 100644 index 000000000..a663a284b --- /dev/null +++ b/hpc/batch_gpu/submit_runtime_imaging_delaunay_nn_a100_hst_fp64_assembly @@ -0,0 +1,64 @@ +#!/bin/bash -l +# +# A100 likelihood-runtime profile: imaging/delaunay_nn x hst x fp64 — PyAutoArray#536 +# A/B leg: **feature**. Single-JIT whole-likelihood cost, the end-to-end companion +# to submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly. +# +# Same private-checkout PYTHONPATH override as the breakdown A/B legs; the +# shared /mnt/ral/jnightin/PyAuto install is untouched. Phase A (--vmap-probe) +# is skipped: only the full single-JIT + vmap timing is wanted here, and this +# cell's probe is known to over-recommend batch 64 (cuFFT batched-plan failure, +# job 342280) — with no probe JSON present the script falls back to the +# VMAP_BATCH table (16), which is the number both legs must share. +# +# control -> 180c8a40 (the branch's merge base = PyAutoArray main) +# feature -> feature/delaunay-nn-constant-split-assembly (04152bb7) +# +# PYAUTO_SIBSON_QUERY_CHUNK is left unset on both legs (the library default +# 4096 is what production uses); the result JSON records the effective chunk. + +#SBATCH -J rt_delnn_asm +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=3:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err + +export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)" +echo "AP_ROOT=$AP_ROOT" +source $AP_ROOT/activate.sh + +AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/PyAutoArray_feature +export PYTHONPATH=$AA_ROOT:$PYTHONPATH + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +# MIG backstop; no-op since 2026-09-05, only bounces under --requeue. +source "$AP_ROOT/hpc/batch_gpu/_gpu_preflight.sh" + +echo "==========================================" +date +echo "Cell: likelihood_runtime/imaging/delaunay_nn" +echo "A/B leg: assembly (feature/delaunay-nn-constant-split-assembly, 04152bb7)" +echo "Sibson chunk: PYAUTO_SIBSON_QUERY_CHUNK=${PYAUTO_SIBSON_QUERY_CHUNK:-unset (library default)}" +printf 'PyAutoArray(AB) %s\n' "$(git -C $AA_ROOT rev-parse HEAD 2>&1)" + +cd $AP_ROOT +python3 -c "import autoarray; print('autoarray.__file__ =', autoarray.__file__)" + +python3 -u scripts/imaging/likelihood_runtime/delaunay_nn.py \ + --instrument hst \ + --config-name hpc_a100_fp64_assembly + +echo "Finished." +date diff --git a/hpc/batch_gpu/submit_runtime_imaging_delaunay_nn_a100_hst_fp64_assembly_control b/hpc/batch_gpu/submit_runtime_imaging_delaunay_nn_a100_hst_fp64_assembly_control new file mode 100644 index 000000000..bfe03ec7c --- /dev/null +++ b/hpc/batch_gpu/submit_runtime_imaging_delaunay_nn_a100_hst_fp64_assembly_control @@ -0,0 +1,64 @@ +#!/bin/bash -l +# +# A100 likelihood-runtime profile: imaging/delaunay_nn x hst x fp64 — PyAutoArray#536 +# A/B leg: **control**. Single-JIT whole-likelihood cost, the end-to-end companion +# to submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly_control. +# +# Same private-checkout PYTHONPATH override as the breakdown A/B legs; the +# shared /mnt/ral/jnightin/PyAuto install is untouched. Phase A (--vmap-probe) +# is skipped: only the full single-JIT + vmap timing is wanted here, and this +# cell's probe is known to over-recommend batch 64 (cuFFT batched-plan failure, +# job 342280) — with no probe JSON present the script falls back to the +# VMAP_BATCH table (16), which is the number both legs must share. +# +# control -> 180c8a40 (the branch's merge base = PyAutoArray main) +# feature -> feature/delaunay-nn-constant-split-assembly (04152bb7) +# +# PYAUTO_SIBSON_QUERY_CHUNK is left unset on both legs (the library default +# 4096 is what production uses); the result JSON records the effective chunk. + +#SBATCH -J rt_delnn_asm_ctl +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=64gb +#SBATCH --time=3:00:00 +#SBATCH -o output/output.%A.out +#SBATCH -e error/error.%A.err + +export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)" +echo "AP_ROOT=$AP_ROOT" +source $AP_ROOT/activate.sh + +AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/PyAutoArray_control +export PYTHONPATH=$AA_ROOT:$PYTHONPATH + +export JAX_PLATFORM_NAME=cuda +export JAX_PLATFORMS=cuda,cpu +export XLA_PYTHON_CLIENT_PREALLOCATE=false +export JAX_ENABLE_X64=True +export NUMBA_CACHE_DIR=/tmp/numba_cache +export MPLCONFIGDIR=/tmp/matplotlib + +nvidia-smi + +# MIG backstop; no-op since 2026-09-05, only bounces under --requeue. +source "$AP_ROOT/hpc/batch_gpu/_gpu_preflight.sh" + +echo "==========================================" +date +echo "Cell: likelihood_runtime/imaging/delaunay_nn" +echo "A/B leg: control (merge base 180c8a40)" +echo "Sibson chunk: PYAUTO_SIBSON_QUERY_CHUNK=${PYAUTO_SIBSON_QUERY_CHUNK:-unset (library default)}" +printf 'PyAutoArray(AB) %s\n' "$(git -C $AA_ROOT rev-parse HEAD 2>&1)" + +cd $AP_ROOT +python3 -c "import autoarray; print('autoarray.__file__ =', autoarray.__file__)" + +python3 -u scripts/imaging/likelihood_runtime/delaunay_nn.py \ + --instrument hst \ + --config-name hpc_a100_fp64_assembly_control + +echo "Finished." +date diff --git a/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly.json b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly.json new file mode 100644 index 000000000..3d959f298 --- /dev/null +++ b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly.json @@ -0,0 +1,75 @@ +{ + "autolens_version": "2026.8.17.1", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 41495 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" + }, + "sibson": { + "PYAUTO_SIBSON_QUERY_CHUNK": null, + "query_chunk": 4096, + "max_cavity_triangles": 32, + "max_neighbors": 32, + "mesh_query_chunk": 4096, + "mesh_max_cavity_triangles": 32, + "mesh_max_neighbors": 32 + }, + "steps": { + "Ray-trace data grid": 0.0001958048902451992, + "Ray-trace mesh grid": 0.0001864529913291335, + "Lens light images (pre-PSF)": 0.00016739601269364358, + "Blurred image (PSF convolution)": 0.0008221165044233203, + "Profile-subtracted image": 0.00011981930583715439, + "Inversion setup (steps 5-8 combined)": 0.02784522119909525, + "Data vector (D)": 0.00036977999843657017, + "Curvature matrix (F)": 0.004851029999554157, + "Regularization matrix (H)": 0.0008001396199688315, + "Regularized reconstruction": 0.030846239486709236, + "Mapped recon + log evidence": 0.00226008640602231 + }, + "total_step_by_step": 0.0684640864143148, + "regularization_matrix_prefix_s": 0.015188775607384741, + "interpolator_prefix_s": 0.01438863598741591, + "split_sibson_prefix_s": 0.018425561604090035, + "regularization_split": { + "Split-point Sibson": 0.004036925616674125, + "Regularization matrix (H, ConstantSplit assembly)": -0.0032367859967052937 + }, + "setup_split": { + "Border relocation": 0.0010862489929422737, + "Triangulation + interpolation": 0.013302386994473636, + "Mapping matrix": -3.101378679275582e-05, + "Blurred mapping matrix (PSF)": 0.008401293191127479 + }, + "vmap_batch": 16, + "steps_vmap_per_call": { + "Inversion setup (steps 5-8 combined)": 0.017416118050459773, + "Regularization matrix (H)": 0.0008017567291972224 + }, + "setup_split_vmap": { + "Border relocation": 6.397715042112395e-05, + "Triangulation + interpolation": 0.00639469081943389, + "Mapping matrix": 8.609752403572219e-05, + "Blurred mapping matrix (PSF)": 0.00832208354404429 + }, + "regularization_matrix_prefix_vmap_per_call_s": 0.007260424699052237, + "interpolator_prefix_vmap_per_call_s": 0.006458667969855014, + "split_sibson_prefix_vmap_per_call_s": 0.008859757317986805, + "regularization_split_vmap": { + "Split-point Sibson": 0.0024010893481317902, + "Regularization matrix (H, ConstantSplit assembly)": -0.0015993326189345678 + } +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly.png b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly.png new file mode 100644 index 000000000..4624d0d7f Binary files /dev/null and b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly.png differ diff --git a/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly_control.json b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly_control.json new file mode 100644 index 000000000..38e9123e7 --- /dev/null +++ b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly_control.json @@ -0,0 +1,75 @@ +{ + "autolens_version": "2026.8.17.1", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 41495 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" + }, + "sibson": { + "PYAUTO_SIBSON_QUERY_CHUNK": null, + "query_chunk": 4096, + "max_cavity_triangles": 32, + "max_neighbors": 32, + "mesh_query_chunk": 4096, + "mesh_max_cavity_triangles": 32, + "mesh_max_neighbors": 32 + }, + "steps": { + "Ray-trace data grid": 0.00014202520251274108, + "Ray-trace mesh grid": 0.00018779991660267115, + "Lens light images (pre-PSF)": 0.00013585828710347415, + "Blurred image (PSF convolution)": 0.0012267312966287136, + "Profile-subtracted image": 0.00015114820562303066, + "Inversion setup (steps 5-8 combined)": 0.03235943019390106, + "Data vector (D)": 0.00036568089853972197, + "Curvature matrix (F)": 0.0048169304151088, + "Regularization matrix (H)": 0.00991666230838746, + "Regularized reconstruction": 0.030680659506469966, + "Mapped recon + log evidence": 0.0022240755846723915 + }, + "total_step_by_step": 0.08220700181555003, + "regularization_matrix_prefix_s": 0.024340781592763962, + "interpolator_prefix_s": 0.014424119284376502, + "split_sibson_prefix_s": 0.018091991101391612, + "regularization_split": { + "Split-point Sibson": 0.0036678718170151097, + "Regularization matrix (H, ConstantSplit assembly)": 0.00624879049137235 + }, + "setup_split": { + "Border relocation": 0.0016336358850821853, + "Triangulation + interpolation": 0.012790483399294316, + "Mapping matrix": -9.57943731918938e-05, + "Blurred mapping matrix (PSF)": 0.012246569199487566 + }, + "vmap_batch": 16, + "steps_vmap_per_call": { + "Inversion setup (steps 5-8 combined)": 0.0176583117005066, + "Regularization matrix (H)": 0.01004674531868659 + }, + "setup_split_vmap": { + "Border relocation": 6.506707577500493e-05, + "Triangulation + interpolation": 0.006311470760556404, + "Mapping matrix": 0.0002295911064720711, + "Blurred mapping matrix (PSF)": 0.010760049194504972 + }, + "regularization_matrix_prefix_vmap_per_call_s": 0.016423283155018, + "interpolator_prefix_vmap_per_call_s": 0.0063765378363314085, + "split_sibson_prefix_vmap_per_call_s": 0.006405580611317419, + "regularization_split_vmap": { + "Split-point Sibson": 2.9042774986010218e-05, + "Regularization matrix (H, ConstantSplit assembly)": 0.010017702543700581 + } +} \ No newline at end of file diff --git a/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly_control.png b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly_control.png new file mode 100644 index 000000000..93b089407 Binary files /dev/null and b/results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly_control.png differ diff --git a/results/delaunay_nn/assembly_bench_a100_fp64.json b/results/delaunay_nn/assembly_bench_a100_fp64.json new file mode 100644 index 000000000..edfb89c86 --- /dev/null +++ b/results/delaunay_nn/assembly_bench_a100_fp64.json @@ -0,0 +1,139 @@ +{ + "label": "A100 fp64, PyAutoArray#536 A/B session", + "backend": "gpu", + "devices": [ + "cuda:0" + ], + "x64": true, + "S": 6000, + "K": 33, + "P": 1500, + "batch": 16, + "repeats": 10, + "tables": "/mnt/ral/jnightin/autolens_profiling_wt/delaunay-nn-constant-split-assembly/results/delaunay_nn/tables_hst_1500.npz", + "duplicate_stats": { + "padded_scatter_entries": 6534000, + "real_contributions": 187242, + "unique_ij_pairs": 29020, + "max_multiplicity": 40, + "mean_multiplicity": 6.452170916609235, + "size_min": 1, + "size_max": 11, + "size_mean": 5.449666666666666 + }, + "xla_flags": "--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0", + "results": [ + { + "variant": "V0", + "backend": "gpu", + "unbatched_ms": 10.494538000784814, + "peak_mem_mib_after_unbatched": 91.29052734375, + "batched_total_ms": 160.48924287315458, + "per_call_vmap_ms": 10.030577679572161, + "peak_mem_mib_after_batched": 669.7373046875, + "h_slogdet": [ + 1.0, + 995.7174023428239 + ], + "h_frobenius": 121.94212479561611 + }, + { + "variant": "V0a_4wide", + "backend": "gpu", + "unbatched_ms": 0.1839539036154747, + "peak_mem_mib_after_unbatched": 669.7373046875, + "batched_total_ms": 0.6146070081740618, + "per_call_vmap_ms": 0.03841293801087886, + "peak_mem_mib_after_batched": 669.7373046875 + }, + { + "variant": "V1_dense_gemm", + "backend": "gpu", + "unbatched_ms": 9.83729655854404, + "peak_mem_mib_after_unbatched": 669.7373046875, + "max_abs_diff": 1.7763568394002505e-15, + "max_rel_diff": 1.380094597460517e-13, + "batched_total_ms": 161.71646094880998, + "per_call_vmap_ms": 10.107278809300624, + "peak_mem_mib_after_batched": 2864.466552734375, + "h_slogdet": [ + 1.0, + 995.7174023299104 + ] + }, + { + "variant": "V1_highest", + "backend": "gpu", + "unbatched_ms": 9.81782702729106, + "peak_mem_mib_after_unbatched": 2864.466552734375, + "max_abs_diff": 1.7763568394002505e-15, + "max_rel_diff": 1.380094597460517e-13, + "batched_total_ms": 161.72618698328733, + "per_call_vmap_ms": 10.107886686455458, + "peak_mem_mib_after_batched": 2864.466552734375 + }, + { + "variant": "V2_bcoo_dense", + "backend": "gpu", + "unbatched_ms": 19.93540208786726, + "peak_mem_mib_after_unbatched": 2864.466552734375, + "max_abs_diff": 1.7763568394002505e-15, + "max_rel_diff": 2.3831650384478404e-13, + "batched_total_ms": 335.31418396160007, + "per_call_vmap_ms": 20.957136497600004, + "peak_mem_mib_after_batched": 2906.764404296875 + }, + { + "variant": "V2b_bcoo_bcoo", + "backend": "gpu", + "unbatched_error": "JaxRuntimeError: UNIMPLEMENTED: Stable sorting of more than 2^31-1 elements is not implemented", + "batched_error": "JaxRuntimeError: UNIMPLEMENTED: Stable sorting of more than 2^31-1 elements is not implemented" + }, + { + "variant": "V3_dedup_segsum", + "backend": "gpu", + "unbatched_ms": 31.788493040949106, + "peak_mem_mib_after_unbatched": 2906.764404296875, + "max_abs_diff": 1.7763568394002505e-15, + "max_rel_diff": 1.2869336668767806e-13, + "batched_total_ms": 626.161597087048, + "per_call_vmap_ms": 39.1350998179405, + "peak_mem_mib_after_batched": 3577.910400390625 + }, + { + "variant": "V5_compact_k12", + "backend": "gpu", + "unbatched_ms": 0.8457755902782083, + "peak_mem_mib_after_unbatched": 3577.910400390625, + "max_abs_diff": 1.7763568394002505e-15, + "max_rel_diff": 4.766330076894545e-13, + "batched_total_ms": 9.260324528440833, + "per_call_vmap_ms": 0.5787702830275521, + "peak_mem_mib_after_batched": 3577.910400390625 + }, + { + "variant": "V5_compact_k16", + "backend": "gpu", + "unbatched_ms": 1.782560022547841, + "peak_mem_mib_after_unbatched": 3577.910400390625, + "max_abs_diff": 1.7763568394002505e-15, + "max_rel_diff": 2.3831650384478404e-13, + "batched_total_ms": 23.340504500083625, + "per_call_vmap_ms": 1.4587815312552266, + "peak_mem_mib_after_batched": 3577.910400390625 + }, + { + "variant": "V6_library", + "backend": "gpu", + "unbatched_ms": 1.2854335363954306, + "peak_mem_mib_after_unbatched": 3577.910400390625, + "max_abs_diff": 1.7763568394002505e-15, + "max_rel_diff": 4.766330076894545e-13, + "batched_total_ms": 13.394007459282875, + "per_call_vmap_ms": 0.8371254662051797, + "peak_mem_mib_after_batched": 3577.910400390625 + } + ], + "autoarray_file": "/mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/PyAutoArray_feature/autoarray/__init__.py", + "autoarray_version": "2026.8.17.1" +} \ No newline at end of file diff --git a/results/delaunay_nn/mesh_points_hst.npy b/results/delaunay_nn/mesh_points_hst.npy new file mode 100644 index 000000000..56141ae25 Binary files /dev/null and b/results/delaunay_nn/mesh_points_hst.npy differ diff --git a/results/delaunay_nn/tables_hst_1500.npz b/results/delaunay_nn/tables_hst_1500.npz new file mode 100644 index 000000000..169a70a69 Binary files /dev/null and b/results/delaunay_nn/tables_hst_1500.npz differ diff --git a/results/notes/delaunay_nn_constant_split_assembly.md b/results/notes/delaunay_nn_constant_split_assembly.md new file mode 100644 index 000000000..9966659e0 --- /dev/null +++ b/results/notes/delaunay_nn_constant_split_assembly.md @@ -0,0 +1,343 @@ +# DelaunayNN ConstantSplit assembly — split-stencil compaction (A100, 2026-09-08) + +PyAutoArray issue #536, branch `feature/delaunay-nn-constant-split-assembly`, commit +`04152bb7`, against its merge base `180c8a40` (PyAutoArray `main`, #534 merged). This is the +follow-up to `delaunay_nn_launch_latency.md` (#532/#533), whose closing section identified +the `ConstantSplit` assembly — not the Sibson loops — as the larger remaining per-call lever +on this cell, at ~10.0 ms per call out of a 16.4 ms params→H prefix at `vmap` 16. + +## What changed + +One function, `pixel_splitted_regularization_matrix_from`'s JAX path in +`autoarray/inversion/regularization/regularization_util.py`: + +1. **The main scatter is compacted to the first `SPLIT_REG_COMPACT_WIDTH = 12` columns.** + The split stencil tables are fixed-shape `(4P, K)` with `K = 33` for `DelaunayNN` + (`SIBSON_MAX_NEIGHBORS` 32 + 1 spare column), while the real occupied width on this cell + is min 1 / mean 5.45 / max 11 over its 6,000 rows. The old code scattered the full padded + `(4P, K, K)` outer product — 6.53 M entries for 187 k real contributions, ~97 % padding, + at a cost quadratic in the padded width. +2. **The `SPLIT_REG_WIDE_ROW_BUDGET = 256` widest rows get a full-width supplement.** The + head × head block the compact pass already scattered is masked to zero and the remaining + head × tail / tail × head / tail × tail blocks go in as one masked `(W, K, K)` scatter — + one kernel launch rather than three block scatters. Rows wider than 12 are therefore + reproduced exactly, not dropped. The cap audit over 101 ensemble geometries + (`delaunay_nn_cap_audit.md`) saw at most 50 above-width rows in one geometry, so 256 is a + ~5× margin. +3. **Beyond the budget the matrix is poisoned with NaN**, matching the Sibson cap convention + (NaN weights → NaN likelihood → the sample is discarded), rather than returning a + silently truncated `H`. When `K <= compact_width` — the `Delaunay` mesh's `K = 4`, and the + adapt-split family — the compaction is a no-op: today's single scatter, no supplement, no + guard. + +## Verdict + +On the A100, the HST / Hilbert-1500 / MGE-60 / `ConstantSplit` **DelaunayNN** imaging +likelihood's **params→H prefix falls 24.34 → 15.19 ms unbatched (−9.15 ms, 1.60×)** and +**16.42 → 7.26 ms per call at `vmap` 16 (−9.16 ms, 2.26×)**. The assembly itself, measured +directly on the real split tables, goes **10.03 → 0.84 ms per call at `vmap` 16 (12.0×)** and +10.49 → 1.29 ms unbatched (8.2×). Whole-likelihood single-JIT runtime goes 75.61 → 66.12 ms +(1.14×) and at `vmap` 16 **50.04 → 40.86 ms per call (1.22×)**. The `29144.581944` pin held +**bit-identically on every leg**. + +**The witness is met on both halves.** + +| Witness | Target | Measured | | +|---|---|---|---| +| "H, ConstantSplit assembly" per call @ `vmap` 16 | < 3 ms (from 10.0) | **0.80** (breakdown H row) / **0.84** (direct bench) | met | +| `regularization_matrix_prefix_s` per call @ `vmap` 16 | < 11 ms (from 16.4) | **7.26** | met | +| `EXPECTED_LOG_EVIDENCE_HST = 29144.581944` at rtol 1e-4 | unchanged | identical to the last printed digit on both legs, `pinned_drift: []` on both | met | + +One honest caveat on the first row, discussed in full below: the `--split-setup` table's +"H, ConstantSplit assembly" cell is a *difference of two independently compiled prefixes* +and on the feature leg it reads **−1.60 ms** per call — negative. That number is an +attribution artifact, not a saving, and is not what the witness is judged on. The two +readings that do carry it — the breakdown's own `Regularization matrix (H)` step row (0.80 ms) +and the direct assembly bench (0.84 ms) — agree with each other to 5 %. + +## Provenance + +| Job | Cell | Leg | Node | Start (2026-09-08 BST) | Elapsed | State | +|---|---|---|---|---|---:|---| +| 342334 | `likelihood_breakdown/imaging/delaunay_nn` | control `180c8a40` | `euclid-ral-gpu-2` | 14:16:38 | 2:36 | COMPLETED | +| 342335 | `likelihood_breakdown/imaging/delaunay_nn` | feature `04152bb7` | `euclid-ral-gpu-2` | 14:19:15 | 2:29 | COMPLETED | +| 342336 | `likelihood_runtime/imaging/delaunay_nn` | control `180c8a40` | `euclid-ral-gpu-2` | 14:21:45 | 1:17 | COMPLETED | +| 342337 | `likelihood_runtime/imaging/delaunay_nn` | feature `04152bb7` | `euclid-ral-gpu-2` | 14:23:03 | 1:22 | COMPLETED | +| 342338 | `misc/delaunay_nn/assembly_bench.py` | feature `04152bb7` | `euclid-ral-gpu-2` | 14:24:26 | 0:33 | COMPLETED | + +- **All five jobs ran back-to-back on the same node inside one 8.4-minute window** + (14:16:38 → 14:24:59), so this is a same-session A/B, not a comparison against a recorded + baseline. +- Breakdown flags `--split-setup --vmap-batch 16`; runtime cell full timing only (no + `--vmap-probe`), `vmap` batch 16 on both legs. `PyAutoLens 2026.8.17.1`, fp64 + (`JAX_ENABLE_X64=True`), dense inversion path, NVIDIA A100 80GB PCIe, `backend = gpu`, + `device = cuda:0` on every leg. +- Recorded `xla_flags` identical on all five job units: + `--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0`. These rows are + therefore comparable with the 2026-09-08 rows in `delaunay_nn_launch_latency.md` and + **not** with the 2026-07-10 rows in `preopt_breakdown_baseline.md`. +- `PYAUTO_SIBSON_QUERY_CHUNK` deliberately unset on every leg; the library default + `query_chunk = 4096` (chosen by #533's sweep) is what both legs ran, recorded under + `sibson` in each result JSON. Caps `max_cavity_triangles = 32`, `max_neighbors = 32`. + +### The A/B differs in PyAutoArray and nothing else + +| Repo | Revision | +|---|---| +| PyAutoArray — control leg | `180c8a40ceaffe1252cc4db86e85461bfe4d036a` | +| PyAutoArray — feature leg | `04152bb776350de589ed4349d07e1a7e2cb1608b` | +| PyAutoNerves | `efe7c04a79f8389bdcccb88b4d6f8781b9b27142` | +| PyAutoFit | `6331b80031b1dedae3ce277a8c2b116b52ddd261` | +| PyAutoGalaxy | `6d216c151c914b2ce79af18fbc0348a6fc5425d7` | +| PyAutoLens | `9468e3e1bc10405c622e7408c04c0760dcb26882` | +| autolens_profiling (`AP_ROOT`) | `3739481432d1ee3322d87158d2881896eb89ba6a` | + +The **shared** `/mnt/ral/jnightin/PyAuto` install was deliberately not touched — the +subhalo-validation runs (`342299`, `342311`) were live on it throughout, on the `ral` +partition. Every leg instead prepends a private PyAutoArray checkout to the `PYTHONPATH` +that `activate.sh` exports: +`/mnt/ral/jnightin/PyAuto_wt/delaunay-nn-constant-split-assembly/PyAutoArray_{control,feature}`, +and each job prints `autoarray.__file__` and the checkout's `git rev-parse HEAD` into its +SLURM log (verified: control leg imported `PyAutoArray_control`, feature leg imported +`PyAutoArray_feature`, on every job). Submits ran from +`/mnt/ral/jnightin/autolens_profiling_wt/delaunay-nn-constant-split-assembly` (detached at +`3739481`). + +`PyAutoFit` on the shared install moved between #533's session and this one +(`2680b32d` → `6331b800`). It is identical across *these* two legs, so this A/B is clean; +the caveat is only for reading absolute rows across the two notes. + +Every `error/*.err` was clean of `Traceback` on the four A/B legs (269 bytes each, the +`nvidia-smi`/preflight banner), and `grep -c "truncated to dtype float32"` was 0 on every +`.out` and `.err`. The bench job's `.err` carries one traceback: the diagnostic variant +`V2b_bcoo_bcoo` fails inside JAX with `UNIMPLEMENTED: Stable sorting of more than 2^31-1 +elements is not implemented`. That is a bench-only variant that never enters the library, +the bench catches it and records it as `FAIL`, and all other rows completed. + +### Pins + +`EXPECTED_LOG_EVIDENCE_HST = 29144.581944` passed on all four A/B legs. The two breakdown +legs printed **identical values to the last digit**: + +| | control (342334) | feature (342335) | +|---|---|---| +| `log_evidence (reference)` — eager `FitImaging` | 29144.581943564488 | 29144.581943564488 | +| `log_evidence (step-by-step)` — the JIT chain | 29144.581943577425 | 29144.581943577425 | +| `log_evidence (inv matrices)` | 29144.5819435774 | 29144.5819435774 | + +Both runtime legs: `pinned_expected = 29144.581943885652`, `pinned_drift: []`. + +So on the production geometry the compacted assembly is not merely within `rtol 1e-4` — the +JIT log evidence is bit-identical to the uncompacted one. That is the expected result: no +row on this cell exceeds the compact width (max size 11 vs `kc = 12`), so the compact pass +alone reproduces every block and the supplement adds exact zeros. The bench's V0-vs-V6 +`max abs diff` of 1.78e-15 (`max rel` 4.77e-13) on the `H` *entries* is GPU scatter +reassociation between two differently-shaped `.at[].add`, not a semantic difference, and it +does not reach the log evidence. + +## Control vs feature — step table + +All values ms per likelihood call. `vmap/16` is `jax.jit(jax.vmap(fn))` over a params pytree +broadcast to batch 16, reported as batch time / 16; only the combined inversion-setup block, +the `--split-setup` prefixes and the params→H prefix are re-timed under `vmap`, so the +remaining rows read `—` there. + +| Step | control | feature | Δ | control vmap/16 | feature vmap/16 | +|---|---:|---:|---:|---:|---:| +| Ray-trace data grid | 0.142 | 0.196 | +0.054 | — | — | +| Ray-trace mesh grid | 0.188 | 0.186 | −0.001 | — | — | +| Lens light images (pre-PSF) | 0.136 | 0.167 | +0.032 | — | — | +| Blurred image (PSF convolution) | 1.227 | 0.822 | −0.405 | — | — | +| Profile-subtracted image | 0.151 | 0.120 | −0.031 | — | — | +| **Inversion setup (steps 5–8 combined)** | **32.359** | **27.845** | **−4.514** | **17.658** | **17.416** | +| Data vector (D) | 0.366 | 0.370 | +0.004 | — | — | +| Curvature matrix (F) | 4.817 | 4.851 | +0.034 | — | — | +| **Regularization matrix (H)** | **9.917** | **0.800** | **−9.117** | **10.047** | **0.802** | +| Regularized reconstruction | 30.681 | 30.846 | +0.166 | — | — | +| Mapped recon + log evidence | 2.224 | 2.260 | +0.036 | — | — | +| **Total step-by-step (unbatched)** | **82.207** | **68.464** | **−13.743** | — | — | + +The `H` row is the whole story: 12.4× unbatched, 12.5× per call at `vmap` 16, and it is the +only row that moves by more than the ±0.4 ms run-to-run noise of the small steps. The +inversion-setup block moves by −4.5 ms unbatched but is flat under `vmap` (17.66 → 17.42), +because the assembly does not live inside it — it lives in the `H` step downstream. + +### Four-way setup split and the prefixes + +| Piece | control | feature | Δ | control vmap/16 | feature vmap/16 | +|---|---:|---:|---:|---:|---:| +| Border relocation | 1.634 | 1.086 | −0.548 | 0.065 | 0.064 | +| Triangulation + interpolation | 12.790 | 13.302 | +0.512 | 6.311 | 6.395 | +| Mapping matrix | −0.096 | −0.031 | +0.065 | 0.230 | 0.086 | +| Blurred mapping matrix (PSF) | 12.247 | 8.401 | −3.846 | 10.760 | 8.322 | +| *Interpolator prefix (params→step 6)* | *14.424* | *14.389* | *−0.035* | *6.377* | *6.459* | +| *Split-Sibson prefix (params→split mappings)* | *18.092* | *18.426* | *+0.334* | *6.406* | *8.860* | +| **params→H prefix (the witness)** | **24.341** | **15.189** | **−9.152 (1.60×)** | **16.423** | **7.260 (2.26×)** | + +The interpolator prefix is flat to 0.2 % on both readings, exactly as it should be: this +change touches nothing upstream of `reg_split_from`. The whole −9.15 ms comes out of the +last interval, the assembly. The "Blurred mapping matrix (PSF)" unbatched sign is the +by-difference artifact documented in `delaunay_nn_breakdown.md` and its 3.8 ms swing is +between two rows that neither leg touches. + +Two sanity anchors against #533's shipped leg, measured twelve hours earlier on the same +node with the same flags: this run's **control** reads 16.423 ms params→H per call at +`vmap` 16 against #533's 16.435 (0.07 % apart) and 10.018 ms assembly against 10.001 +(0.2 % apart). The control leg reproduces the predecessor to within a fifth of a percent, +which is the strongest available statement that the two notes' rows are on one scale. + +### Split-point Sibson vs ConstantSplit assembly — and why the feature cell goes negative + +| Piece | control | feature | control vmap/16 | feature vmap/16 | +|---|---:|---:|---:|---:| +| Split-point Sibson | 3.668 | 4.037 | 0.029 | 2.401 | +| H, ConstantSplit assembly | 6.249 | −3.237 | **10.018** | **−1.599** | + +**The feature leg's negative assembly cell is an attribution artifact, not a saving of more +than 100 %.** This row is `regularization_matrix_prefix − split_sibson_prefix`, a difference +of two *independently compiled* programs. The split-Sibson prefix must **materialise** the +split mappings, sizes and weights as its outputs; the params→H prefix consumes them straight +into `reg_split_from` and never writes them out. Once the assembly costs only ~0.8 ms, that +materialisation overhead (~1.6 ms per call at `vmap` 16 here) is larger than the interval +being measured, and the subtraction goes below zero. #533's note recorded the same +phenomenon one boundary earlier, where its "Split-point Sibson" row read −2.510; the caveat +is the one in `delaunay_walk_early_exit.md`, that a prefix difference is only meaningful when +it is large against the compile-boundary slack. + +Three independent readings of the same quantity at `vmap` 16 — and only the first is a +prefix difference: + +| reading | control | feature | +|---|---:|---:| +| `--split-setup` differenced cell | 10.018 | −1.599 (artifact) | +| breakdown `Regularization matrix (H)` step row | 10.047 | 0.802 | +| direct assembly bench, real tables (below) | 10.031 | 0.837 | + +The last two agree to 4 %, on two different machines-worth of code paths, and they are what +the verdict rests on. The differenced cell is retained here only because dropping a row that +misbehaves would hide the mechanism. + +## Direct assembly bench (job 342338, `scripts/misc/delaunay_nn/assembly_bench.py`) + +The assembly function timed on its own, on the real `(4P, K) = (6000, 33)` HST split stencil +tables committed at `results/delaunay_nn/tables_hst_1500.npz` (`P = 1500` mesh pixels; +occupied sizes min 1 / mean 5.45 / max 11; 6,534,000 padded scatter entries for 187,242 real +contributions into 29,020 unique `(i, j)` pairs, mean multiplicity 6.45). `V0` is the +pre-#536 full padded scatter, forced from the same source by `compact_width = K`; every +other variant is checked against it. fp64, 3 warm-ups + median of 10, `block_until_ready`. + +| variant | unbatched ms | per call @ `vmap` 16 | max abs diff vs V0 | max rel diff | +|---|---:|---:|---:|---:| +| **V0** (pre-#536, full `(4P, 33, 33)` scatter) | **10.495** | **10.031** | — | — | +| V0a_4wide (the `Delaunay` mesh's `K = 4`, for scale) | 0.184 | 0.038 | — | — | +| V1_dense_gemm | 9.837 | 10.107 | 1.78e-15 | 1.38e-13 | +| V1_highest | 9.818 | 10.108 | 1.78e-15 | 1.38e-13 | +| V2_bcoo_dense | 19.935 | 20.957 | 1.78e-15 | 2.38e-13 | +| V2b_bcoo_bcoo | FAIL | FAIL | — | — | +| V3_dedup_segsum | 31.788 | 39.135 | 1.78e-15 | 1.29e-13 | +| V5_compact_k12 (compaction only, no supplement) | 0.846 | 0.579 | 1.78e-15 | 4.77e-13 | +| V5_compact_k16 (compaction only, no supplement) | 1.783 | 1.459 | 1.78e-15 | 2.38e-13 | +| **V6_library** (shipped: `kc = 12` + wide-row supplement) | **1.285** | **0.837** | **1.78e-15** | **4.77e-13** | + +**V0 → V6_library is 8.2× unbatched and 12.0× per call at `vmap` 16**, and it lands within +5 % of the breakdown's own `H` row on both readings. Three things this table settles: + +- **The width, not the algorithm, was the cost.** Every clever reformulation loses: a dense + GEMM (V1) is flat at 9.8 ms because it is still `K`-wide; a BCOO route (V2) is 2× worse; + deduplicate-then-segment-sum (V3) is 3× worse and 4× worse batched. Only narrowing the + scatter helps — and it helps *faster* than the `O(K²)` entry count predicts. 33 → 12 is + 7.6× by entry count and 17.3× measured (10.031 → 0.579); 33 → 4, the `Delaunay` mesh's own + width, is 68× by entry count and 264× measured (10.031 → 0.038). The extra factor is the + scatter's own behaviour, not arithmetic: a narrower block means fewer atomic conflicts per + destination and a smaller working set. `O(K²)` is therefore a floor on the saving, not an + estimate of it. +- **The safety guarantee costs ~0.26 ms per call.** V5_compact_k12 is the compaction with + no supplement and no guard, at 0.579 ms; V6_library adds the 256-row full-width supplement + and the overflow check for 0.837 ms. On this geometry the supplement contributes exactly + zero to the matrix (no row exceeds 12), so that 0.26 ms buys only correctness on the tail + geometries the cap audit found — which is the right trade, but it should be recorded as a + price rather than assumed free. +- **Width 12 is the right pick.** V5 at 16 costs 1.459 ms per call against 0.579 at 12 — + 2.5× — for margin the supplement already provides exactly. + +`peak MiB` in the bench JSON is a process high-water mark that only ratchets upward across +variants, so it does not attribute memory per variant and is not quoted here. + +## Whole-likelihood runtime cell (`likelihood_runtime/imaging/delaunay_nn`) + +| | control | feature | Δ | +|---|---:|---:|---:| +| Full pipeline (single JIT) | 75.614 | 66.117 | **−9.50 (1.14×)** | +| `vmap` batch 16, per call | 50.044 | 40.858 | **−9.19 (1.22×)** | + +This is what a production evaluation pays. Both savings track the params→H prefix drop +(−9.15 / −9.16 ms) essentially exactly — unlike #533, where the unbatched saving exceeded the +prefix drop by fusion slack. That is the signature of a change local to one late step: the +work removed is removed whole, and nothing upstream re-fuses around it. + +At `vmap` 16 the whole likelihood is now **40.86 ms per call**, down from 52.62 ms at the end +of #533 and 58.02 ms before it — 1.42× for the two changes together on the batched production +number. Those two earlier figures are from #533's session rather than this one; the +within-session statement is this note's own 50.04 → 40.86. + +## Where the remaining per-call cost is + +At `vmap` 16 on the feature leg, the 7.26 ms params→H prefix decomposes as: + +| piece | per call @ `vmap` 16 | share of params→H | comparison | +|---|---:|---:|---| +| Sibson (locate + both interpolation passes) | 6.459 | 89 % | Delaunay's whole interpolator prefix: 5.07 ms | +| **H, ConstantSplit assembly** | **0.80–0.84** | **11 %** | Delaunay's H row: 0.07 ms | + +(6.459 + 0.84 = 7.30 against the measured 7.26, so the decomposition closes to within 0.6 %.) + +**This inverts #533's closing finding.** There, after Phase A, ~60 % of the per-call params→H +prefix was the 33-wide assembly and ~40 % was Sibson, and the note concluded that "the bigger +per-call lever is a ConstantSplit-assembly rework, not the cavity early exit — by roughly 8×". +That rework is now done, and the ratio is 11 : 89 the other way. The assembly is 12× closer +to Delaunay's 0.07 ms H row than it was, and there is no second order of magnitude left in +it: 0.80 ms is 2.0 % of the 40.86 ms whole likelihood, so even taking it to zero would be +worth less than half of what this change was worth. + +### Phase B re-costed against these numbers + +Phase B as planned is the cavity `fori_loop` → `lax.while_loop` early exit. It targets the +Sibson share only, which is now the whole of the remaining prefix: + +- **Its headroom at `vmap` 16 is 6.46 ms per call** (up from a 6.4 ms headroom out of a + 16.4 ms prefix — the headroom did not grow, the prefix around it shrank). The observed + cavity occupancy on the production audit is 11 of the 32 fixed trips (`max_cavity = 11`, + `max_neighbors = 13` in the workspace_test gate), so the early exit removes the tail of a + 32-trip loop that is already latency-amortised at chunk 4096. The plan's own estimate was + ~20 % of the remaining Sibson cost, i.e. **≈ 1.3 ms per call at `vmap` 16** and ≈ 3 ms + unbatched. +- Against the new whole-likelihood number that ≈ 1.3 ms is **≈ 3.2 % of a production + evaluation**, up from ≈ 2.5 % when #533 costed it, purely because the denominator fell + from 52.6 to 40.9 ms. +- **Phase B is now the largest remaining lever on the params→H prefix**, by ~8× over what is + left of the assembly — the exact reverse of the ordering #533 recorded, and the reason its + recommendation should not be read forward without this note. + +Phase C (loop-free k-ring cavity, which changes fp summation order) also targets the Sibson +share and should be re-costed against the same 6.46 ms rather than against #533's figures. + +Outside the params→H prefix, the standing rock is unchanged and is now much the largest +single row: **Regularized reconstruction at 30.85 ms unbatched**, the dense solve, untouched +by any of this work. + +## Artifacts + +- `results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly_control.{json,png}` +- `results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_assembly.{json,png}` +- `results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly_control.json` +- `results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly.json` +- `results/delaunay_nn/assembly_bench_a100_fp64.json` +- `results/delaunay_nn/{tables_hst_1500.npz,mesh_points_hst.npy}` (the bench's inputs) +- `hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_assembly{,_control}` +- `hpc/batch_gpu/submit_runtime_imaging_delaunay_nn_a100_hst_fp64_assembly{,_control}` +- `hpc/batch_gpu/submit_delaunay_nn_assembly_bench_a100` + +The 2026-09-08 `..._launch_latency` rows in `delaunay_nn_launch_latency.md` are left +untouched; this note's control leg supersedes that note's rows for any comparison against +this branch (its merge base moved from `d7c96762` to `180c8a40`), and the two control legs +agree to 0.2 %, so the row sets are in practice on one scale. diff --git a/results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly.json b/results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly.json new file mode 100644 index 000000000..a5ff96a05 --- /dev/null +++ b/results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly.json @@ -0,0 +1,30 @@ +{ + "autolens_version": "2026.8.17.1", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 35333 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" + }, + "full_pipeline_single_jit": 0.06611713680904359, + "vmap": { + "batch_size": 16, + "batch_time": 0.6537213676143437, + "per_call": 0.040857585475896484, + "speedup_vs_single_jit": 1.6 + }, + "pinned_expected": 29144.581943885652, + "pinned_drift": [] +} \ No newline at end of file diff --git a/results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly_control.json b/results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly_control.json new file mode 100644 index 000000000..f3cfdc7a3 --- /dev/null +++ b/results/runtime/imaging/delaunay_nn/delaunay_nn_hpc_a100_fp64_assembly_control.json @@ -0,0 +1,30 @@ +{ + "autolens_version": "2026.8.17.1", + "device": { + "backend": "gpu", + "device": "cuda:0", + "xla_flags": "--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0", + "omp_num_threads": null, + "cpu_count": 124, + "nvidia_smi": "NVIDIA A100 80GB PCIe, 35331 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" + }, + "full_pipeline_single_jit": 0.07561378970276564, + "vmap": { + "batch_size": 16, + "batch_time": 0.8007008988177404, + "per_call": 0.05004380617610878, + "speedup_vs_single_jit": 1.5 + }, + "pinned_expected": 29144.581943885652, + "pinned_drift": [] +} \ No newline at end of file diff --git a/scripts/misc/delaunay_nn/README.md b/scripts/misc/delaunay_nn/README.md index 83229b52b..dfb81727f 100644 --- a/scripts/misc/delaunay_nn/README.md +++ b/scripts/misc/delaunay_nn/README.md @@ -47,3 +47,46 @@ actual Hilbert meshes ray-traced through a mass-model ensemble. The current cap decision and measured tail distributions are recorded in [`results/notes/delaunay_nn_cap_audit.md`](../../../results/notes/delaunay_nn_cap_audit.md). + +## ConstantSplit assembly bench (`assembly_bench.py`) + +`assembly_bench.py` times the *regularization* half of the DelaunayNN cell — +`reg_split_from` plus `pixel_splitted_regularization_matrix_from` — which the +launch-latency work left as the largest remaining per-call cost +([`results/notes/delaunay_nn_launch_latency.md`](../../../results/notes/delaunay_nn_launch_latency.md)). +It runs the real `(6000, 33)` HST split stencil tables, checked in at +`results/delaunay_nn/tables_hst_1500.npz`, through every candidate formulation: + +| variant | what it measures | +|---|---| +| `V0` | the pre-#536 path: the full padded `(4P, K, K)` scatter (`compact_width=K`); the reference every other variant is diffed against | +| `V0a_4wide` | the same scatter on 4-wide `Delaunay`-style tables, as the width calibration | +| `V1_dense_gemm` / `V1_highest` | dense `B^T diag(s) B` | +| `V2_bcoo_dense` / `V2b_bcoo_bcoo` | sparse BCOO assembly | +| `V3_dedup_segsum` | dedup sort + `segment_sum` | +| `V5_compact_k` | a fixed compaction to `N` columns with no wide-row supplement — the width sweep that priced the compaction | +| `V6_library` | **the shipped library path**: default `SPLIT_REG_COMPACT_WIDTH` compaction plus the wide-row supplement | + +```bash +python scripts/misc/delaunay_nn/assembly_bench.py +python scripts/misc/delaunay_nn/assembly_bench.py --pixels 300 --repeats 3 # smoke +``` + +Defaults read `results/delaunay_nn/tables_hst_1500.npz` and write +`results/delaunay_nn/assembly_bench.json`. The tables are checked in (330 KB); +rebuild them from the HST breakdown cell with + +```bash +python scripts/misc/delaunay_nn/extract_mesh_points.py \ + --out results/delaunay_nn/mesh_points_hst.npy +python scripts/misc/delaunay_nn/build_tables.py \ + --points results/delaunay_nn/mesh_points_hst.npy \ + --out results/delaunay_nn/tables_hst_1500.npz +``` + +`V6_library` must reproduce `V0` exactly on these tables (max abs diff 0.0 on +CPU): every real stencil row is narrower than the compact width, so the +supplement contributes only zeros. The tail geometries that do exceed it are +audited by `autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn_caps.py`, +whose per-geometry count of above-width split rows is the margin the wide-row +budget is set against. diff --git a/scripts/misc/delaunay_nn/assembly_bench.py b/scripts/misc/delaunay_nn/assembly_bench.py new file mode 100644 index 000000000..0c584cf88 --- /dev/null +++ b/scripts/misc/delaunay_nn/assembly_bench.py @@ -0,0 +1,564 @@ +"""Benchmark alternative formulations of the DelaunayNN ConstantSplit +regularization-matrix assembly (PyAutoArray ``reg_split_from`` + +``pixel_splitted_regularization_matrix_from``). + +Diagnostic only: nothing here is imported by the library. Inputs are the real +(4P, K) split stencil tables built once from the HST DelaunayNN cell's +source-plane mesh (``build_tables.py``), loaded from +``results/delaunay_nn/tables_hst_1500.npz`` by default. + +``V0`` is the pre-#536 assembly (the full padded ``(4P, K, K)`` scatter, forced +by ``compact_width=K``) and is the reference every other variant is checked +against; ``V6_library`` is the shipped library path with its default compaction +(``SPLIT_REG_COMPACT_WIDTH`` + the wide-row supplement), so the bench measures +what production actually runs. + +Every variant is timed as a device-argument ``jax.jit`` function (unbatched) +and as ``jax.jit(jax.vmap(...))`` over a batch of 16, fp64, with +``block_until_ready`` and 3 warm-ups + median of N repeats. + +Run from the autolens_profiling root: + +```bash +python scripts/misc/delaunay_nn/assembly_bench.py +``` +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import time +import traceback +from pathlib import Path + +import numpy as np + + +def _profiling_root() -> Path: + for parent in Path(__file__).resolve().parents: + if (parent / "ruff.toml").exists(): + return parent + raise RuntimeError("autolens_profiling root (ruff.toml) not found") + + +ROOT = _profiling_root() +DEFAULT_TABLES = ROOT / "results" / "delaunay_nn" / "tables_hst_1500.npz" +DEFAULT_OUT = ROOT / "results" / "delaunay_nn" / "assembly_bench.json" + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +parser = argparse.ArgumentParser() +parser.add_argument("--tables", default=str(DEFAULT_TABLES), help=".npz from build_tables.py") +parser.add_argument("--out", default=str(DEFAULT_OUT), help="JSON output path") +parser.add_argument( + "--pixels", type=int, default=0, help="truncate to this many mesh pixels (smoke)" +) +parser.add_argument("--batch", type=int, default=16) +parser.add_argument("--repeats", type=int, default=10) +parser.add_argument("--skip", default="", help="comma-separated variant names to skip") +parser.add_argument("--label", default="", help="free-text label recorded in the JSON") +parser.add_argument( + "--compact-widths", + default="12,16", + help="comma-separated fixed compaction widths for the V5 diagnostic (no supplement)", +) +args = parser.parse_args() + +SKIP = {s.strip() for s in args.skip.split(",") if s.strip()} + +os.environ.setdefault("JAX_ENABLE_X64", "True") + +import jax +import jax.numpy as jnp + +jax.config.update("jax_enable_x64", True) + +from autoarray.inversion.regularization import regularization_util as reg_util + +BACKEND = jax.default_backend() +print(f"backend : {BACKEND}") +print(f"devices : {jax.devices()}") +print(f"x64 enabled : {jax.config.jax_enable_x64}") + +try: + import autoarray + + print(f"autoarray : {autoarray.__file__}") +except Exception: # pragma: no cover + pass + + +# --------------------------------------------------------------------------- +# Tables +# --------------------------------------------------------------------------- + +data = np.load(args.tables) +mappings_np = data["mappings"].astype(np.int32) # (S, K) +sizes_np = data["sizes"].astype(np.int32) # (S,) +weights_np = data["weights"].astype(np.float64) # (S, K) + +if args.pixels: + keep_pixels = int(args.pixels) + keep_rows = 4 * keep_pixels + mappings_np = mappings_np[:keep_rows].copy() + sizes_np = sizes_np[:keep_rows].copy() + weights_np = weights_np[:keep_rows].copy() + # Drop stencil entries that point at dropped pixels; renormalise the row. + bad = mappings_np >= keep_pixels + mappings_np = np.where(bad, -1, mappings_np) + weights_np = np.where(bad, 0.0, weights_np) + # Compact each row so the surviving entries are contiguous within size. + for r in range(mappings_np.shape[0]): + k = int(sizes_np[r]) + m = mappings_np[r, :k] + w = weights_np[r, :k] + good = m >= 0 + m2, w2 = m[good], w[good] + n = m2.shape[0] + mappings_np[r, :] = -1 + weights_np[r, :] = 0.0 + mappings_np[r, :n] = m2 + weights_np[r, :n] = w2 / (w2.sum() if w2.sum() != 0 else 1.0) + sizes_np[r] = n + +S, K = mappings_np.shape +P = S // 4 +BATCH = args.batch +REPEATS = args.repeats + +print(f"tables : S={S} K={K} P={P}") +print(f"sizes : min={sizes_np.min()} max={sizes_np.max()} mean={sizes_np.mean():.2f}") + + +def duplicate_stats(mappings, sizes): + """Unique (i, j) pairs vs total contributions in the (P, P) scatter.""" + valid = mappings != -1 + rows_i = [] + rows_j = [] + total = 0 + for r in range(mappings.shape[0]): + m = mappings[r][valid[r]] + total += m.size * m.size + ii, jj = np.meshgrid(m, m, indexing="ij") + rows_i.append(ii.ravel()) + rows_j.append(jj.ravel()) + ii = np.concatenate(rows_i) + jj = np.concatenate(rows_j) + keys = ii.astype(np.int64) * P + jj.astype(np.int64) + uniq, counts = np.unique(keys, return_counts=True) + return { + "padded_scatter_entries": int(S * K * K), + "real_contributions": int(total), + "unique_ij_pairs": int(uniq.size), + "max_multiplicity": int(counts.max()), + "mean_multiplicity": float(counts.mean()), + "size_min": int(sizes.min()), + "size_max": int(sizes.max()), + "size_mean": float(sizes.mean()), + } + + +DUP = duplicate_stats(mappings_np, sizes_np) +print("duplicate stats:", DUP) + + +# --------------------------------------------------------------------------- +# Device inputs +# --------------------------------------------------------------------------- + +mappings = jnp.asarray(mappings_np) +sizes = jnp.asarray(sizes_np) +weights = jnp.asarray(weights_np) +reg_w = jnp.full((P,), 1.0, dtype=jnp.float64) # ConstantSplit(coefficient=1.0) + +# Batch: mappings/sizes are the same connectivity, weights perturbed per lane +# (under a production vmap every lane re-runs Sibson, so all four are batched). +rng = np.random.default_rng(0) +pert = 1.0 + 1e-3 * rng.standard_normal((BATCH, 1, 1)) +weights_b = jnp.asarray(weights_np[None, :, :] * pert) +mappings_b = jnp.broadcast_to(mappings, (BATCH,) + mappings.shape) +sizes_b = jnp.broadcast_to(sizes, (BATCH,) + sizes.shape) +reg_w_b = jnp.broadcast_to(reg_w, (BATCH,) + reg_w.shape) +mappings_b = jnp.asarray(mappings_b) +sizes_b = jnp.asarray(sizes_b) +reg_w_b = jnp.asarray(reg_w_b) + +# 4-wide (Delaunay-style) calibration tables: 3 real columns + 1 spare. +K4 = 4 +mappings4_np = -np.ones((S, K4), dtype=np.int32) +weights4_np = np.zeros((S, K4), dtype=np.float64) +sizes4_np = np.minimum(sizes_np, 3).astype(np.int32) +for r in range(S): + k = int(sizes4_np[r]) + mappings4_np[r, :k] = mappings_np[r, :k] + w = weights_np[r, :k] + tot = w.sum() + weights4_np[r, :k] = w / (tot if tot != 0 else 1.0) +mappings4 = jnp.asarray(mappings4_np) +sizes4 = jnp.asarray(sizes4_np) +weights4 = jnp.asarray(weights4_np) +mappings4_b = jnp.asarray(np.broadcast_to(mappings4_np, (BATCH, S, K4))) +sizes4_b = jnp.asarray(np.broadcast_to(sizes4_np, (BATCH, S))) +weights4_b = jnp.asarray(weights4_np[None, :, :] * pert) + + +# --------------------------------------------------------------------------- +# Variants +# --------------------------------------------------------------------------- + + +def _split(mappings, sizes, weights): + """``reg_split_from`` JAX branch — shared prologue of every variant.""" + return reg_util.reg_split_from( + splitted_mappings=mappings, + splitted_sizes=sizes, + splitted_weights=weights, + xp=jnp, + ) + + +def v0(mappings, sizes, weights, reg_weights): + """Pre-#536 path: reg_split_from + the full padded (4P, K, K) scatter. + + ``compact_width=K`` disables the library's compaction, so this is the + uncompacted scatter the compaction replaced, and the reference result. + """ + m, s, w = _split(mappings, sizes, weights) + return reg_util.pixel_splitted_regularization_matrix_from( + regularization_weights=reg_weights, + splitted_mappings=m, + splitted_sizes=s, + splitted_weights=w, + xp=jnp, + compact_width=m.shape[1], + ) + + +def v6(mappings, sizes, weights, reg_weights): + """The shipped library path: default compaction + wide-row supplement.""" + m, s, w = _split(mappings, sizes, weights) + return reg_util.pixel_splitted_regularization_matrix_from( + regularization_weights=reg_weights, + splitted_mappings=m, + splitted_sizes=s, + splitted_weights=w, + xp=jnp, + ) + + +def _dense_b(m, w, reg_weights): + """Dense split-mapping matrix B (S, P) and the per-row scale s (S,).""" + n_rows = m.shape[0] + n_pix = reg_weights.shape[0] + valid = m != -1 + m_fixed = jnp.where(valid, m, 0) + w_fixed = jnp.where(valid, w, 0.0) + rows = jnp.arange(n_rows, dtype=m_fixed.dtype)[:, None] + rows = jnp.broadcast_to(rows, m_fixed.shape) + b = jnp.zeros((n_rows, n_pix), dtype=w.dtype).at[rows, m_fixed].add(w_fixed) + scale = reg_weights[jnp.arange(n_rows) // 4] ** 2.0 + return b, scale + + +def _v1(mappings, sizes, weights, reg_weights, precision=None): + m, s, w = _split(mappings, sizes, weights) + b, scale = _dense_b(m, w, reg_weights) + h = jnp.matmul(b.T, scale[:, None] * b, precision=precision) + return h + jnp.eye(reg_weights.shape[0], dtype=w.dtype) * 1e-8 + + +def v1(mappings, sizes, weights, reg_weights): + """Dense GEMM: H = B^T diag(s) B.""" + return _v1(mappings, sizes, weights, reg_weights) + + +def v1_highest(mappings, sizes, weights, reg_weights): + return _v1(mappings, sizes, weights, reg_weights, precision=jax.lax.Precision.HIGHEST) + + +def v2(mappings, sizes, weights, reg_weights): + """BCOO sparse B^T @ dense (scaled) B.""" + from jax.experimental import sparse + + m, s, w = _split(mappings, sizes, weights) + n_rows = m.shape[0] + n_pix = reg_weights.shape[0] + valid = m != -1 + m_fixed = jnp.where(valid, m, 0) + w_fixed = jnp.where(valid, w, 0.0) + rows = jnp.broadcast_to(jnp.arange(n_rows, dtype=m_fixed.dtype)[:, None], m_fixed.shape) + idx = jnp.stack([rows.ravel(), m_fixed.ravel()], axis=-1) + scale = reg_weights[jnp.arange(n_rows) // 4] ** 2.0 + b_sp = sparse.BCOO((w_fixed.ravel(), idx), shape=(n_rows, n_pix)) + b_dense = jnp.zeros((n_rows, n_pix), dtype=w.dtype).at[rows, m_fixed].add(w_fixed) + h = b_sp.T @ (scale[:, None] * b_dense) + return h + jnp.eye(n_pix, dtype=w.dtype) * 1e-8 + + +def v2b(mappings, sizes, weights, reg_weights): + """BCOO sparse-sparse: (sqrt(s) B)^T @ (sqrt(s) B) via bcoo_dot_general.""" + from jax.experimental import sparse + + m, s, w = _split(mappings, sizes, weights) + n_rows = m.shape[0] + n_pix = reg_weights.shape[0] + valid = m != -1 + m_fixed = jnp.where(valid, m, 0) + w_fixed = jnp.where(valid, w, 0.0) + scale = jnp.sqrt(reg_weights[jnp.arange(n_rows) // 4] ** 2.0) + w_scaled = w_fixed * scale[:, None] + rows = jnp.broadcast_to(jnp.arange(n_rows, dtype=m_fixed.dtype)[:, None], m_fixed.shape) + idx_t = jnp.stack([m_fixed.ravel(), rows.ravel()], axis=-1) # B^T indices + bt_sp = sparse.BCOO((w_scaled.ravel(), idx_t), shape=(n_pix, n_rows)) + idx = jnp.stack([rows.ravel(), m_fixed.ravel()], axis=-1) + b_sp = sparse.BCOO((w_scaled.ravel(), idx), shape=(n_rows, n_pix)) + h = sparse.bcoo_dot_general( + bt_sp, + b_sp, + dimension_numbers=(((1,), (0,)), ((), ())), + ) + h = h.todense() if hasattr(h, "todense") else h + return h + jnp.eye(n_pix, dtype=w.dtype) * 1e-8 + + +def v3(mappings, sizes, weights, reg_weights): + """Dedup-then-scatter: sort the (i, j) triples, segment_sum, one scatter.""" + m, s, w = _split(mappings, sizes, weights) + n_rows = m.shape[0] + n_pix = reg_weights.shape[0] + valid = m != -1 + m_fixed = jnp.where(valid, m, 0) + w_fixed = jnp.where(valid, w, 0.0) + scale = reg_weights[jnp.arange(n_rows) // 4] ** 2.0 + outer = (w_fixed[:, :, None] * w_fixed[:, None, :]) * scale[:, None, None] + ii = jnp.broadcast_to(m_fixed[:, :, None], outer.shape) + jj = jnp.broadcast_to(m_fixed[:, None, :], outer.shape) + keys = (ii.astype(jnp.int32) * n_pix + jj.astype(jnp.int32)).ravel() + vals = outer.ravel() + order = jnp.argsort(keys) + keys_s = keys[order] + vals_s = vals[order] + new = jnp.concatenate( + [jnp.ones((1,), dtype=jnp.int32), (keys_s[1:] != keys_s[:-1]).astype(jnp.int32)] + ) + seg = jnp.cumsum(new) - 1 + n_seg = keys_s.shape[0] + summed = jax.ops.segment_sum(vals_s, seg, num_segments=n_seg, indices_are_sorted=True) + # First key of each segment. + seg_keys = jnp.zeros((n_seg,), dtype=keys_s.dtype).at[seg].max(keys_s) + flat = jnp.zeros((n_pix * n_pix,), dtype=w.dtype).at[seg_keys].add(summed) + return flat.reshape((n_pix, n_pix)) + jnp.eye(n_pix, dtype=w.dtype) * 1e-8 + + +def make_v5(kc: int): + """Fixed stencil compaction to ``kc`` columns, with no wide-row supplement. + + Numerically identical to V0 whenever every post-split size <= kc, which the + real HST tables satisfy (max size 11 + 1 insertion). Kept as the width sweep + that priced the compaction: it isolates the cost of the main scatter alone, + where ``V6_library`` adds the supplement that makes the result exact in the + tail. + """ + + def v5(mappings, sizes, weights, reg_weights): + m, s, w = _split(mappings, sizes, weights) + m = m[:, :kc] + w = w[:, :kc] + return reg_util.pixel_splitted_regularization_matrix_from( + regularization_weights=reg_weights, + splitted_mappings=m, + splitted_sizes=jnp.minimum(s, kc), + splitted_weights=w, + xp=jnp, + compact_width=kc, + ) + + return v5 + + +# --------------------------------------------------------------------------- +# Timing harness +# --------------------------------------------------------------------------- + + +def _block(x): + for leaf in jax.tree_util.tree_leaves(x): + if hasattr(leaf, "block_until_ready"): + leaf.block_until_ready() + return x + + +def peak_mem_bytes(): + try: + stats = jax.local_devices()[0].memory_stats() + if stats is None: + return None + return int(stats.get("peak_bytes_in_use", 0)) + except Exception: + return None + + +def time_fn(fn, arrays, repeats=REPEATS, warmups=3): + compiled = jax.jit(fn) + result = _block(compiled(*arrays)) + for _ in range(warmups - 1): + _block(compiled(*arrays)) + times = [] + for _ in range(repeats): + t0 = time.perf_counter() + _block(compiled(*arrays)) + times.append(time.perf_counter() - t0) + return statistics.median(times) * 1e3, result + + +def run_variant(name, fn, single_args, batch_args, check_ref=None, batched=True): + entry = {"variant": name, "backend": BACKEND} + try: + ms, result = time_fn(fn, single_args) + entry["unbatched_ms"] = ms + entry["peak_mem_mib_after_unbatched"] = ( + (peak_mem_bytes() or 0) / 2**20 if peak_mem_bytes() is not None else None + ) + if check_ref is not None: + ref = np.asarray(check_ref, dtype=np.float64) + got = np.asarray(result, dtype=np.float64) + diff = np.abs(got - ref) + denom = np.abs(ref) + entry["max_abs_diff"] = float(diff.max()) + nz = denom > 0 + entry["max_rel_diff"] = float((diff[nz] / denom[nz]).max()) if nz.any() else 0.0 + print(f" {name:<16} unbatched {ms:9.3f} ms", flush=True) + except Exception as exc: # pragma: no cover + entry["unbatched_error"] = f"{type(exc).__name__}: {exc}" + print(f" {name:<16} unbatched FAILED: {exc}", flush=True) + traceback.print_exc() + result = None + + if batched: + try: + ms_b, _ = time_fn(jax.vmap(fn), batch_args) + entry["batched_total_ms"] = ms_b + entry["per_call_vmap_ms"] = ms_b / BATCH + entry["peak_mem_mib_after_batched"] = ( + (peak_mem_bytes() or 0) / 2**20 if peak_mem_bytes() is not None else None + ) + print(f" {name:<16} vmap{BATCH:<3d} {ms_b / BATCH:9.3f} ms/call", flush=True) + except Exception as exc: # pragma: no cover + entry["batched_error"] = f"{type(exc).__name__}: {exc}" + print(f" {name:<16} vmap{BATCH:<3d} FAILED: {exc}", flush=True) + return entry, result + + +single_args = (mappings, sizes, weights, reg_w) +batch_args = (mappings_b, sizes_b, weights_b, reg_w_b) +single_args4 = (mappings4, sizes4, weights4, reg_w) +batch_args4 = (mappings4_b, sizes4_b, weights4_b, reg_w_b) + +results = [] + +print("\n--- V0 (baseline) ---", flush=True) +entry, h_ref = run_variant("V0", v0, single_args, batch_args) +results.append(entry) +h_ref_np = np.asarray(h_ref, dtype=np.float64) if h_ref is not None else None +if h_ref_np is not None: + print( + f" H: shape {h_ref_np.shape} nnz {int((h_ref_np != 0).sum())} " + f"frobenius {np.linalg.norm(h_ref_np):.6e}" + ) + sign, logdet = np.linalg.slogdet(h_ref_np) + print(f" slogdet(H) = ({sign}, {logdet:.9f})") + results[-1]["h_slogdet"] = [float(sign), float(logdet)] + results[-1]["h_frobenius"] = float(np.linalg.norm(h_ref_np)) + +VARIANTS = [ + ("V0a_4wide", v0, single_args4, batch_args4, None), + ("V1_dense_gemm", v1, single_args, batch_args, h_ref_np), + ("V1_highest", v1_highest, single_args, batch_args, h_ref_np), + ("V2_bcoo_dense", v2, single_args, batch_args, h_ref_np), + ("V2b_bcoo_bcoo", v2b, single_args, batch_args, h_ref_np), + ("V3_dedup_segsum", v3, single_args, batch_args, h_ref_np), +] + +for _kc in [int(x) for x in args.compact_widths.split(",") if x.strip()]: + VARIANTS.append((f"V5_compact_k{_kc}", make_v5(_kc), single_args, batch_args, h_ref_np)) + +VARIANTS.append(("V6_library", v6, single_args, batch_args, h_ref_np)) + +for name, fn, sa, ba, ref in VARIANTS: + if name in SKIP: + print(f"\n--- {name} (skipped) ---", flush=True) + results.append({"variant": name, "backend": BACKEND, "skipped": True}) + continue + print(f"\n--- {name} ---", flush=True) + entry, res = run_variant(name, fn, sa, ba, check_ref=ref) + if name == "V1_dense_gemm" and res is not None and h_ref_np is not None: + try: + s1, l1 = np.linalg.slogdet(np.asarray(res, dtype=np.float64)) + entry["h_slogdet"] = [float(s1), float(l1)] + except Exception: + pass + results.append(entry) + + +# --------------------------------------------------------------------------- +# Report +# --------------------------------------------------------------------------- + +header = ( + f"{'variant':<18}{'backend':<9}{'unbatched ms':>14}{'per-call@vmap' + str(BATCH):>18}" + f"{'max abs diff':>16}{'max rel diff':>16}{'peak MiB':>12}" +) +print("\n" + "=" * len(header)) +print(header) +print("=" * len(header)) +for r in results: + if r.get("skipped"): + print(f"{r['variant']:<18}{r['backend']:<9}{'skipped':>14}") + continue + ub = r.get("unbatched_ms") + pc = r.get("per_call_vmap_ms") + ad = r.get("max_abs_diff") + rd = r.get("max_rel_diff") + pm = r.get("peak_mem_mib_after_batched") or r.get("peak_mem_mib_after_unbatched") + print( + f"{r['variant']:<18}{r['backend']:<9}" + f"{(f'{ub:.3f}' if ub is not None else 'FAIL'):>14}" + f"{(f'{pc:.3f}' if pc is not None else 'FAIL'):>18}" + f"{(f'{ad:.3e}' if ad is not None else '-'):>16}" + f"{(f'{rd:.3e}' if rd is not None else '-'):>16}" + f"{(f'{pm:.0f}' if pm is not None else '-'):>12}" + ) +print("=" * len(header)) + +payload = { + "label": args.label, + "backend": BACKEND, + "devices": [str(d) for d in jax.devices()], + "x64": bool(jax.config.jax_enable_x64), + "S": int(S), + "K": int(K), + "P": int(P), + "batch": BATCH, + "repeats": REPEATS, + "tables": os.path.abspath(args.tables), + "duplicate_stats": DUP, + "xla_flags": os.environ.get("XLA_FLAGS"), + "results": results, +} +try: + import autoarray + + payload["autoarray_file"] = autoarray.__file__ + payload["autoarray_version"] = getattr(autoarray, "__version__", None) +except Exception: + pass + +with open(args.out, "w") as f: + json.dump(payload, f, indent=2) +print(f"\nwrote {args.out}") diff --git a/scripts/misc/delaunay_nn/build_tables.py b/scripts/misc/delaunay_nn/build_tables.py new file mode 100644 index 000000000..9ead84b60 --- /dev/null +++ b/scripts/misc/delaunay_nn/build_tables.py @@ -0,0 +1,79 @@ +"""Build the real (S, K) DelaunayNN ConstantSplit stencil tables from a mesh point set. + +Runs the NumPy Sibson path (``scipy_delaunay_nn``) on the source-plane mesh +points, then applies ``InterpolatorDelaunayNN._mappings_sizes_weights_split``'s +spare-column hstack, giving the exact (4P, 33) tables ConstantSplit consumes. + +The checked-in tables used by ``assembly_bench.py`` were built from the HST +breakdown cell's 1,500-vertex source-plane mesh, run from the autolens_profiling +root: + +```bash +python scripts/misc/delaunay_nn/extract_mesh_points.py \ + --out results/delaunay_nn/mesh_points_hst.npy +python scripts/misc/delaunay_nn/build_tables.py \ + --points results/delaunay_nn/mesh_points_hst.npy \ + --out results/delaunay_nn/tables_hst_1500.npz +``` +""" + +import argparse + +import numpy as np + +parser = argparse.ArgumentParser() +parser.add_argument("--points", required=True) +parser.add_argument("--out", required=True) +parser.add_argument( + "--pixels", type=int, default=0, help="subsample the point set to this many pixels" +) +parser.add_argument("--seed", type=int, default=0) +args = parser.parse_args() + +from autoarray.inversion.mesh.interpolator.sibson import scipy_delaunay_nn + +pts = np.load(args.points).astype(np.float64) +if args.pixels and args.pixels < pts.shape[0]: + rng = np.random.default_rng(args.seed) + idx = np.sort(rng.choice(pts.shape[0], size=args.pixels, replace=False)) + pts = pts[idx] +print("points:", pts.shape) + +# The data-grid query set is irrelevant to the split tables; keep it tiny. +query = pts[:4].copy() + +out = scipy_delaunay_nn(points_np=pts, query_points_np=query, areas_factor=0.5) +( + points, + simplices_padded, + mappings, + sizes, + weights, + split_points, + splitted_mappings, + splitted_sizes, + splitted_weights, + cavity_sizes, + overflow, + degenerate, + split_cavity_sizes, + split_overflow, + split_degenerate, +) = out + +print("split tables:", splitted_mappings.shape, splitted_weights.shape) +print("split overflow:", int(np.sum(split_overflow)), "degenerate:", int(np.sum(split_degenerate))) +print("nan weights rows:", int(np.sum(np.isnan(splitted_weights).any(axis=1)))) +print( + f"split sizes: min {splitted_sizes.min()} max {splitted_sizes.max()} " + f"mean {splitted_sizes.mean():.2f}" +) + +# InterpolatorDelaunayNN._mappings_sizes_weights_split: one spare padded column. +row_count = splitted_mappings.shape[0] +m = np.hstack([splitted_mappings.astype(np.int32), -np.ones((row_count, 1), dtype=np.int32)]) +w = np.hstack([splitted_weights.astype(np.float64), np.zeros((row_count, 1), dtype=np.float64)]) +s = splitted_sizes.astype(np.int32) + +np.savez_compressed(args.out, mappings=m, sizes=s, weights=w, points=points) +print("wrote", args.out, m.shape, w.shape, s.shape) diff --git a/scripts/misc/delaunay_nn/extract_mesh_points.py b/scripts/misc/delaunay_nn/extract_mesh_points.py new file mode 100644 index 000000000..21497f8c8 --- /dev/null +++ b/scripts/misc/delaunay_nn/extract_mesh_points.py @@ -0,0 +1,115 @@ +"""Extract the real source-plane DelaunayNN mesh grid from the HST breakdown cell. + +Reuses the minimum of scripts/imaging/likelihood_breakdown/delaunay_nn.py needed +to reach ``relocated_mesh_grid`` (1500 Hilbert vertices traced to the source +plane and border-relocated). Writes it to --out as .npy. + +Run from the autolens_profiling root. +""" + +import argparse +import sys +from pathlib import Path + +import numpy as np + +parser = argparse.ArgumentParser() +parser.add_argument("--out", required=True) +args = parser.parse_args() + +root = Path.cwd() +sys.path.insert(0, str(root)) +sys.path.insert(0, str(root / "scripts" / "misc")) + +import autofit as af +import autolens as al +import jax.numpy as jnp +from simulators.imaging import INSTRUMENTS + +from _adapt_image_util import adapt_image_for_dataset + +instrument = "hst" +pixel_scale = INSTRUMENTS[instrument]["pixel_scale"] +dataset_path = Path("dataset") / "imaging" / instrument + +dataset = al.Imaging.from_fits( + data_path=dataset_path / "data.fits", + psf_path=dataset_path / "psf.fits", + noise_map_path=dataset_path / "noise_map.fits", + pixel_scales=pixel_scale, +) + +mask_radius = 3.5 +mask = al.Mask2D.circular( + shape_native=dataset.shape_native, + pixel_scales=dataset.pixel_scales, + radius=mask_radius, +) +dataset = dataset.apply_mask(mask=mask) +dataset = dataset.apply_over_sampling(over_sample_size_lp=4, over_sample_size_pixelization=1) +over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from( + grid=dataset.grid, + sub_size_list=[4, 2, 1], + radial_list=[0.3, 0.6], + centre_list=[(0.0, 0.0)], +) +dataset = dataset.apply_over_sampling( + over_sample_size_lp=over_sample_size, over_sample_size_pixelization=1 +) + +adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) + +n_mesh_vertices = 1500 +image_mesh = al.image_mesh.Hilbert(pixels=n_mesh_vertices, weight_power=1.0, weight_floor=0.0) +image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from( + mask=dataset.mask, adapt_data=adapt_image +) +print("mesh vertices placed:", image_plane_mesh_grid.shape[0]) + +lens_bulge = al.model_util.mge_model_from( + mask_radius=mask_radius, total_gaussians=60, centre_prior_is_uniform=True +) +mass = af.Model(al.mp.Isothermal) +mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005) +mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005) +mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05) +_ell = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0) +mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=_ell[0], sigma=0.01) +mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=_ell[1], sigma=0.01) +shear = af.Model(al.mp.ExternalShear) +shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005) +shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005) +lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear) + +mesh = al.mesh.DelaunayNN(pixels=n_mesh_vertices, areas_factor=0.5, zeroed_pixels=0) +regularization = al.reg.ConstantSplit(coefficient=1.0) +pixelization = al.Pixelization(mesh=mesh, regularization=regularization) +source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization) +model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) + +instance = model.instance_from_vector(vector=model.physical_values_from_prior_medians) +tracer = al.Tracer(galaxies=list(instance.galaxies)) + +from autoarray.inversion.mesh.border_relocator import BorderRelocator + +border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1) +traced_source_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp)[-1] +traced_mesh_source = tracer.traced_grid_2d_list_from( + grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp +)[-1] +relocated_mesh_grid = border_relocator.relocated_mesh_grid_from( + grid=traced_source_grid, mesh_grid=traced_mesh_source +) + +pts = np.asarray(relocated_mesh_grid.array, dtype=np.float64) +print("relocated mesh grid:", pts.shape, pts.dtype) +np.save(args.out, pts) +print( + "areas_factor:", + mesh.areas_factor, + "max_neighbors:", + mesh.max_neighbors, + "max_cavity_triangles:", + mesh.max_cavity_triangles, +) +print("wrote", args.out)