Skip to content

Commit 9eec803

Browse files
committed
init
1 parent a309b25 commit 9eec803

64 files changed

Lines changed: 4150 additions & 0 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
.DS_Store
2+
*.csv
3+
*.json
14
# Byte-compiled / optimized / DLL files
25
__pycache__/
36
*.py[codz]

README.md

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# Proposal
2+
3+
Breaking the Speed Limit: Fast Statistical Models with Python 3.14, Numba, and JAX
4+
5+
## Section Type
6+
7+
30 minute talk
8+
9+
## Abstract
10+
11+
Data scientists and domain experts often face a dilemma: we understand the models, and we use Python, but we aren't C++ or Rust engineers. We need code that is quick to write, easy to work with, and still fast enough to run on large, real‑world datasets. How do we choose the right tool without getting lost in low‑level details?
12+
13+
With the new free‑threaded build and experimental JIT in Python 3.14, combined with tools like Numba and JAX, we finally have a realistic way to push back against the "Python is slow" stereotype. In this talk, we'll use two concrete workloads to illustrate this modern stack: complex iterative loops (k‑means) and massive data parallelism (permutation test). The focus is on computational patterns rather than statistical theory.
14+
15+
We'll compare plain NumPy, Python 3.14 (with free-threaded or JIT configurations), Numba, and JAX across varying data scales, highlighting the trade-offs in runtime, memory, debuggability, and developer experience. Along the way, we'll also demonstrate how AI coding tools can serve as a copilot, helping to translate clear mathematical code into high-performance kernels without requiring deep compiler expertise.
16+
17+
-----
18+
19+
20+
# Outline
21+
22+
## Introduction (3 minutes)
23+
24+
* **Background and motivation:** The tension between development speed, execution speed, and safety/debuggability: we want code that is fast to write, fast to run, and still easy to reason about.
25+
* **Evaluation matrix:**
26+
27+
* **Runtime:** Behavior across different data sizes.
28+
* **Memory:** Extra temporaries and **copies vs shared data** in parallel code.
29+
* **Debuggability:** How easy it is to troubleshoot or handle numerical issues.
30+
* **Developer effort:** How far you have to move away from "plain NumPy".
31+
* **Two compute patterns** (no background in statistics required):
32+
33+
* **k‑means for the iterative pattern:** A loop‑heavy, sequential algorithm where each step depends on the previous step. Good for testing how different runtimes handle *tight Python loops* over arrays.
34+
* **Permutation Test for the parallel pattern:** A "run the same function thousands of times on different resamples" scenario. Embarrassingly parallel computation with a lot of big array math.
35+
* **Data scales:**
36+
37+
* **Toy datasets:** tens of thousands of points (great for illustration and debugging, runs on a laptop).
38+
* **Large-scale datasets:** tens of millions of data points or gene expression matrices (typical research scale, usually needs a server).
39+
40+
41+
## Loop‑heavy k‑means (12 minutes)
42+
43+
* **Baseline:** k‑means clustering implemented in plain NumPy, highlighting the iterative pattern: assign points to clusters, then update centroids.
44+
* **Numba:** Add `@njit` to the assignment and centroid‑update functions, fix unsupported features, and show how compiling the inner loops changes performance and code structure.
45+
* **Python 3.14 with JIT (GIL build):** Take a more loop‑heavy k‑means variant using explicit Python `for` loops and run it on a JIT‑enabled CPython 3.14 build. Demonstrate that the JIT accelerates native Python control flow but has minimal impact on the NumPy‑heavy baseline, where most work is already in C.
46+
* **JAX:** Rewrite k‑means in a functional style using JAX arrays, `jax.lax.scan` for the refinement steps, and `jax.jit` for compilation.
47+
* **Discussion of trade‑offs:**
48+
* Show where the CPython JIT helps and where it doesn’t.
49+
* Compare speedups and costs: type restrictions, debugging experience, and how much each approach diverges from the simple NumPy baseline.
50+
51+
52+
## Parallel Permutation Test (10 minutes)
53+
54+
* **Baseline:** Implement a permutation test using `multiprocessing` (e.g. `ProcessPoolExecutor` or `multiprocessing.Pool`) on a standard GIL build. Highlight how large arrays are serialized and effectively copied into multiple processes.
55+
* **Free‑threaded Python 3.14:** Switch to `concurrent.futures.ThreadPoolExecutor` on the free‑threaded (no‑GIL) CPython 3.14 build, where all threads share a single in‑process copy of the data while running permutations in parallel.
56+
* **Numba:** Package the permutation test into a compiled kernel using `@njit(parallel=True)` and `prange` to parallelize across resamples. Discuss how Numba's own thread pool gives parallel speedups while keeping everything in one process and one shared array, independent of whether CPython is free‑threaded.
57+
* **JAX:** Express the permutation test in JAX and use `jax.vmap` (with `jax.jit`) to vectorize over permutations and batch them into a single compiled computation, optionally on a GPU/TPU.
58+
* **Discussion of trade‑offs:**
59+
60+
* Show how **naive multiprocessing** tends to create multiple large copies of the data, while free‑threaded threads, Numba, and JAX all operate on **one shared dataset** (in RAM or on device).
61+
* Compare performance and memory behavior between these approaches, and discuss where each one wins (e.g., simplicity vs maximum speedup).
62+
* Briefly touch on randomization and reproducibility: how random initialization and parallelism can change the order of operations and random draws, and simple strategies to solve this.
63+
* Compare developer effort and style shifts.
64+
65+
66+
## Developer Experience (2 minutes)
67+
68+
Practical tips for working with these tools on numerical code:
69+
70+
* **Start simple:** Build and test your model in plain NumPy (or NumPy plus a small Numba kernel) to validate correctness before chasing performance.
71+
* **Profile first:** Use basic profiling to identify hotspots and verify that JIT/Numba/JAX changes actually improve the end‑to‑end runtime and memory behavior.
72+
* **AI as copilot, not driver:** Use AI tools to help translate statistical formulas to NumPy code and further Numba or JAX variants, but keep human control over correctness, numerical stability, and performance checks.
73+
74+
75+
## Conclusion (1 minute)
76+
77+
Wrap up with a compact "take‑home" decision guide:
78+
79+
* When a **small refactor + Numba** is the best trade‑off for loop‑heavy numerical work.
80+
* When the **CPython 3.14 JIT** is "good enough" as a low‑effort speedup for Python‑heavy glue code, and when the **free‑threaded build** is worth using to replace `multiprocessing` with threads for parallel workloads.
81+
* When it's worth adopting JAX's functional style for **maximum speed and scalability**, especially with accelerators.
82+
* How AI tools can accelerate refactoring and exploration, but should not replace human judgment.
83+
84+
The goal is not to crown a single winner, but to give an honest picture of **what you gain for each extra unit of complexity**, so attendees can choose the right toolchain for their own numerical workloads.
85+
86+
-------
87+
88+
# Notes
89+
90+
## What motivates us to submit this proposal
91+
92+
We are PhD candidates in Biostatistics working with high‑dimensional genomic data. This talk is based on our experience with statistical modeling in Python, particularly in the context of large‑scale genomic data analysis, where runtime and memory efficiency are crucial.
93+
94+
Unlike software engineers, statisticians prioritize **correctness, reproducibility, and ease of collaboration** over raw performance. However, as datasets grow larger, the need for efficient computation becomes unavoidable. Historically, the answer was "rewrite the hot path in C/C++". This is the current situation for most statisticians in our research area: they rewrite their R code using Rcpp, which not only requires a deep knowledge of C++ but also results in a fragile package dependency. Is there a better way?
95+
96+
With the advent of Python 3.14's free‑threaded build and experimental JIT, along with mature tools like Numba and JAX, we now have a more accessible path to high performance without leaving Python. This talk aims to share practical insights and benchmarks that help statisticians and data scientists make informed decisions about optimizing their Python code.
97+
98+
In our research group, we have successfully applied these techniques to accelerate various statistical models, including EM algorithms and moment-based estimators, demonstrating significant speedups while maintaining code clarity and reproducibility. We believe that our experiences can provide valuable guidance to others facing similar challenges in the Python ecosystem. Thus, we are motivated to share our findings and help the broader community navigate the evolving landscape of high-performance Python for statistical computing.
99+
100+
While the examples come from Biostatistics, the core ideas are general to anyone doing numerical Python. We focus on **runtime behavior, memory layout, and developer workflow**, not on the underlying statistical theory. Developers of numerical computing packages can also benefit from this talk by gaining insight into the user's perspective on performance trade-offs in Python. Therefore, we believe this talk will resonate with a wide audience in the Python community.
101+
102+
## Previous Experience
103+
104+
We have given talks at PyCon HK 2025, focusing on functional programming in Python and new features in recent Python releases (especially Python 3.14). Although the YouTube video is not available yet, here are our slides for reference:
105+
106+
1. [PyCon HK 2025: Functional Programming in Python](https://lucajiang.github.io/functional_python/)
107+
2. [PyCon HK 2025: Shall We Upgrade? Navigating Python's Rapid Evolution](https://lucajiang.github.io/new_in_python/)
108+
109+
---
110+
111+
## Repository layout
112+
113+
- **`docs/`** — Reference notes (Python 3.14, Numba, JAX, k-means, permutation tests, benchmarking, related talks).
114+
- **`experiments/`** — Runnable benchmarks: `setup/` (deps), `kmeans/`, `permutation_test/`, `devex/`, `visualization/`, `results/`. See each subdirectory `README.md`.
115+
- **`slides/`** — Reveal.js HTML deck for the talk ([`slides/index.html`](slides/index.html)); see [`slides/README.md`](slides/README.md).

docs/01-python314-features.md

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
# Python 3.14:Free-threaded 与实验性 JIT
2+
3+
本文档汇总与演讲「Breaking the Speed Limit」相关的 **Python 3.14** 运行时特性。不仅提供结论,更深入探讨其底层设计机制,为现场回答高级 Python 工程师的技术提问提供充足弹药。
4+
5+
## 1. Free-threaded 构建(PEP 703 / PEP 779)
6+
7+
### 1.1 核心机制:移除 GIL 带来的底层重构
8+
9+
- **内存分配器(mimalloc)**:为了在多线程环境下保证无锁的内存分配效率,Free-threaded 构建默认集成了 `mimalloc`。由于去除了 GIL,多个线程会并发分配/释放内存,如果继续使用旧版的 `pymalloc` 将会导致严重的锁争用。
10+
- **偏置引用计数(Biased Reference Counting, BRC)与延迟引用计数(Deferred Reference Counting)**
11+
- 传统 CPython 每个对象的赋值/销毁都会原子的增减 refcount,这在无 GIL 时会造成可怕的 CPU 缓存行伪共享(False Sharing)。
12+
- BRC 允许将引用计数偏向"创建对象的那个线程",通过本地线程进行操作,只有跨线程访问时才使用原子的原子操作(atomic instructions)。
13+
- 对于常驻对象(如单例、内置类型或顶层函数),采用延迟引用计数或不追踪(Immortalization, PEP 683),进一步降低锁争用。
14+
15+
### 1.2 行为与性能代价
16+
17+
- 在 free-threaded 模式下,**单线程**代码相对传统 GIL 构建通常有 **约 5–10%** 的性能损失。原因在于即便有 BRC,依然会有不可避免的轻量级锁检查和内存屏障(Memory Barriers)开销。
18+
- 本地自适应解释器(PEP 659 的 Specializing Adaptive Interpreter)在无 GIL 环境下面临并发覆写字节码的问题。在 Python 3.13/3.14 中,这需要精细的锁或 RCU (Read-Copy-Update) 机制来保证类型专化的线程安全。
19+
20+
### 1.3 实验中的表现(在 GIL 构建上的天花板)
21+
22+
本仓库的置换检验实验在标准 GIL 构建(Python 3.11.6, macOS 15.1, Apple Silicon 8 核心)上运行。即便在 GIL 下,`ThreadPoolExecutor` 仍然取得了 **~1.4× 加速**(n=10k, R=10000,详见 [`perm_scaling.png`](../experiments/results/v2/perm_scaling.png))。原因:NumPy 的 `.sum()``.permutation()` 在 C 层释放 GIL,使线程可以重叠。**这就是 GIL 构建上线程加速的上限。** 在 free-threaded 构建上,同一段 Python 代码应该继续向 8 核心靠拢——这是演讲最具说服力的「升级即收益」论点。
23+
24+
---
25+
26+
## 2. 实验性 JIT(PEP 744 / PEP 774)
27+
28+
### 2.1 Copy-and-Patch 技术解析
29+
30+
Python 3.14 搭载的实验性 JIT 并非传统的追踪式 JIT(如 PyPy 的 Tracing JIT)或基于完整 LLVM 编译的 JIT,而是基于 **Copy-and-Patch** 技术的模板 JIT(Template JIT)。
31+
32+
- **Stencil(模板生成)**:在 CPython 的构建阶段(而非运行阶段),使用 LLVM 对 C 语言写好的解释器指令(Opcode)进行编译,提取出机器码模板(Stencils)。这正是 PEP 774 的核心——不再要求用户的运行环境有 LLVM,而是直接在 CPython 源码树中预编译机器码片段。
33+
- **运行时 Patch**:当一段 Python 代码变"热"(Hot)时,JIT 引擎只需把预设的机器码模板像"拼图"一样拷贝到内存可执行页中,把函数指针和偏移量作为参数(Patch)填入即可。
34+
- **优势**:编译开销极低(因为只是内存拷贝和重定位),预热极快,完全不需要运行时进行复杂的寄存器分配和指令选择。
35+
- **劣势**:因为 stencil 在构建时生成,**无法进行跨 opcode 的优化(例如 loop hoisting、公共子表达式消除)**,因此对于数值代码,与 Numba / JAX 的完整 LLVM 编译差距依然显著。
36+
37+
### 2.2 演讲实验的启示:什么能被加速?
38+
39+
- JIT 对于 **纯 Python 密集循环**、条件分支(控制流)有着显著加速。
40+
- **NumPy/C 扩展盲区**:如果在 CPython JIT 下运行高度 NumPy 向量化的代码,JIT **毫无作用**。因为运行时大部分时间在 `numpy` 的 C 库中。
41+
- 在我们的 k-means 实验中,`kmeans_loops.py`(N=2000, d=10, k=5)在标准 CPython 3.11 下单次完整跑需 **0.84 s**;对照 Numba 版本的 **0.010 s**(N=100k,80 倍大)。这一巨大差距正是 Python 循环解释开销的量化,也是 CPython JIT 值得去优化的地方——但即便 3.14 JIT 把纯循环加速 3×,距 Numba 也还有一个数量级。**这是向观众传达的真实界限**:JIT 加速的是 Python 解释器层面的开销,而不是改变算法的渐进复杂度或 C 核心的执行速度。
42+
43+
---
44+
45+
## 3. 陷阱:JIT 与 Free-threaded 的正交与互斥
46+
47+
**Python 3.14** 的实验性阶段,你需要向观众诚实指出:
48+
49+
- **现阶段,JIT 与 Free-threading 是互相独立的,甚至部分版本构建存在排斥。**
50+
- 并发修改字节码(JIT 编译时替换指令)在 Free-threaded 下需要非常复杂的同步原语,因此很多时候我们是在 **单线程 JIT****多线程无 JIT** 之间做选择。
51+
- 演讲建议:将两者解耦讲解。K-means 用于讲 JIT 带来的"原生 Python"收益;Permutation Test 用于讲 Free-threaded 解决的多进程共享内存痛点。
52+
- 验证方式:`python -VV``sys._is_gil_enabled()``sys.flags.jit`。3.15+ 可能会放宽此约束。
53+
54+
---
55+
56+
## 4. Free-threaded 与科学计算栈的真实交互
57+
58+
在 Permutation Test 的实验中,我们会发现 Free-threaded 版本的线程池性能非常优异,但有两个容易踩坑的点:
59+
60+
1. **Thread-Safety 并不免费**:即便释放了 GIL,如果在 Python 循环里依然对全局状态(如往 `list.append` 结果)频繁访问,会导致 CPython 内部的微观锁争用,从而拖慢速度。正确的做法是:**为每个线程分配独立的写入 Buffer,最后统一 merge**(本仓库 [`permtest_freethreaded.py`](../experiments/permutation_test/permtest_freethreaded.py) 即按此设计)。
61+
2. **NumPy 的内部锁**:NumPy 在执行向量化运算时,内部的内存分配也可能遭遇 C 层面的锁。好在现代 NumPy 和底层 BLAS(如 OpenBLAS / MKL)对多线程有着不同程度的支持和隔离。
62+
63+
---
64+
65+
## 5. 专家视角结论
66+
67+
在准备这篇演讲时,应该让观众明确:Python 3.14 的到来并不意味着我们可以盲目写纯 Python 循环。它降低了"偶尔写出慢代码"的惩罚,并为构建高性能的多线程 Python 库(如在单进程内共享巨大矩阵的统计工具)提供了基础设施。
68+
69+
**实测数据要点(Apple Silicon 8 核, Python 3.11, NumPy 1.24, Numba 0.57, JAX 0.4)**
70+
71+
- 即使在 GIL 上,`ThreadPoolExecutor` 在 NumPy-heavy 置换检验里仍能取得 1.4× 加速(见 [`perm_scaling.png`](../experiments/results/v2/perm_scaling.png))。无 GIL 构建应进一步扩大此差距。
72+
- `multiprocessing` 在 R=10000 时的子进程 RSS 加总 **757 MB**(8 worker × 10k-length float64 array + 解释器基础内存),对照同规模线程池的 **~1.4 MB**[`perm_memory.png`](../experiments/results/v2/perm_memory.png) 把这一差距做成了一张直观的对数刻度条形图。
73+
- 在 CPU 上运行基于 `jax.vmap` 的置换检验(R=10000)耗时 **~75 s**——比纯 NumPy 慢 37×。JAX 在加速器上才能发挥,它在 CPU 上并非默认选项。

0 commit comments

Comments
 (0)