Add Nx CPU benchmarks and refresh docs (#1)

* Add nx benchmarks for cpu

* Bump package version to 0.2.0 and elixir version to 1.19

* Fix clippy bug

Author: goodhamgupta

README.md

@@ -35,14 +35,18 @@ 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
 
   # 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<f32[2]>
 ```
+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
 

assets/benchmark_d_len.png

@@ -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
+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

assets/benchmark_data.csv

@@ -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):")