Add multi-pathway benchmark tests (bare metal, Jupyter, jupyter-book)

Author: mmcky

.github/workflows/ci.yml

@@ -25,9 +25,35 @@ jobs:
           pip install -U "jax[cuda13]"
           pip install numpyro
           python scripts/test-jax-install.py
-      - name: Run Hardware Benchmarks
+      # === Benchmark Tests (Bare Metal, Jupyter, Jupyter-Book) ===
+      - name: Run Hardware Benchmarks (Bare Metal)
         shell: bash -l {0}
-        run: python scripts/benchmark-hardware.py
+        run: |
+          echo "=== Bare Metal Python Script Execution ==="
+          python scripts/benchmark-hardware.py
+      - name: Run Jupyter Notebook Benchmark (via nbconvert)
+        shell: bash -l {0}
+        run: |
+          echo "=== Jupyter Kernel Execution ==="
+          jupyter nbconvert --to notebook --execute scripts/benchmark-jupyter.ipynb --output benchmark-jupyter-executed.ipynb
+          echo "Notebook executed successfully"
+      - name: Run Jupyter-Book Benchmark
+        shell: bash -l {0}
+        run: |
+          echo "=== Jupyter-Book Execution ==="
+          # Build just the benchmark file using jupyter-book
+          mkdir -p benchmark_test
+          cp scripts/benchmark-jupyterbook.md benchmark_test/
+          # Create minimal _config.yml
+          echo "title: Benchmark Test" > benchmark_test/_config.yml
+          echo "execute:" >> benchmark_test/_config.yml
+          echo "  execute_notebooks: force" >> benchmark_test/_config.yml
+          # Create minimal _toc.yml
+          echo "format: jb-book" > benchmark_test/_toc.yml
+          echo "root: benchmark-jupyterbook" >> benchmark_test/_toc.yml
+          # Build
+          jb build benchmark_test --path-output benchmark_build/
+          echo "Jupyter-Book build completed successfully"
       - name: Install latex dependencies
         run: |
           sudo apt-get -qq update

scripts/benchmark-jupyter.ipynb

@@ -0,0 +1,156 @@
+---
+jupytext:
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+kernelspec:
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+---
+
+# JAX Performance Benchmark - Jupyter Book Execution
+
+This file tests JAX performance when executed through Jupyter Book's notebook execution.
+Compare results with direct script and nbconvert execution.
+
+```{code-cell} ipython3
+import time
+import platform
+import os
+
+print("=" * 60)
+print("JUPYTER BOOK EXECUTION BENCHMARK")
+print("=" * 60)
+print(f"Platform: {platform.platform()}")
+print(f"Python: {platform.python_version()}")
+print(f"CPU Count: {os.cpu_count()}")
+```
+
+```{code-cell} ipython3
+# Import JAX and check devices
+import jax
+import jax.numpy as jnp
+
+devices = jax.devices()
+default_backend = jax.default_backend()
+has_gpu = any('cuda' in str(d).lower() or 'gpu' in str(d).lower() for d in devices)
+
+print(f"JAX devices: {devices}")
+print(f"Default backend: {default_backend}")
+print(f"GPU Available: {has_gpu}")
+```
+
+```{code-cell} ipython3
+# Define JIT-compiled function
+@jax.jit
+def matmul(a, b):
+    return jnp.dot(a, b)
+
+print("matmul function defined with @jax.jit")
+```
+
+```{code-cell} ipython3
+# Benchmark 1: Small matrix (1000x1000) - includes JIT compilation
+print("\n" + "=" * 60)
+print("BENCHMARK 1: Small Matrix (1000x1000)")
+print("=" * 60)
+
+n = 1000
+key = jax.random.PRNGKey(0)
+A = jax.random.normal(key, (n, n))
+B = jax.random.normal(key, (n, n))
+
+# Warm-up run (includes compilation)
+start = time.perf_counter()
+C = matmul(A, B).block_until_ready()
+warmup_time = time.perf_counter() - start
+print(f"Warm-up (includes JIT compile): {warmup_time:.3f} seconds")
+
+# Compiled run
+start = time.perf_counter()
+C = matmul(A, B).block_until_ready()
+compiled_time = time.perf_counter() - start
+print(f"Compiled execution: {compiled_time:.3f} seconds")
+```
+
+```{code-cell} ipython3
+# Benchmark 2: Large matrix (3000x3000) - triggers recompilation
+print("\n" + "=" * 60)
+print("BENCHMARK 2: Large Matrix (3000x3000)")
+print("=" * 60)
+
+n = 3000
+A = jax.random.normal(key, (n, n))
+B = jax.random.normal(key, (n, n))
+
+# Warm-up run (recompilation for new size)
+start = time.perf_counter()
+C = matmul(A, B).block_until_ready()
+warmup_time = time.perf_counter() - start
+print(f"Warm-up (recompile for new size): {warmup_time:.3f} seconds")
+
+# Compiled run
+start = time.perf_counter()
+C = matmul(A, B).block_until_ready()
+compiled_time = time.perf_counter() - start
+print(f"Compiled execution: {compiled_time:.3f} seconds")
+```
+
+```{code-cell} ipython3
+# Benchmark 3: Element-wise operations (50M elements)
+print("\n" + "=" * 60)
+print("BENCHMARK 3: Element-wise Operations (50M elements)")
+print("=" * 60)
+
+@jax.jit
+def elementwise_ops(x):
+    return jnp.cos(x**2) + jnp.sin(x)
+
+x = jax.random.normal(key, (50_000_000,))
+
+# Warm-up
+start = time.perf_counter()
+y = elementwise_ops(x).block_until_ready()
+warmup_time = time.perf_counter() - start
+print(f"Warm-up (includes JIT compile): {warmup_time:.3f} seconds")
+
+# Compiled
+start = time.perf_counter()
+y = elementwise_ops(x).block_until_ready()
+compiled_time = time.perf_counter() - start
+print(f"Compiled execution: {compiled_time:.3f} seconds")
+```
+
+```{code-cell} ipython3
+# Benchmark 4: Multiple small operations (simulates lecture cells)
+print("\n" + "=" * 60)
+print("BENCHMARK 4: Multiple Small Operations (lecture simulation)")
+print("=" * 60)
+
+total_start = time.perf_counter()
+
+# Simulate multiple cell executions with different operations
+for i, size in enumerate([100, 500, 1000, 2000, 3000]):
+    @jax.jit
+    def compute(a, b):
+        return jnp.dot(a, b) + jnp.sum(a)
+    
+    A = jax.random.normal(key, (size, size))
+    B = jax.random.normal(key, (size, size))
+    
+    start = time.perf_counter()
+    result = compute(A, B).block_until_ready()
+    elapsed = time.perf_counter() - start
+    print(f"  Size {size}x{size}: {elapsed:.3f} seconds")
+
+total_time = time.perf_counter() - total_start
+print(f"\nTotal time for all operations: {total_time:.3f} seconds")
+```
+
+```{code-cell} ipython3
+print("\n" + "=" * 60)
+print("JUPYTER BOOK EXECUTION BENCHMARK COMPLETE")
+print("=" * 60)
+```