|
| 1 | +# bench/svd_bench.exs |
| 2 | +# |
| 3 | +# Benchmarks `Nx.LinAlg.svd/1` across two MLX-backed Nx stacks — EMLX and Emily |
| 4 | +# (ausimian/emily) — and, within each, across: |
| 5 | +# |
| 6 | +# - device: :cpu vs :gpu (GPU skipped automatically if Metal is unavailable) |
| 7 | +# - compiler: :defn_evaluator — `compiler: Nx.Defn.Evaluator`, dispatching op-by-op |
| 8 | +# onto whichever backend the input already lives on |
| 9 | +# (EMLX.Backend or Emily.Backend) |
| 10 | +# :defn_compiled — `compiler: EMLX`, which lowers the whole graph to |
| 11 | +# a single fused native call |
| 12 | +# :defn_native — `compiler: Emily.Compiler, native: true`, Emily's |
| 13 | +# equivalent single-NIF fused replay |
| 14 | +# |
| 15 | +# Every lane is built *once* per `{n, device}` via `Nx.Defn.compile/3` and that |
| 16 | +# one closure is replayed across all Benchee-measured iterations — the |
| 17 | +# realistic "compile once, run many" pattern (see `EMLX`'s moduledoc), not a |
| 18 | +# fresh retrace (and, pre-`native_lowerable_block?/2`, dispatch-key rehash) on |
| 19 | +# every call, which is what repeated `Nx.Defn.jit_apply/3` would measure. |
| 20 | +# |
| 21 | +# This isolates two independent costs: library (EMLX vs Emily) and dispatch overhead |
| 22 | +# (per-op NIF round-trips vs one fused native call), on both CPU and GPU/Metal. |
| 23 | +# |
| 24 | +# This script is standalone — it does not need to run inside the `emlx` Mix project. |
| 25 | +# `Mix.install/2` pulls in EMLX from this checkout (via a `path:` dependency) and |
| 26 | +# Emily from Hex (Apple Silicon only), plus Benchee for the actual measurement. |
| 27 | +# |
| 28 | +# Run with (from anywhere): |
| 29 | +# elixir bench/svd_bench.exs |
| 30 | +# |
| 31 | +# Tunables (env vars): |
| 32 | +# EMLX_SVD_SIZES=128,256,512,1024 (default: 128,256,512,1024) |
| 33 | +# EMLX_SVD_TIME=2 (Benchee measurement time per lane, in seconds) |
| 34 | +# EMLX_SVD_WARMUP_TIME=1 (Benchee warmup time per lane, in seconds; also |
| 35 | +# primes the defn_compiled/defn_native JIT caches) |
| 36 | + |
| 37 | +apple_silicon? = |
| 38 | + match?({:unix, :darwin}, :os.type()) and |
| 39 | + :erlang.system_info(:system_architecture) |
| 40 | + |> to_string() |
| 41 | + |> String.starts_with?("aarch64") |
| 42 | + |
| 43 | +# Emily's precompiled NIFs are Apple Silicon macOS only (see its README's |
| 44 | +# Prerequisites section), so it's only added to the install list there — on any |
| 45 | +# other host this script benchmarks EMLX alone. |
| 46 | +deps = |
| 47 | + [ |
| 48 | + {:emlx, path: Path.expand("..", __DIR__)}, |
| 49 | + # Pin the same `nx` source as `emlx`'s own mix.exs (git main), overriding |
| 50 | + # Emily's Hex-sourced `~> 0.12` requirement, which main satisfies. |
| 51 | + {:nx, github: "elixir-nx/nx", branch: "main", sparse: "nx", override: true}, |
| 52 | + {:benchee, "~> 1.3"} |
| 53 | + ] ++ if(apple_silicon?, do: [{:emily, "~> 0.7"}], else: []) |
| 54 | + |
| 55 | +Mix.install(deps) |
| 56 | + |
| 57 | +# MLX's CPU backend JIT-compiles fused kernels via `popen`/`pclose`, which fails with |
| 58 | +# `ECHILD` under the BEAM's default SIGCHLD=SIG_IGN disposition — see |
| 59 | +# `EMLX.Application`'s moduledoc. `elixir` scripts own this trade-off just like our |
| 60 | +# test suite does. |
| 61 | +:os.set_signal(:sigchld, :default) |
| 62 | + |
| 63 | +defmodule Bench.SVD do |
| 64 | + # The benchmarked computation is always exactly `Nx.LinAlg.svd/1` |
| 65 | + # (`full_matrices?: true`) — every lane only varies the `Nx.Defn.compile/3` |
| 66 | + # `opts` (`:compiler`, `:device`, `:native`). Each lane's runner is built |
| 67 | + # *once* via `compile/3` (see `build/2`) and replayed across every |
| 68 | + # Benchee-measured iteration, mirroring the "compile once, run many" pattern |
| 69 | + # `EMLX`'s moduledoc recommends — not a fresh retrace per iteration, which |
| 70 | + # is the worst case `Nx.Defn.jit_apply/3` would otherwise measure. |
| 71 | + # |
| 72 | + # The compiled closure is stashed in `:persistent_term` (keyed by `key`) |
| 73 | + # rather than returned directly, because Benchee runs each job in its own |
| 74 | + # process: closing over the closure itself would make the *spawn* copy it |
| 75 | + # into that process, and `Nx.Defn.compile/3`'s closure keeps a reference to |
| 76 | + # the full (un-lowered) `Nx.Block` computation graph — including any |
| 77 | + # `default_expr` fallback, e.g. `Nx.LinAlg.svd/1`'s ~100-iteration |
| 78 | + # Jacobi-rotation algorithm, a DAG whose nodes heavily share |
| 79 | + # sub-expressions across iterations. Erlang's term-copy path doesn't |
| 80 | + # preserve that sharing, so flattening the DAG into a tree to copy it |
| 81 | + # blows up combinatorially and the spawn effectively never finishes. |
| 82 | + # `:persistent_term` values are read by reference (no copy), so fetching |
| 83 | + # the closure from inside the already-spawned job process sidesteps this |
| 84 | + # entirely. |
| 85 | + def build(key, template, opts) do |
| 86 | + compiled = Nx.Defn.compile(&Nx.LinAlg.svd(&1, full_matrices?: true), [template], opts) |
| 87 | + :persistent_term.put(key, compiled) |
| 88 | + key |
| 89 | + end |
| 90 | + |
| 91 | + # The `Nx.sum`/`Nx.to_number` calls aren't part of the compiled closure; |
| 92 | + # they just force materialization of all three outputs so a lazily-elided, |
| 93 | + # fused-but-unused output can't skew the timing. |
| 94 | + def run(key, a) do |
| 95 | + {u, s, vt} = :persistent_term.get(key).(a) |
| 96 | + Nx.to_number(Nx.sum(u)) + Nx.to_number(Nx.sum(s)) + Nx.to_number(Nx.sum(vt)) |
| 97 | + end |
| 98 | +end |
| 99 | + |
| 100 | +sizes = |
| 101 | + case System.get_env("EMLX_SVD_SIZES") do |
| 102 | + nil -> [128, 256, 512, 1024] |
| 103 | + csv -> csv |> String.split(",") |> Enum.map(&String.to_integer(String.trim(&1))) |
| 104 | + end |
| 105 | + |
| 106 | +parse_seconds = fn env, default -> {f, ""} = Float.parse(System.get_env(env, default)); f end |
| 107 | + |
| 108 | +bench_time = parse_seconds.("EMLX_SVD_TIME", "2") |
| 109 | +warmup_time = parse_seconds.("EMLX_SVD_WARMUP_TIME", "1") |
| 110 | + |
| 111 | +emily_available? = apple_silicon? and Code.ensure_loaded?(Emily.Backend) |
| 112 | + |
| 113 | +devices = |
| 114 | + case EMLX.NIF.command_queue_new(:gpu) do |
| 115 | + {:ok, _} -> |
| 116 | + [:cpu, :gpu] |
| 117 | + |
| 118 | + {:error, reason} -> |
| 119 | + IO.puts("==> GPU unavailable (#{inspect(reason)}), benchmarking CPU only.\n") |
| 120 | + [:cpu] |
| 121 | + end |
| 122 | + |
| 123 | +engines = [:emlx] ++ if(emily_available?, do: [:emily], else: []) |
| 124 | + |
| 125 | +unless emily_available? do |
| 126 | + IO.puts("==> Emily unavailable on this host (Apple Silicon macOS only), skipping.\n") |
| 127 | +end |
| 128 | + |
| 129 | +# Deterministic input, generated once on the host and transferred per-{engine,device} |
| 130 | +# below — keeps the matrix identical across every combination. |
| 131 | +random_matrix = fn n -> |
| 132 | + key = Nx.Random.key(42) |
| 133 | + {t, _key} = Nx.Random.uniform(key, shape: {n, n}, type: :f32) |
| 134 | + t |
| 135 | +end |
| 136 | + |
| 137 | +backend_for = fn |
| 138 | + :emlx, device -> {EMLX.Backend, device: device} |
| 139 | + :emily, device -> {Emily.Backend, device: device} |
| 140 | +end |
| 141 | + |
| 142 | +compiled_opts_for = fn |
| 143 | + :emlx, device -> [compiler: EMLX, device: device] |
| 144 | + :emily, device -> [compiler: Emily.Compiler, native: true, device: device] |
| 145 | +end |
| 146 | + |
| 147 | +# `emlx defn_evaluator` on `:gpu` has no native `full_matrices?: true` kernel |
| 148 | +# (see `EMLX.Backend.block/4`'s `:gpu` clause), so `Nx.Defn.Evaluator` runs the |
| 149 | +# `default_expr` fallback (a ~100-iteration Jacobi-rotation algorithm) fully |
| 150 | +# eagerly, one NIF round-trip per op. That's *correct* but tens of seconds |
| 151 | +# slow even at `n=128` and grows with `n`, so it's skipped by default — a |
| 152 | +# single stuck-looking lane otherwise dominates the whole suite's wall time. |
| 153 | +# Set `EMLX_SVD_INCLUDE_SLOW_GPU_EVALUATOR=1` to include it anyway. |
| 154 | +include_slow_gpu_evaluator? = System.get_env("EMLX_SVD_INCLUDE_SLOW_GPU_EVALUATOR") == "1" |
| 155 | + |
| 156 | +# Each lane builds its `Nx.Defn.compile/3` closure once (for `n`'s template) |
| 157 | +# and returns a runner that replays it — the benchmarked function itself |
| 158 | +# always lives in `Bench.SVD`, never as an inline anonymous function. |
| 159 | +# (`defn_evaluator` has no comparable per-call compile cost to amortize, but |
| 160 | +# building it the same way keeps lane-construction code uniform.) |
| 161 | +lanes_for = fn engine, device, n -> |
| 162 | + native_key = {:svd_bench, engine, device, n, :native} |
| 163 | + template = Nx.template({n, n}, {:f, 32}) |
| 164 | + Bench.SVD.build(native_key, template, compiled_opts_for.(engine, device)) |
| 165 | + |
| 166 | + native_lane = [ |
| 167 | + {(if engine == :emlx, do: :defn_compiled, else: :defn_native), |
| 168 | + &Bench.SVD.run(native_key, &1)} |
| 169 | + ] |
| 170 | + |
| 171 | + if engine == :emlx and device == :gpu and not include_slow_gpu_evaluator? do |
| 172 | + native_lane |
| 173 | + else |
| 174 | + evaluator_key = {:svd_bench, engine, device, n, :defn_evaluator} |
| 175 | + Bench.SVD.build(evaluator_key, template, compiler: Nx.Defn.Evaluator) |
| 176 | + [{:defn_evaluator, &Bench.SVD.run(evaluator_key, &1)} | native_lane] |
| 177 | + end |
| 178 | +end |
| 179 | + |
| 180 | +IO.puts("==> SVD benchmark (EMLX vs Emily)") |
| 181 | +IO.puts(" sizes: #{Enum.join(sizes, ", ")}") |
| 182 | +IO.puts(" time: #{bench_time}s per lane (warmup: #{warmup_time}s)") |
| 183 | +IO.puts(" devices: #{Enum.join(devices, ", ")}") |
| 184 | +IO.puts(" engines: #{Enum.join(engines, ", ")}") |
| 185 | + |
| 186 | +unless include_slow_gpu_evaluator? do |
| 187 | + IO.puts( |
| 188 | + " (skipping emlx defn_evaluator on :gpu — eager Jacobi fallback, tens of " <> |
| 189 | + "seconds+ per call; set EMLX_SVD_INCLUDE_SLOW_GPU_EVALUATOR=1 to include it)" |
| 190 | + ) |
| 191 | +end |
| 192 | + |
| 193 | +IO.puts("") |
| 194 | + |
| 195 | +suites = |
| 196 | + for n <- sizes, device <- devices, into: %{} do |
| 197 | + IO.puts("--- n=#{n} device=#{device} " <> String.duplicate("-", 40)) |
| 198 | + |
| 199 | + jobs = |
| 200 | + for engine <- engines, |
| 201 | + {mode, runner} <- lanes_for.(engine, device, n), |
| 202 | + reduce: %{} do |
| 203 | + jobs -> |
| 204 | + a = Nx.backend_transfer(random_matrix.(n), backend_for.(engine, device)) |
| 205 | + label = "#{engine} #{mode}" |
| 206 | + Map.put(jobs, label, fn -> runner.(a) end) |
| 207 | + end |
| 208 | + |
| 209 | + suite = |
| 210 | + Benchee.run(jobs, |
| 211 | + time: bench_time, |
| 212 | + warmup: warmup_time, |
| 213 | + memory_time: 0, |
| 214 | + print: [benchmarking: false, fast_warning: false], |
| 215 | + formatters: [Benchee.Formatters.Console] |
| 216 | + ) |
| 217 | + |
| 218 | + {{n, device}, suite} |
| 219 | + end |
| 220 | + |
| 221 | +# Cross-device speedup for the fused/native lane of each engine — Benchee only |
| 222 | +# compares jobs within a single run, so CPU vs GPU (two separate runs) is |
| 223 | +# summarized here from each run's own statistics. |
| 224 | +median_ms = fn suite, job_name -> |
| 225 | + case suite && Enum.find(suite.scenarios, &(&1.job_name == job_name)) do |
| 226 | + nil -> nil |
| 227 | + scenario -> scenario.run_time_data.statistics.median / 1_000_000 |
| 228 | + end |
| 229 | +end |
| 230 | + |
| 231 | +if :cpu in devices and :gpu in devices do |
| 232 | + IO.puts("\n=== CPU vs GPU speedup (fused/native lane, median ms) ===\n") |
| 233 | + |
| 234 | + Enum.each(engines, fn engine -> |
| 235 | + mode = if engine == :emlx, do: "defn_compiled", else: "defn_native" |
| 236 | + job_name = "#{engine} #{mode}" |
| 237 | + |
| 238 | + Enum.each(sizes, fn n -> |
| 239 | + cpu_ms = median_ms.(suites[{n, :cpu}], job_name) |
| 240 | + gpu_ms = median_ms.(suites[{n, :gpu}], job_name) |
| 241 | + |
| 242 | + case {cpu_ms, gpu_ms} do |
| 243 | + {c, g} when is_number(c) and is_number(g) and g > 0 -> |
| 244 | + speedup = Float.round(c / g, 2) |
| 245 | + |
| 246 | + IO.puts( |
| 247 | + " #{engine} n=#{n}: cpu=#{Float.round(c, 3)} ms, gpu=#{Float.round(g, 3)} ms, speedup=#{speedup}x" |
| 248 | + ) |
| 249 | + |
| 250 | + _ -> |
| 251 | + :ok |
| 252 | + end |
| 253 | + end) |
| 254 | + end) |
| 255 | +end |
0 commit comments