diff --git a/README.md b/README.md index a92e65a..54709cf 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ end ### System Requirements - **macOS (Apple Silicon)**: Uses Accelerate framework (no additional setup needed) -- **Linux x86_64**: Requires OpenBLAS and AVX2 support +- **Linux x86_64**: Requires OpenBLAS; AVX2 is recommended for SIMD speedups ```bash # Ubuntu/Debian sudo apt-get install libopenblas-dev @@ -43,6 +43,10 @@ end # Fedora/RHEL sudo dnf install openblas-devel ``` + For source builds that should use AVX2, set: + ```bash + export RUSTFLAGS="-C target-feature=+avx2" + ``` ### Rust Toolchain @@ -73,6 +77,7 @@ docs = Nx.tensor([ scores = ExMaxsimCpu.maxsim_scores(query, docs) # => #Nx.Tensor ``` +Note: MaxSim expects L2-normalized embeddings. Normalize query and docs along the token axis before scoring. ### Variable-Length Documents @@ -138,34 +143,50 @@ For additional performance on Intel CPUs, you can enable libxsmm: ## Benchmarks -ExMaxsimCpu delivers exceptional performance, **even outperforming GPU-accelerated alternatives**: +ExMaxsimCpu delivers strong CPU performance, especially on small-to-moderate batch sizes: + +### Benchmark Environment + +- OS: macOS 26.2 (25C56) +- CPU: Apple M4 Pro (12 physical / 12 logical cores) +- Architecture: arm64 +- Elixir: 1.19.5 (OTP 28) +- BLAS: Accelerate +- Env: OPENBLAS_NUM_THREADS=1, RAYON_NUM_THREADS=8 +- Nx CPU backend: Torchx (CPU) ### Performance Summary | Implementation | Typical Latency | vs ExMaxsimCpu | |---------------|-----------------|----------------| -| **ExMaxsimCpu** (BLAS+SIMD) | 0.09 - 0.28 ms | — | -| Nx + Torchx MPS (Apple GPU) | 0.86 - 4.43 ms | ~14x slower | -| Nx BinaryBackend (pure Elixir) | 426 - 5,842 ms | ~13,000x slower | +| **ExMaxsimCpu** (BLAS+SIMD) | 0.08 - 0.28 ms | — | +| Nx CPU backend (Torchx) | 0.20 - 0.89 ms | ~4x slower | +| Nx + Torchx MPS (Apple GPU) | 1.09 - 3.24 ms | ~15x slower | +| Nx BinaryBackend (unaccelerated) | 386 - 5,038 ms | ~13,800x slower | -> **Key insight:** ExMaxsimCpu is faster than GPU acceleration! The combination of optimized BLAS (Accelerate/OpenBLAS) and hand-tuned SIMD (AVX2/NEON) outperforms even Apple's Metal Performance Shaders. +Notes: +- Nx BinaryBackend is an unaccelerated baseline; EXLA/Torchx CPU will be faster. +- Nx CPU backend results (EXLA or Torchx) are included when available. +- MPS timings are for small shapes where transfer overhead can dominate. +- Results are machine- and shape-dependent; run your own benchmarks for production sizing. -### Three-Way Performance Comparison +### Performance Comparison ![Three-Way Comparison](assets/benchmark_three_way.png) -*Benchmark on Apple M4 Pro, showing ExMaxsimCpu (blue) vs Nx+MPS GPU (green) vs Nx BinaryBackend (orange) on a log scale.* +*Benchmark on Apple Silicon (this run), showing ExMaxsimCpu (blue), Nx BinaryBackend (orange), and optional Nx CPU/MPS backends on a log scale.* ### Speedup Analysis | Comparison | Speedup Factor | |------------|----------------| -| vs Nx BinaryBackend | **4,000x - 19,000x** | -| vs Nx + Torchx MPS (GPU) | **6x - 29x** | +| vs Nx BinaryBackend | **4,100x - 24,200x** | +| vs Nx CPU backend (Torchx) | **2.1x - 6.2x** | +| vs Nx + Torchx MPS (GPU) | **7.5x - 24.6x** | ![Speedup Comparison](assets/benchmark_speedup.png) -### Why is CPU faster than GPU? +### Why can CPU win on these shapes? 1. **Optimized BLAS**: Uses Apple Accelerate (or OpenBLAS on Linux) which is highly tuned for the CPU architecture 2. **SIMD Instructions**: Hand-optimized AVX2/NEON code for max-reduction operations @@ -190,7 +211,7 @@ ExMaxsimCpu delivers exceptional performance, **even outperforming GPU-accelerat ### Running Benchmarks -Generate benchmark data (includes MPS GPU comparison on Apple Silicon): +Generate benchmark data (includes optional MPS and Nx CPU backends when available): ```bash OPENBLAS_NUM_THREADS=1 mix run bench/generate_plots.exs ``` @@ -200,7 +221,7 @@ Generate plots: uv run bench/plot_benchmarks.py ``` -> **Note:** For MPS benchmarks, ensure `torchx` is installed. The benchmark automatically detects GPU availability. +> **Note:** For MPS benchmarks, ensure `torchx` is installed. For Nx CPU backends, install `torchx` or add `exla` to your dev deps. The benchmark auto-detects availability. ## Platform Support diff --git a/assets/benchmark_d_len.png b/assets/benchmark_d_len.png index fc9efaf..3dfcf75 100644 Binary files a/assets/benchmark_d_len.png and b/assets/benchmark_d_len.png differ diff --git a/assets/benchmark_data.csv b/assets/benchmark_data.csv index f6d514f..d7debe7 100644 --- a/assets/benchmark_data.csv +++ b/assets/benchmark_data.csv @@ -1,11 +1,11 @@ -config,param_value,ex_time_ms,nx_time_ms,mps_time_ms,mps_transfer_time_ms,speedup_vs_nx,speedup_vs_mps -n_docs,10,0.1234,426.714,0.862,1.257,3457.97,6.99 -n_docs,25,0.1808,1076.1877,4.4306,2.1413,5952.37,24.51 -n_docs,50,0.1732,2176.558,2.5852,2.7917,12566.73,14.93 -n_docs,100,0.2842,5397.9557,2.6614,3.1407,18993.51,9.36 -d_len,32,0.0929,1037.171,2.593,2.7057,11164.38,27.91 -d_len,64,0.1801,2149.108,1.7498,2.1703,11932.86,9.72 -d_len,128,0.246,5842.7707,1.6214,1.865,23751.1,6.59 -dim,64,0.178,1197.7157,2.3052,2.2873,6728.74,12.95 -dim,128,0.2202,2253.2173,2.6466,1.6743,10232.59,12.02 -dim,256,0.1771,4542.2847,2.8658,3.8633,25648.13,16.18 \ No newline at end of file +config,param_value,ex_time_ms,nx_time_ms,nx_cpu_time_ms,nx_cpu_backend,mps_time_ms,mps_transfer_time_ms,speedup_vs_nx,speedup_vs_nx_cpu,speedup_vs_mps +n_docs,10,0.0936,386.2173,0.198,Torchx (CPU),1.087,1.9057,4126.25,2.12,11.61 +n_docs,25,0.1231,970.5887,0.427,Torchx (CPU),1.376,2.182,7884.55,3.47,11.18 +n_docs,50,0.1502,1990.481,0.5047,Torchx (CPU),2.1504,2.3657,13252.2,3.36,14.32 +n_docs,100,0.2561,4990.392,0.8927,Torchx (CPU),3.2372,3.5997,19486.11,3.49,12.64 +d_len,32,0.0836,992.1037,0.489,Torchx (CPU),2.0586,2.6113,11867.27,5.85,24.62 +d_len,64,0.1499,2025.9087,0.64,Torchx (CPU),2.2548,2.2373,13515.07,4.27,15.04 +d_len,128,0.2753,5037.7807,0.7853,Torchx (CPU),2.0598,2.2373,18299.24,2.85,7.48 +dim,64,0.084,1069.5737,0.5213,Torchx (CPU),1.7228,2.34,12733.02,6.21,20.51 +dim,128,0.1691,2114.0027,0.6673,Torchx (CPU),2.227,2.3307,12501.49,3.95,13.17 +dim,256,0.1661,4027.811,0.758,Torchx (CPU),2.882,3.2097,24249.31,4.56,17.35 \ No newline at end of file diff --git a/assets/benchmark_dim.png b/assets/benchmark_dim.png index 1c53c57..b41381b 100644 Binary files a/assets/benchmark_dim.png and b/assets/benchmark_dim.png differ diff --git a/assets/benchmark_n_docs.png b/assets/benchmark_n_docs.png index 66af783..3554c95 100644 Binary files a/assets/benchmark_n_docs.png and b/assets/benchmark_n_docs.png differ diff --git a/assets/benchmark_speedup.png b/assets/benchmark_speedup.png index e3a8f7b..68ab1b7 100644 Binary files a/assets/benchmark_speedup.png and b/assets/benchmark_speedup.png differ diff --git a/assets/benchmark_summary.png b/assets/benchmark_summary.png index f53c12b..42e1d88 100644 Binary files a/assets/benchmark_summary.png and b/assets/benchmark_summary.png differ diff --git a/assets/benchmark_three_way.png b/assets/benchmark_three_way.png index e2f6ab7..6a3c72a 100644 Binary files a/assets/benchmark_three_way.png and b/assets/benchmark_three_way.png differ diff --git a/bench/benchmark.exs b/bench/benchmark.exs index f12e4bc..8be44ae 100644 --- a/bench/benchmark.exs +++ b/bench/benchmark.exs @@ -10,12 +10,47 @@ IO.puts("======================\n") # Generate random normalized vectors defmodule BenchHelper do + # Suppress compile warnings for optional backends + @compile {:no_warn_undefined, Torchx} + @compile {:no_warn_undefined, EXLA} + def random_normalized(shape) do key = Nx.Random.key(System.unique_integer([:positive])) {tensor, _key} = Nx.Random.normal(key, shape: shape, type: :f32) norms = Nx.sqrt(Nx.sum(Nx.pow(tensor, 2), axes: [-1], keep_axes: true)) Nx.divide(tensor, norms) end + + def torchx_cpu_backend do + with true <- Code.ensure_loaded?(Torchx), + {:ok, _apps} <- Application.ensure_all_started(:torchx) do + {Torchx.Backend, device: :cpu} + else + _ -> nil + end + end + + def exla_cpu_backend do + with true <- Code.ensure_loaded?(EXLA), + {:ok, _apps} <- Application.ensure_all_started(:exla) do + {EXLA.Backend, client: :host} + else + _ -> nil + end + end + + def cpu_backend do + case exla_cpu_backend() do + nil -> + case torchx_cpu_backend() do + nil -> nil + backend -> {"Torchx (CPU)", backend} + end + + backend -> + {"EXLA (CPU)", backend} + end + end end defmodule NxReference do @@ -44,6 +79,14 @@ end IO.puts("Generating test data...") query = BenchHelper.random_normalized({32, 128}) docs = BenchHelper.random_normalized({100, 64, 128}) +cpu_backend = BenchHelper.cpu_backend() + +if cpu_backend do + {label, _backend} = cpu_backend + IO.puts("✓ Nx CPU backend available: #{label}") +else + IO.puts("⚠ Nx CPU backend not available (install torchx or exla)") +end # Correctness check IO.puts("\nCorrectness check:") @@ -59,18 +102,29 @@ IO.puts(" Status: #{if max_diff < 1.0e-4, do: "✓ PASS", else: "✗ FAIL"}") # Benchmark comparison IO.puts("\nPerformance comparison (32 query tokens, 128 dim, 64 doc tokens, 100 docs):") -Benchee.run( - %{ - "ExMaxsimCpu (BLAS+SIMD)" => fn -> - ExMaxsimCpu.maxsim_scores(query, docs) - end, - "Pure Nx Reference" => fn -> - NxReference.maxsim_scores(query, docs) - end - }, - warmup: 1, - time: 3 -) +benchmarks = %{ + "ExMaxsimCpu (BLAS+SIMD)" => fn -> + ExMaxsimCpu.maxsim_scores(query, docs) + end, + "Nx BinaryBackend" => fn -> + NxReference.maxsim_scores(query, docs) |> Nx.to_binary() + end +} + +benchmarks = + if cpu_backend do + {label, backend} = cpu_backend + query_cpu = Nx.backend_transfer(query, backend) + docs_cpu = Nx.backend_transfer(docs, backend) + + Map.put(benchmarks, "Nx #{label}", fn -> + NxReference.maxsim_scores(query_cpu, docs_cpu) |> Nx.to_binary() + end) + else + benchmarks + end + +Benchee.run(benchmarks, warmup: 1, time: 3) # Larger scale benchmark (ExMaxsimCpu only - Nx reference is too slow) IO.puts("\n\nLarger scale benchmark (ExMaxsimCpu only):") diff --git a/bench/generate_plots.exs b/bench/generate_plots.exs index 1dc4b45..9edd968 100644 --- a/bench/generate_plots.exs +++ b/bench/generate_plots.exs @@ -13,6 +13,7 @@ IO.puts("==========================================\n") defmodule BenchHelper do # Suppress compile warnings for optional Torchx @compile {:no_warn_undefined, Torchx} + @compile {:no_warn_undefined, EXLA} def random_normalized(shape, backend \\ Nx.BinaryBackend) do key = Nx.Random.key(System.unique_integer([:positive])) @@ -51,6 +52,37 @@ defmodule BenchHelper do end end + def torchx_cpu_backend do + with true <- Code.ensure_loaded?(Torchx), + {:ok, _apps} <- Application.ensure_all_started(:torchx) do + {Torchx.Backend, device: :cpu} + else + _ -> nil + end + end + + def exla_cpu_backend do + with true <- Code.ensure_loaded?(EXLA), + {:ok, _apps} <- Application.ensure_all_started(:exla) do + {EXLA.Backend, client: :host} + else + _ -> nil + end + end + + def cpu_backend do + case exla_cpu_backend() do + nil -> + case torchx_cpu_backend() do + nil -> nil + backend -> {"Torchx (CPU)", backend} + end + + backend -> + {"EXLA (CPU)", backend} + end + end + defp mps_available? do # Use Torchx's built-in device availability check try do @@ -97,6 +129,7 @@ end # Check for MPS availability mps_backend = BenchHelper.torchx_mps_backend() +cpu_backend = BenchHelper.cpu_backend() if mps_backend do IO.puts("✓ Torchx MPS backend available - will benchmark GPU acceleration") @@ -105,8 +138,29 @@ else IO.puts(" (Install torchx and run on Apple Silicon for MPS benchmarks)") end +if cpu_backend do + {cpu_label, _backend} = cpu_backend + IO.puts("✓ #{cpu_label} backend available - will benchmark Nx CPU backend") +else + IO.puts("⚠ EXLA/Torchx CPU backend not available - skipping Nx CPU benchmarks") + IO.puts(" (Install exla or torchx to benchmark Nx CPU backends)") +end + IO.puts("") +# One-time warmup to avoid NIF load cost in the first measurement +IO.puts("Warming up NIF...") +warmup_query = BenchHelper.random_normalized({2, 2}) +warmup_docs = BenchHelper.random_normalized({1, 2, 2}) +_ = ExMaxsimCpu.maxsim_scores(warmup_query, warmup_docs) + +if cpu_backend do + {_label, backend} = cpu_backend + warmup_query_cpu = Nx.backend_transfer(warmup_query, backend) + warmup_docs_cpu = Nx.backend_transfer(warmup_docs, backend) + _ = NxReference.maxsim_scores(warmup_query_cpu, warmup_docs_cpu) |> Nx.to_binary() +end + # Benchmark configurations # Only include values where we can run comparisons without timeout configs = [ @@ -149,6 +203,24 @@ results = :ok end, 3) + # Measure Nx CPU backend (EXLA/Torchx) if available + {nx_cpu_time, nx_cpu_backend_label} = + if cpu_backend do + {label, backend} = cpu_backend + query_cpu = Nx.backend_transfer(query_bin, backend) + docs_cpu = Nx.backend_transfer(docs_bin, backend) + + time = BenchHelper.measure_time(fn -> + result = NxReference.maxsim_scores(query_cpu, docs_cpu) + _ = Nx.to_binary(result) + :ok + end, 3) + + {time, label} + else + {nil, nil} + end + # Measure Nx with MPS backend (if available) {mps_time, mps_transfer_time} = if mps_backend do @@ -183,19 +255,32 @@ results = # Calculate speedups speedup_vs_nx = nx_time / ex_time speedup_vs_mps = if mps_time, do: mps_time / ex_time, else: nil + speedup_vs_nx_cpu = if nx_cpu_time, do: nx_cpu_time / ex_time, else: nil # Print results mps_str = if mps_time, do: ", MPS: #{Float.round(mps_time, 2)}ms", else: "" - IO.puts("ExMaxsim: #{Float.round(ex_time, 2)}ms, Nx: #{Float.round(nx_time, 2)}ms#{mps_str}") + cpu_str = + if nx_cpu_time do + ", Nx CPU (#{nx_cpu_backend_label}): #{Float.round(nx_cpu_time, 2)}ms" + else + "" + end + + IO.puts( + "ExMaxsim: #{Float.round(ex_time, 2)}ms, Nx: #{Float.round(nx_time, 2)}ms#{cpu_str}#{mps_str}" + ) %{ config: config.name, param_value: value, ex_time_ms: ex_time, nx_time_ms: nx_time, + nx_cpu_time_ms: nx_cpu_time, + nx_cpu_backend: nx_cpu_backend_label, mps_time_ms: mps_time, mps_transfer_time_ms: mps_transfer_time, speedup_vs_nx: speedup_vs_nx, + speedup_vs_nx_cpu: speedup_vs_nx_cpu, speedup_vs_mps: speedup_vs_mps } end @@ -203,16 +288,20 @@ results = # Write CSV csv_path = "assets/benchmark_data.csv" -csv_header = "config,param_value,ex_time_ms,nx_time_ms,mps_time_ms,mps_transfer_time_ms,speedup_vs_nx,speedup_vs_mps" +csv_header = + "config,param_value,ex_time_ms,nx_time_ms,nx_cpu_time_ms,nx_cpu_backend,mps_time_ms,mps_transfer_time_ms,speedup_vs_nx,speedup_vs_nx_cpu,speedup_vs_mps" csv_content = [csv_header] ++ Enum.map(results, fn r -> mps_time = if r.mps_time_ms, do: Float.round(r.mps_time_ms, 4), else: "" mps_transfer = if r.mps_transfer_time_ms, do: Float.round(r.mps_transfer_time_ms, 4), else: "" + nx_cpu_time = if r.nx_cpu_time_ms, do: Float.round(r.nx_cpu_time_ms, 4), else: "" + nx_cpu_backend = r.nx_cpu_backend || "" + speedup_nx_cpu = if r.speedup_vs_nx_cpu, do: Float.round(r.speedup_vs_nx_cpu, 2), else: "" speedup_mps = if r.speedup_vs_mps, do: Float.round(r.speedup_vs_mps, 2), else: "" - "#{r.config},#{r.param_value},#{Float.round(r.ex_time_ms, 4)},#{Float.round(r.nx_time_ms, 4)},#{mps_time},#{mps_transfer},#{Float.round(r.speedup_vs_nx, 2)},#{speedup_mps}" + "#{r.config},#{r.param_value},#{Float.round(r.ex_time_ms, 4)},#{Float.round(r.nx_time_ms, 4)},#{nx_cpu_time},#{nx_cpu_backend},#{mps_time},#{mps_transfer},#{Float.round(r.speedup_vs_nx, 2)},#{speedup_nx_cpu},#{speedup_mps}" end) File.write!(csv_path, Enum.join(csv_content, "\n")) @@ -223,6 +312,16 @@ IO.puts("\n=== Summary ===") avg_speedup_nx = Enum.sum(Enum.map(results, & &1.speedup_vs_nx)) / length(results) IO.puts("Average speedup vs Nx (BinaryBackend): #{Float.round(avg_speedup_nx, 0)}x") +if cpu_backend do + {cpu_label, _backend} = cpu_backend + cpu_results = Enum.filter(results, & &1.speedup_vs_nx_cpu) + if length(cpu_results) > 0 do + avg_speedup_nx_cpu = + Enum.sum(Enum.map(cpu_results, & &1.speedup_vs_nx_cpu)) / length(cpu_results) + IO.puts("Average speedup vs Nx CPU (#{cpu_label}): #{Float.round(avg_speedup_nx_cpu, 1)}x") + end +end + if mps_backend do mps_results = Enum.filter(results, & &1.speedup_vs_mps) if length(mps_results) > 0 do diff --git a/bench/plot_benchmarks.py b/bench/plot_benchmarks.py index f462ebf..5431c0a 100644 --- a/bench/plot_benchmarks.py +++ b/bench/plot_benchmarks.py @@ -6,7 +6,7 @@ # ] # /// """ -Generate benchmark comparison plots for ExMaxsimCpu vs Nx vs Nx+MPS. +Generate benchmark comparison plots for ExMaxsimCpu vs Nx vs Nx+MPS (and Nx CPU backends if available). Run after: mix run bench/generate_plots.exs Usage: uv run bench/plot_benchmarks.py @@ -29,6 +29,7 @@ EX_COLOR = '#6366f1' # Indigo - ExMaxsimCpu NX_COLOR = '#f97316' # Orange - Pure Nx MPS_COLOR = '#10b981' # Emerald - Nx + MPS +NX_CPU_COLOR = '#0ea5e9' # Sky - Nx CPU backend SPEEDUP_COLOR = '#8b5cf6' # Purple @@ -36,6 +37,8 @@ def load_data(csv_path): """Load benchmark data from CSV.""" data = {'n_docs': [], 'd_len': [], 'dim': []} has_mps = False + has_nx_cpu = False + nx_cpu_label = None with open(csv_path, 'r') as f: reader = csv.DictReader(f) @@ -45,19 +48,32 @@ def load_data(csv_path): 'param': int(row['param_value']), 'ex_time': float(row['ex_time_ms']), 'nx_time': float(row['nx_time_ms']), + 'nx_cpu_time': float(row['nx_cpu_time_ms']) if row.get('nx_cpu_time_ms') else None, + 'nx_cpu_backend': row.get('nx_cpu_backend') or None, 'mps_time': float(row['mps_time_ms']) if row['mps_time_ms'] else None, 'mps_transfer_time': float(row['mps_transfer_time_ms']) if row.get('mps_transfer_time_ms') else None, 'speedup_nx': float(row['speedup_vs_nx']), + 'speedup_nx_cpu': float(row['speedup_vs_nx_cpu']) if row.get('speedup_vs_nx_cpu') else None, 'speedup_mps': float(row['speedup_vs_mps']) if row['speedup_vs_mps'] else None, } data[config].append(entry) if entry['mps_time'] is not None: has_mps = True + if entry['nx_cpu_time'] is not None: + has_nx_cpu = True + if entry['nx_cpu_backend'] and not nx_cpu_label: + nx_cpu_label = entry['nx_cpu_backend'] - return data, has_mps + meta = { + 'has_mps': has_mps, + 'has_nx_cpu': has_nx_cpu, + 'nx_cpu_label': nx_cpu_label, + } + return data, meta -def plot_time_comparison(data, config_name, param_label, output_path, has_mps): + +def plot_time_comparison(data, config_name, param_label, output_path, has_mps, has_nx_cpu, nx_cpu_label): """Create a bar chart comparing execution times.""" fig, ax = plt.subplots(figsize=(12, 6)) @@ -65,18 +81,23 @@ def plot_time_comparison(data, config_name, param_label, output_path, has_mps): ex_times = [d['ex_time'] for d in data] nx_times = [d['nx_time'] for d in data] mps_times = [d['mps_time'] if d['mps_time'] else 0 for d in data] + nx_cpu_times = [d['nx_cpu_time'] if d['nx_cpu_time'] else 0 for d in data] x = np.arange(len(params)) + series = [ + ('ExMaxsimCpu (BLAS+SIMD)', ex_times, EX_COLOR), + ] + if has_nx_cpu and any(nx_cpu_times): + series.append((nx_cpu_label or 'Nx CPU Backend', nx_cpu_times, NX_CPU_COLOR)) if has_mps and any(mps_times): - width = 0.25 - bars1 = ax.bar(x - width, ex_times, width, label='ExMaxsimCpu (BLAS+SIMD)', color=EX_COLOR, alpha=0.8) - bars2 = ax.bar(x, mps_times, width, label='Nx + Torchx MPS (GPU)', color=MPS_COLOR, alpha=0.8) - bars3 = ax.bar(x + width, nx_times, width, label='Nx BinaryBackend', color=NX_COLOR, alpha=0.8) - else: - width = 0.35 - bars1 = ax.bar(x - width/2, ex_times, width, label='ExMaxsimCpu (BLAS+SIMD)', color=EX_COLOR, alpha=0.8) - bars3 = ax.bar(x + width/2, nx_times, width, label='Nx BinaryBackend', color=NX_COLOR, alpha=0.8) + series.append(('Nx + Torchx MPS (GPU)', mps_times, MPS_COLOR)) + series.append(('Nx BinaryBackend', nx_times, NX_COLOR)) + + width = 0.8 / len(series) + offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, len(series)) + for offset, (label, times, color) in zip(offsets, series): + ax.bar(x + offset, times, width, label=label, color=color, alpha=0.8) ax.set_xlabel(param_label) ax.set_ylabel('Time (ms)') @@ -92,7 +113,7 @@ def plot_time_comparison(data, config_name, param_label, output_path, has_mps): print(f"Saved: {output_path}") -def plot_speedup_chart(all_data, output_path, has_mps): +def plot_speedup_chart(all_data, output_path, has_mps, has_nx_cpu): """Create a combined speedup chart.""" n_cols = 3 fig, axes = plt.subplots(1, n_cols, figsize=(14, 5)) @@ -108,46 +129,42 @@ def plot_speedup_chart(all_data, output_path, has_mps): params = [str(d['param']) for d in data] speedups_nx = [d['speedup_nx'] for d in data] speedups_mps = [d['speedup_mps'] if d['speedup_mps'] else 0 for d in data] + speedups_nx_cpu = [d['speedup_nx_cpu'] if d['speedup_nx_cpu'] else 0 for d in data] x = np.arange(len(params)) + series = [ + ('vs Nx Binary', speedups_nx, NX_COLOR), + ] + if has_nx_cpu and any(speedups_nx_cpu): + series.append(('vs Nx CPU', speedups_nx_cpu, NX_CPU_COLOR)) if has_mps and any(speedups_mps): - width = 0.35 - bars1 = ax.bar(x - width/2, speedups_nx, width, label='vs Nx Binary', color=NX_COLOR, alpha=0.8) - bars2 = ax.bar(x + width/2, speedups_mps, width, label='vs Nx MPS', color=MPS_COLOR, alpha=0.8) - - # Add value labels - for bar, val in zip(bars1, speedups_nx): - ax.annotate(f'{val:.0f}x', - xy=(bar.get_x() + bar.get_width()/2, bar.get_height()), - xytext=(0, 3), textcoords='offset points', - ha='center', va='bottom', fontsize=8, fontweight='bold') - for bar, val in zip(bars2, speedups_mps): + series.append(('vs Nx MPS', speedups_mps, MPS_COLOR)) + + width = 0.8 / len(series) + offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, len(series)) + for offset, (series_label, values, color) in zip(offsets, series): + bars = ax.bar(x + offset, values, width, label=series_label, color=color, alpha=0.8) + for bar, val in zip(bars, values): if val > 0: - ax.annotate(f'{val:.1f}x', + ax.annotate(f'{val:.0f}x', xy=(bar.get_x() + bar.get_width()/2, bar.get_height()), xytext=(0, 3), textcoords='offset points', ha='center', va='bottom', fontsize=8, fontweight='bold') - else: - bars = ax.bar(x, speedups_nx, color=NX_COLOR, alpha=0.8) - for bar, val in zip(bars, speedups_nx): - ax.annotate(f'{val:.0f}x', - xy=(bar.get_x() + bar.get_width()/2, bar.get_height()), - xytext=(0, 3), textcoords='offset points', - ha='center', va='bottom', fontsize=9, fontweight='bold') ax.set_xlabel(label) ax.set_ylabel('Speedup (x times faster)') ax.set_xticks(x) ax.set_xticklabels(params) - if has_mps: + if has_mps or has_nx_cpu: ax.legend(loc='upper left', fontsize=8) title = 'ExMaxsimCpu Speedup' + title += ' vs Nx (BinaryBackend)' + if has_nx_cpu: + title += ' and Nx (CPU)' if has_mps: - title += ' vs Nx (BinaryBackend) and Nx (MPS GPU)' - else: - title += ' vs Nx (BinaryBackend)' + title += ' and Nx (MPS GPU)' fig.suptitle(title, fontsize=14, fontweight='bold') plt.tight_layout() @@ -156,7 +173,7 @@ def plot_speedup_chart(all_data, output_path, has_mps): print(f"Saved: {output_path}") -def plot_summary(all_data, output_path, has_mps): +def plot_summary(all_data, output_path, has_mps, has_nx_cpu, nx_cpu_label): """Create a summary visualization.""" fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) @@ -166,12 +183,17 @@ def plot_summary(all_data, output_path, has_mps): ex_times = [d['ex_time'] for d in data] nx_times = [d['nx_time'] for d in data] mps_times = [d['mps_time'] if d['mps_time'] else None for d in data] + nx_cpu_times = [d['nx_cpu_time'] if d['nx_cpu_time'] else None for d in data] ax1.plot(params, ex_times, 'o-', color=EX_COLOR, linewidth=2, markersize=8, label='ExMaxsimCpu (BLAS+SIMD)') if has_mps and any(mps_times): mps_valid = [(p, t) for p, t in zip(params, mps_times) if t is not None] if mps_valid: ax1.plot([p for p, _ in mps_valid], [t for _, t in mps_valid], 'd-', color=MPS_COLOR, linewidth=2, markersize=8, label='Nx + Torchx MPS (GPU)') + if has_nx_cpu and any(nx_cpu_times): + cpu_valid = [(p, t) for p, t in zip(params, nx_cpu_times) if t is not None] + if cpu_valid: + ax1.plot([p for p, _ in cpu_valid], [t for _, t in cpu_valid], '^-', color=NX_CPU_COLOR, linewidth=2, markersize=8, label=nx_cpu_label or 'Nx CPU Backend') ax1.plot(params, nx_times, 's-', color=NX_COLOR, linewidth=2, markersize=8, label='Nx BinaryBackend') ax1.set_xlabel('Number of Documents') @@ -216,10 +238,10 @@ def plot_summary(all_data, output_path, has_mps): print(f"Saved: {output_path}") -def plot_three_way_comparison(all_data, output_path, has_mps): - """Create a comprehensive three-way comparison chart.""" - if not has_mps: - print(f"Skipping three-way comparison (no MPS data)") +def plot_three_way_comparison(all_data, output_path, has_mps, has_nx_cpu, nx_cpu_label): + """Create a comprehensive comparison chart.""" + if not has_mps and not has_nx_cpu: + print("Skipping multi-way comparison (no MPS/CPU backend data)") return fig, ax = plt.subplots(figsize=(12, 7)) @@ -229,25 +251,35 @@ def plot_three_way_comparison(all_data, output_path, has_mps): ex_times = [] nx_times = [] mps_times = [] + nx_cpu_times = [] for config in ['n_docs', 'd_len', 'dim']: for d in all_data[config]: - if d['mps_time']: + if d['mps_time'] or d['nx_cpu_time']: all_configs.append(f"{config}={d['param']}") ex_times.append(d['ex_time']) nx_times.append(d['nx_time']) - mps_times.append(d['mps_time']) + mps_times.append(d['mps_time'] if d['mps_time'] else 0) + nx_cpu_times.append(d['nx_cpu_time'] if d['nx_cpu_time'] else 0) if not all_configs: print(f"Skipping three-way comparison (no complete data)") return x = np.arange(len(all_configs)) - width = 0.25 + series = [ + ('ExMaxsimCpu (BLAS+SIMD)', ex_times, EX_COLOR), + ] + if has_nx_cpu and any(nx_cpu_times): + series.append((nx_cpu_label or 'Nx CPU Backend', nx_cpu_times, NX_CPU_COLOR)) + if has_mps and any(mps_times): + series.append(('Nx + Torchx MPS (GPU)', mps_times, MPS_COLOR)) + series.append(('Nx BinaryBackend', nx_times, NX_COLOR)) - bars1 = ax.bar(x - width, ex_times, width, label='ExMaxsimCpu (BLAS+SIMD)', color=EX_COLOR, alpha=0.8) - bars2 = ax.bar(x, mps_times, width, label='Nx + Torchx MPS (GPU)', color=MPS_COLOR, alpha=0.8) - bars3 = ax.bar(x + width, nx_times, width, label='Nx BinaryBackend', color=NX_COLOR, alpha=0.8) + width = 0.8 / len(series) + offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, len(series)) + for offset, (label, times, color) in zip(offsets, series): + ax.bar(x + offset, times, width, label=label, color=color, alpha=0.8) ax.set_xlabel('Configuration') ax.set_ylabel('Time (ms)') @@ -255,7 +287,7 @@ def plot_three_way_comparison(all_data, output_path, has_mps): ax.set_xticklabels(all_configs, rotation=45, ha='right') ax.set_yscale('log') ax.legend(loc='upper left') - ax.set_title('Three-Way Performance Comparison: ExMaxsimCpu vs Nx (MPS) vs Nx (Binary)') + ax.set_title('Performance Comparison: ExMaxsimCpu vs Nx (Binary/MPS/CPU)') plt.tight_layout() plt.savefig(output_path, dpi=150, bbox_inches='tight') @@ -274,32 +306,41 @@ def main(): return print("Loading benchmark data...") - data, has_mps = load_data(csv_path) + data, meta = load_data(csv_path) + has_mps = meta['has_mps'] + has_nx_cpu = meta['has_nx_cpu'] + nx_cpu_label = meta['nx_cpu_label'] if has_mps: - print("MPS data detected - generating three-way comparison plots") + print("MPS data detected - generating comparison plots") else: - print("No MPS data - generating two-way comparison plots") + print("No MPS data - generating comparison plots") print("\nGenerating plots...") # Individual comparison plots - plot_time_comparison(data['n_docs'], 'n_docs', 'Number of Documents', - assets_dir / 'benchmark_n_docs.png', has_mps) - plot_time_comparison(data['d_len'], 'd_len', 'Document Length (tokens)', - assets_dir / 'benchmark_d_len.png', has_mps) - plot_time_comparison(data['dim'], 'dim', 'Embedding Dimension', - assets_dir / 'benchmark_dim.png', has_mps) + plot_time_comparison( + data['n_docs'], 'n_docs', 'Number of Documents', + assets_dir / 'benchmark_n_docs.png', has_mps, has_nx_cpu, nx_cpu_label + ) + plot_time_comparison( + data['d_len'], 'd_len', 'Document Length (tokens)', + assets_dir / 'benchmark_d_len.png', has_mps, has_nx_cpu, nx_cpu_label + ) + plot_time_comparison( + data['dim'], 'dim', 'Embedding Dimension', + assets_dir / 'benchmark_dim.png', has_mps, has_nx_cpu, nx_cpu_label + ) # Combined speedup chart - plot_speedup_chart(data, assets_dir / 'benchmark_speedup.png', has_mps) + plot_speedup_chart(data, assets_dir / 'benchmark_speedup.png', has_mps, has_nx_cpu) # Summary plot - plot_summary(data, assets_dir / 'benchmark_summary.png', has_mps) + plot_summary(data, assets_dir / 'benchmark_summary.png', has_mps, has_nx_cpu, nx_cpu_label) - # Three-way comparison (if MPS data available) - if has_mps: - plot_three_way_comparison(data, assets_dir / 'benchmark_three_way.png', has_mps) + # Multi-way comparison (if MPS/CPU data available) + if has_mps or has_nx_cpu: + plot_three_way_comparison(data, assets_dir / 'benchmark_three_way.png', has_mps, has_nx_cpu, nx_cpu_label) print("\nAll plots generated successfully!") diff --git a/lib/ex_maxsim_cpu.ex b/lib/ex_maxsim_cpu.ex index 3be1f22..1f6ebe0 100644 --- a/lib/ex_maxsim_cpu.ex +++ b/lib/ex_maxsim_cpu.ex @@ -40,6 +40,10 @@ defmodule ExMaxsimCpu do - `RAYON_NUM_THREADS`: Control Rayon parallelism (default: number of CPUs) - `OPENBLAS_NUM_THREADS`: Set to 1 to avoid oversubscription with Rayon + + ## Input Requirements + + - Embeddings must be L2-normalized per token for correct MaxSim scores """ alias ExMaxsimCpu.Nif diff --git a/mix.exs b/mix.exs index db1cf90..364d1d5 100644 --- a/mix.exs +++ b/mix.exs @@ -1,14 +1,14 @@ defmodule ExMaxsimCpu.MixProject do use Mix.Project - @version "0.1.0" + @version "0.2.0" @source_url "https://github.com/mixedbread-ai/maxsim-cpu" def project do [ app: :ex_maxsim_cpu, version: @version, - elixir: "~> 1.14", + elixir: "~> 1.19", start_permanent: Mix.env() == :prod, deps: deps(), description: description(), @@ -30,7 +30,7 @@ defmodule ExMaxsimCpu.MixProject do defp deps do [ {:rustler, "~> 0.37.0"}, - {:nx, "~> 0.7", optional: true}, + {:nx, "~> 0.7"}, {:torchx, "~> 0.7", optional: true}, {:ex_doc, "~> 0.31", only: :dev, runtime: false}, {:benchee, "~> 1.0", only: :dev} diff --git a/native/maxsim_cpu/src/core/simd.rs b/native/maxsim_cpu/src/core/simd.rs index dd431e4..cb0aebb 100644 --- a/native/maxsim_cpu/src/core/simd.rs +++ b/native/maxsim_cpu/src/core/simd.rs @@ -1,18 +1,18 @@ //! SIMD-optimized operations for MaxSim. //! //! Provides platform-specific SIMD implementations: -//! - AVX2 for x86_64 +//! - AVX2 for x86_64 (when enabled at compile time) //! - NEON for aarch64 //! - Scalar fallback for other platforms -#[cfg(target_arch = "x86_64")] +#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] use std::arch::x86_64::*; #[cfg(target_arch = "aarch64")] use std::arch::aarch64::*; /// Find maximum value in a slice using AVX2 SIMD with prefetching. -#[cfg(target_arch = "x86_64")] +#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] #[inline] pub fn simd_max(slice: &[f32]) -> f32 { if slice.len() < 8 { @@ -138,7 +138,10 @@ pub fn simd_max(slice: &[f32]) -> f32 { } /// Scalar fallback for platforms without SIMD support. -#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] +#[cfg(not(any( + all(target_arch = "x86_64", target_feature = "avx2"), + target_arch = "aarch64" +)))] #[inline] pub fn simd_max(slice: &[f32]) -> f32 { slice.iter().copied().fold(f32::NEG_INFINITY, f32::max) diff --git a/native/maxsim_cpu/src/lib.rs b/native/maxsim_cpu/src/lib.rs index 7f27153..3be55c5 100644 --- a/native/maxsim_cpu/src/lib.rs +++ b/native/maxsim_cpu/src/lib.rs @@ -19,6 +19,8 @@ fn load(_env: Env, _info: Term) -> bool { rustler::init!("Elixir.ExMaxsimCpu.Nif", load = load); +type DocData<'a> = Vec<(usize, Cow<'a, [f32]>)>; + /// Convert a binary to &[f32], handling alignment. /// Returns Cow::Borrowed if aligned, Cow::Owned if copy was needed. fn binary_to_f32<'a>(bin: &'a Binary) -> Result, &'static str> { @@ -137,10 +139,10 @@ fn maxsim_scores_nif<'a>( #[rustler::nif(schedule = "DirtyCpu")] fn maxsim_scores_variable_nif<'a>( env: Env<'a>, - query_bin: Binary, + query_bin: Binary<'a>, q_len: usize, dim: usize, - doc_bins: Vec, + doc_bins: Vec>, doc_lens: Vec, ) -> NifResult> { // Wrap in catch_unwind to prevent panics from crashing BEAM @@ -163,7 +165,7 @@ fn maxsim_scores_variable_nif<'a>( let query = binary_to_f32(&query_bin).map_err(|e| e.to_string())?; // Convert document binaries and validate sizes - let doc_data: Result)>, String> = doc_bins + let doc_data: Result, String> = doc_bins .iter() .zip(doc_lens.iter()) .enumerate() @@ -179,8 +181,7 @@ fn maxsim_scores_variable_nif<'a>( } let data = binary_to_f32(bin).map_err(|e| e.to_string())?; - // Convert Cow to owned Vec for lifetime reasons - Ok((len, data.into_owned())) + Ok((len, data)) }) .collect(); @@ -189,7 +190,7 @@ fn maxsim_scores_variable_nif<'a>( // Create slice references for the algorithm let doc_slices: Vec<(usize, &[f32])> = doc_data .iter() - .map(|(len, data)| (*len, data.as_slice())) + .map(|(len, data)| (*len, data.as_ref())) .collect(); // Compute scores