perf: comrade performance benchmark

benchmark/Project.toml

@@ -8,4 +8,4 @@ PrettyTables = "3"
 julia = "1.11"
 
 [workspace]
-projects = ["nn", "misc", "polybench"]
+projects = ["nn", "misc", "polybench", "comrade"]

benchmark/comrade/CondaPkg.toml

@@ -0,0 +1,2 @@
+[deps]
+python = "<3.11,>=3.9,<4"

benchmark/comrade/Project.toml

@@ -0,0 +1,41 @@
+[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"
+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"
+PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
+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 = {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 = "main"}
+VLBILikelihoods = {url = "https://github.com/ptiede/VLBILikelihoods.jl", rev = "main"}
+VLBISkyModels = {url = "https://github.com/EHTJulia/VLBISkyModels.jl", rev = "ptiede-reactnfft"}
+
+[compat]
+julia = "1.11"

benchmark/comrade/comimager.jl

@@ -0,0 +1,124 @@
+# 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{eltype(mimg)}())
+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(::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())
+
+    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(IntegSeg(), Normal(T(0.0), T(0.2)));
+            LM=IIDSitePrior(IntegSeg(), Normal(T(0.0), T(1.0))),
+        ),
+        gp=ArrayPrior(
+            IIDSitePrior(IntegSeg(), 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

benchmark/comrade/common.jl

@@ -0,0 +1,86 @@
+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)
+
+    Ts = if backend == "TPU"
+        Float32
+    else
+        Float64
+    end
+
+    x = Reactant.to_rarray(randn(rng, Ts, 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

benchmark/comrade/runbenchmarks.jl

@@ -0,0 +1,56 @@
+# Comrade Benchmarks Runner
+# This script runs all Comrade benchmarks and stores results to a JSON file
+# Load dependencies
+using Reactant
+using Chairmarks: @b
+using Random: Random
+using Printf: @sprintf
+
+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))
+
+backend = get_backend()
+
+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)
+
+    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
+
+results = run_all_benchmarks(backend)
+
+save_results(results, joinpath(@__DIR__, "results"), "comrade", backend)
+pretty_print_results(results, "comrade", backend)