7.6。使用 Numba Rewrite Pass 獲得樂趣和優化 · Numba 0.44 中文文檔

# 7.6。使用 Numba Rewrite Pass 獲得樂趣和優化 > 原文： [http://numba.pydata.org/numba-doc/latest/developer/rewrites.html](http://numba.pydata.org/numba-doc/latest/developer/rewrites.html) ## 7.6.1。概述本節介紹中間表示（IR）重寫，以及如何使用它們來實現優化。正如前面在“[階段 6a：重寫類型 IR](architecture.html#rewrite-typed-ir) ”中所討論的那樣，重寫 Numba IR 允許我們執行在較低 LLVM 級別執行將更加困難的優化。與 Numba 類型和降低子系統類似，重寫子系統是用戶可擴展的。這種可擴展性為 Numba 提供了支持各種特定領域優化（DSO）的可能性。其余小節詳細介紹了實現重寫的機制，使用重寫注冊表注冊重寫，并提供添加新重寫的示例，以及數組表達式優化傳遞的內部結構。最后，我們回顧一下示例中公開的一些用例，并回顧開發人員應該注意的任何問題。 ## 7.6.2。重寫通行證重寫通道有一個簡單的 [`match()`](#Rewrite.match "Rewrite.match") 和 [`apply()`](#Rewrite.apply "Rewrite.apply") 接口。匹配和重寫之間的劃分遵循如何在聲明性域特定語言（DSL）中定義術語重寫。在這樣的 DSL 中，可以按如下方式編寫重寫： ```py <match> => <replacement> ``` `<match>`和`<replacement>`符號表示 IR 術語表達式，其中左側表示要匹配的模式，右側表示 IR 術語構造函數以在匹配時構建。每當重寫與 IR 模式匹配時，左側的任何自由變量都綁定在自定義環境中。應用時，重寫使用模式匹配環境綁定右側的任何自由變量。由于 Python 不常用于聲明性容量，因此 Numba 使用對象狀態來處理匹配和應用程序步驟之間的信息傳輸。 ### 7.6.2.1。 [`Rewrite`](#Rewrite "Rewrite") 基類 ```py class Rewrite ``` [`Rewrite`](#Rewrite "Rewrite") 類只是為 Numba 重寫定義了一個抽象基類。開發人員應將重寫定義為此基類型的子類，重載 [`match()`](#Rewrite.match "Rewrite.match") 和 [`apply()`](#Rewrite.apply "Rewrite.apply") 方法。 ```py pipeline ``` pipeline 屬性包含當前正在編譯正在考慮重寫的函數的`numba.compiler.Pipeline`實例。 ```py __init__(self, pipeline, *args, **kws) ``` 用于重寫的基礎構造函數只是將其參數存儲到同名的屬性中。除非用于調試或測試，否則重寫應僅由`RewriteRegistry.apply()`方法中的`RewriteRegistry`構造，并且構造接口應保持穩定（盡管管道通常包含所知的所有內容）。 ```py match(self, block, typemap, callmap) ``` [`match()`](#Rewrite.match "Rewrite.match") 方法除 _self_ 外還有四個參數： * _func_ir_ ：這是被重寫函數的`numba.ir.FunctionIR`實例。 * _ 塊 _：這是`numba.ir.Block`的一個實例。匹配方法應迭代`numba.ir.Block.body`成員中包含的指令。 * _typemap_ ：這是一個 Python [`dict`](https://docs.python.org/3/library/stdtypes.html#dict "(in Python v3.7)") 實例，它從 IR 中的符號名稱（表示為字符串）映射到 Numba 類型。 * _callmap_ ：這是另一個 [`dict`](https://docs.python.org/3/library/stdtypes.html#dict "(in Python v3.7)") 實例從呼叫（表示為`numba.ir.Expr`實例）映射到它們對應的呼叫站點類型簽名，表示為`numba.typing.templates.Signature`實例。 [`match()`](#Rewrite.match "Rewrite.match") 方法應返回 [`bool`](https://docs.python.org/3/library/functions.html#bool "(in Python v3.7)") 結果。 [`True`](https://docs.python.org/3/library/constants.html#True "(in Python v3.7)") 結果應表明發現了一個或多個匹配， [`apply()`](#Rewrite.apply "Rewrite.apply") 方法將返回新的替換`numba.ir.Block`實例。 [`False`](https://docs.python.org/3/library/constants.html#False "(in Python v3.7)") 結果應表明未找到匹配項，隨后對 [`apply()`](#Rewrite.apply "Rewrite.apply") 的調用將返回未定義或無效的結果。 ```py apply(self) ``` 只有在成功調用 [`match()`](#Rewrite.match "Rewrite.match") 后才能調用 [`apply()`](#Rewrite.apply "Rewrite.apply") 方法。此方法不需要 _self_ 以外的其他參數，并應返回替換`numba.ir.Block`實例。如上所述，調用 [`apply()`](#Rewrite.apply "Rewrite.apply") 的行為是未定義的，除非 [`match()`](#Rewrite.match "Rewrite.match") 已被調用并返回 [`True`](https://docs.python.org/3/library/constants.html#True "(in Python v3.7)") 。 ### 7.6.2.2。子類 [`Rewrite`](#Rewrite "Rewrite") 在進入任何 [`Rewrite`](#Rewrite "Rewrite") 子類必須具有的重載方法的期望之前，讓我們回過頭來回顧一下這里發生的事情。通過提供可擴展的編譯器，Numba 向用戶定義的代碼生成器開放，這些代碼生成器可能不完整，或者更糟，不正確。當代碼生成器出錯時，它可能導致異常的程序行為或提??前終止。用戶定義的重寫增加了一個新的復雜程度，因為它們不僅必須生成正確的代碼，而且它們生成的代碼應該確保編譯器不會卡在匹配/應用循環中。編譯器的非終止將直接導致用戶函數調用的非終止。有幾種方法可以幫助確保重寫終止： * _ 鍵入 _：重寫通常應嘗試分解復合類型，并避免編寫新類型。如果重寫與特定類型匹配，則將表達式類型更改為較低級別類型將確保在應用重寫后它們不會長期匹配。 * _ 特殊指令 _：重寫可以合成自定義運算符或在目標 IR 中使用特殊函數。此技術再次生成不再位于原始匹配域內的代碼，并且重寫將終止。在下面的“[案例研究：數組表達式](#case-study-array-expressions)”小節中，我們將看到數組表達式重寫器如何使用這兩種技術。 ### 7.6.2.3。重載 [`Rewrite.match()`](#Rewrite.match "Rewrite.match") 每個重寫開發人員都應盡量讓 [`match()`](#Rewrite.match "Rewrite.match") 的實現盡快返回 [`False`](https://docs.python.org/3/library/constants.html#False "(in Python v3.7)") 值。 Numba 是一個即時編譯器，添加編譯時間最終會增加用戶的運行時間。當重寫為給定塊返回 [`False`](https://docs.python.org/3/library/constants.html#False "(in Python v3.7)") 時，注冊表將不再使用該重寫處理該塊，并且編譯器更接近于繼續降低。這種對及時性的需求必須與收集必要信息以進行重寫相匹配。重寫開發人員應該很樂意為其子類添加動態屬性，然后使用這些新屬性來指導替換基本塊的構造。 ### 7.6.2.4。重載 [`Rewrite.apply()`](#Rewrite.apply "Rewrite.apply") [`apply()`](#Rewrite.apply "Rewrite.apply") 方法應返回替換`numba.ir.Block`實例以替換包含重寫匹配的基本塊。如上所述，由 [`apply()`](#Rewrite.apply "Rewrite.apply") 方法構建的 IR 應該保留用戶代碼的語義，但也試圖避免為相同的重寫或重寫集生成另一個匹配。 ## 7.6.3。重寫注冊表如果要在重寫過程中包含重寫，則應將其注冊到重寫注冊表。 `numba.rewrites`模塊提供抽象基類和類裝飾器，用于掛鉤到 Numba 重寫子系統。以下說明了新重寫的存根定義： ```py from numba import rewrites @rewrites.register_rewrite class MyRewrite(rewrites.Rewrite): def match(self, block, typemap, calltypes): raise NotImplementedError("FIXME") def apply(self): raise NotImplementedError("FIXME") ``` 開發人員應注意，使用如上所示的類裝飾器將在導入時注冊重寫。在編譯開始之前，開發人員有責任確保加載擴展。 ## 7.6.4。案例研究：數組表達式本小節更深入地介紹了數組表達式重寫器。數組表達式重寫器及其大部分支持功能可在`numba.npyufunc.array_exprs`模塊中找到。重寫傳遞本身在`RewriteArrayExprs`類中實現。除了重寫器之外，`array_exprs`模塊還包括用于降低數組表達式的函數`_lower_array_expr()`。整體優化過程如下： * `RewriteArrayExprs.match()`：重寫過程查找兩個或多個形成數組表達式的數組操作。 * `RewriteArrayExprs.apply()`：一旦找到數組表達式，重寫器就會用一種新的 IR 表達式`arrayexpr`替換各個數組操作。 * `numba.npyufunc.array_exprs._lower_array_expr()`：在降低期間，只要找到`arrayexpr` IR 表達式，代碼生成器就會調用`_lower_array_expr()`。下面給出關于優化的每個步驟的更多細節。 ### 7.6.4.1。 `RewriteArrayExprs.match()`方法數組表達式優化傳遞首先查找數組操作，包括對支持的`ufunc`和用戶定義的 [`DUFunc`](../reference/jit-compilation.html#numba.DUFunc "numba.DUFunc") 的調用。 Numba IR 遵循靜態單指派（SSA）語言的約定，這意味著搜索數組運算符首先要查找賦值指令。當重寫傳遞調用`RewriteArrayExprs.match()`方法時，它首先檢查它是否可以簡單地拒絕基本塊。如果方法確定塊是匹配的候選者，則它在重寫對象中設置以下狀態變量： * _crnt_block_ ：當前匹配的基本塊。 * _typemap_ ：匹配函數的 _typemap_ 。 * _ 匹配 _：引用數組表達式的變量名列表。 * _array_assigns_ ：從賦值變量名稱到定義給定變量的實際賦值指令的映射。 * _const_assigns_ ：從賦值變量名到定義常量變量的常量值表達式的映射。此時，匹配方法迭代迭代輸入基本塊中的賦值指令。對于每個賦值指令，匹配器會查找以下兩種情況之一： * 數組操作：如果賦值指令的右側是表達式，并且該表達式的結果是數組類型，則匹配器檢查表達式是已知數組操作還是對通用函數的調用。如果找到數組運算符，匹配器將左側變量名和整個指令存儲在 _array_assigns_ 成員中。最后，匹配器測試以查看陣列操作的任何操作數是否也已被識別為其他陣列操作的目標。如果一個或多個操作數也是數組操作的目標，則匹配器還會將左側變量名稱附加到 _ 匹配 _ 成員。 * 常量：常量（甚至標量）可以是數組操作的操作數。不必擔心數組表達式的常量分離，匹配器在 _const_assigns_ 成員中存儲常量名稱和值。匹配方法的結束只是檢查非空 _ 匹配 _ 列表，如果有一個或多個匹配則返回 [`True`](https://docs.python.org/3/library/constants.html#True "(in Python v3.7)") ， [`False`](https://docs.python.org/3/library/constants.html#False "(in Python v3.7)") 當 _ 匹配 _ 為空時。 ### 7.6.4.2。 `RewriteArrayExprs.apply()`方法當`RewriteArrayExprs.match()`找到一個或匹配的數組表達式時，重寫過程將調用`RewriteArrayExprs.apply()`。 apply 方法有兩次通過。第一遍迭代所找到的匹配，并根據舊基本塊中的指令構建映射到新基本塊中的新指令。第二遍迭代舊基本塊中的指令，復制未被重寫改變的指令，以及替換或刪除第一遍所識別的指令。 `RewriteArrayExprs._handle_matches()`實現重寫的代碼生成部分的第一遍。對于每個匹配，此方法構建一個特殊的 IR 表達式，其中包含數組表達式的表達式樹。為了計算表達式樹的葉子，`_handle_matches()`方法迭代所識別的根操作的操作數。如果操作數是另一個數組操作，則將其轉換為表達式子樹。如果操作數是常量，`_handle_matches()`將復制常量值。否則，操作數被標記為由數組表達式使用。當該方法構建數組表達式節點時，它構建從舊指令到新指令（ _replace_map_ ）的映射，以及可能已移動的變量集（ _used_vars_ ）和變量應該完全刪除（ _dead_vars_ ）。這三個數據結構返回到調用`RewriteArrayExprs.apply()`方法。 `RewriteArrayExprs.apply()`方法的剩余部分迭代舊基本塊中的指令。對于每條指令，此方法根據`RewriteArrayExprs._handle_matches()`的結果替換，刪除或復制該指令。以下列表描述了優化如何處理單個指令： * 當指令是賦值時，`apply()`檢查它是否在替換指令映射中。當在指令映射中找到賦值指令時，`apply()`必須檢查替換指令是否也在替換映射中。優化器繼續此檢查，直到它到達 [`None`](https://docs.python.org/3/library/constants.html#None "(in Python v3.7)") 值或不在替換映射中的指令。刪除具有 [`None`](https://docs.python.org/3/library/constants.html#None "(in Python v3.7)") 的替換的指令。更換非 [`None`](https://docs.python.org/3/library/constants.html#None "(in Python v3.7)") 替換的說明。不在替換映射中的賦值指令被附加到新的基本塊而不進行任何更改。 * 當指令是刪除指令時，重寫檢查它是否刪除可能仍被稍后的數組表達式使用的變量，或者它是否刪除了死變量。刪除已使用變量的指令被添加到延遲刪除指令的映射中，`apply()`使用它來移動它們超過該變量的任何用途。循環復制刪除非死變量的指令，并忽略死變量的刪除指令（有效地從基本塊中刪除它們）。 * 所有其他指令都附加到新的基本塊。最后，`apply()`方法返回用于降低的新基本塊。 ### 7.6.4.3。 `_lower_array_expr()`功能如果我們只是重寫，那么編譯器的降低階段就會失敗，抱怨它不知道如何降低`arrayexpr`操作。我們首先在編譯器實例化`RewriteArrayExprs`類時將降低函數掛鉤到目標上下文中。只要遇到`arrayexr`運算符，此掛鉤就會導致降低傳遞調用`_lower_array_expr()`。這個功能有兩個步驟： * 合成實現數組表達式的 Python 函數：這個新的 Python 函數本質上就像一個 Numpy `ufunc`，在廣播數組參數中的標量值上返回表達式的結果。降低功能通過從數組表達式樹轉換為 Python AST 來實現這一點。 * 將合成 Python 函數編譯到內核中：此時，降低函數依賴于現有代碼來降低 ufunc 和 DUFunc 內核，在定義如何降低對合成函數的調用之后調用`numba.targets.numpyimpl.numpy_ufunc_kernel()`。最終結果類似于 Numba 對象模式中的循環提升。 ## 7.6.5。結論和注意事項我們已經了解了如何在 Numba 中實現重寫，從界面開始，以實際優化結束。本節的重點是： * 在編寫好的插件時，匹配器應該盡快獲得 go / no-go 結果。 * 重寫應用程序部分的計算成本可能更高，但仍應生成不會在編譯器中導致無限循環的代碼。 * 我們使用對象狀態將任何匹配結果傳遞給重寫應用程序傳遞。