From 509f729e5db002b5fdf7d1bb05070e6722c555cb Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Wed, 11 Feb 2026 15:45:08 -0500 Subject: [PATCH 01/24] Re add changes --- benchmark/Comrade/CondaPkg.toml | 2 + benchmark/Comrade/Project.toml | 57 ++++++++++++++ benchmark/Comrade/Sampling/main.py | 87 +++++++++++++++++++++ benchmark/Comrade/Sampling/pyproject.toml | 11 +++ benchmark/Comrade/comimager.jl | 93 +++++++++++++++++++++++ benchmark/Comrade/common.jl | 88 +++++++++++++++++++++ benchmark/Comrade/export_to_jax.jl | 18 +++++ benchmark/Comrade/runbenchmarks.jl | 67 ++++++++++++++++ 8 files changed, 423 insertions(+) create mode 100644 benchmark/Comrade/CondaPkg.toml create mode 100644 benchmark/Comrade/Project.toml create mode 100644 benchmark/Comrade/Sampling/main.py create mode 100644 benchmark/Comrade/Sampling/pyproject.toml create mode 100644 benchmark/Comrade/comimager.jl create mode 100644 benchmark/Comrade/common.jl create mode 100644 benchmark/Comrade/export_to_jax.jl create mode 100644 benchmark/Comrade/runbenchmarks.jl diff --git a/benchmark/Comrade/CondaPkg.toml b/benchmark/Comrade/CondaPkg.toml new file mode 100644 index 0000000000..2f0958226c --- /dev/null +++ b/benchmark/Comrade/CondaPkg.toml @@ -0,0 +1,2 @@ +[deps] +python = "<3.11,>=3.9,<4" diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml new file mode 100644 index 0000000000..275eb0a47f --- /dev/null +++ b/benchmark/Comrade/Project.toml @@ -0,0 +1,57 @@ +[deps] +AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" +Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" +Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" +Comrade = "99d987ce-9a1e-4df8-bc0b-1ea019aa547b" +ComradeBase = "6d8c423b-a35f-4ef1-850c-862fe21f82c4" +CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" +NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Pyehtim = "3d61700d-6e5b-419a-8e22-9c066cf00468" +PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff" +VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029" +VLBILikelihoods = "90db92cd-0007-4c0a-8e51-dbf0782ce592" +VLBISkyModels = "d6343c73-7174-4e0f-bb64-562643efbeca" + +[sources.Comrade] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/Comrade.jl" + +[sources.ComradeBase] +rev = "main" +url = "https://github.com/ptiede/ComradeBase.jl" + +[sources.NFFT] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/NFFT.jl" + +[sources.Reactant] +path = "../../" + +[sources.TransformVariables] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/TransformVariables.jl" + +[sources.VLBIImagePriors] +rev = "ptiede-reactantperf" +url = "https://github.com/ptiede/VLBIImagePriors.jl" + +[sources.VLBILikelihoods] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/VLBILikelihoods.jl" + +[sources.VLBISkyModels] +rev = "ptiede-reactnfft" +url = "https://github.com/EHTJulia/VLBISkyModels.jl" diff --git a/benchmark/Comrade/Sampling/main.py b/benchmark/Comrade/Sampling/main.py new file mode 100644 index 0000000000..f380948e01 --- /dev/null +++ b/benchmark/Comrade/Sampling/main.py @@ -0,0 +1,87 @@ +import sys + +sys.path.insert(0, "Serialized/Fwd") +sys.path.insert(0, "Serialized/Bwd") + + +import jax +jax.config.update('jax_enable_x64', True) +import jax.numpy as jnp +import blackjax +import enzyme_ad +from functools import partial + +import logdensityof as lg +import gl as gl + +from logdensityof import run_logdensityof +from gl import run_gl + +lg_inputs = lg.load_inputs() +gl_inputs = gl.load_inputs() + +tpost = lg_inputs[:-1] +xr = lg_inputs[-1] + + + + +jlr = jax.jit(run_logdensityof) + +jtpost0 = jnp.array(tpost[0]) +jtpost1 = jnp.array(tpost[1]) +jtpost2 = jnp.array(tpost[2]) +jtpost3 = jnp.array(tpost[3]) +jtpost4 = jnp.array(tpost[4]) +jxr = jnp.array(xr) + +out = jlr(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, xr) + +run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) +run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) + + +@jax.custom_vjp +def f(x): + out = run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x) + return out[0] + +def f_fwd(x): + j = run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x)[0] + return f(x), (j,) + +def f_bwd(res, g): + j = res[0] + return (g * j,) + +f.defvjp(f_fwd, f_bwd) + +logdensity = lambda x: f(**x) + +inv_mass_matrix = jnp.ones(len(jxr)) +initial_position = {"x": jxr} + +rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) + +# adaptation +warmup = blackjax.window_adaptation(blackjax.nuts, logdensity, progress_bar=True) +rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3) +(state, parameters), _ = warmup.run(warmup_key, initial_position, num_steps=1000) + + +def inference_loop(rng_key, kernel, init, nsamples): + @jax.jit + def step(state, rng_key): + state, _ = kernel(rng_key, state) + return state, state + + keys = jax.random.split(rng_key, nsamples) + _, states = jax.lax.scan(step, init, keys) + return states + +# inference loop +rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) +kernel = blackjax.nuts(logdensity, **parameters).step +states = inference_loop(sample_key, kernel, state, nsamples=1000) + + diff --git a/benchmark/Comrade/Sampling/pyproject.toml b/benchmark/Comrade/Sampling/pyproject.toml new file mode 100644 index 0000000000..387857e6ff --- /dev/null +++ b/benchmark/Comrade/Sampling/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "sampling" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +requires-python = ">=3.10, <3.12" +dependencies = [ + "blackjax>=1.3", + "ipython>=8.38.0", + "jax[cuda12]==0.5.0", +] diff --git a/benchmark/Comrade/comimager.jl b/benchmark/Comrade/comimager.jl new file mode 100644 index 0000000000..a4536740f1 --- /dev/null +++ b/benchmark/Comrade/comimager.jl @@ -0,0 +1,93 @@ + +# ENV["JULIA_DEBUG"] = "Reactant" + + +const dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" +const dataf = Base.download(dataurl) + + +# TODO Make Distributions package that is compatible with Reactant +Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) +function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) + λ = rate(d) + z = log(λ) - λ * x + @trace if x < 0 + out = oftype(z, -Inf) + else + out = z + end + return z +end + +# THis is very much needed or else `@compile` hangs +function Distributions.logpdf(d::Distributions.DiagNormal, x::Reactant.AnyTracedRVector) + l = VLBILikelihoods._unnormed_logpdf_μΣ(d.μ, d.Σ.diag, x) + n = VLBILikelihoods._gaussnorm(d.μ, d.Σ.diag) + return l + n +end + +function sky(θ, metadata) + (; z, ρs, σ) = θ + (; srf, grid, mimg) = metadata + x = genfield(StationaryRandomField(MarkovPS(ρs), srf), z) + x .*= σ + mx = maximum(x) + bmimg = baseimage(mimg) + rast = @. exp(x - mx) * bmimg + rast ./= sum(rast) + return ContinuousImage(rast, grid, DeltaPulse{3}()) +end + + +function build_post(fov, npix, dataf) + + # @testset "Comrade Integration Imaging" begin + obs = ehtim.obsdata.load_uvfits(dataf) + obsavg = obs.add_fractional_noise(0.02) + dvis = extract_table(obsavg, Visibilities()) + + npix = npix + fovx = fov + fovy = fov + + # Now let's form our cache's. First, we have our usual image cache which is needed to numerically + # compute the visibilities. + grd = imagepixels(fovx, fovy, npix, npix) + pl = StationaryRandomFieldPlan(grd) + mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(25.0))), grd) + skymeta = (; srf=pl, grid=grd, mimg=mimg) + + ρs = ntuple(Returns(Uniform(0.01, max(size(grd)...))), 3) + zprior = std_dist(pl) + prior = (z=zprior, ρs=ρs, σ=Exponential(1.0)) + + skymr = SkyModel(sky, prior, grd; metadata=skymeta, algorithm=VLBISkyModels.ReactantAlg()) # Need to do this so that we allocate proper Reactant arrays for internal stuff + + g(x) = exp(complex(x.lg, x.gp)) + G = SingleStokesGain(g) + + intpr = ( + lg=ArrayPrior( + IIDSitePrior(ScanSeg(), Normal(0.0, 0.2)); + LM=IIDSitePrior(ScanSeg(), Normal(0.0, 1.0)), + ), + gp=ArrayPrior( + IIDSitePrior(ScanSeg(), DiagonalVonMises(0.0, inv(π^2))); + refant=SEFDReference(0.0), + phase=true, + ), + ) + intmodel = InstrumentModel(G, intpr) + + postr = VLBIPosterior(skymr, intmodel, dvis) + tpostr = asflat(postr) + + return tpostr +end + +if abspath(PROGRAM_FILE) == @__FILE__ + tpostr = build_post(μas2rad(200.0), 64, dataf) + results = Dict() + backend = get_backend() + run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr, "forward", "test") +end \ No newline at end of file diff --git a/benchmark/Comrade/common.jl b/benchmark/Comrade/common.jl new file mode 100644 index 0000000000..d86622a60c --- /dev/null +++ b/benchmark/Comrade/common.jl @@ -0,0 +1,88 @@ +using Printf: @sprintf +using Reactant: Reactant, @compile +using Enzyme: Enzyme, Const +using Random: Random + +include("../utils.jl") + +logdensityofref(tpostr, xr) = logdensityof(tpostr[], xr) +function gradref(tpostr, xr) + derivs, val = (Enzyme.gradient(ReverseWithPrimal, logdensityofref, Ref(tpostr), xr)) + return last(derivs), val +end + +function run_comrade_benchmark!( + results::Dict, + benchmark_name::String, + backend::String, + tpostr, +) + # TODO which Enzyme passes do I want to enable and disable + run_benchmark!(results, benchmark_name, backend, tpostr, "forward", "Default") + run_benchmark!(results, benchmark_name, backend, tpostr, "backward", "Default") + return nothing +end + +function run_benchmark!( + results::Dict, + benchmark_name::String, + backend::String, + tpost, + fwd_or_bwd::String, + tag::String +) + + if !haskey(results, "TFLOP/s") + results["TFLOP/s"] = Dict{String,Float64}() + end + if !haskey(results, "Runtime (s)") + results["Runtime (s)"] = Dict{String,Float64}() + end + + prim_or_rev = fwd_or_bwd == "forward" ? "primal" : "reverse" + full_benchmark_name = string(benchmark_name, "/", prim_or_rev, "/", backend, "/", tag) + @assert !haskey(results["Runtime (s)"], full_benchmark_name) "Benchmark already \ + exists: \ + $(full_benchmark_name)" + @assert !haskey(results["TFLOP/s"], full_benchmark_name) "Benchmark already exists: \ + $(full_benchmark_name)" + + rng = Random.default_rng() # don't use any other rng + Random.seed!(rng, 0) + x = Reactant.to_rarray(randn(rng, dimension(tpost))) + + fn = if fwd_or_bwd == "forward" + logdensityof + elseif fwd_or_bwd == "backward" + gradref + else + error("Unknown fwd_or_bwd: $(fwd_or_bwd)") + end + + prof_result = Reactant.Profiler.profile_with_xprof( + fn,tpost, x; nrepeat=10, warmup=3 + ) + + + results["Runtime (s)"][full_benchmark_name] = + prof_result.profiling_result.runtime_ns / 1e9 + results["TFLOP/s"][full_benchmark_name] = + if prof_result.profiling_result.flops_data === nothing + -1 + else + prof_result.profiling_result.flops_data.RawFlopsRate / 1e12 + end + + GC.gc(true) + + print_stmt = @sprintf( + "%100s : %.5gs %.5g TFLOP/s", + full_benchmark_name, + results["Runtime (s)"][full_benchmark_name], + results["TFLOP/s"][full_benchmark_name] + ) + @info print_stmt + return nothing +end + + diff --git a/benchmark/Comrade/export_to_jax.jl b/benchmark/Comrade/export_to_jax.jl new file mode 100644 index 0000000000..1316226bf6 --- /dev/null +++ b/benchmark/Comrade/export_to_jax.jl @@ -0,0 +1,18 @@ +using Pkg; +Pkg.activate(@__DIR__); +ENV["XLA_REACTANT_GPU_PREALLOCATE"] = "false" +using Reactant, NPZ, PythonCall + +include("comimager.jl") + +logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( + logdensityof, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Fwd") +) + +println("Exported to: $logdensity_python_file_path") + +grad_logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( + gl, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Bwd") +) + +println("Exported to: $grad_logdensity_python_file_path") diff --git a/benchmark/Comrade/runbenchmarks.jl b/benchmark/Comrade/runbenchmarks.jl new file mode 100644 index 0000000000..d31e71e939 --- /dev/null +++ b/benchmark/Comrade/runbenchmarks.jl @@ -0,0 +1,67 @@ +# Comrade Benchmarks Runner +# This script runs all Comrade benchmarks and stores results to a JSON file + +using Pkg + +@static if VERSION ≥ v"1.10-" && VERSION < v"1.11" + Pkg.add([ + PackageSpec(; name="ComradeBase", rev="main"), + PackageSpec(; name="Comrade", rev="ptiede-reactant"), + PackageSpec(; name="VLBISkyModels", rev="ptiede-reactnfft"), + PackageSpec(; name="VLBILikelihoods", rev="ptiede-reactant"), + PackageSpec(; name="VLBIImagePriors", rev="ptiede-reactantperf"), + PackageSpec(; + url="https://github.com/ptiede/TransformVariables.jl", rev="ptiede-reactant" + ), + PackageSpec(; + url="https://github.com/ptiede/NFFT.jl", rev="ptiede-reactant" + ) + ]) + Pkg.develop(; path=joinpath(@__DIR__, "../../")) +end + + +# Load dependencies +using Reactant +using Chairmarks: @b +using Random: Random +using Printf: @sprintf + + + +using Reactant +using LinearAlgebra +using AbstractFFTs + +using VLBISkyModels +using VLBILikelihoods +using Comrade +using Distributions +using VLBIImagePriors +using LogExpFunctions +import TransformVariables as TV + +using Downloads +using Distributions +using Enzyme + +using Pyehtim +using Test + +include("common.jl") + +@info sprint(io -> versioninfo(io; verbose=true)) + +include("comimager.jl") + +function run_all_benchmarks(backend::String) + results = Dict{String,Dict{String,Float64}}() + + tpostr = build_post(μas2rad(200.0), 64, dataf) + run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) + + return results +end + + +res = run_all_benchmarks(get_backend()) \ No newline at end of file From 71d8f99a6515313d36e79e066d047094fbe45657 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 13:34:42 -0500 Subject: [PATCH 02/24] Update benchmarking --- benchmark/Comrade/Project.toml | 1 + benchmark/Comrade/comimager.jl | 6 ------ benchmark/Comrade/runbenchmarks.jl | 10 ++++++++++ 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml index 275eb0a47f..060ee7704b 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/Comrade/Project.toml @@ -12,6 +12,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" diff --git a/benchmark/Comrade/comimager.jl b/benchmark/Comrade/comimager.jl index a4536740f1..7a245b1822 100644 --- a/benchmark/Comrade/comimager.jl +++ b/benchmark/Comrade/comimager.jl @@ -1,10 +1,4 @@ -# ENV["JULIA_DEBUG"] = "Reactant" - - -const dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" -const dataf = Base.download(dataurl) - # TODO Make Distributions package that is compatible with Reactant Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) diff --git a/benchmark/Comrade/runbenchmarks.jl b/benchmark/Comrade/runbenchmarks.jl index d31e71e939..ef597cd536 100644 --- a/benchmark/Comrade/runbenchmarks.jl +++ b/benchmark/Comrade/runbenchmarks.jl @@ -57,9 +57,19 @@ include("comimager.jl") function run_all_benchmarks(backend::String) results = Dict{String,Dict{String,Float64}}() + dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" + dataf = Base.download(dataurl) + tpostr = build_post(μas2rad(200.0), 64, dataf) run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) + + tpostr = build_post(μas2rad(200.0), 128, dataf) + run_comrade_benchmark!(results, "Comrade EHT Imaging 128 x 128", backend, tpostr) + + tpostr = build_post(μas2rad(200.0), 256, dataf) + run_comrade_benchmark!(results, "Comrade EHT Imaging 256 x 256", backend, tpostr) + return results end From 07f247d4ddb22068425c55efd861d040f471d0e3 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 13:44:30 -0500 Subject: [PATCH 03/24] Removed sampling from benchmarks --- benchmark/Comrade/Sampling/main.py | 87 ----------------------- benchmark/Comrade/Sampling/pyproject.toml | 11 --- benchmark/Comrade/export_to_jax.jl | 18 ----- 3 files changed, 116 deletions(-) delete mode 100644 benchmark/Comrade/Sampling/main.py delete mode 100644 benchmark/Comrade/Sampling/pyproject.toml delete mode 100644 benchmark/Comrade/export_to_jax.jl diff --git a/benchmark/Comrade/Sampling/main.py b/benchmark/Comrade/Sampling/main.py deleted file mode 100644 index f380948e01..0000000000 --- a/benchmark/Comrade/Sampling/main.py +++ /dev/null @@ -1,87 +0,0 @@ -import sys - -sys.path.insert(0, "Serialized/Fwd") -sys.path.insert(0, "Serialized/Bwd") - - -import jax -jax.config.update('jax_enable_x64', True) -import jax.numpy as jnp -import blackjax -import enzyme_ad -from functools import partial - -import logdensityof as lg -import gl as gl - -from logdensityof import run_logdensityof -from gl import run_gl - -lg_inputs = lg.load_inputs() -gl_inputs = gl.load_inputs() - -tpost = lg_inputs[:-1] -xr = lg_inputs[-1] - - - - -jlr = jax.jit(run_logdensityof) - -jtpost0 = jnp.array(tpost[0]) -jtpost1 = jnp.array(tpost[1]) -jtpost2 = jnp.array(tpost[2]) -jtpost3 = jnp.array(tpost[3]) -jtpost4 = jnp.array(tpost[4]) -jxr = jnp.array(xr) - -out = jlr(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, xr) - -run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) -run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) - - -@jax.custom_vjp -def f(x): - out = run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x) - return out[0] - -def f_fwd(x): - j = run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x)[0] - return f(x), (j,) - -def f_bwd(res, g): - j = res[0] - return (g * j,) - -f.defvjp(f_fwd, f_bwd) - -logdensity = lambda x: f(**x) - -inv_mass_matrix = jnp.ones(len(jxr)) -initial_position = {"x": jxr} - -rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) - -# adaptation -warmup = blackjax.window_adaptation(blackjax.nuts, logdensity, progress_bar=True) -rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3) -(state, parameters), _ = warmup.run(warmup_key, initial_position, num_steps=1000) - - -def inference_loop(rng_key, kernel, init, nsamples): - @jax.jit - def step(state, rng_key): - state, _ = kernel(rng_key, state) - return state, state - - keys = jax.random.split(rng_key, nsamples) - _, states = jax.lax.scan(step, init, keys) - return states - -# inference loop -rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) -kernel = blackjax.nuts(logdensity, **parameters).step -states = inference_loop(sample_key, kernel, state, nsamples=1000) - - diff --git a/benchmark/Comrade/Sampling/pyproject.toml b/benchmark/Comrade/Sampling/pyproject.toml deleted file mode 100644 index 387857e6ff..0000000000 --- a/benchmark/Comrade/Sampling/pyproject.toml +++ /dev/null @@ -1,11 +0,0 @@ -[project] -name = "sampling" -version = "0.1.0" -description = "Add your description here" -readme = "README.md" -requires-python = ">=3.10, <3.12" -dependencies = [ - "blackjax>=1.3", - "ipython>=8.38.0", - "jax[cuda12]==0.5.0", -] diff --git a/benchmark/Comrade/export_to_jax.jl b/benchmark/Comrade/export_to_jax.jl deleted file mode 100644 index 1316226bf6..0000000000 --- a/benchmark/Comrade/export_to_jax.jl +++ /dev/null @@ -1,18 +0,0 @@ -using Pkg; -Pkg.activate(@__DIR__); -ENV["XLA_REACTANT_GPU_PREALLOCATE"] = "false" -using Reactant, NPZ, PythonCall - -include("comimager.jl") - -logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( - logdensityof, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Fwd") -) - -println("Exported to: $logdensity_python_file_path") - -grad_logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( - gl, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Bwd") -) - -println("Exported to: $grad_logdensity_python_file_path") From 173877a65543a312d56fd9838a8487ceba7ac1c8 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 14:40:21 -0500 Subject: [PATCH 04/24] Add JSON --- benchmark/Comrade/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml index 060ee7704b..6685c8302d 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/Comrade/Project.toml @@ -12,6 +12,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" From ad940ee31947b646466ec59c5838ef5e34eb70b1 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 14:42:12 -0500 Subject: [PATCH 05/24] Remove JSON3 --- benchmark/Comrade/Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml index 6685c8302d..b0aa841164 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/Comrade/Project.toml @@ -13,7 +13,6 @@ Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" -JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" From 7247a09fa28900791bf86a8bc7866ae1a644a9f4 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Wed, 11 Feb 2026 15:45:08 -0500 Subject: [PATCH 06/24] Re add changes --- benchmark/Comrade/CondaPkg.toml | 2 + benchmark/Comrade/Project.toml | 57 ++++++++++++++ benchmark/Comrade/Sampling/main.py | 87 +++++++++++++++++++++ benchmark/Comrade/Sampling/pyproject.toml | 11 +++ benchmark/Comrade/comimager.jl | 93 +++++++++++++++++++++++ benchmark/Comrade/common.jl | 88 +++++++++++++++++++++ benchmark/Comrade/export_to_jax.jl | 18 +++++ benchmark/Comrade/runbenchmarks.jl | 67 ++++++++++++++++ 8 files changed, 423 insertions(+) create mode 100644 benchmark/Comrade/CondaPkg.toml create mode 100644 benchmark/Comrade/Project.toml create mode 100644 benchmark/Comrade/Sampling/main.py create mode 100644 benchmark/Comrade/Sampling/pyproject.toml create mode 100644 benchmark/Comrade/comimager.jl create mode 100644 benchmark/Comrade/common.jl create mode 100644 benchmark/Comrade/export_to_jax.jl create mode 100644 benchmark/Comrade/runbenchmarks.jl diff --git a/benchmark/Comrade/CondaPkg.toml b/benchmark/Comrade/CondaPkg.toml new file mode 100644 index 0000000000..2f0958226c --- /dev/null +++ b/benchmark/Comrade/CondaPkg.toml @@ -0,0 +1,2 @@ +[deps] +python = "<3.11,>=3.9,<4" diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml new file mode 100644 index 0000000000..275eb0a47f --- /dev/null +++ b/benchmark/Comrade/Project.toml @@ -0,0 +1,57 @@ +[deps] +AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" +Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" +Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" +Comrade = "99d987ce-9a1e-4df8-bc0b-1ea019aa547b" +ComradeBase = "6d8c423b-a35f-4ef1-850c-862fe21f82c4" +CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" +NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Pyehtim = "3d61700d-6e5b-419a-8e22-9c066cf00468" +PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff" +VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029" +VLBILikelihoods = "90db92cd-0007-4c0a-8e51-dbf0782ce592" +VLBISkyModels = "d6343c73-7174-4e0f-bb64-562643efbeca" + +[sources.Comrade] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/Comrade.jl" + +[sources.ComradeBase] +rev = "main" +url = "https://github.com/ptiede/ComradeBase.jl" + +[sources.NFFT] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/NFFT.jl" + +[sources.Reactant] +path = "../../" + +[sources.TransformVariables] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/TransformVariables.jl" + +[sources.VLBIImagePriors] +rev = "ptiede-reactantperf" +url = "https://github.com/ptiede/VLBIImagePriors.jl" + +[sources.VLBILikelihoods] +rev = "ptiede-reactant" +url = "https://github.com/ptiede/VLBILikelihoods.jl" + +[sources.VLBISkyModels] +rev = "ptiede-reactnfft" +url = "https://github.com/EHTJulia/VLBISkyModels.jl" diff --git a/benchmark/Comrade/Sampling/main.py b/benchmark/Comrade/Sampling/main.py new file mode 100644 index 0000000000..f380948e01 --- /dev/null +++ b/benchmark/Comrade/Sampling/main.py @@ -0,0 +1,87 @@ +import sys + +sys.path.insert(0, "Serialized/Fwd") +sys.path.insert(0, "Serialized/Bwd") + + +import jax +jax.config.update('jax_enable_x64', True) +import jax.numpy as jnp +import blackjax +import enzyme_ad +from functools import partial + +import logdensityof as lg +import gl as gl + +from logdensityof import run_logdensityof +from gl import run_gl + +lg_inputs = lg.load_inputs() +gl_inputs = gl.load_inputs() + +tpost = lg_inputs[:-1] +xr = lg_inputs[-1] + + + + +jlr = jax.jit(run_logdensityof) + +jtpost0 = jnp.array(tpost[0]) +jtpost1 = jnp.array(tpost[1]) +jtpost2 = jnp.array(tpost[2]) +jtpost3 = jnp.array(tpost[3]) +jtpost4 = jnp.array(tpost[4]) +jxr = jnp.array(xr) + +out = jlr(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, xr) + +run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) +run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) + + +@jax.custom_vjp +def f(x): + out = run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x) + return out[0] + +def f_fwd(x): + j = run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x)[0] + return f(x), (j,) + +def f_bwd(res, g): + j = res[0] + return (g * j,) + +f.defvjp(f_fwd, f_bwd) + +logdensity = lambda x: f(**x) + +inv_mass_matrix = jnp.ones(len(jxr)) +initial_position = {"x": jxr} + +rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) + +# adaptation +warmup = blackjax.window_adaptation(blackjax.nuts, logdensity, progress_bar=True) +rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3) +(state, parameters), _ = warmup.run(warmup_key, initial_position, num_steps=1000) + + +def inference_loop(rng_key, kernel, init, nsamples): + @jax.jit + def step(state, rng_key): + state, _ = kernel(rng_key, state) + return state, state + + keys = jax.random.split(rng_key, nsamples) + _, states = jax.lax.scan(step, init, keys) + return states + +# inference loop +rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) +kernel = blackjax.nuts(logdensity, **parameters).step +states = inference_loop(sample_key, kernel, state, nsamples=1000) + + diff --git a/benchmark/Comrade/Sampling/pyproject.toml b/benchmark/Comrade/Sampling/pyproject.toml new file mode 100644 index 0000000000..387857e6ff --- /dev/null +++ b/benchmark/Comrade/Sampling/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "sampling" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +requires-python = ">=3.10, <3.12" +dependencies = [ + "blackjax>=1.3", + "ipython>=8.38.0", + "jax[cuda12]==0.5.0", +] diff --git a/benchmark/Comrade/comimager.jl b/benchmark/Comrade/comimager.jl new file mode 100644 index 0000000000..a4536740f1 --- /dev/null +++ b/benchmark/Comrade/comimager.jl @@ -0,0 +1,93 @@ + +# ENV["JULIA_DEBUG"] = "Reactant" + + +const dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" +const dataf = Base.download(dataurl) + + +# TODO Make Distributions package that is compatible with Reactant +Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) +function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) + λ = rate(d) + z = log(λ) - λ * x + @trace if x < 0 + out = oftype(z, -Inf) + else + out = z + end + return z +end + +# THis is very much needed or else `@compile` hangs +function Distributions.logpdf(d::Distributions.DiagNormal, x::Reactant.AnyTracedRVector) + l = VLBILikelihoods._unnormed_logpdf_μΣ(d.μ, d.Σ.diag, x) + n = VLBILikelihoods._gaussnorm(d.μ, d.Σ.diag) + return l + n +end + +function sky(θ, metadata) + (; z, ρs, σ) = θ + (; srf, grid, mimg) = metadata + x = genfield(StationaryRandomField(MarkovPS(ρs), srf), z) + x .*= σ + mx = maximum(x) + bmimg = baseimage(mimg) + rast = @. exp(x - mx) * bmimg + rast ./= sum(rast) + return ContinuousImage(rast, grid, DeltaPulse{3}()) +end + + +function build_post(fov, npix, dataf) + + # @testset "Comrade Integration Imaging" begin + obs = ehtim.obsdata.load_uvfits(dataf) + obsavg = obs.add_fractional_noise(0.02) + dvis = extract_table(obsavg, Visibilities()) + + npix = npix + fovx = fov + fovy = fov + + # Now let's form our cache's. First, we have our usual image cache which is needed to numerically + # compute the visibilities. + grd = imagepixels(fovx, fovy, npix, npix) + pl = StationaryRandomFieldPlan(grd) + mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(25.0))), grd) + skymeta = (; srf=pl, grid=grd, mimg=mimg) + + ρs = ntuple(Returns(Uniform(0.01, max(size(grd)...))), 3) + zprior = std_dist(pl) + prior = (z=zprior, ρs=ρs, σ=Exponential(1.0)) + + skymr = SkyModel(sky, prior, grd; metadata=skymeta, algorithm=VLBISkyModels.ReactantAlg()) # Need to do this so that we allocate proper Reactant arrays for internal stuff + + g(x) = exp(complex(x.lg, x.gp)) + G = SingleStokesGain(g) + + intpr = ( + lg=ArrayPrior( + IIDSitePrior(ScanSeg(), Normal(0.0, 0.2)); + LM=IIDSitePrior(ScanSeg(), Normal(0.0, 1.0)), + ), + gp=ArrayPrior( + IIDSitePrior(ScanSeg(), DiagonalVonMises(0.0, inv(π^2))); + refant=SEFDReference(0.0), + phase=true, + ), + ) + intmodel = InstrumentModel(G, intpr) + + postr = VLBIPosterior(skymr, intmodel, dvis) + tpostr = asflat(postr) + + return tpostr +end + +if abspath(PROGRAM_FILE) == @__FILE__ + tpostr = build_post(μas2rad(200.0), 64, dataf) + results = Dict() + backend = get_backend() + run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr, "forward", "test") +end \ No newline at end of file diff --git a/benchmark/Comrade/common.jl b/benchmark/Comrade/common.jl new file mode 100644 index 0000000000..d86622a60c --- /dev/null +++ b/benchmark/Comrade/common.jl @@ -0,0 +1,88 @@ +using Printf: @sprintf +using Reactant: Reactant, @compile +using Enzyme: Enzyme, Const +using Random: Random + +include("../utils.jl") + +logdensityofref(tpostr, xr) = logdensityof(tpostr[], xr) +function gradref(tpostr, xr) + derivs, val = (Enzyme.gradient(ReverseWithPrimal, logdensityofref, Ref(tpostr), xr)) + return last(derivs), val +end + +function run_comrade_benchmark!( + results::Dict, + benchmark_name::String, + backend::String, + tpostr, +) + # TODO which Enzyme passes do I want to enable and disable + run_benchmark!(results, benchmark_name, backend, tpostr, "forward", "Default") + run_benchmark!(results, benchmark_name, backend, tpostr, "backward", "Default") + return nothing +end + +function run_benchmark!( + results::Dict, + benchmark_name::String, + backend::String, + tpost, + fwd_or_bwd::String, + tag::String +) + + if !haskey(results, "TFLOP/s") + results["TFLOP/s"] = Dict{String,Float64}() + end + if !haskey(results, "Runtime (s)") + results["Runtime (s)"] = Dict{String,Float64}() + end + + prim_or_rev = fwd_or_bwd == "forward" ? "primal" : "reverse" + full_benchmark_name = string(benchmark_name, "/", prim_or_rev, "/", backend, "/", tag) + @assert !haskey(results["Runtime (s)"], full_benchmark_name) "Benchmark already \ + exists: \ + $(full_benchmark_name)" + @assert !haskey(results["TFLOP/s"], full_benchmark_name) "Benchmark already exists: \ + $(full_benchmark_name)" + + rng = Random.default_rng() # don't use any other rng + Random.seed!(rng, 0) + x = Reactant.to_rarray(randn(rng, dimension(tpost))) + + fn = if fwd_or_bwd == "forward" + logdensityof + elseif fwd_or_bwd == "backward" + gradref + else + error("Unknown fwd_or_bwd: $(fwd_or_bwd)") + end + + prof_result = Reactant.Profiler.profile_with_xprof( + fn,tpost, x; nrepeat=10, warmup=3 + ) + + + results["Runtime (s)"][full_benchmark_name] = + prof_result.profiling_result.runtime_ns / 1e9 + results["TFLOP/s"][full_benchmark_name] = + if prof_result.profiling_result.flops_data === nothing + -1 + else + prof_result.profiling_result.flops_data.RawFlopsRate / 1e12 + end + + GC.gc(true) + + print_stmt = @sprintf( + "%100s : %.5gs %.5g TFLOP/s", + full_benchmark_name, + results["Runtime (s)"][full_benchmark_name], + results["TFLOP/s"][full_benchmark_name] + ) + @info print_stmt + return nothing +end + + diff --git a/benchmark/Comrade/export_to_jax.jl b/benchmark/Comrade/export_to_jax.jl new file mode 100644 index 0000000000..1316226bf6 --- /dev/null +++ b/benchmark/Comrade/export_to_jax.jl @@ -0,0 +1,18 @@ +using Pkg; +Pkg.activate(@__DIR__); +ENV["XLA_REACTANT_GPU_PREALLOCATE"] = "false" +using Reactant, NPZ, PythonCall + +include("comimager.jl") + +logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( + logdensityof, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Fwd") +) + +println("Exported to: $logdensity_python_file_path") + +grad_logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( + gl, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Bwd") +) + +println("Exported to: $grad_logdensity_python_file_path") diff --git a/benchmark/Comrade/runbenchmarks.jl b/benchmark/Comrade/runbenchmarks.jl new file mode 100644 index 0000000000..d31e71e939 --- /dev/null +++ b/benchmark/Comrade/runbenchmarks.jl @@ -0,0 +1,67 @@ +# Comrade Benchmarks Runner +# This script runs all Comrade benchmarks and stores results to a JSON file + +using Pkg + +@static if VERSION ≥ v"1.10-" && VERSION < v"1.11" + Pkg.add([ + PackageSpec(; name="ComradeBase", rev="main"), + PackageSpec(; name="Comrade", rev="ptiede-reactant"), + PackageSpec(; name="VLBISkyModels", rev="ptiede-reactnfft"), + PackageSpec(; name="VLBILikelihoods", rev="ptiede-reactant"), + PackageSpec(; name="VLBIImagePriors", rev="ptiede-reactantperf"), + PackageSpec(; + url="https://github.com/ptiede/TransformVariables.jl", rev="ptiede-reactant" + ), + PackageSpec(; + url="https://github.com/ptiede/NFFT.jl", rev="ptiede-reactant" + ) + ]) + Pkg.develop(; path=joinpath(@__DIR__, "../../")) +end + + +# Load dependencies +using Reactant +using Chairmarks: @b +using Random: Random +using Printf: @sprintf + + + +using Reactant +using LinearAlgebra +using AbstractFFTs + +using VLBISkyModels +using VLBILikelihoods +using Comrade +using Distributions +using VLBIImagePriors +using LogExpFunctions +import TransformVariables as TV + +using Downloads +using Distributions +using Enzyme + +using Pyehtim +using Test + +include("common.jl") + +@info sprint(io -> versioninfo(io; verbose=true)) + +include("comimager.jl") + +function run_all_benchmarks(backend::String) + results = Dict{String,Dict{String,Float64}}() + + tpostr = build_post(μas2rad(200.0), 64, dataf) + run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) + + return results +end + + +res = run_all_benchmarks(get_backend()) \ No newline at end of file From d5ec57e5d43bbe012082beacc873f1480ec9bb02 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 13:34:42 -0500 Subject: [PATCH 07/24] Update benchmarking --- benchmark/Comrade/Project.toml | 1 + benchmark/Comrade/comimager.jl | 6 ------ benchmark/Comrade/runbenchmarks.jl | 10 ++++++++++ 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml index 275eb0a47f..060ee7704b 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/Comrade/Project.toml @@ -12,6 +12,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" diff --git a/benchmark/Comrade/comimager.jl b/benchmark/Comrade/comimager.jl index a4536740f1..7a245b1822 100644 --- a/benchmark/Comrade/comimager.jl +++ b/benchmark/Comrade/comimager.jl @@ -1,10 +1,4 @@ -# ENV["JULIA_DEBUG"] = "Reactant" - - -const dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" -const dataf = Base.download(dataurl) - # TODO Make Distributions package that is compatible with Reactant Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) diff --git a/benchmark/Comrade/runbenchmarks.jl b/benchmark/Comrade/runbenchmarks.jl index d31e71e939..ef597cd536 100644 --- a/benchmark/Comrade/runbenchmarks.jl +++ b/benchmark/Comrade/runbenchmarks.jl @@ -57,9 +57,19 @@ include("comimager.jl") function run_all_benchmarks(backend::String) results = Dict{String,Dict{String,Float64}}() + dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" + dataf = Base.download(dataurl) + tpostr = build_post(μas2rad(200.0), 64, dataf) run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) + + tpostr = build_post(μas2rad(200.0), 128, dataf) + run_comrade_benchmark!(results, "Comrade EHT Imaging 128 x 128", backend, tpostr) + + tpostr = build_post(μas2rad(200.0), 256, dataf) + run_comrade_benchmark!(results, "Comrade EHT Imaging 256 x 256", backend, tpostr) + return results end From 8c48bcd4aa5b733e1d1ed451537667e6e0c6ebe7 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 13:44:30 -0500 Subject: [PATCH 08/24] Removed sampling from benchmarks --- benchmark/Comrade/Sampling/main.py | 87 ----------------------- benchmark/Comrade/Sampling/pyproject.toml | 11 --- benchmark/Comrade/export_to_jax.jl | 18 ----- 3 files changed, 116 deletions(-) delete mode 100644 benchmark/Comrade/Sampling/main.py delete mode 100644 benchmark/Comrade/Sampling/pyproject.toml delete mode 100644 benchmark/Comrade/export_to_jax.jl diff --git a/benchmark/Comrade/Sampling/main.py b/benchmark/Comrade/Sampling/main.py deleted file mode 100644 index f380948e01..0000000000 --- a/benchmark/Comrade/Sampling/main.py +++ /dev/null @@ -1,87 +0,0 @@ -import sys - -sys.path.insert(0, "Serialized/Fwd") -sys.path.insert(0, "Serialized/Bwd") - - -import jax -jax.config.update('jax_enable_x64', True) -import jax.numpy as jnp -import blackjax -import enzyme_ad -from functools import partial - -import logdensityof as lg -import gl as gl - -from logdensityof import run_logdensityof -from gl import run_gl - -lg_inputs = lg.load_inputs() -gl_inputs = gl.load_inputs() - -tpost = lg_inputs[:-1] -xr = lg_inputs[-1] - - - - -jlr = jax.jit(run_logdensityof) - -jtpost0 = jnp.array(tpost[0]) -jtpost1 = jnp.array(tpost[1]) -jtpost2 = jnp.array(tpost[2]) -jtpost3 = jnp.array(tpost[3]) -jtpost4 = jnp.array(tpost[4]) -jxr = jnp.array(xr) - -out = jlr(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, xr) - -run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) -run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, jxr) - - -@jax.custom_vjp -def f(x): - out = run_logdensityof(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x) - return out[0] - -def f_fwd(x): - j = run_gl(jtpost0, jtpost1, jtpost2, jtpost3, jtpost4, x)[0] - return f(x), (j,) - -def f_bwd(res, g): - j = res[0] - return (g * j,) - -f.defvjp(f_fwd, f_bwd) - -logdensity = lambda x: f(**x) - -inv_mass_matrix = jnp.ones(len(jxr)) -initial_position = {"x": jxr} - -rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) - -# adaptation -warmup = blackjax.window_adaptation(blackjax.nuts, logdensity, progress_bar=True) -rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3) -(state, parameters), _ = warmup.run(warmup_key, initial_position, num_steps=1000) - - -def inference_loop(rng_key, kernel, init, nsamples): - @jax.jit - def step(state, rng_key): - state, _ = kernel(rng_key, state) - return state, state - - keys = jax.random.split(rng_key, nsamples) - _, states = jax.lax.scan(step, init, keys) - return states - -# inference loop -rng_key, sample_key = jax.random.split(jax.random.PRNGKey(0)) -kernel = blackjax.nuts(logdensity, **parameters).step -states = inference_loop(sample_key, kernel, state, nsamples=1000) - - diff --git a/benchmark/Comrade/Sampling/pyproject.toml b/benchmark/Comrade/Sampling/pyproject.toml deleted file mode 100644 index 387857e6ff..0000000000 --- a/benchmark/Comrade/Sampling/pyproject.toml +++ /dev/null @@ -1,11 +0,0 @@ -[project] -name = "sampling" -version = "0.1.0" -description = "Add your description here" -readme = "README.md" -requires-python = ">=3.10, <3.12" -dependencies = [ - "blackjax>=1.3", - "ipython>=8.38.0", - "jax[cuda12]==0.5.0", -] diff --git a/benchmark/Comrade/export_to_jax.jl b/benchmark/Comrade/export_to_jax.jl deleted file mode 100644 index 1316226bf6..0000000000 --- a/benchmark/Comrade/export_to_jax.jl +++ /dev/null @@ -1,18 +0,0 @@ -using Pkg; -Pkg.activate(@__DIR__); -ENV["XLA_REACTANT_GPU_PREALLOCATE"] = "false" -using Reactant, NPZ, PythonCall - -include("comimager.jl") - -logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( - logdensityof, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Fwd") -) - -println("Exported to: $logdensity_python_file_path") - -grad_logdensity_python_file_path = Reactant.Serialization.export_to_enzymejax( - gl, tpostr, xr; output_dir=joinpath(@__DIR__, "Sampling/Serialized/Bwd") -) - -println("Exported to: $grad_logdensity_python_file_path") From ecae34cb6460d0abc9b5e333bf6fce0de49f3605 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 14:40:21 -0500 Subject: [PATCH 09/24] Add JSON --- benchmark/Comrade/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml index 060ee7704b..6685c8302d 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/Comrade/Project.toml @@ -12,6 +12,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" From a35f1621f72de0e4b86ca5566077ab0190bb67a0 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 14:42:12 -0500 Subject: [PATCH 10/24] Remove JSON3 --- benchmark/Comrade/Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml index 6685c8302d..b0aa841164 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/Comrade/Project.toml @@ -13,7 +13,6 @@ Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" -JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" From 91ce070401bed018e67594de6da4ecdf6000228e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 12 Feb 2026 14:42:11 -0600 Subject: [PATCH 11/24] chore: minor cleanup --- benchmark/Project.toml | 2 +- benchmark/{Comrade => comrade}/CondaPkg.toml | 0 benchmark/{Comrade => comrade}/Project.toml | 42 ++++++------------- benchmark/{Comrade => comrade}/comimager.jl | 0 benchmark/{Comrade => comrade}/common.jl | 17 ++------ .../{Comrade => comrade}/runbenchmarks.jl | 30 ++----------- 6 files changed, 21 insertions(+), 70 deletions(-) rename benchmark/{Comrade => comrade}/CondaPkg.toml (100%) rename benchmark/{Comrade => comrade}/Project.toml (65%) rename benchmark/{Comrade => comrade}/comimager.jl (100%) rename benchmark/{Comrade => comrade}/common.jl (91%) rename benchmark/{Comrade => comrade}/runbenchmarks.jl (63%) diff --git a/benchmark/Project.toml b/benchmark/Project.toml index a0a1922ee1..d3d6b7ccd5 100644 --- a/benchmark/Project.toml +++ b/benchmark/Project.toml @@ -8,4 +8,4 @@ PrettyTables = "3" julia = "1.11" [workspace] -projects = ["nn", "misc", "polybench"] +projects = ["nn", "misc", "polybench", "comrade"] diff --git a/benchmark/Comrade/CondaPkg.toml b/benchmark/comrade/CondaPkg.toml similarity index 100% rename from benchmark/Comrade/CondaPkg.toml rename to benchmark/comrade/CondaPkg.toml diff --git a/benchmark/Comrade/Project.toml b/benchmark/comrade/Project.toml similarity index 65% rename from benchmark/Comrade/Project.toml rename to benchmark/comrade/Project.toml index b0aa841164..f85af27d06 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/comrade/Project.toml @@ -26,33 +26,15 @@ VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029" VLBILikelihoods = "90db92cd-0007-4c0a-8e51-dbf0782ce592" VLBISkyModels = "d6343c73-7174-4e0f-bb64-562643efbeca" -[sources.Comrade] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/Comrade.jl" - -[sources.ComradeBase] -rev = "main" -url = "https://github.com/ptiede/ComradeBase.jl" - -[sources.NFFT] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/NFFT.jl" - -[sources.Reactant] -path = "../../" - -[sources.TransformVariables] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/TransformVariables.jl" - -[sources.VLBIImagePriors] -rev = "ptiede-reactantperf" -url = "https://github.com/ptiede/VLBIImagePriors.jl" - -[sources.VLBILikelihoods] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/VLBILikelihoods.jl" - -[sources.VLBISkyModels] -rev = "ptiede-reactnfft" -url = "https://github.com/EHTJulia/VLBISkyModels.jl" +[sources] +Comrade = {url = "https://github.com/ptiede/Comrade.jl", rev = "ptiede-reactant"} +ComradeBase = {url = "https://github.com/ptiede/ComradeBase.jl", rev = "main"} +NFFT = {url = "https://github.com/ptiede/NFFT.jl", rev = "ptiede-reactant"} +Reactant = {path = "../.."} +TransformVariables = {url = "https://github.com/ptiede/TransformVariables.jl", rev = "ptiede-reactant"} +VLBIImagePriors = {url = "https://github.com/ptiede/VLBIImagePriors.jl", rev = "ptiede-reactantperf"} +VLBILikelihoods = {url = "https://github.com/ptiede/VLBILikelihoods.jl", rev = "ptiede-reactant"} +VLBISkyModels = {url = "https://github.com/EHTJulia/VLBISkyModels.jl", rev = "ptiede-reactnfft"} + +[compat] +julia = "1.11" diff --git a/benchmark/Comrade/comimager.jl b/benchmark/comrade/comimager.jl similarity index 100% rename from benchmark/Comrade/comimager.jl rename to benchmark/comrade/comimager.jl diff --git a/benchmark/Comrade/common.jl b/benchmark/comrade/common.jl similarity index 91% rename from benchmark/Comrade/common.jl rename to benchmark/comrade/common.jl index d86622a60c..3a3d3c1f33 100644 --- a/benchmark/Comrade/common.jl +++ b/benchmark/comrade/common.jl @@ -12,10 +12,7 @@ function gradref(tpostr, xr) end function run_comrade_benchmark!( - results::Dict, - benchmark_name::String, - backend::String, - tpostr, + results::Dict, benchmark_name::String, backend::String, tpostr ) # TODO which Enzyme passes do I want to enable and disable run_benchmark!(results, benchmark_name, backend, tpostr, "forward", "Default") @@ -27,11 +24,10 @@ function run_benchmark!( results::Dict, benchmark_name::String, backend::String, - tpost, + tpost, fwd_or_bwd::String, - tag::String + tag::String, ) - if !haskey(results, "TFLOP/s") results["TFLOP/s"] = Dict{String,Float64}() end @@ -59,10 +55,7 @@ function run_benchmark!( error("Unknown fwd_or_bwd: $(fwd_or_bwd)") end - prof_result = Reactant.Profiler.profile_with_xprof( - fn,tpost, x; nrepeat=10, warmup=3 - ) - + prof_result = Reactant.Profiler.profile_with_xprof(fn, tpost, x; nrepeat=10, warmup=3) results["Runtime (s)"][full_benchmark_name] = prof_result.profiling_result.runtime_ns / 1e9 @@ -84,5 +77,3 @@ function run_benchmark!( @info print_stmt return nothing end - - diff --git a/benchmark/Comrade/runbenchmarks.jl b/benchmark/comrade/runbenchmarks.jl similarity index 63% rename from benchmark/Comrade/runbenchmarks.jl rename to benchmark/comrade/runbenchmarks.jl index ef597cd536..d3c54dcb63 100644 --- a/benchmark/Comrade/runbenchmarks.jl +++ b/benchmark/comrade/runbenchmarks.jl @@ -1,35 +1,11 @@ # Comrade Benchmarks Runner # This script runs all Comrade benchmarks and stores results to a JSON file - -using Pkg - -@static if VERSION ≥ v"1.10-" && VERSION < v"1.11" - Pkg.add([ - PackageSpec(; name="ComradeBase", rev="main"), - PackageSpec(; name="Comrade", rev="ptiede-reactant"), - PackageSpec(; name="VLBISkyModels", rev="ptiede-reactnfft"), - PackageSpec(; name="VLBILikelihoods", rev="ptiede-reactant"), - PackageSpec(; name="VLBIImagePriors", rev="ptiede-reactantperf"), - PackageSpec(; - url="https://github.com/ptiede/TransformVariables.jl", rev="ptiede-reactant" - ), - PackageSpec(; - url="https://github.com/ptiede/NFFT.jl", rev="ptiede-reactant" - ) - ]) - Pkg.develop(; path=joinpath(@__DIR__, "../../")) -end - - # Load dependencies using Reactant using Chairmarks: @b using Random: Random using Printf: @sprintf - - -using Reactant using LinearAlgebra using AbstractFFTs @@ -52,6 +28,8 @@ include("common.jl") @info sprint(io -> versioninfo(io; verbose=true)) +backend = get_backend() + include("comimager.jl") function run_all_benchmarks(backend::String) @@ -63,7 +41,6 @@ function run_all_benchmarks(backend::String) tpostr = build_post(μas2rad(200.0), 64, dataf) run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) - tpostr = build_post(μas2rad(200.0), 128, dataf) run_comrade_benchmark!(results, "Comrade EHT Imaging 128 x 128", backend, tpostr) @@ -73,5 +50,6 @@ function run_all_benchmarks(backend::String) return results end +res = run_all_benchmarks(backend) -res = run_all_benchmarks(get_backend()) \ No newline at end of file +save_results(results, joinpath(@__DIR__, "results"), "comrade", backend) From 0b50907b0d1893409252c87b2751ea20a53a0f90 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 12 Feb 2026 15:02:32 -0600 Subject: [PATCH 12/24] chore: fmt --- benchmark/comrade/comimager.jl | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/benchmark/comrade/comimager.jl b/benchmark/comrade/comimager.jl index 7a245b1822..4b0f845386 100644 --- a/benchmark/comrade/comimager.jl +++ b/benchmark/comrade/comimager.jl @@ -1,5 +1,4 @@ - # TODO Make Distributions package that is compatible with Reactant Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) @@ -32,7 +31,6 @@ function sky(θ, metadata) return ContinuousImage(rast, grid, DeltaPulse{3}()) end - function build_post(fov, npix, dataf) # @testset "Comrade Integration Imaging" begin @@ -55,7 +53,9 @@ function build_post(fov, npix, dataf) zprior = std_dist(pl) prior = (z=zprior, ρs=ρs, σ=Exponential(1.0)) - skymr = SkyModel(sky, prior, grd; metadata=skymeta, algorithm=VLBISkyModels.ReactantAlg()) # Need to do this so that we allocate proper Reactant arrays for internal stuff + skymr = SkyModel( + sky, prior, grd; metadata=skymeta, algorithm=VLBISkyModels.ReactantAlg() + ) # Need to do this so that we allocate proper Reactant arrays for internal stuff g(x) = exp(complex(x.lg, x.gp)) G = SingleStokesGain(g) @@ -83,5 +83,7 @@ if abspath(PROGRAM_FILE) == @__FILE__ tpostr = build_post(μas2rad(200.0), 64, dataf) results = Dict() backend = get_backend() - run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr, "forward", "test") -end \ No newline at end of file + run_comrade_benchmark!( + results, "Comrade EHT Imaging 64 x 64", backend, tpostr, "forward", "test" + ) +end From eedd9ded3def23923c5127a343f3343a7a6985b4 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 18:44:37 -0500 Subject: [PATCH 13/24] Enable Float32 support --- benchmark/Comrade/Project.toml | 3 +- benchmark/Comrade/comimager.jl | 81 ++++++++++++++++++++++++------ benchmark/Comrade/common.jl | 9 +++- benchmark/Comrade/runbenchmarks.jl | 6 ++- 4 files changed, 81 insertions(+), 18 deletions(-) diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml index b0aa841164..981b5e9c9a 100644 --- a/benchmark/Comrade/Project.toml +++ b/benchmark/Comrade/Project.toml @@ -1,6 +1,7 @@ [deps] AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" @@ -46,7 +47,7 @@ rev = "ptiede-reactant" url = "https://github.com/ptiede/TransformVariables.jl" [sources.VLBIImagePriors] -rev = "ptiede-reactantperf" +rev = "main" url = "https://github.com/ptiede/VLBIImagePriors.jl" [sources.VLBILikelihoods] diff --git a/benchmark/Comrade/comimager.jl b/benchmark/Comrade/comimager.jl index 7a245b1822..4934b492a8 100644 --- a/benchmark/Comrade/comimager.jl +++ b/benchmark/Comrade/comimager.jl @@ -2,6 +2,12 @@ # TODO Make Distributions package that is compatible with Reactant Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) + + +Distributions.minimum(::Exponential{T}) where {T<:AbstractFloat} = zero(T) +Distributions.maximum(::Exponential{T}) where {T<:AbstractFloat} = convert(T, Inf) + + function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) λ = rate(d) z = log(λ) - λ * x @@ -13,8 +19,19 @@ function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) return z end +# SO DUMB but Distributions doesn't support RNG from Flat32 vonmises distributions +function Distributions._rand!(rng::Random.AbstractRNG, d::DiagonalVonMises, x::AbstractVector{<:Float32}) + dv = Distributions.product_distribution(Distributions.VonMises.(Float64.(d.μ), Float64.(d.κ))) + x64 = rand(rng, dv) + x .= Float32.(x64) + return x +end + + +const MyDiagNormal{T} = MvNormal{T,Distributions.PDMats.PDiagMat{T,Vector{T}},Vector{T}} + # THis is very much needed or else `@compile` hangs -function Distributions.logpdf(d::Distributions.DiagNormal, x::Reactant.AnyTracedRVector) +function Distributions.logpdf(d::MyDiagNormal, x::Reactant.AnyTracedRVector) l = VLBILikelihoods._unnormed_logpdf_μΣ(d.μ, d.Σ.diag, x) n = VLBILikelihoods._gaussnorm(d.μ, d.Σ.diag) return l + n @@ -29,31 +46,67 @@ function sky(θ, metadata) bmimg = baseimage(mimg) rast = @. exp(x - mx) * bmimg rast ./= sum(rast) - return ContinuousImage(rast, grid, DeltaPulse{3}()) + return ContinuousImage(rast, grid, DeltaPulse()) end +function convert_table(T, dvis) + dt = datatable(dvis) + + @reset dt.baseline.U = T.(dt.baseline.U) + @reset dt.baseline.V = T.(dt.baseline.V) + @reset dt.baseline.Ti = T.(dt.baseline.Ti) + @reset dt.baseline.Fr = T.(dt.baseline.Fr) + @reset dt.measurement = Complex{T}.(dt.measurement) + @reset dt.noise = T.(dt.noise) + dvisT = Comrade.rebuild(dvis, dt) + Td = Comrade.datumtype(dvisT) + config = arrayconfig(dvisT) + confT = Comrade.EHTArrayConfiguration(; + bandwidth=T(config.bandwidth), + tarr = config.tarr, + scans=config.scans, + mjd=config.mjd, + ra=T(config.ra), + dec=T(config.dec), + source = config.source, + timetype = config.timetype, + datatable = config.datatable, + ) + + return EHTObservationTable{Td}(dvisT.measurement, dvisT.noise, confT) +end + +function build_post(fov, npix, dataf, backend) -function build_post(fov, npix, dataf) + if backend == "TPU" + T = Float32 + else + T = Float64 + end # @testset "Comrade Integration Imaging" begin obs = ehtim.obsdata.load_uvfits(dataf) obsavg = obs.add_fractional_noise(0.02) - dvis = extract_table(obsavg, Visibilities()) + dvis0 = extract_table(obsavg, Visibilities()) + + dvis = convert_table(T, dvis0) + + npix = npix - fovx = fov - fovy = fov + fovx = T(fov) + fovy = T(fov) # Now let's form our cache's. First, we have our usual image cache which is needed to numerically # compute the visibilities. grd = imagepixels(fovx, fovy, npix, npix) pl = StationaryRandomFieldPlan(grd) - mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(25.0))), grd) + mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(T(25.0)))), grd) skymeta = (; srf=pl, grid=grd, mimg=mimg) - ρs = ntuple(Returns(Uniform(0.01, max(size(grd)...))), 3) + ρs = ntuple(Returns(Uniform(T(0.01), T(max(size(grd)...)))), 3) zprior = std_dist(pl) - prior = (z=zprior, ρs=ρs, σ=Exponential(1.0)) + prior = (z=zprior, ρs=ρs, σ=Exponential(T(1.0))) skymr = SkyModel(sky, prior, grd; metadata=skymeta, algorithm=VLBISkyModels.ReactantAlg()) # Need to do this so that we allocate proper Reactant arrays for internal stuff @@ -62,12 +115,12 @@ function build_post(fov, npix, dataf) intpr = ( lg=ArrayPrior( - IIDSitePrior(ScanSeg(), Normal(0.0, 0.2)); - LM=IIDSitePrior(ScanSeg(), Normal(0.0, 1.0)), + IIDSitePrior(ScanSeg(), Normal(T(0.0), T(0.2))); + LM=IIDSitePrior(ScanSeg(), Normal(T(0.0), T(1.0))), ), gp=ArrayPrior( - IIDSitePrior(ScanSeg(), DiagonalVonMises(0.0, inv(π^2))); - refant=SEFDReference(0.0), + IIDSitePrior(ScanSeg(), DiagonalVonMises(T(0.0), T(inv(π^2)))); + refant=SEFDReference(T(0.0)), phase=true, ), ) @@ -80,7 +133,7 @@ function build_post(fov, npix, dataf) end if abspath(PROGRAM_FILE) == @__FILE__ - tpostr = build_post(μas2rad(200.0), 64, dataf) + tpostr = build_post(T(μas2rad(200.0)), 64, dataf) results = Dict() backend = get_backend() run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr, "forward", "test") diff --git a/benchmark/Comrade/common.jl b/benchmark/Comrade/common.jl index d86622a60c..2f47ddede5 100644 --- a/benchmark/Comrade/common.jl +++ b/benchmark/Comrade/common.jl @@ -49,7 +49,14 @@ function run_benchmark!( rng = Random.default_rng() # don't use any other rng Random.seed!(rng, 0) - x = Reactant.to_rarray(randn(rng, dimension(tpost))) + + if backend == "TPU" + T = Float32 + else + T = Float64 + end + + x = Reactant.to_rarray(randn(rng, T, dimension(tpost))) fn = if fwd_or_bwd == "forward" logdensityof diff --git a/benchmark/Comrade/runbenchmarks.jl b/benchmark/Comrade/runbenchmarks.jl index ef597cd536..1caf72344b 100644 --- a/benchmark/Comrade/runbenchmarks.jl +++ b/benchmark/Comrade/runbenchmarks.jl @@ -9,7 +9,7 @@ using Pkg PackageSpec(; name="Comrade", rev="ptiede-reactant"), PackageSpec(; name="VLBISkyModels", rev="ptiede-reactnfft"), PackageSpec(; name="VLBILikelihoods", rev="ptiede-reactant"), - PackageSpec(; name="VLBIImagePriors", rev="ptiede-reactantperf"), + PackageSpec(; name="VLBIImagePriors", rev="main"), PackageSpec(; url="https://github.com/ptiede/TransformVariables.jl", rev="ptiede-reactant" ), @@ -32,6 +32,7 @@ using Printf: @sprintf using Reactant using LinearAlgebra using AbstractFFTs +using Accessors: @set, @reset using VLBISkyModels using VLBILikelihoods @@ -56,11 +57,12 @@ include("comimager.jl") function run_all_benchmarks(backend::String) results = Dict{String,Dict{String,Float64}}() + bkend = get_backend() dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" dataf = Base.download(dataurl) - tpostr = build_post(μas2rad(200.0), 64, dataf) + tpostr = build_post(μas2rad(200.0), 64, dataf, bkend) run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) From 39337c18d8f8bb081eb2194e606ee003f34f7cc6 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 20:09:14 -0500 Subject: [PATCH 14/24] Add Comrade benchmark --- benchmark/Comrade/CondaPkg.toml | 2 - benchmark/Comrade/Project.toml | 59 ------------ benchmark/Comrade/comimager.jl | 140 ----------------------------- benchmark/Comrade/common.jl | 95 -------------------- benchmark/Comrade/runbenchmarks.jl | 79 ---------------- benchmark/comrade/comimager.jl | 87 +++++++++++++----- benchmark/comrade/common.jl | 9 +- benchmark/comrade/runbenchmarks.jl | 8 +- 8 files changed, 77 insertions(+), 402 deletions(-) delete mode 100644 benchmark/Comrade/CondaPkg.toml delete mode 100644 benchmark/Comrade/Project.toml delete mode 100644 benchmark/Comrade/comimager.jl delete mode 100644 benchmark/Comrade/common.jl delete mode 100644 benchmark/Comrade/runbenchmarks.jl diff --git a/benchmark/Comrade/CondaPkg.toml b/benchmark/Comrade/CondaPkg.toml deleted file mode 100644 index 2f0958226c..0000000000 --- a/benchmark/Comrade/CondaPkg.toml +++ /dev/null @@ -1,2 +0,0 @@ -[deps] -python = "<3.11,>=3.9,<4" diff --git a/benchmark/Comrade/Project.toml b/benchmark/Comrade/Project.toml deleted file mode 100644 index 981b5e9c9a..0000000000 --- a/benchmark/Comrade/Project.toml +++ /dev/null @@ -1,59 +0,0 @@ -[deps] -AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" -Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" -Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" -Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" -Comrade = "99d987ce-9a1e-4df8-bc0b-1ea019aa547b" -ComradeBase = "6d8c423b-a35f-4ef1-850c-862fe21f82c4" -CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" -Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" -JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" -LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" -NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" -NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" -Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -Pyehtim = "3d61700d-6e5b-419a-8e22-9c066cf00468" -PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" -TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff" -VLBIImagePriors = "b1ba175b-8447-452c-b961-7db2d6f7a029" -VLBILikelihoods = "90db92cd-0007-4c0a-8e51-dbf0782ce592" -VLBISkyModels = "d6343c73-7174-4e0f-bb64-562643efbeca" - -[sources.Comrade] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/Comrade.jl" - -[sources.ComradeBase] -rev = "main" -url = "https://github.com/ptiede/ComradeBase.jl" - -[sources.NFFT] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/NFFT.jl" - -[sources.Reactant] -path = "../../" - -[sources.TransformVariables] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/TransformVariables.jl" - -[sources.VLBIImagePriors] -rev = "main" -url = "https://github.com/ptiede/VLBIImagePriors.jl" - -[sources.VLBILikelihoods] -rev = "ptiede-reactant" -url = "https://github.com/ptiede/VLBILikelihoods.jl" - -[sources.VLBISkyModels] -rev = "ptiede-reactnfft" -url = "https://github.com/EHTJulia/VLBISkyModels.jl" diff --git a/benchmark/Comrade/comimager.jl b/benchmark/Comrade/comimager.jl deleted file mode 100644 index 4934b492a8..0000000000 --- a/benchmark/Comrade/comimager.jl +++ /dev/null @@ -1,140 +0,0 @@ - - -# TODO Make Distributions package that is compatible with Reactant -Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) - - -Distributions.minimum(::Exponential{T}) where {T<:AbstractFloat} = zero(T) -Distributions.maximum(::Exponential{T}) where {T<:AbstractFloat} = convert(T, Inf) - - -function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) - λ = rate(d) - z = log(λ) - λ * x - @trace if x < 0 - out = oftype(z, -Inf) - else - out = z - end - return z -end - -# SO DUMB but Distributions doesn't support RNG from Flat32 vonmises distributions -function Distributions._rand!(rng::Random.AbstractRNG, d::DiagonalVonMises, x::AbstractVector{<:Float32}) - dv = Distributions.product_distribution(Distributions.VonMises.(Float64.(d.μ), Float64.(d.κ))) - x64 = rand(rng, dv) - x .= Float32.(x64) - return x -end - - -const MyDiagNormal{T} = MvNormal{T,Distributions.PDMats.PDiagMat{T,Vector{T}},Vector{T}} - -# THis is very much needed or else `@compile` hangs -function Distributions.logpdf(d::MyDiagNormal, x::Reactant.AnyTracedRVector) - l = VLBILikelihoods._unnormed_logpdf_μΣ(d.μ, d.Σ.diag, x) - n = VLBILikelihoods._gaussnorm(d.μ, d.Σ.diag) - return l + n -end - -function sky(θ, metadata) - (; z, ρs, σ) = θ - (; srf, grid, mimg) = metadata - x = genfield(StationaryRandomField(MarkovPS(ρs), srf), z) - x .*= σ - mx = maximum(x) - bmimg = baseimage(mimg) - rast = @. exp(x - mx) * bmimg - rast ./= sum(rast) - return ContinuousImage(rast, grid, DeltaPulse()) -end - -function convert_table(T, dvis) - dt = datatable(dvis) - - @reset dt.baseline.U = T.(dt.baseline.U) - @reset dt.baseline.V = T.(dt.baseline.V) - @reset dt.baseline.Ti = T.(dt.baseline.Ti) - @reset dt.baseline.Fr = T.(dt.baseline.Fr) - @reset dt.measurement = Complex{T}.(dt.measurement) - @reset dt.noise = T.(dt.noise) - dvisT = Comrade.rebuild(dvis, dt) - Td = Comrade.datumtype(dvisT) - config = arrayconfig(dvisT) - confT = Comrade.EHTArrayConfiguration(; - bandwidth=T(config.bandwidth), - tarr = config.tarr, - scans=config.scans, - mjd=config.mjd, - ra=T(config.ra), - dec=T(config.dec), - source = config.source, - timetype = config.timetype, - datatable = config.datatable, - ) - - return EHTObservationTable{Td}(dvisT.measurement, dvisT.noise, confT) -end - -function build_post(fov, npix, dataf, backend) - - if backend == "TPU" - T = Float32 - else - T = Float64 - end - - # @testset "Comrade Integration Imaging" begin - obs = ehtim.obsdata.load_uvfits(dataf) - obsavg = obs.add_fractional_noise(0.02) - dvis0 = extract_table(obsavg, Visibilities()) - - dvis = convert_table(T, dvis0) - - - - npix = npix - fovx = T(fov) - fovy = T(fov) - - # Now let's form our cache's. First, we have our usual image cache which is needed to numerically - # compute the visibilities. - grd = imagepixels(fovx, fovy, npix, npix) - pl = StationaryRandomFieldPlan(grd) - mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(T(25.0)))), grd) - skymeta = (; srf=pl, grid=grd, mimg=mimg) - - ρs = ntuple(Returns(Uniform(T(0.01), T(max(size(grd)...)))), 3) - zprior = std_dist(pl) - prior = (z=zprior, ρs=ρs, σ=Exponential(T(1.0))) - - skymr = SkyModel(sky, prior, grd; metadata=skymeta, algorithm=VLBISkyModels.ReactantAlg()) # Need to do this so that we allocate proper Reactant arrays for internal stuff - - g(x) = exp(complex(x.lg, x.gp)) - G = SingleStokesGain(g) - - intpr = ( - lg=ArrayPrior( - IIDSitePrior(ScanSeg(), Normal(T(0.0), T(0.2))); - LM=IIDSitePrior(ScanSeg(), Normal(T(0.0), T(1.0))), - ), - gp=ArrayPrior( - IIDSitePrior(ScanSeg(), DiagonalVonMises(T(0.0), T(inv(π^2)))); - refant=SEFDReference(T(0.0)), - phase=true, - ), - ) - intmodel = InstrumentModel(G, intpr) - - postr = VLBIPosterior(skymr, intmodel, dvis) - tpostr = asflat(postr) - - return tpostr -end - -if abspath(PROGRAM_FILE) == @__FILE__ - tpostr = build_post(T(μas2rad(200.0)), 64, dataf) - results = Dict() - backend = get_backend() - run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr, "forward", "test") -end \ No newline at end of file diff --git a/benchmark/Comrade/common.jl b/benchmark/Comrade/common.jl deleted file mode 100644 index 2f47ddede5..0000000000 --- a/benchmark/Comrade/common.jl +++ /dev/null @@ -1,95 +0,0 @@ -using Printf: @sprintf -using Reactant: Reactant, @compile -using Enzyme: Enzyme, Const -using Random: Random - -include("../utils.jl") - -logdensityofref(tpostr, xr) = logdensityof(tpostr[], xr) -function gradref(tpostr, xr) - derivs, val = (Enzyme.gradient(ReverseWithPrimal, logdensityofref, Ref(tpostr), xr)) - return last(derivs), val -end - -function run_comrade_benchmark!( - results::Dict, - benchmark_name::String, - backend::String, - tpostr, -) - # TODO which Enzyme passes do I want to enable and disable - run_benchmark!(results, benchmark_name, backend, tpostr, "forward", "Default") - run_benchmark!(results, benchmark_name, backend, tpostr, "backward", "Default") - return nothing -end - -function run_benchmark!( - results::Dict, - benchmark_name::String, - backend::String, - tpost, - fwd_or_bwd::String, - tag::String -) - - if !haskey(results, "TFLOP/s") - results["TFLOP/s"] = Dict{String,Float64}() - end - if !haskey(results, "Runtime (s)") - results["Runtime (s)"] = Dict{String,Float64}() - end - - prim_or_rev = fwd_or_bwd == "forward" ? "primal" : "reverse" - full_benchmark_name = string(benchmark_name, "/", prim_or_rev, "/", backend, "/", tag) - @assert !haskey(results["Runtime (s)"], full_benchmark_name) "Benchmark already \ - exists: \ - $(full_benchmark_name)" - @assert !haskey(results["TFLOP/s"], full_benchmark_name) "Benchmark already exists: \ - $(full_benchmark_name)" - - rng = Random.default_rng() # don't use any other rng - Random.seed!(rng, 0) - - if backend == "TPU" - T = Float32 - else - T = Float64 - end - - x = Reactant.to_rarray(randn(rng, T, dimension(tpost))) - - fn = if fwd_or_bwd == "forward" - logdensityof - elseif fwd_or_bwd == "backward" - gradref - else - error("Unknown fwd_or_bwd: $(fwd_or_bwd)") - end - - prof_result = Reactant.Profiler.profile_with_xprof( - fn,tpost, x; nrepeat=10, warmup=3 - ) - - - results["Runtime (s)"][full_benchmark_name] = - prof_result.profiling_result.runtime_ns / 1e9 - results["TFLOP/s"][full_benchmark_name] = - if prof_result.profiling_result.flops_data === nothing - -1 - else - prof_result.profiling_result.flops_data.RawFlopsRate / 1e12 - end - - GC.gc(true) - - print_stmt = @sprintf( - "%100s : %.5gs %.5g TFLOP/s", - full_benchmark_name, - results["Runtime (s)"][full_benchmark_name], - results["TFLOP/s"][full_benchmark_name] - ) - @info print_stmt - return nothing -end - - diff --git a/benchmark/Comrade/runbenchmarks.jl b/benchmark/Comrade/runbenchmarks.jl deleted file mode 100644 index 1caf72344b..0000000000 --- a/benchmark/Comrade/runbenchmarks.jl +++ /dev/null @@ -1,79 +0,0 @@ -# Comrade Benchmarks Runner -# This script runs all Comrade benchmarks and stores results to a JSON file - -using Pkg - -@static if VERSION ≥ v"1.10-" && VERSION < v"1.11" - Pkg.add([ - PackageSpec(; name="ComradeBase", rev="main"), - PackageSpec(; name="Comrade", rev="ptiede-reactant"), - PackageSpec(; name="VLBISkyModels", rev="ptiede-reactnfft"), - PackageSpec(; name="VLBILikelihoods", rev="ptiede-reactant"), - PackageSpec(; name="VLBIImagePriors", rev="main"), - PackageSpec(; - url="https://github.com/ptiede/TransformVariables.jl", rev="ptiede-reactant" - ), - PackageSpec(; - url="https://github.com/ptiede/NFFT.jl", rev="ptiede-reactant" - ) - ]) - Pkg.develop(; path=joinpath(@__DIR__, "../../")) -end - - -# Load dependencies -using Reactant -using Chairmarks: @b -using Random: Random -using Printf: @sprintf - - - -using Reactant -using LinearAlgebra -using AbstractFFTs -using Accessors: @set, @reset - -using VLBISkyModels -using VLBILikelihoods -using Comrade -using Distributions -using VLBIImagePriors -using LogExpFunctions -import TransformVariables as TV - -using Downloads -using Distributions -using Enzyme - -using Pyehtim -using Test - -include("common.jl") - -@info sprint(io -> versioninfo(io; verbose=true)) - -include("comimager.jl") - -function run_all_benchmarks(backend::String) - results = Dict{String,Dict{String,Float64}}() - bkend = get_backend() - - dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" - dataf = Base.download(dataurl) - - tpostr = build_post(μas2rad(200.0), 64, dataf, bkend) - run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) - - - tpostr = build_post(μas2rad(200.0), 128, dataf) - run_comrade_benchmark!(results, "Comrade EHT Imaging 128 x 128", backend, tpostr) - - tpostr = build_post(μas2rad(200.0), 256, dataf) - run_comrade_benchmark!(results, "Comrade EHT Imaging 256 x 256", backend, tpostr) - - return results -end - - -res = run_all_benchmarks(get_backend()) \ No newline at end of file diff --git a/benchmark/comrade/comimager.jl b/benchmark/comrade/comimager.jl index 4b0f845386..1ba9b4709f 100644 --- a/benchmark/comrade/comimager.jl +++ b/benchmark/comrade/comimager.jl @@ -1,6 +1,10 @@ # TODO Make Distributions package that is compatible with Reactant Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) + +Distributions.minimum(::Exponential{T}) where {T<:AbstractFloat} = zero(T) +Distributions.maximum(::Exponential{T}) where {T<:AbstractFloat} = convert(T, Inf) + function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) λ = rate(d) z = log(λ) - λ * x @@ -12,8 +16,22 @@ function Distributions.logpdf(d::Exponential, x::Reactant.TracedRNumber) return z end +# SO DUMB but Distributions doesn't support RNG from Flat32 vonmises distributions +function Distributions._rand!( + rng::Random.AbstractRNG, d::DiagonalVonMises, x::AbstractVector{<:Float32} +) + dv = Distributions.product_distribution( + Distributions.VonMises.(Float64.(d.μ), Float64.(d.κ)) + ) + x64 = rand(rng, dv) + x .= Float32.(x64) + return x +end + +const MyDiagNormal{T} = MvNormal{T,Distributions.PDMats.PDiagMat{T,Vector{T}},Vector{T}} + # THis is very much needed or else `@compile` hangs -function Distributions.logpdf(d::Distributions.DiagNormal, x::Reactant.AnyTracedRVector) +function Distributions.logpdf(d::MyDiagNormal, x::Reactant.AnyTracedRVector) l = VLBILikelihoods._unnormed_logpdf_μΣ(d.μ, d.Σ.diag, x) n = VLBILikelihoods._gaussnorm(d.μ, d.Σ.diag) return l + n @@ -28,30 +46,64 @@ function sky(θ, metadata) bmimg = baseimage(mimg) rast = @. exp(x - mx) * bmimg rast ./= sum(rast) - return ContinuousImage(rast, grid, DeltaPulse{3}()) + return ContinuousImage(rast, grid, DeltaPulse()) end -function build_post(fov, npix, dataf) +function convert_table(T, dvis) + dt = datatable(dvis) + + @reset dt.baseline.U = T.(dt.baseline.U) + @reset dt.baseline.V = T.(dt.baseline.V) + @reset dt.baseline.Ti = T.(dt.baseline.Ti) + @reset dt.baseline.Fr = T.(dt.baseline.Fr) + @reset dt.measurement = Complex{T}.(dt.measurement) + @reset dt.noise = T.(dt.noise) + dvisT = Comrade.rebuild(dvis, dt) + Td = Comrade.datumtype(dvisT) + config = arrayconfig(dvisT) + confT = Comrade.EHTArrayConfiguration(; + bandwidth=T(config.bandwidth), + tarr=config.tarr, + scans=config.scans, + mjd=config.mjd, + ra=T(config.ra), + dec=T(config.dec), + source=config.source, + timetype=config.timetype, + datatable=config.datatable, + ) + + return EHTObservationTable{Td}(dvisT.measurement, dvisT.noise, confT) +end + +function build_post(fov, npix, dataf, backend) + if backend == "TPU" + T = Float32 + else + T = Float64 + end # @testset "Comrade Integration Imaging" begin obs = ehtim.obsdata.load_uvfits(dataf) obsavg = obs.add_fractional_noise(0.02) - dvis = extract_table(obsavg, Visibilities()) + dvis0 = extract_table(obsavg, Visibilities()) + + dvis = convert_table(T, dvis0) npix = npix - fovx = fov - fovy = fov + fovx = T(fov) + fovy = T(fov) # Now let's form our cache's. First, we have our usual image cache which is needed to numerically # compute the visibilities. grd = imagepixels(fovx, fovy, npix, npix) pl = StationaryRandomFieldPlan(grd) - mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(25.0))), grd) + mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(T(25.0)))), grd) skymeta = (; srf=pl, grid=grd, mimg=mimg) - ρs = ntuple(Returns(Uniform(0.01, max(size(grd)...))), 3) + ρs = ntuple(Returns(Uniform(T(0.01), T(max(size(grd)...)))), 3) zprior = std_dist(pl) - prior = (z=zprior, ρs=ρs, σ=Exponential(1.0)) + prior = (z=zprior, ρs=ρs, σ=Exponential(T(1.0))) skymr = SkyModel( sky, prior, grd; metadata=skymeta, algorithm=VLBISkyModels.ReactantAlg() @@ -62,12 +114,12 @@ function build_post(fov, npix, dataf) intpr = ( lg=ArrayPrior( - IIDSitePrior(ScanSeg(), Normal(0.0, 0.2)); - LM=IIDSitePrior(ScanSeg(), Normal(0.0, 1.0)), + IIDSitePrior(IntegSeg(), Normal(T(0.0), T(0.2))); + LM=IIDSitePrior(IntegSeg(), Normal(T(0.0), T(1.0))), ), gp=ArrayPrior( - IIDSitePrior(ScanSeg(), DiagonalVonMises(0.0, inv(π^2))); - refant=SEFDReference(0.0), + IIDSitePrior(IntegSeg(), DiagonalVonMises(T(0.0), T(inv(π^2)))); + refant=SEFDReference(T(0.0)), phase=true, ), ) @@ -78,12 +130,3 @@ function build_post(fov, npix, dataf) return tpostr end - -if abspath(PROGRAM_FILE) == @__FILE__ - tpostr = build_post(μas2rad(200.0), 64, dataf) - results = Dict() - backend = get_backend() - run_comrade_benchmark!( - results, "Comrade EHT Imaging 64 x 64", backend, tpostr, "forward", "test" - ) -end diff --git a/benchmark/comrade/common.jl b/benchmark/comrade/common.jl index 3a3d3c1f33..0e69d86ef7 100644 --- a/benchmark/comrade/common.jl +++ b/benchmark/comrade/common.jl @@ -45,7 +45,14 @@ function run_benchmark!( rng = Random.default_rng() # don't use any other rng Random.seed!(rng, 0) - x = Reactant.to_rarray(randn(rng, dimension(tpost))) + + Ts = if backend == "TPU" + Float32 + else + Float64 + end + + x = Reactant.to_rarray(randn(rng, Ts, dimension(tpost))) fn = if fwd_or_bwd == "forward" logdensityof diff --git a/benchmark/comrade/runbenchmarks.jl b/benchmark/comrade/runbenchmarks.jl index d3c54dcb63..192adbb0e8 100644 --- a/benchmark/comrade/runbenchmarks.jl +++ b/benchmark/comrade/runbenchmarks.jl @@ -8,7 +8,7 @@ using Printf: @sprintf using LinearAlgebra using AbstractFFTs - +using Accessors: @set, @reset using VLBISkyModels using VLBILikelihoods using Comrade @@ -38,13 +38,13 @@ function run_all_benchmarks(backend::String) dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" dataf = Base.download(dataurl) - tpostr = build_post(μas2rad(200.0), 64, dataf) + tpostr = build_post(μas2rad(200.0), 64, dataf, backend) run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) - tpostr = build_post(μas2rad(200.0), 128, dataf) + tpostr = build_post(μas2rad(200.0), 128, dataf, backend) run_comrade_benchmark!(results, "Comrade EHT Imaging 128 x 128", backend, tpostr) - tpostr = build_post(μas2rad(200.0), 256, dataf) + tpostr = build_post(μas2rad(200.0), 256, dataf, backend) run_comrade_benchmark!(results, "Comrade EHT Imaging 256 x 256", backend, tpostr) return results From 9ee71d2680eb49810c1faee4b2c8f121f435cf80 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 21:14:33 -0500 Subject: [PATCH 15/24] Fix --- benchmark/comrade/runbenchmarks.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark/comrade/runbenchmarks.jl b/benchmark/comrade/runbenchmarks.jl index 192adbb0e8..27ea1b01f8 100644 --- a/benchmark/comrade/runbenchmarks.jl +++ b/benchmark/comrade/runbenchmarks.jl @@ -52,4 +52,4 @@ end res = run_all_benchmarks(backend) -save_results(results, joinpath(@__DIR__, "results"), "comrade", backend) +save_results(res, joinpath(@__DIR__, "results"), "comrade", backend) From 3221f4886f066dd60cd4c628f0be3bd3ec3c7e49 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 22:05:49 -0500 Subject: [PATCH 16/24] Print out the function --- benchmark/comrade/common.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/benchmark/comrade/common.jl b/benchmark/comrade/common.jl index 0e69d86ef7..d8a8222d0b 100644 --- a/benchmark/comrade/common.jl +++ b/benchmark/comrade/common.jl @@ -54,6 +54,9 @@ function run_benchmark!( x = Reactant.to_rarray(randn(rng, Ts, dimension(tpost))) + @info typeof(x) + @info typeof(tpost) + fn = if fwd_or_bwd == "forward" logdensityof elseif fwd_or_bwd == "backward" @@ -64,6 +67,8 @@ function run_benchmark!( prof_result = Reactant.Profiler.profile_with_xprof(fn, tpost, x; nrepeat=10, warmup=3) + @code_hlo fn(tpost, x) + results["Runtime (s)"][full_benchmark_name] = prof_result.profiling_result.runtime_ns / 1e9 results["TFLOP/s"][full_benchmark_name] = From b7cebbc2021540c9a1162d4533b5037a95adc1dd Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 22:49:34 -0500 Subject: [PATCH 17/24] Fixed accidental type promotion --- benchmark/comrade/Project.toml | 2 +- benchmark/comrade/comimager.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmark/comrade/Project.toml b/benchmark/comrade/Project.toml index f85af27d06..ff2c85e3d4 100644 --- a/benchmark/comrade/Project.toml +++ b/benchmark/comrade/Project.toml @@ -32,7 +32,7 @@ ComradeBase = {url = "https://github.com/ptiede/ComradeBase.jl", rev = "main"} NFFT = {url = "https://github.com/ptiede/NFFT.jl", rev = "ptiede-reactant"} Reactant = {path = "../.."} TransformVariables = {url = "https://github.com/ptiede/TransformVariables.jl", rev = "ptiede-reactant"} -VLBIImagePriors = {url = "https://github.com/ptiede/VLBIImagePriors.jl", rev = "ptiede-reactantperf"} +# VLBIImagePriors = {url = "https://github.com/ptiede/VLBIImagePriors.jl", rev = "ptiede-reactantperf"} VLBILikelihoods = {url = "https://github.com/ptiede/VLBILikelihoods.jl", rev = "ptiede-reactant"} VLBISkyModels = {url = "https://github.com/EHTJulia/VLBISkyModels.jl", rev = "ptiede-reactnfft"} diff --git a/benchmark/comrade/comimager.jl b/benchmark/comrade/comimager.jl index 1ba9b4709f..c7482d3565 100644 --- a/benchmark/comrade/comimager.jl +++ b/benchmark/comrade/comimager.jl @@ -46,7 +46,7 @@ function sky(θ, metadata) bmimg = baseimage(mimg) rast = @. exp(x - mx) * bmimg rast ./= sum(rast) - return ContinuousImage(rast, grid, DeltaPulse()) + return ContinuousImage(rast, grid, DeltaPulse{eltype(mimg)}()) end function convert_table(T, dvis) From ac73fbd9a2f5ea9f0726af91a3500d9c91ddcaef Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 23:09:00 -0500 Subject: [PATCH 18/24] Update to use correct branches --- benchmark/comrade/Project.toml | 6 +++--- benchmark/comrade/common.jl | 5 ----- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/benchmark/comrade/Project.toml b/benchmark/comrade/Project.toml index ff2c85e3d4..36f9ff1cea 100644 --- a/benchmark/comrade/Project.toml +++ b/benchmark/comrade/Project.toml @@ -32,9 +32,9 @@ ComradeBase = {url = "https://github.com/ptiede/ComradeBase.jl", rev = "main"} NFFT = {url = "https://github.com/ptiede/NFFT.jl", rev = "ptiede-reactant"} Reactant = {path = "../.."} TransformVariables = {url = "https://github.com/ptiede/TransformVariables.jl", rev = "ptiede-reactant"} -# VLBIImagePriors = {url = "https://github.com/ptiede/VLBIImagePriors.jl", rev = "ptiede-reactantperf"} -VLBILikelihoods = {url = "https://github.com/ptiede/VLBILikelihoods.jl", rev = "ptiede-reactant"} -VLBISkyModels = {url = "https://github.com/EHTJulia/VLBISkyModels.jl", rev = "ptiede-reactnfft"} +VLBIImagePriors = {url = "https://github.com/ptiede/VLBIImagePriors.jl", rev = "main"} +VLBILikelihoods = {url = "https://github.com/ptiede/VLBILikelihoods.jl", rev = "main"} +VLBISkyModels = {url = "https://github.com/EHTJulia/VLBISkyModels.jl", rev = "main"} [compat] julia = "1.11" diff --git a/benchmark/comrade/common.jl b/benchmark/comrade/common.jl index d8a8222d0b..0e69d86ef7 100644 --- a/benchmark/comrade/common.jl +++ b/benchmark/comrade/common.jl @@ -54,9 +54,6 @@ function run_benchmark!( x = Reactant.to_rarray(randn(rng, Ts, dimension(tpost))) - @info typeof(x) - @info typeof(tpost) - fn = if fwd_or_bwd == "forward" logdensityof elseif fwd_or_bwd == "backward" @@ -67,8 +64,6 @@ function run_benchmark!( prof_result = Reactant.Profiler.profile_with_xprof(fn, tpost, x; nrepeat=10, warmup=3) - @code_hlo fn(tpost, x) - results["Runtime (s)"][full_benchmark_name] = prof_result.profiling_result.runtime_ns / 1e9 results["TFLOP/s"][full_benchmark_name] = From b1808ec9e00d3fe2585bca7712928e9e56ef52b3 Mon Sep 17 00:00:00 2001 From: Paul Tiede Date: Thu, 12 Feb 2026 23:16:32 -0500 Subject: [PATCH 19/24] Oops --- benchmark/comrade/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark/comrade/Project.toml b/benchmark/comrade/Project.toml index 36f9ff1cea..133a90a72b 100644 --- a/benchmark/comrade/Project.toml +++ b/benchmark/comrade/Project.toml @@ -34,7 +34,7 @@ Reactant = {path = "../.."} TransformVariables = {url = "https://github.com/ptiede/TransformVariables.jl", rev = "ptiede-reactant"} VLBIImagePriors = {url = "https://github.com/ptiede/VLBIImagePriors.jl", rev = "main"} VLBILikelihoods = {url = "https://github.com/ptiede/VLBILikelihoods.jl", rev = "main"} -VLBISkyModels = {url = "https://github.com/EHTJulia/VLBISkyModels.jl", rev = "main"} +VLBISkyModels = {url = "https://github.com/EHTJulia/VLBISkyModels.jl", rev = "ptiede-reactnfft"} [compat] julia = "1.11" From e66d75f9a9ed247cdbc1134fad9d7b5bcf2edd5c Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 13 Feb 2026 07:30:29 -0600 Subject: [PATCH 20/24] perf: comrade runner --- .github/workflows/benchmark.yml | 1 + benchmark/comrade/runbenchmarks.jl | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 4bca11adc1..0e81e4fa4d 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -63,6 +63,7 @@ jobs: - nn - misc - polybench + - comrade os: - linux-x86-n2-32 - linux-x86-ct6e-180-4tpu diff --git a/benchmark/comrade/runbenchmarks.jl b/benchmark/comrade/runbenchmarks.jl index 27ea1b01f8..709dcc985c 100644 --- a/benchmark/comrade/runbenchmarks.jl +++ b/benchmark/comrade/runbenchmarks.jl @@ -50,6 +50,7 @@ function run_all_benchmarks(backend::String) return results end -res = run_all_benchmarks(backend) +results = run_all_benchmarks(backend) -save_results(res, joinpath(@__DIR__, "results"), "comrade", backend) +save_results(results, joinpath(@__DIR__, "results"), "comrade", backend) +pretty_print_results(results, "comrade", backend) From db12b20f8c62c5b4cd18e87ffa5900b4b00223d8 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 13 Feb 2026 07:33:29 -0600 Subject: [PATCH 21/24] ci: skip fetch --- .github/workflows/benchmark.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 0e81e4fa4d..af611d497d 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -166,3 +166,4 @@ jobs: fail-on-alert: false auto-push: ${{ github.event_name != 'pull_request' }} max-items-in-chart: 50 + skip-fetch-gh-pages: true From 22106b3bbce0778d29b8e3bd5d79b056479cd58a Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 13 Feb 2026 07:37:59 -0600 Subject: [PATCH 22/24] ci: smart selection of specific benchmarks to run --- .github/workflows/benchmark.yml | 104 +++++++++++++++++++++++--------- 1 file changed, 74 insertions(+), 30 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index af611d497d..c06f88652f 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -8,7 +8,7 @@ permissions: on: schedule: - - cron: '0 3 * * *' # Nightly at 3am UTC + - cron: "0 3 * * *" # Nightly at 3am UTC workflow_dispatch: # Manual trigger pull_request: @@ -26,44 +26,95 @@ concurrency: cancel-in-progress: true jobs: - check-benchmark-changes: - if: github.event_name == 'pull_request' + determine-suites: runs-on: linux-x86-n2-32 container: - image: 'ghcr.io/enzymead/reactant-docker-images@sha256:7004a6ebbdd77bd047900b2bffc542e8576864056dc27a9c94d30666d6f7ea01' + image: "ghcr.io/enzymead/reactant-docker-images@sha256:7004a6ebbdd77bd047900b2bffc542e8576864056dc27a9c94d30666d6f7ea01" outputs: - benchmark_changed: ${{ steps.filter.outputs.benchmark }} + suites: ${{ steps.compute.outputs.suites }} + any_suite: ${{ steps.compute.outputs.any_suite }} steps: - uses: actions/checkout@v6 + with: + fetch-depth: 0 - uses: dorny/paths-filter@v3 + if: github.event_name == 'pull_request' id: filter with: filters: | - benchmark: - - 'benchmark/**' + common: + # TODO: reenable + # - '.github/workflows/benchmark.yml' + # - 'benchmark/Project.toml' + - 'ext/**' + - 'lib/**' + - 'src/**' + - 'Project.toml' + - 'benchmark/*.jl' + - name: Compute suites to run + id: compute + shell: bash + run: | + # Discover all benchmark suites (subdirs of benchmark/ with a runbenchmarks.jl) + ALL_SUITES=$(find benchmark -maxdepth 2 -name runbenchmarks.jl -path 'benchmark/*/runbenchmarks.jl' \ + | sed 's|benchmark/\(.*\)/runbenchmarks.jl|\1|' | sort | jq -Rsc 'split("\n") | map(select(. != ""))') + echo "Discovered suites: ${ALL_SUITES}" + + # Run all suites for schedule, workflow_dispatch, or 'run benchmarks' label + if [[ "${{ github.event_name }}" != "pull_request" ]]; then + echo "suites=${ALL_SUITES}" >> $GITHUB_OUTPUT + echo "any_suite=true" >> $GITHUB_OUTPUT + exit 0 + fi + + HAS_LABEL="${{ contains(join(github.event.pull_request.labels.*.name, ','), 'run benchmarks') }}" + if [[ "${HAS_LABEL}" == "true" ]]; then + echo "suites=${ALL_SUITES}" >> $GITHUB_OUTPUT + echo "any_suite=true" >> $GITHUB_OUTPUT + exit 0 + fi + + # For PRs, check if common files changed (run all) or only specific suites + COMMON="${{ steps.filter.outputs.common }}" + if [[ "${COMMON}" == "true" ]]; then + echo "suites=${ALL_SUITES}" >> $GITHUB_OUTPUT + echo "any_suite=true" >> $GITHUB_OUTPUT + exit 0 + fi + + # Fetch the base branch for comparison + git config --global --add safe.directory "${GITHUB_WORKSPACE}" + git fetch origin "${{ github.event.pull_request.base.ref }}" + + # Check which suite directories had changes using git diff + SUITES="[]" + for suite in $(echo "${ALL_SUITES}" | jq -r '.[]'); do + if git diff --quiet "origin/${{ github.event.pull_request.base.ref }}" -- "benchmark/${suite}/"; then + echo "Suite ${suite}: no changes" + else + echo "Suite ${suite}: changed" + SUITES=$(echo "${SUITES}" | jq -c ". + [\"${suite}\"]") + fi + done + + if [[ "${SUITES}" == "[]" ]]; then + echo "any_suite=false" >> $GITHUB_OUTPUT + else + echo "any_suite=true" >> $GITHUB_OUTPUT + fi + echo "suites=${SUITES}" >> $GITHUB_OUTPUT benchmark: timeout-minutes: 90 - needs: [check-benchmark-changes] - if: | - !cancelled() && ( - (github.ref == 'refs/heads/main' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')) || - (github.event_name == 'pull_request' && ( - contains(join(github.event.pull_request.labels.*.name, ','), 'run benchmarks') || - needs.check-benchmark-changes.outputs.benchmark_changed == 'true' - )) - ) + needs: [determine-suites] + if: needs.determine-suites.outputs.any_suite == 'true' runs-on: ${{ matrix.os }} container: image: ${{ contains(matrix.os, 'linux') && 'ghcr.io/enzymead/reactant-docker-images@sha256:7004a6ebbdd77bd047900b2bffc542e8576864056dc27a9c94d30666d6f7ea01' || '' }} strategy: fail-fast: false matrix: - suite: - - nn - - misc - - polybench - - comrade + suite: ${{ fromJson(needs.determine-suites.outputs.suites) }} os: - linux-x86-n2-32 - linux-x86-ct6e-180-4tpu @@ -93,15 +144,8 @@ jobs: benchmark-aggregate: runs-on: ubuntu-latest - if: | - !cancelled() && ( - (github.ref == 'refs/heads/main' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')) || - (github.event_name == 'pull_request' && ( - contains(join(github.event.pull_request.labels.*.name, ','), 'run benchmarks') || - needs.check-benchmark-changes.outputs.benchmark_changed == 'true' - )) - ) && needs.benchmark.result == 'success' - needs: [check-benchmark-changes, benchmark] + if: ${{ !cancelled() && needs.benchmark.result == 'success' }} + needs: [determine-suites, benchmark] steps: - uses: actions/checkout@v6 - uses: julia-actions/setup-julia@v2 From d6fbd824a4e90cd87d6cc9fac486c93f291300c0 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 13 Feb 2026 08:02:57 -0600 Subject: [PATCH 23/24] fix: missing dep --- benchmark/comrade/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmark/comrade/Project.toml b/benchmark/comrade/Project.toml index 133a90a72b..3ab90dc435 100644 --- a/benchmark/comrade/Project.toml +++ b/benchmark/comrade/Project.toml @@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d" NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" +PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Pyehtim = "3d61700d-6e5b-419a-8e22-9c066cf00468" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" From eb3fe75bd8fc253ea6b034fef426c0bb7ed88ced Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 13 Feb 2026 10:36:21 -0600 Subject: [PATCH 24/24] fix: naming --- benchmark/comrade/comimager.jl | 14 +++----------- benchmark/comrade/runbenchmarks.jl | 16 ++++++++-------- 2 files changed, 11 insertions(+), 19 deletions(-) diff --git a/benchmark/comrade/comimager.jl b/benchmark/comrade/comimager.jl index c7482d3565..547e58fbdd 100644 --- a/benchmark/comrade/comimager.jl +++ b/benchmark/comrade/comimager.jl @@ -1,4 +1,3 @@ - # TODO Make Distributions package that is compatible with Reactant Distributions.logpdf(d::Uniform, x::Reactant.TracedRNumber) = oftype(x, -log(d.b - d.a)) @@ -76,14 +75,7 @@ function convert_table(T, dvis) return EHTObservationTable{Td}(dvisT.measurement, dvisT.noise, confT) end -function build_post(fov, npix, dataf, backend) - if backend == "TPU" - T = Float32 - else - T = Float64 - end - - # @testset "Comrade Integration Imaging" begin +function build_post(::Type{T}, fov, npix, dataf) where {T} obs = ehtim.obsdata.load_uvfits(dataf) obsavg = obs.add_fractional_noise(0.02) dvis0 = extract_table(obsavg, Visibilities()) @@ -94,8 +86,8 @@ function build_post(fov, npix, dataf, backend) fovx = T(fov) fovy = T(fov) - # Now let's form our cache's. First, we have our usual image cache which is needed to numerically - # compute the visibilities. + # Now let's form our cache's. First, we have our usual image cache which is needed + # to numerically compute the visibilities. grd = imagepixels(fovx, fovy, npix, npix) pl = StationaryRandomFieldPlan(grd) mimg = intensitymap(modify(Gaussian(), Stretch(μas2rad(T(25.0)))), grd) diff --git a/benchmark/comrade/runbenchmarks.jl b/benchmark/comrade/runbenchmarks.jl index 709dcc985c..b1910380da 100644 --- a/benchmark/comrade/runbenchmarks.jl +++ b/benchmark/comrade/runbenchmarks.jl @@ -38,14 +38,14 @@ function run_all_benchmarks(backend::String) dataurl = "https://de.cyverse.org/anon-files/iplant/home/shared/commons_repo/curated/EHTC_M87pol2017_Nov2023/hops_data/April06/SR2_M87_2017_096_lo_hops_ALMArot.uvfits" dataf = Base.download(dataurl) - tpostr = build_post(μas2rad(200.0), 64, dataf, backend) - run_comrade_benchmark!(results, "Comrade EHT Imaging 64 x 64", backend, tpostr) - - tpostr = build_post(μas2rad(200.0), 128, dataf, backend) - run_comrade_benchmark!(results, "Comrade EHT Imaging 128 x 128", backend, tpostr) - - tpostr = build_post(μas2rad(200.0), 256, dataf, backend) - run_comrade_benchmark!(results, "Comrade EHT Imaging 256 x 256", backend, tpostr) + T = backend == "TPU" ? Float32 : Float64 + + for sz in (64, 128, 256) + tpostr = build_post(T, μas2rad(200.0), 64, dataf) + run_comrade_benchmark!( + results, "Comrade EHT Imaging $(sz) x $(sz) [$(T)]", backend, tpostr + ) + end return results end