内存规划¶
受众：对自定义 ExecuTorch 程序运行内存区域感兴趣的后端集成商和嵌入式开发人员。
概述¶
MemoryPlanning 是执行 ExecuTorch 程序之前执行并执行的最后一个操作。在此过程中，ExecuTorch 会计算每个可变张量的大小和寿命，并规划它们在固定大小的内存区域中的位置。ExportedProgram
具体来说，有三个与内存规划相关的过程：
SpecPropPass计算图中每个张量（inputs、intermediates 或 outputs）的 TensorSpec。张量规范最重要的字段是张量形状的符号表达式，其中初始符号集来自输入张量的维度，中间张量形状的符号表达式通过张量运算传播。用户可以将维度标记为动态或静态，当维度是动态的时，用户需要使用 ValueRange 对维度进行注释。SymShapeEvalPass将符号表达式计算为具有上限的具体整数。有两种方法可以进行上限特化： HintBasedSymShapeEval（即将弃用）是计算上限的旧方法。它不查看元件的 ValueRange，而是使用示例输入的形状来替换所有元件。我们称其为“基于提示”，因为示例输入的形状只是输入形状在运行时可能是什么的提示，仅用于跟踪。ValueRangeBasedSymShapeEval 是执行 UpperBoundMemory 规划的推荐方法。它实际上会查看 symbols 的 ValueRange 并对范围进行推理以获得真正的上限。MemoryPlanningPass给定所有张量的实际内存规划是否获得具有具体整数形状的 TensorSpec。
算法¶
ExecuTorch 为内存规划算法提供了两个开箱即用的选项，但如果提供的选项不适合或不足以满足其使用案例，用户可以定义自己的选项。
朴素算法只是将所有张量连接在一个线性内存块中，而不考虑内存重用。它用作总内存消耗的上限，并用作基线。
Greedy 算法尝试根据最佳条件重用已分配的内存。具体说来： 当没有分配的内存的生命周期与我们尝试为其进行内存规划的当前张量不重叠时，我们会分配一个与当前张量具有相同大小和生命周期的新内存缓冲区。当有一个或多个分配的内存缓冲区，其生命周期与当前张量重叠时，我们选择与当前张量大小最接近的缓冲区，以减少内存碎片。最后，我们在内存中线性分配这些内存缓冲区。
方法输入和输出¶
这提供了不内存规划程序输入和输出的选项。如果未规划 IO，则用户将需要提供数据缓冲区以在运行时支持这些值。例：MemoryPlanningPass
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 进行内存规划，而是在运行时为 input 和 output 提供相同的缓冲区，以避免复制。
自定义内存计划¶
用户可以编写自定义内存计划以利用多个内存位置（如 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,
)
)
)
尝试编写自定义内存规划算法的用户应首先查看贪婪算法的实现。