CPython 效能將提升 5 倍?faster-python 專案 PEP 659 原始碼級解讀

語言: CN / TW / HK

在 2021 年早些時候,Python 作者 Guido van Rossum 被微軟返聘繼續進行 CPython 相關工作,他們提出了一個 faster-python 計劃,計劃在 4 年內將 CPython 的效能提升 5 倍,整個專案被開放在 GitHub 的 faster-cpython Group,通過 Activity 可見該專案的一部分 ideas 已經有了相應的程式碼實現和驗證。

本文將就該專案中的一個重要提案 PEP 659 進行解讀和原始碼分析,並從中學習在 bytecode level 進行虛擬機器效能優化的思路和手段,希望能對大家有所啟發。

提案解讀

PEP 659 創建於 2021 年 4 月,全稱為 Specializing Adaptive Interpreter,這裡有兩個關鍵詞:Specializing 和 Adaptive,這裡可以簡單理解為對 特定位置 的程式碼進行適配(Adaptive), 替換 為特殊的程式碼(Specializing)從而提高特定位置操作的執行速度。比如通過觀察發現某個查詢 dict 的程式碼在多次執行過程中 dict 沒有變動,那麼我們可以針對這段程式碼進行優化,將 dict entry 的 index 直接快取起來,這樣在下次查詢時就避免了 hashtable 查詢的過程從而提高效能,這裡的觀察就對應到 Adaptive,替換程式碼的過程則對應到 Specializing。

上面的例子並不準確,只是幫助大家對 Specializing Adaptive Interpreter 有一個初步的印象,下面我們將摘錄提案中的關鍵語句進行解讀。

首先要明確的一點是,PEP 659 並不是一個 JIT 方案,因為它的初衷在於讓那些無法直接使用 PyPy 等包含 JIT Compiler 的使用者也能享受到 faster CPython 的紅利。例如在 iOS 平臺下,使用者程序受限於 codesign 動態建立的可執行的內碼表在缺頁中斷時會因為未包含合法簽名而被拒絕,因此無法直接使用包含 JIT Compiler 的 Python 虛擬機器。

看到這裡可能有些人會擔心,不使用 JIT 單純從虛擬機器層面進行優化的空間和收益如何呢?在 PEP 659 中作者也給出了一些解釋:

Specialization is typically done in the context of a JIT compiler, but research shows specialization in an interpreter can boost performance significantly, even outperforming a naive compiler.

即研究發現僅僅在 Interpreter 層面進行 Specialization 優化也可以獲得顯著的效能提升,效能收益甚至可以超過一些初級的 JIT 方案,作者在這裡還引用了一篇自己之前的論文,感興趣的同學可以自行去 PEP 659 提案的參考文獻部分檢視。

到這裡我們也就明確了 PEP 659 不包含 JIT Compiler,簡單地說就是它不生成程式碼,它只是程式碼的搬運工,我們需要窮舉所有可能的優化情況,並且提前準備好程式碼,在觀察到匹配的優化條件時將位元組碼進行替換,當發現不滿足優化條件時還必須能夠優雅的退回到優化前的程式碼以保證程式的正確性。

為了能更好的窮舉優化情況和切換程式碼,需要選擇合適的優化粒度,提案原文是:

By using adaptive and speculative specialization at the granularity of individual virtual machine instructions, we get a faster interpreter that also generates profiling information for more sophisticated optimizations in the future.

即在虛擬機器指令層面進行優化,而不是像 JIT 那樣在一個區域或者函式維度進行 優化,這樣可以針對特定指令進行細分,例如在 CPython 中獲取 globals 和 builtins 都是通過 LOAD_GLOBAL 指令,首先在 globals 中查詢,查詢失敗後再 fallback 到 builtins 中查詢,在這裡可能的情況只有 2 種,因此我們可以為虛擬機器新增兩條指令 LOAD_GLOBAL_MODULE 和 LOAD_GLOBAL_BUILTIN,當發現某段位元組碼中的 LOAD_GLOBAL 一直在查詢 globals 時,我們可以將其優化為前者,反之優化為後者,同時可以對 globals 和 builtins dict 的 entry index 進行快取避免重複訪問 dict 的 hashtable,當發現不滿足優化條件(例如查詢失敗,或是 dict 被修改)時再回滾到 LOAD_GLOBAL 指令保證程式的正確性。

上述從 LOAD_GLOBAL 到 LOAD_GLOBAL_MODULE / LOAD_GLOBAL_BUILTIN 的過程實際上就是 PEP 標題中的 Specializing,而選擇將指令替換為 LOAD_GLOBAL_MODULE 還是 LOAD_GLOBAL_BUILTIN 的過程其實就是 Adaptive,它的職責是觀察特定程式碼中的指令的執行情況,以為其選擇正確的優化指令,觀察的過程也是虛擬機器程式碼執行的過程,因此在這裡還需要額外引入一個 Adaptive 指令 LOAD_GLOBAL_ADAPATIVE 用來執行觀察和替換邏輯。

雖然說 Specializing 能夠通過減少判斷和增加快取來提速,但從原始指令到 Adaptive 指令,從 Adaptive 觀察得出 Specializing 指令的過程也是有損耗的,因此需要基於一定的策略進行優化,而不是無腦的嘗試優化所有程式碼中的指令,原文中提到:

Typical optimizations for virtual machines are expensive, so a long "warm up" time is required to gain confidence that the cost of optimization is justified.

即對虛擬機器的 Specializing 優化是昂貴的,這種昂貴體現在我們需要在位元組碼執行的大迴圈中插入優化程式碼,甚至在每一次迴圈都需要額外的處理邏輯,而很多程式碼可能都只會執行一次,如果對他們進行優化純屬浪費時間,因此我們需要找出那些需要優化的程式碼,這樣程式碼的一個主要特點是被高頻呼叫,我們可以對每個 PyCodeObject (CPython 中儲存位元組碼和環境資訊的物件) 增加計數器,只有執行次數超過某個閾值才執行優化邏輯,這個過程被稱之為 warm up。

在虛擬機器層面,當某個位元組碼物件 PyCodeObject 被執行(可以簡單理解為一段 Python 程式碼對應的位元組碼被執行)或是發生絕對跳轉時,程式碼物件的 co_warmup 計數器會被累加,當達到閾值後這段位元組碼中所有可優化指令都會被替換為 Adaptive 指令進行觀察和特殊化,此後這些指令只會在 Adaptive 和 Specializing 之間轉換,當 Specializing 指令的優化條件被破壞時,我們不會立即回滾到 Adaptive 指令,而是允許一定次數的 miss 以防止出現優化和去優化之間的顛簸,同樣的當去優化回滾到 Adaptive 指令後我們也會暫停觀察和優化邏輯,讓指令按照原始邏輯執行一段時間,這個過程被稱為 deferred,整個過程的狀態圖如下:

到這裡我們已經對 PEP 659 的原理和大致工作過程有了瞭解,但為了高效能的實現這個優化器還有很多細節需要探討,例如哪些指令需要優化,如何優雅的替換指令和還原,如何設計指令快取等,這些具體問題單純從理論層面分析太過枯燥和晦澀,因此接下來我們將結合 Python 3.11 中 PEP 659 的原始碼實現來進行講解。

原始碼分析

實際上 PEP 659 的基礎設施和部分優化指令已經在 CPython 3.11 的分支上有所體現,我們這裡以 LOAD_GLOBAL 改造為例對上述流程進行細緻分析,之所以選擇 LOAD_GLOBAL 是因為它在虛擬機器層面的操作較為簡單,只涉及到兩個名稱空間的查詢,優化和去優化判斷也比較直接,足夠清晰的說明問題又不會因為指令的晦澀給大家的閱讀帶來困難。

整個優化過程主要包括 Warmup, Adaptive, Specializing & Deoptimize 這幾個階段,下面我們將就各階段的功能和核心程式碼做一些分析和講解。

Warmup

如上文所述,warmup 解決的問題是找到真正高頻執行的程式碼,避免優化那些此後再也不會被執行到的程式碼,這裡我們不需要計算程式碼執行的頻率,只需要設定一個閾值,並對特定的位元組碼物件 PyCodeObject 的執行次數計數,當達到閾值時認為完成了 warmup,因此在這裡我們為 PyCodeObject 引入了一個新欄位 co_warmup:

/* Bytecode object */
struct PyCodeObject {
PyObject_HEAD
// The hottest fields (in the eval loop) are grouped here at the top.
PyObject *co_consts; /* list (constants used) */
PyObject *co_names; /* list of strings (names used) */
// ...
int co_warmup; /* Warmup counter for quickening */
// ...


/* Quickened instructions and cache, or NULL
This should be treated as opaque by all code except the specializer and
interpreter. */
union _cache_or_instruction *co_quickened;
};

在 PyCodeObject 物件被建立時,co_warmup 的初始值被設定為 QUICKENING_WARMUP_DELAY,它是一個負值,每當 PyCodeObject 被執行或是在程式碼內發生絕對跳轉時 co_warmup 會 +1,當達到閾值 0 後則進入優化邏輯:

#define QUICKENING_WARMUP_DELAY 8
/* We want to compare to zero for efficiency, so we offset values accordingly */
#define QUICKENING_INITIAL_WARMUP_VALUE (-QUICKENING_WARMUP_DELAY)


static void
init_code(PyCodeObject *co, struct _PyCodeConstructor *con) {
// ...
co->co_warmup = QUICKENING_INITIAL_WARMUP_VALUE;
co->co_quickened = NULL;
}

Adaptive

當 co_warmup 被累加到 0 後,會走到 _Py_Quicken 函式執行優化邏輯,由於涉及到位元組碼調整,官方在這裡為原始位元組碼 co_code 拷貝了一份副本存放在 quickened 中,此後的所有修改都發生在 quickened 中:

/* Bytecode object */
struct PyCodeObject {
PyObject_HEAD
// The hottest fields (in the eval loop) are grouped here at the top.
PyObject *co_consts; /* list (constants used) */
PyObject *co_names; /* list of strings (names used) */
_Py_CODEUNIT *co_firstinstr; /* Pointer to first instruction, used for quickening. */
// ...
PyObject *co_code; /* instruction opcodes */

/* Quickened instructions and cache, or NULL
This should be treated as opaque by all code except the specializer and
interpreter. */
union _cache_or_instruction *co_quickened;
};


typedef uint16_t _Py_CODEUNIT;


int
_Py_Quicken(PyCodeObject *code) {
if (code->co_quickened) {
return 0;
}
Py_ssize_t size = PyBytes_GET_SIZE(code->co_code);
int instr_count = (int)(size/sizeof(_Py_CODEUNIT));
if (instr_count > MAX_SIZE_TO_QUICKEN) {
code->co_warmup = QUICKENING_WARMUP_COLDEST;
return 0;
}
int entry_count = entries_needed(code->co_firstinstr, instr_count);
SpecializedCacheOrInstruction *quickened = allocate(entry_count, instr_count);
if (quickened == NULL) {
return -1;
}
_Py_CODEUNIT *new_instructions = first_instruction(quickened);
memcpy(new_instructions, code->co_firstinstr, size);
optimize(quickened, instr_count);
code->co_quickened = quickened;
code->co_firstinstr = new_instructions;
return 0;
}

看到上面的程式碼你可能會疑惑為什麼 quickened 包含了除去位元組碼以外的額外空間,實際上僅僅通過指令 替換是難以實現優化的,我們還需要對指令操作的物件儘可能進行快取,以 LOAD_GLOBAL 為例,最壞情況我們需要查 2 次字典,第一次查 globals,第二次查 builtins,在 CPython 中字典底層是 hashtable,因此在發生 hash 碰撞時它的複雜度是大於 O(1) 的。為了優化變動不頻繁的字典,我們可以將 key 對應的 hashtable index 進行快取,顯然這些快取和指令是一一對應的,因此 quickened 的額外空間的作用就是儲存優化後指令的額外資訊。

quickened 是一個雙向陣列,左邊存放是 cache,右邊存放位元組碼,在陣列的最左端額外分配了一個 cache entry 用來儲存 cache count,通過 cache count 我們可以快速定位到位元組碼陣列,前面看到的 first_instruction 就是通過 cache_count 從 quickened 定位到 instr 0 的:

/* We layout the quickened data as a bi-directional array:
* Instructions upwards, cache entries downwards.
* first_instr is aligned to a SpecializedCacheEntry.
* The nth instruction is located at first_instr[n]
* The nth cache is located at ((SpecializedCacheEntry *)first_instr)[-1-n]
* The first (index 0) cache entry is reserved for the count, to enable finding
* the first instruction from the base pointer.
* The cache_count argument must include space for the count.
* We use the SpecializedCacheOrInstruction union to refer to the data
* to avoid type punning.


Layout of quickened data, each line 8 bytes for M cache entries and N instructions:


<cache_count> <---- co->co_quickened
<cache M-1>
<cache M-2>
...
<cache 0>
<instr 0> <instr 1> <instr 2> <instr 3> <--- co->co_first_instr
<instr 4> <instr 5> <instr 6> <instr 7>
...
<instr N-1>
*/

每條被優化的指令都需要獨立的 cache entries,我們需要一種手段 O(1) 的從指令索引到 cache,在 Python 3 中每條指令包括 8 bit 的 opcode 和 8 bit 的 operand,官方選擇使用 offset + operand 構建二級索引,這裡的 offset 用來確定索引的塊範圍(這裡有點像頁表搜尋,offset 代表 PAGE,operand 代表 PAGEOFF), operand 用來修正 offset,被索引覆蓋的 operand 則存放在 cache 中,上述設計在 PEP 659 中的原文如下:

For instructions that need specialization data, the operand in the quickened array will serve as a partial index, along with the offset of the instruction, to find the first specialization data entry for that instruction.

對於 LOAD_GLOBAL 而言,需要快取的主要是 dict 的版本資訊以及 key index,此外還有一些優化器所需的額外資訊,具體如下:

  1. 字典鍵值對的版本 dk_version,由於這裡同時涉及到 globals 和 builtins,我們需要快取兩個 dk_version,每個 dk_version 都是一個 uint32_t,所以加起來 dk_version 需要一個 8B,也就是一個 cache entry;

  1. key 對應的 index,它是一個 uint16_t,由於我們已經佔用了一個 cache entry,這裡只能再使用一個額外的 cache entry;

  1. 由於原來的 operand 被用來做 partial index,我們還需要額外的一個 uint8_t 來儲存 original_oparg,這個可以和 2 中的 uint16_t 合併儲存;

  1. 在提案解讀部分我們提到在 Adaptive, Specilization 和 Deoptimize 之間轉換需要一個計數器緩衝來避免顛簸,需要一個 counter,官方在這裡使用了一個 uint8_t。

綜合上述分析,LOAD_GLOBAL 需要兩個 cache entry,第一個儲存 orignal_oparg + counter + index,第二個儲存 globals 和 builtins 的 dk_version:

typedef struct {
uint8_t original_oparg;
uint8_t counter;
uint16_t index;
} _PyAdaptiveEntry;


typedef struct {
uint32_t module_keys_version;
uint32_t builtin_keys_version;
} _PyLoadGlobalCache;

分析到這裡其實實施 Adaptive 策略已經水到渠成了,對於 LOAD_GLOBAL 而言無外乎查詢 globals 和 builtins,當某個 LOAD_GLOBAL 在 Adaptive 邏輯中命中 globals 時,我們將其優化為 LOAD_GLOBAL_MODULE,快取好 index 和 module_keys_version 即可;當某個 LOAD_GLOBAL 沒有命中 globals 時,我們需要同時快取 module_keys_version 和 builtin_keys_version,這是因為當 globals 發生變化後可能導致下次 LOAD_GLOBAL 不再命中 builtins,在這裡我們將其優化為 LOAD_GLOBAL_BUILTIN,這個優化選擇和快取的過程實際上就是 Adaptive。

Specializing & Deoptimize

如上文所述,位元組碼物件 PyCodeObject 在經過了 Warmup 和 Adaptive 過程之後其中可優化指令均已變成了 Specialized 形式,例如 LOAD_GLOBAL 已經全部以 LOAD_GLOBAL_MODULE 或 LOAD_GLOBAL_BUILTIN 的形式存在(這裡先不考慮 Deoptimize),用大白話說就是位元組碼中的指令完成了對當前環境的最優適配,下面我們來看看這些適配的程式碼是怎樣的。

其實經過上面的分析,相信大家已經把適配程式碼猜個八九不離十了,下面我們以較為複雜的 LOAD_GLOBAL_BUILTIN 為例分析一下原始碼:

TARGET(LOAD_GLOBAL_BUILTIN) {
// ...
PyDictObject *mdict = (PyDictObject *)GLOBALS();
PyDictObject *bdict = (PyDictObject *)BUILTINS();
SpecializedCacheEntry *caches = GET_CACHE();
_PyAdaptiveEntry *cache0 = &caches[0].adaptive;
_PyLoadGlobalCache *cache1 = &caches[-1].load_global;
DEOPT_IF(mdict->ma_keys->dk_version != cache1->module_keys_version, LOAD_GLOBAL);
DEOPT_IF(bdict->ma_keys->dk_version != cache1->builtin_keys_version, LOAD_GLOBAL);
PyDictKeyEntry *ep = DK_ENTRIES(bdict->ma_keys) + cache0->index;
PyObject *res = ep->me_value;
DEOPT_IF(res == NULL, LOAD_GLOBAL);
STAT_INC(LOAD_GLOBAL, hit);
Py_INCREF(res);
PUSH(res);
DISPATCH();
}

我們先忽略 DEOPT_IF 看一下主邏輯,首先取出當前命令的 cache entries,第一個 entry 記錄了 index,第二個 entry 記錄了 globals 和 builtins 的 dk_version,在快取命中的情況下我們只需要從 builtins dict 的 hashtable[index] 取出鍵值對,將它的 value 返回即可,相比於原來先查 globals dict 再查 builtins dict 著實快了不少。

但是先不要高興太早,該少的減少其實一個都少不了,我們知道只有 globals 找不到的時候才會去找 builtins,如果 globals 變了那快取必然會失效,此外我們的 index 快取前提也是 builtins dict 不能改變,綜合上述兩點我們必須先確認兩個字典的 dk_version 都沒變才能繼續執行,這其實也是 Deoptimize 的觸發條件之一,這正是 DEOPT_IF 所做的事情:

#define DEOPT_IF(cond, instname) if (cond) { goto instname ## _miss; }
#define JUMP_TO_INSTRUCTION(op) goto PREDICT_ID(op)
#define ADAPTIVE_CACHE_BACKOFF 64


static inline void
cache_backoff(_PyAdaptiveEntry *entry) {
entry->counter = ADAPTIVE_CACHE_BACKOFF;
}


LOAD_GLOBAL_miss:
{
STAT_INC(LOAD_GLOBAL, miss);
_PyAdaptiveEntry *cache = &GET_CACHE()->adaptive;
cache->counter--;
if (cache->counter == 0) {
next_instr[-1] = _Py_MAKECODEUNIT(LOAD_GLOBAL_ADAPTIVE, _Py_OPARG(next_instr[-1]));
STAT_INC(LOAD_GLOBAL, deopt);
cache_backoff(cache);
}
oparg = cache->original_oparg;
STAT_DEC(LOAD_GLOBAL, unquickened);
JUMP_TO_INSTRUCTION(LOAD_GLOBAL);
}

DEOPT_IF 做的事情是跳到指令的 cache miss 分支,對於 LOAD_GLOBAL 是 LOAD_GLOBAL_miss,miss 分支做的事情非常固定,首先可以明確的是它會 fallback 到常規的 LOAD_GLOBAL 分支(底部的 JUMP_TO_INSTRUCTION),但它永遠不會將指令回滾到 LOAD_GLOBAL,只會在 Adaptive 和 Specialized 指令之間轉移。為了避免快取顛簸,上面的程式碼會在 counter 遞減到 0 時才退化到 Adaptive 指令,在此之前會一直嘗試先訪問快取。

到這裡整個優化和去優化過程我們就分析完了,整個流程還是比較複雜的,但具有明顯的套路,只要構建好 Cache 和 Adaptive 基礎設施,對其他指令的改造就比較容易了。通過這種方式進行指令加速的一個難點是如何設計快取以減少分支,因為 DEOPT_IF 的存在我們很難直接減少前置的條件判斷,只能將其轉化為更高效的形式。

關注我們,每週 3 篇移動技術實踐&乾貨給你思考!