内存规划¶
受众： 对自定义 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,
)
)
)
尝试编写自定义内存规划算法的用户应该从查看 贪心算法的实现 开始。
