記憶體規劃¶
目標對象：對自訂 ExecuTorch 程式運作的記憶體區域感興趣的後端整合人員和嵌入式開發人員。
概觀¶
MemoryPlanning 是在取得 ExportedProgram 並進行發送到 ExecuTorch 程式之前的最後一個動作。在此過程中，ExecuTorch 會取得每個可變張量的大小和生命週期，並在固定大小的記憶體領域中規劃它們的位置。
具體而言，有三個與記憶體規劃相關的 Pass：
SpecPropPass為圖 (graph) 中的每個張量 (輸入、中間值或輸出) 計算一個 TensorSpec。tensor spec 中最重要的欄位是張量形狀的符號表示式，其中符號的初始集合來自輸入張量的維度，中間張量形狀的符號表示式則透過張量運算傳播。維度可以由使用者標記為動態或靜態，當維度為動態時，使用者需要使用 ValueRange 註解該維度。SymShapeEvalPass將符號表示式評估為具有上限的具體整數。有兩種方法可以進行上限特化：HintBasedSymShapeEval（即將棄用）是評估上限的舊方法。它不查看符號的 ValueRange，而是使用範例輸入的形狀來替換所有符號。我們稱之為「基於提示」，因為範例輸入的形狀只是輸入形狀在運行時可能為何的提示，僅用於追蹤。ValueRangeBasedSymShapeEval 是建議的UpperBoundMemory 規劃方法。它實際上會查看符號的 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,
)
)
)
嘗試編寫自訂記憶體規劃演算法的使用者應該首先查看貪婪演算法的實現。
