在 ONNX Runtime 推理引擎中,ExecutionPlan(执行计划)、AllocationPlan(内存分配计划) 和 MemoryPattern(内存模式) 构成了一个完整的三层优化体系。这三个组件按照严格的构建顺序协同工作,通过静态分析计算图,分别确定最优的算子执行顺序、内存分配策略和内存布局方案。本文将深入剖析这三个组件的生成算法、构建顺序及其协同工作机制。
目录
-
1. ExecutionPlan 执行计划生成原理 -
• 1.1 执行计划的核心作用与基础地位 -
• 1.2 生成流程与关键组件 -
2. AllocationPlan 内存分配计划生成机制 -
• 2.1 基于执行计划的内存策略制定 -
• 2.2 OrtValueIndex:张量的唯一标识 -
• 2.3 核心数据结构与分配类型 -
• 2.4 AllocationPlan 生成算法详解 -
3. MemoryPattern 内存模式生成与优化 -
• 3.1 内存模式的概念与价值 -
• 3.2 内存跟踪机制的原理与流程 -
• 3.3 OrtValuePatternPlanner 的核心作用 -
• 3.4 内存模式的构建过程 -
• 3.5 内存模式的使用机制与优化效果 -
4. 三层优化体系的协同工作流程 -
• 4.1 构建阶段的严格依赖关系 -
• 4.2 运行时的高效协同执行
1 ExecutionPlan 执行计划生成原理
1.1 执行计划的核心作用与基础地位
ExecutionPlan 是整个优化体系的基础,它定义了模型中算子的执行顺序和资源分配。作为第一个生成的计划,它为后续的内存优化提供了必要的结构信息。
主要职责:
-
• 拓扑排序:确保算子按数据依赖关系正确执行 -
• 设备分配:将算子分配到合适的执行提供者 -
• 子图划分:根据设备边界优化计算图结构 -
• Kernel选择:为每个算子选择最优的实现Kernel
1.2 生成流程与关键组件
执行计划生成入口:
执行计划的生成由 SequentialPlanner 协调,这是整个优化流程的起点。
class SequentialPlanner {
public:
Status CreatePlan(const GraphViewer& graph_viewer,
const ExecutionProviders& providers,
const KernelCreateInfoMap& kernel_create_info_map,
SequentialExecutionPlan& plan) {
PlannerImpl planner(graph_viewer, providers, kernel_create_info_map, plan);
return planner.CreatePlan();
}
};
代码说明:SequentialPlanner 是执行计划生成的入口类,它接收计算图、执行提供者和Kernel注册信息作为输入,输出完整的 SequentialExecutionPlan。这个类协调了整个执行计划的生成过程,包括拓扑排序、设备分配和Kernel选择等关键步骤。PlannerImpl 是实际的实现类,封装了所有核心算法。
执行计划生成流程:
ExecutionPlan 的生成是一个系统化的过程,为后续的内存优化奠定基础。
关键生成步骤:
-
1. 拓扑排序与依赖分析
拓扑排序确保算子按照数据依赖关系正确执行,这是后续内存优化的前提。
std::vector<NodeIndex> GetTopologicalOrder() {
std::vector<NodeIndex> order;
std::queue<NodeIndex> ready_nodes;
// 计算每个节点的入度
std::unordered_map<NodeIndex, size_t> in_degree;
for (auto node_index : graph_viewer_.GetNodeIndices()) {
in_degree[node_index] = graph_viewer_.GetNode(node_index)->GetInputEdgesCount();
if (in_degree[node_index] == 0) {
ready_nodes.push(node_index);
}
}
// 生成拓扑排序
while (!ready_nodes.empty()) {
NodeIndex current = ready_nodes.front();
ready_nodes.pop();
order.push_back(current);
for (auto output_edge : graph_viewer_.GetNode(current)->GetOutputEdges()) {
NodeIndex successor = output_edge.GetNode().Index();
if (--in_degree[successor] == 0) {
ready_nodes.push(successor);
}
}
}
return order;
}
代码说明:
拓扑排序算法使用 Kahn 算法,通过维护入度表和就绪队列来生成无环图的线性执行顺序。这确保了每个节点的输入在其执行前都已就绪,为后续的内存生命周期分析提供了基础。算法的时间复杂度为 O(V+E),其中 V 是节点数,E 是边数。
-
2. 设备分配与Kernel选择
设备分配基于各执行提供者的能力评估,为每个算子选择最优的执行设备和Kernel实现。
Status AssignDevicesAndKernels() {
for (auto node_index : topological_order_) {
auto node = graph_viewer_.GetNode(node_index);
// 为节点选择最优执行提供者
auto best_provider = SelectBestExecutionProvider(node);
// 获取对应的Kernel实现
const KernelCreateInfo* kernel_info =
kernel_registry_manager_.GetKernelCreateInfo(node, best_provider);
if (!kernel_info) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
"No kernel found for node: ", node->Name());
}
// 注册Kernel到执行计划
plan_.AddKernel(node_index, kernel_info, best_provider);
}
return Status::OK();
}
代码说明:
设备分配过程遍历拓扑排序后的节点序列,为每个节点选择最优的执行提供者(如 CPU、CUDA、TensorRT 等),并查找对应的Kernel实现。选择策略综合考虑了算子的计算特性、设备的能力限制和数据传输开销。如果找不到合适的Kernel实现,会返回错误状态。
2 AllocationPlan 内存分配计划生成机制
2.1 基于执行计划的内存策略制定
AllocationPlan 是优化体系的第二层,它基于 ExecutionPlan 提供的执行顺序信息,制定详细的内存分配策略。这一阶段的核心任务是将执行顺序转换为内存使用策略。
2.2 OrtValueIndex:张量的唯一标识
在 ONNX Runtime 中,OrtValueIndex 是用于唯一标识张量的整数索引。这个索引系统在整个内存管理流程中起着关键作用:
OrtValueNameIdxMap 的核心实现:
class OrtValueNameIdxMap {
int next_idx_ = 0;
InlinedHashMap<std::string, int> map_;
InlinedHashMap<int, std::string> idx_name_map_;
public:
int Add(const std::string& name) {
const int idx = next_idx_;
auto p = map_.emplace(name, idx);
if (p.second) {
idx_name_map_[idx] = name;
next_idx_++;
return idx;
}
return p.first->second;
}
common::Status GetIdx(std::string_view name, int& idx) const {
auto it = map_.find(std::string(name));
if (it == map_.end()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Could not find OrtValue with name '", name, "'");
}
idx = it->second;
return common::Status::OK();
}
};
代码说明:OrtValueNameIdxMap 维护了张量名称与整数索引之间的双向映射关系。通过 Add 方法为每个唯一的张量名称分配连续的整数索引,GetIdx 方法提供名称到索引的查找功能。这种设计既保证了查找效率,又便于使用数组进行高效的内存管理。
OrtValueIndex 的关键特性:
|
|
|
|
|---|---|---|
| 唯一性 |
|
|
| 双向映射 |
|
|
| 连续性 |
|
|
| 稳定性 |
|
|
2.3 核心数据结构与分配类型
基础数据结构定义:
AllocationPlan 的核心数据结构包含了分配类型、设备位置、复用缓冲区索引和生命周期跟踪信息。
struct AllocPlanPerValue {
AllocKind alloc_kind;
OrtDevice location;
OrtValueIndex reused_buffer;
ProgramCounter program_counter;
enum AllocKind {
kAllocate, // 新分配 - 常规中间结果
kReuse, // 复用现有 - 生命周期不重叠
kPreExisting, // 外部传入 - 计算图输入
kAllocateStatically, // 静态分配 - 模型权重
kAllocateOutput // 输出分配 - 计算图输出
};
};
代码说明:AllocPlanPerValue 结构体为每个张量定义了完整的内存分配策略。alloc_kind 指定分配类型,location 确定设备位置,reused_buffer 记录复用关系,program_counter 跟踪生命周期。这种精细化的设计使得 ONNX Runtime 能够为每个张量制定最优的内存管理策略。
ProgramCounter 生命周期跟踪:
ProgramCounter 负责精确记录每个张量在计算图中的开始和结束使用位置,这是内存复用决策的基础。
class ProgramCounter {
std::vector<size_t> starts_;
std::vector<size_t> ends_;
public:
void AddStart(size_t start) {
ORT_ENFORCE(starts_.size() == ends_.size());
starts_.push_back(start);
}
void AddEnd(size_t end) {
ORT_ENFORCE(starts_.size() == ends_.size() + 1);
ends_.push_back(end);
}
size_t GetStart(int index) const { return starts_[index]; }
size_t GetEnd(int index) const { return ends_[index]; }
};
代码说明:ProgramCounter 通过 starts_ 和 ends_ 数组分别记录每个张量的开始和结束使用步骤。AddStart 和 AddEnd 方法确保生命周期记录的完整性,GetStart 和 GetEnd 方法提供查询接口。这种设计使得生命周期分析能够精确到具体的执行步骤,为内存复用提供准确的时间窗口信息。
2.4 AllocationPlan 生成算法详解
完整的生成流程:
AllocationPlan 的生成依赖于 ExecutionPlan 提供的执行顺序信息,是一个多阶段的复杂过程。
算法核心步骤实现:
-
1. 引用计数统计与初始化
引用计数统计基于 ExecutionPlan 的执行顺序,准确记录每个张量被后续节点引用的次数。
void PlannerImpl::ComputeReuseCount() {
// 初始化所有张量的使用计数
use_counts_.resize(all_values_size_, 0);
// 处理计算图输入:确保不被复用
for (auto graph_input : graph_viewer_.GetInputs()) {
OrtValueIndex index = Index(graph_input->Name());
use_counts_[index] = std::numeric_limits<int>::max(); // 标记为永久使用
}
// 基于ExecutionPlan的执行顺序统计引用
for (auto node_index : execution_plan_.execution_order) {
auto pnode = graph_viewer_.GetNode(node_index);
// 处理节点输入:增加使用计数
for (auto input : pnode->InputDefs()) {
if (input->Exists()) {
OrtValueIndex input_index = Index(input->Name());
use_counts_[input_index]++;
}
}
// 处理节点输出:设置初始计数
for (auto output : pnode->OutputDefs()) {
if (output->Exists()) {
OrtValueIndex output_index = Index(output->Name());
use_counts_[output_index] = 1; // 初始计数为1
}
}
}
}
代码说明:
引用计数统计遍历 ExecutionPlan 中的执行顺序,为每个张量计算被后续节点引用的次数。输入张量的引用计数会递增,输出张量设置初始计数。计算图输入被标记为永久使用以避免被复用。这个过程建立了张量使用关系的完整图谱,为后续的生命周期分析奠定基础。
-
2. 生命周期分析与 ProgramCounter 构建
基于 ExecutionPlan 的执行步骤,精确记录每个张量的开始和结束使用位置。
void PlannerImpl::BuildProgramCounters() {
size_t current_step = 0;
for (auto node_index : execution_plan_.execution_order) {
auto pnode = graph_viewer_.GetNode(node_index);
// 处理节点输出:记录开始位置
for (auto output : pnode->OutputDefs()) {
if (output->Exists()) {
OrtValueIndex output_index = Index(output->Name());
program_counters_[output_index].AddStart(current_step);
}
}
// 处理节点输入:记录结束位置
for (auto input : pnode->InputDefs()) {
if (input->Exists()) {
OrtValueIndex input_index = Index(input->Name());
// 减少使用计数,检查是否为最后一次使用
if (--use_counts_[input_index] == 0) {
program_counters_[input_index].AddEnd(current_step);
}
}
}
current_step++;
}
// 处理计算图输出:在整个执行期间保持活跃
for (auto graph_output : graph_viewer_.GetOutputs()) {
OrtValueIndex output_index = Index(graph_output->Name());
program_counters_[output_index].AddEnd(current_step); // 最后一步结束
}
}
代码说明:
生命周期分析按照执行步骤顺序进行,输出张量在创建节点记录开始位置,输入张量在最后一次使用时记录结束位置。通过引用计数递减到零的判断,准确捕捉每个张量的完整生命周期。计算图输出在整个执行期间保持活跃,直到最后一步才结束。
3 MemoryPattern 内存模式生成与优化
3.1 内存模式的概念与价值
MemoryPattern 是优化体系的第三层,也是最高级的优化阶段。它基于 AllocationPlan 提供的内存分配策略,通过分析内存使用模式,生成最优的内存布局方案。
核心价值:
-
• 预分配优化:避免运行时频繁的内存分配/释放操作 -
• 布局优化:通过连续内存布局减少碎片,提高缓存局部性 -
• 性能提升:将运行时分配转换为简单的指针偏移操作
3.2 内存跟踪机制的原理与流程
内存跟踪机制是 MemoryPattern 生成的基础,它通过精确记录每次内存分配和释放操作,构建完整的内存使用轨迹。
内存跟踪的核心原理:
内存跟踪基于以下关键观察:对于固定形状的模型,其内存分配模式在多次运行间是稳定的。通过捕获这种稳定模式,可以预分配内存并优化布局。
跟踪数据的核心内容:
|
|
|
|
|---|---|---|
| 分配时间点 |
|
|
| 释放时间点 |
|
|
| 分配大小 |
|
|
| 设备位置 |
|
|
| 张量索引 | OrtValueIndex
|
|
内存跟踪的完整流程:
内存跟踪的调用机制:
跟踪机制通过重写内存分配接口来实现,在以下时机自动调用:
-
1. ExecutionFrame初始化时:跟踪预分配的内存 -
2. 节点输出内存分配时:跟踪动态分配的内存 -
3. 内存复用发生时:跟踪复用关系的变化 -
4. 内存释放时:跟踪生命周期结束
跟踪数据的存储结构:
struct MemoryTraceRecord {
OrtValueIndex tensor_index;
size_t allocation_step;
size_t free_step;
size_t size;
OrtDevice device;
AllocKind alloc_type;
};
这种详细的跟踪数据为后续的内存模式生成提供了坚实的基础。
3.3 OrtValuePatternPlanner 的核心作用
OrtValuePatternPlanner 是内存模式生成的核心组件,它负责:
-
• 收集分配信息:通过 TraceAllocation和TraceFree记录完整的内存使用轨迹 -
• 分析使用模式:识别内存分配的规律和峰值需求 -
• 生成布局方案:创建最优的内存偏移映射表
OrtValuePatternPlanner 的初始化:
OrtValuePatternPlanner::OrtValuePatternPlanner(
const ExecutionPlanBase& execution_plan,
bool trace_using_counters) : execution_planner_(execution_plan) {
planner_map_.reserve(execution_plan.GetAllLocations().size());
for (auto& location : execution_plan.GetAllLocations()) {
planner_map_.emplace(location, trace_using_counters);
}
}
代码说明:OrtValuePatternPlanner 在构造时根据执行计划中的设备位置信息,为每个设备创建独立的 MemoryPatternPlanner。planner_map_ 维护了设备到对应规划器的映射关系。这种按设备分别优化的策略确保了不同设备上的内存特性得到充分考虑,比如 GPU 内存的对齐要求和 CPU 内存的缓存友好性。
3.4 内存模式的构建过程
内存跟踪机制:
内存模式的构建依赖于详细的内存分配跟踪数据。
common::Status OrtValuePatternPlanner::TraceAllocation(
int ort_value_idx, size_t size) {
const auto& location = execution_planner_.GetLocation(ort_value_idx);
auto it = planner_map_.find(location);
if (it == planner_map_.end()) {
return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT);
}
// 记录分配信息
it->second.TraceAllocation(ort_value_idx, size);
return common::Status::OK();
}
common::Status OrtValuePatternPlanner::TraceFree(int ort_value_index) {
const auto& location = execution_planner_.GetLocation(ort_value_index);
auto it = planner_map_.find(location);
if (it == planner_map_.end()) {
return common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT);
}
// 记录释放信息
it->second.TraceFree(ort_value_index);
return common::Status::OK();
}
代码说明:TraceAllocation 和 TraceFree 方法在每次内存分配和释放时被调用,记录张量索引、分配大小和设备位置。这些方法首先根据张量索引查找对应的设备位置,然后将跟踪信息转发给相应设备的 MemoryPatternPlanner。这些数据用于分析内存使用模式和计算峰值需求。
模式生成流程:
模式生成实现:
common::Status OrtValuePatternPlanner::GeneratePatterns(
MemoryPatternGroup& out) {
out.locations.reserve(planner_map_.size());
out.patterns.reserve(planner_map_.size());
// 为每个设备生成对应的内存模式
for (auto& it : planner_map_) {
out.locations.push_back(it.first);
out.patterns.push_back(it.second.GenerateMemPattern());
}
return common::Status::OK();
}
代码说明:GeneratePatterns 方法遍历所有设备的模式规划器,生成对应的内存模式。每个设备的模式包含了该设备上所有张量的内存布局信息,最终组合成完整的 MemoryPatternGroup。这种方法确保了多设备场景下的内存优化能够分别进行,避免设备间的相互干扰。
内存模式的数据结构:
struct MemoryPatternGroup {
std::vector<OrtDevice> locations;
std::vector<MemoryPattern> patterns;
const MemoryPattern* GetPattern(const OrtDevice& location) const {
for (size_t i = 0; i < locations.size(); ++i) {
if (locations[i] == location) {
return &patterns[i];
}
}
return nullptr;
}
};
struct MemoryPattern {
size_t peak_size;
std::unordered_map<OrtValueIndex, size_t> offsets;
size_t GetOffset(OrtValueIndex idx) const {
auto it = offsets.find(idx);
if (it != offsets.end()) {
return it->second;
}
return SIZE_MAX;
}
};
代码说明:MemoryPatternGroup 管理多个设备的内存模式,通过 GetPattern 方法提供设备到模式的查找功能。MemoryPattern 描述单个设备上的内存布局,peak_size 表示峰值内存需求,offsets 映射表记录每个张量在预分配块中的偏移量。GetOffset 方法提供安全的偏移量查询,避免访问不存在的张量。
3.5 内存模式的使用机制与优化效果
启用内存模式:
SessionOptions session_options;
session_options.enable_mem_pattern = true;
session_options.enable_mem_reuse = true;
代码说明:
通过设置 SessionOptions 中的 enable_mem_pattern 和 enable_mem_reuse 标志来启用内存模式优化。这些设置会在会话初始化时触发内存模式的生成和应用。enable_mem_pattern 控制是否使用预分配的内存块,enable_mem_reuse 控制是否启用内存复用策略。
执行时的内存分配:
当内存模式启用时,ExecutionFrame 会使用预分配的内存块:
Status ExecutionFrame::AllocateWithMemoryPattern(OrtValueIndex index) {
const auto& pattern = memory_pattern_group_.GetPattern(alloc_plan.location);
if (!pattern) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No memory pattern for device");
}
size_t offset = pattern->GetOffset(index);
if (offset == SIZE_MAX) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No offset found for tensor");
}
// 在预分配块中按偏移量获取内存
void* memory_ptr = static_cast<char*>(preallocated_block_) + offset;
// 创建 OrtValue 并设置内存指针
return CreateOrtValueWithPreallocatedMemory(index, memory_ptr);
}
代码说明:
在内存模式启用的情况下,内存分配不再是动态的 malloc 调用,而是简单的指针偏移计算。AllocateWithMemoryPattern 方法首先查找对应设备的内存模式,然后根据张量索引获取预计算的偏移量,最后通过指针算术在预分配的内存块中定位具体的内存位置。这种方法完全消除了运行时内存分配的开销,大大提高了性能。
优化效果对比:
|
|
|
|
|
|---|---|---|---|
| 内存分配次数 |
|
|
|
| 内存碎片 |
|
|
|
| 分配性能 |
|
|
|
| 确定性 |
|
|
|
适用场景分析:
|
|
|
|
|---|---|---|
| 静态形状推理 |
|
|
| 动态形状推理 |
|
|
| 训练模式 |
|
|
| 边缘设备 |
|
|
4 三层优化体系的协同工作流程
4.1 构建阶段的严格依赖关系
三个组件的构建存在严格的依赖关系,形成完整的优化流水线。
4.2 运行时的高效协同执行
在推理执行过程中,三个组件通过紧密配合实现极致性能。
执行器与优化组件的协同:
Status SequentialExecutor::Execute(const SessionState& session_state,
const std::vector<OrtValue>& feeds,
std::vector<OrtValue>& fetches) {
// 初始化ExecutionFrame,加载所有优化计划
ExecutionFrame frame(feeds, fetches, session_state);
// 按照ExecutionPlan顺序执行节点
for (auto node_index : session_state.GetExecutionPlan().execution_order) {
// 根据AllocationPlan和MemoryPattern准备内存
ORT_RETURN_IF_ERROR(frame.PrepareForNode(node_index));
// 执行计算Kernel
ORT_RETURN_IF_ERROR(ExecuteKernel(node_index, frame));
// 根据AllocationPlan回收内存
ORT_RETURN_IF_ERROR(frame.CompleteNode(node_index));
}
return Status::OK();
}
代码说明:SequentialExecutor 按照 ExecutionPlan 的执行顺序遍历节点,在每個节点执行前后调用 ExecutionFrame 的内存管理方法。PrepareForNode 方法根据 AllocationPlan 和 MemoryPattern 为节点准备输入输出内存,ExecuteKernel 执行实际的计算,CompleteNode 根据 AllocationPlan 回收不再需要的内存。这些方法内部会使用所有优化组件的策略。
内存分配决策的运行时执行:
ExecutionFrame 在运行时综合运用三个优化组件的信息:
Status ExecutionFrame::GetOrCreateMLValue(OrtValueIndex index,
OrtValue*& ort_value) {
const auto& alloc_plan = allocation_plan_[index];
// 优先使用MemoryPattern的预分配内存
if (memory_pattern_group_ && alloc_plan.alloc_kind == AllocKind::kAllocate) {
return AllocateWithMemoryPattern(index);
}
// fallback 到常规分配策略
switch (alloc_plan.alloc_kind) {
case AllocKind::kReuse: {
OrtValueIndex reused_index = alloc_plan.reused_buffer;
return ReuseBuffer(reused_index, index);
}
case AllocKind::kAllocate: {
return AllocateNewBuffer(index);
}
case AllocKind::kPreExisting: {
return UsePreExistingBuffer(index);
}
case AllocKind::kAllocateStatically: {
return UseStaticBuffer(index);
}
case AllocKind::kAllocateOutput: {
return AllocateOutputBuffer(index);
}
default:
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unknown allocation kind");
}
}
代码说明:GetOrCreateMLValue 方法根据 AllocationPlan 中定义的分配策略,优先使用 MemoryPattern 的预分配内存,其次是内存复用,最后才是新分配。这种优先级确保了最优的性能表现。方法支持所有分配类型,包括新分配、内存复用、外部内存、静态分配和输出分配,为不同的张量需求提供针对性的内存管理策略。
结论
ONNX Runtime 的三层优化体系通过 ExecutionPlan、AllocationPlan 和 MemoryPattern 的紧密协作,实现了从计算调度到内存管理的全方位优化。
体系架构优势:
-
1. 分层优化:三个组件各司其职,分别解决执行顺序、分配策略和内存布局问题 -
2. 严格依赖:构建顺序确保每个阶段都能基于前序阶段的优化结果 -
3. 渐进优化:从基础执行到高级内存布局,优化级别逐步提升 -
4. 场景适配:支持从动态形状的弹性管理到静态形状的极致优化
关键技术突破:
-
• 执行计划的基础性:为后续优化提供准确的执行顺序和依赖关系 -
• 分配计划的智能性:基于生命周期分析实现精准的内存复用 -
• 内存模式的极致性:通过预分配和布局优化达到接近手写Kernel的性能
未来发展方向:
随着AI模型和硬件平台的不断发展,这个三层优化体系将继续演进:
-
• 跨设备协同:优化多设备间的内存共享和数据传输 -
• 动态适应性:增强对动态形状和条件分支的优化能力 -
• 学习型优化:引入机器学习技术预测和优化内存使用模式 -
• 异构集成:更好地支持新兴硬件和专用加速器
深入理解这个三层优化体系的构建顺序和协同机制,对于充分发挥 ONNX Runtime 性能潜力、诊断优化推理应用具有至关重要的价值。

