記憶體規劃¶
受眾： 對自定義 ExecuTorch 程式執行所在的記憶體區域感興趣的後端整合者和嵌入式開發者。
概述¶
記憶體規劃 (MemoryPlanning) 是在獲取 ExportedProgram 並將其生成 ExecuTorch 程式之前採取的最後一步操作。在此過程中，ExecuTorch 會獲取每個可變張量的大小和生命週期，並規劃它們在固定大小記憶體區域中的位置。
具體來說，有三個與記憶體規劃相關的 Pass：
SpecPropPass計算圖中每個張量（輸入、中間張量或輸出）的 TensorSpec。張量規範 (Tensor spec) 最重要的欄位是張量形狀的符號表達式，其中初始符號集來自輸入張量的維度，中間張量形狀的符號表達式透過張量操作進行傳播。使用者可以將維度標記為動態或靜態，並且當維度為動態時，使用者需要使用 ValueRange 註解該維度。SymShapeEvalPass將符號表達式評估為具體的整數及其上界。有兩種方法進行上界特化：HintBasedSymShapeEval（即將棄用）是舊的評估上界的方式。它不檢視符號的 ValueRange，而是使用示例輸入的形狀來替換所有符號。我們稱其為“基於提示”，因為示例輸入的形狀只是執行時輸入形狀可能是什麼的提示，僅用於追蹤。ValueRangeBasedSymShapeEval 是進行上界記憶體規劃的推薦方式。它會實際檢視符號的 ValueRange，並對範圍進行推斷以獲得實際的上界。MemoryPlanningPass在所有張量都獲得帶有具體整數形狀的 TensorSpec 後，執行實際的記憶體規劃。
演算法¶
ExecuTorch 開箱即用地提供了兩種記憶體規劃演算法選項，但如果提供的選項不合適或不足，使用者可以自定義自己的演算法。
樸素演算法簡單地將所有張量串聯線上性記憶體塊中，不考慮記憶體複用。它作為總記憶體消耗的上界，並作為基準。
貪心演算法嘗試根據最佳適配準則複用已分配的記憶體。具體來說：當不存在一個已分配的記憶體，其生命週期與當前我們嘗試進行記憶體規劃的張量不重疊時，我們分配一個新的記憶體緩衝區，其大小和生命週期與當前張量相同。當存在一個或多個已分配的記憶體緩衝區，其生命週期與當前張量重疊時，我們選擇大小與當前張量最接近的緩衝區，以減少記憶體碎片。最後，我們線上性記憶體中分配這些記憶體緩衝區。
方法輸入與輸出¶
MemoryPlanningPass 提供了不對程式輸入和輸出進行記憶體規劃的選項。如果 IO 未規劃，則使用者將被要求在執行時提供資料緩衝區來儲存這些值。示例
program = edge_program.to_executorch(
exir.ExecutorchBackendConfig(
memory_planning_pass=MemoryPlanningPass(
alloc_graph_input=False, # Inputs will not be memory planned, the data_ptr for input tensors after model load will be nullptr
alloc_graph_output=True, # Outputs will be memory planned, the data_ptr for output tensors after model load will be in the `planned_memory`.
)
)
)
一個常見的設定是對於模型輸出作為後續推理輸入的模型。在這種情況下，通常最好不對 IO 進行記憶體規劃，而是在執行時為輸入和輸出提供相同的緩衝區，以避免複製。
自定義記憶體規劃¶
使用者可以編寫自定義記憶體規劃，以利用多個記憶體位置（如 SRAM 和 DRAM），將特定節點的輸出放置在特定位置，甚至改變規劃演算法本身。以下示例展示瞭如何複用提供的規劃演算法，但具有多個層級並將特定運算元的輸出放置在特定記憶體區域。
class CustomPoolMemoryPlanningPass(MemoryPlanningPass):
def run(self, graph_module: GraphModule, graph_signature: Optional[ExportGraphSignature]) -> PassResult:
for subgm in graph_module.modules():
if not isinstance(subgm, GraphModule):
continue
for node in subgm.graph.nodes:
# mem_id = 1 placeholder and outputs of mul
# mem_id = 2 for outputs of add
# parent class will copy spec will to alloc nodes
if node.op == "placeholder":
node.meta["spec"].mem_id = 1
continue
if node.op != "call_function":
continue
if node.target == torch.ops.aten.add.out:
node.meta["spec"].mem_id = 2
elif node.target == torch.ops.aten.mul.out:
node.meta["spec"].mem_id = 1
return super().run(graph_module, graph_signature)
然後，在降低到 ExecuTorch 時，你可以按以下方式使用你的自定義規劃
program = edge_program.to_executorch(
exir.ExecutorchBackendConfig(
memory_planning_pass=CustomPoolMemoryPlanningPass(
memory_planning_algo=greedy,
)
)
)
嘗試編寫自定義記憶體規劃演算法的使用者應該從檢視 貪心演算法的實現 開始。
