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 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..3ab90dc435 --- /dev/null +++ b/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" diff --git a/benchmark/comrade/comimager.jl b/benchmark/comrade/comimager.jl new file mode 100644 index 0000000000..547e58fbdd --- /dev/null +++ b/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 diff --git a/benchmark/comrade/common.jl b/benchmark/comrade/common.jl new file mode 100644 index 0000000000..0e69d86ef7 --- /dev/null +++ b/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 diff --git a/benchmark/comrade/runbenchmarks.jl b/benchmark/comrade/runbenchmarks.jl new file mode 100644 index 0000000000..b1910380da --- /dev/null +++ b/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)