企业官网建设流程全解析

PyTorch Static Runtime 静态运行时：面向 CPU 推理的 TorchScript 优化执行引擎

【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch

导读

Static Runtime 是 PyTorch 提供的一个面向 CPU 推理场景的优化运行时，可作为 TorchScript JIT 解释器的"即插即用"替代方案（drop-in replacement），在 C++ 与 Python 中均可使用。它的核心思路是通过**内存规划（Memory Planning）**把中间张量的分配合并到单块连续内存（slab）中，从而大幅降低推理过程中的动态分配开销。本文以官方文档 torch/csrc/jit/runtime/static/README.md 为骨架，结合仓库源码（impl.h、impl.cpp、ops.h、init.cpp）与测试用例（test/test_static_runtime.py），深入讲解其适用场景、线程模型、内存规划原理、算子注册机制与核心数据结构。读完本文，你将掌握如何用 Static Runtime 替换 JIT 解释器、如何配置其选项以获得最优推理性能，并理解其底层实现原理。

这是一个实验性功能（experimental feature）。在使用前请确认你的模型与运行环境满足下文所述的前提条件。

一、为什么需要 Static Runtime：适用场景与前提假设

Static Runtime 专门为 CPU 推理优化，其价值体现在以下三个条件下最为明显（对应官方文档的适用场景）：

1. 模型几乎不含控制流（very little control flow）。Static Runtime 在加载时对整个图做静态分析与规划，控制流（prim::If、prim::Loop、prim::fork）虽受支持，但会引入额外的运行时决策与子块管理开销。
2. PyTorch 的运行时开销（如张量创建）占模型运行时间不可忽视的比例。尤其是当张量分配消耗大量时间时，Static Runtime 的帮助最大——中间张量的内存被合并进单块 slab，推理过程中绝大多数动态分配被消除。
3. 推理性能极其重要（inference performance is extremely important）。

使用前提假设（Assumptions）

官方文档明确了当前使用该功能的一组假设：

• 仅推理执行（Inference only），且仅 CPU；
• 静态输入 dtype（Static input dtypes）；
• 静态输入 shape：运行时虽然支持动态 shape，但过度的动态 shape 会显著降低性能（原因见下文"标准尺寸调整"一节）。

此外，从源码实现看还有一条硬性约束：impl.cpp 中的canEnableStaticRuntime会检查图是否可被 Static Runtime 支持，其注释明确指出Graph must be frozen——在构造StaticModule时若传入的模块尚未冻结，构造函数会自动调用freeze_module完成冻结（见impl.cpp第 337-340 行），以保证图中没有可变的 Python 状态干扰静态分析。

二、线程模型与两种执行模式

Static Runtime 支持两种执行模式（对应文档 "Threading model" 章节），两者的差别在于并发模型不同。

模式 1：单线程执行（仅算子内部并行）

单线程执行，除算子内并行（intra-op parallelism）外不做任何线程级并行。此时有两种等价写法：

// m 是 TorchScript 模块 auto runtime = StaticRuntime(m, opts); auto output = runtime.run(args, kwargs);

或先调用PrepareForStaticRuntime对图做预处理，再构造运行时：

auto mod = PrepareForStaticRuntime(m); auto runtime = StaticRuntime(mod, opts); auto output = runtime.run(args, kwargs);

模式 2：数据并行式多线程（每线程一个运行时实例）

类似于数据并行：同一模型在不同线程上同时处理不同输入。此时必须先用PrepareForStaticRuntime准备图。一个模型对应一个StaticModule（或官方术语 InferenceModule）实例，每个运行线程对应一个StaticRuntime实例。为避免频繁构造StaticRuntime，建议用同步栈（如boost::lockfree::stack）缓存所有运行时实例：

// 初始化阶段：每个模型只做一次 auto mod = PrepareForStaticRuntime(m); // 128 对大多数场景足够，可按需调整 boost::lockfree::stack<std::shared_ptr<StaticRuntime>, boost::lockfree::fixed_sized<true>> pool(128); // 推理阶段 std::shared_ptr<StaticRuntime> runtime = nullptr; pool.pop(runtime); if (!runtime) { runtime = std::make_shared<StaticRuntime>(mod, opts); } auto output = runtime->run(args, kwargs); pool.push(runtime);

在 impl.h 的类注释中可以看到与之对应的StaticModule/StaticRuntime用法（StaticRuntime(*module)构造并自行管理内存）。两种模式共同的硬性约束为官方文档加粗强调的：

In both modes,StaticRuntimemay not be used after its associatedStaticModuleis destructed!（两种模式下，StaticRuntime都不能在其关联的StaticModule析构之后继续使用！）

原因在于StaticRuntime内部以非拥有（unowned）引用的方式指向StaticModule持有的共享数据结构（见下文"结构图与生命周期"）。

三、内存规划（Memory Planning）：性能优化的核心

Static Runtime 的加速本质来自其内存规划器（Memory Planner）做的两件事（对应文档 "Memory Planning" 章节）：

1. 合并内部张量存储的分配（coalesce internal allocations for tensor storage）；
2. 通过静态分析确定如何高效复用内存（static analysis to figure out how to efficiently reuse memory）。

3.1 标准尺寸调整（Standard Resizing）

Static Runtime 会在第一次推理迭代时记录它见到的每个受管中间张量（managed intermediate tensor）所需的空间。一个中间张量成为"受管"（managed）需满足两个条件：

1. 产生它的算子具有 out variant。Out variant 是对算子的包装，概念上把算子签名从Tensor some_op(const Tensor& some_arg)变换为void some_op(Tensor& output, const Tensor& some_arg)。Out variant 通过REGISTER_OPERATOR_FUNCTOR宏注册，详见下文"注册算子"。
2. 该张量不对图输出产生别名（does not alias a graph output）。输出张量由内存规划器单独处理，见"受管输出张量"。

在上述算法下，静态分析把张量归入若干StorageGroup（存储组）。同一存储组内的张量共享内存，而两个张量能够进入同一存储组的前提是它们的生命周期（lifetime）不重叠。

在后续迭代中，Static Runtime 在每次 run 开始时分配张量缓冲区，分配的内存大小为：

sum([max(tensor.size()) for tensor in storage_groups])

即各存储组内最大张量尺寸之和。如果后续运行中某张量需要超过已分配空间，就会发生一次动态分配——这正是动态 shape 会降低性能的原因。在标准尺寸调整策略下，Static Runtime 会在迭代结束时记录每个存储组内新的最大张量尺寸，并在下一次迭代分配可能更大的缓冲区。

3.2 受管输出张量（Managed Output Tensors）

StaticRuntime可以通过StaticModuleOptions中的manage_output_tensors选项选择性地管理输出张量。开启该选项后，输出张量的分配会被合并在一起。注意：包含输出张量的缓冲区与包含中间张量的缓冲区是分离的——前者需要存活到推理 run 结束之后（因为调用方还要使用输出），而后者需要在 run 结束时被释放。

底层实现上，每个返回的Tensor中都保存了一个指向输出 arena 的引用计数指针（refcounted pointer），arena 被显式销毁。

3.3StaticModuleOptions：内存相关配置项

StaticModuleOptions定义于 impl.h，其中与内存规划直接相关的字段如下（全部带默认值）：

在 impl.cpp 的OptimizeGraph中可以看到这些选项如何被消费：若enable_tensorexpr_fusion为真且提供了sample_inputs，则先执行performTensorExprFusion，随后依次执行Inline、ConstantPropagation、Canonicalize、RemoveTensorMutation等 pass 完成图优化。

四、注册算子（Registering Ops）：三种执行模式

Static Runtime 对图中的每个torch::jit::Node选择执行实现，共有三种模式（对应文档 "Registering Ops" 章节）：

1. Out variants：返回张量且输出可被内存规划器管理的算子。通过 ops.h 中的REGISTER_OPERATOR_FUNCTOR宏注册，宏的使用方式如下：

REGISTER_OPERATOR_FUNCTOR( aten::op_name, aten_op_name, // 该宏会生成一个 struct，此字段为其命名 [](torch::jit::Node* n) -> SROperator { // 该机制让我们可以支持 schema 的子集 if (n->matches(some_schema)) { return some_overload; } else if (n->matches(another_schema)) { return another_overload; } return nullptr; })

其中SROperator是std::function<void(ProcessedNode*)>的类型别名（定义见 impl.h），ProcessedNode的细节见下文"实现细节"。

2. Native functions：与 out variant 类似，区别在于输出无法被管理——因为算子的返回类型不是张量，或者它是视图算子（view op，返回张量别名而非新张量）。通过REGISTER_NATIVE_OPERATOR_FUNCTOR宏注册（ops.h），用法与REGISTER_OPERATOR_FUNCTOR相同。

3. JIT fallback：Static Runtime 对该算子没有实现，因此回退选用 JIT 解释器使用的实现。

加载模型时的算子选择流程

对图中的每个torch::jit::Node，按以下顺序决策：

1. 若注册了 out variant，把节点传给生成SROperator的函数；结果非nullptr则使用该算子；
2. 若注册了 native function，把节点传给生成SROperator的函数；结果非nullptr则使用该算子；
3. 使用 JIT 实现；若 JIT 实现也不存在，Static Runtime 会抛出异常。

从源码看，这一选择逻辑落在ProcessedFunction上（impl.h），它用枚举Kind { kOutVariant, kNativeFunction, kInterpreterFallback }标记每个节点最终选定的执行方式，ProcessedFunction在StaticModule构造时初始化。

五、实现细节（Implementation Details）

5.1 核心数据结构与生命周期

官方文档给出了核心数据结构的引用关系图（箭头A→B表示A存储了对B的引用；若是非拥有引用，A不能活得比B（或B间接引用的任何对象）更久；若是拥有引用，则A与B生命周期相同）：

IValue array◄────────────────┐─────────────────────────────────────────┐ ▲ │ Owns │ Owns │ │ ┌───────────────────────────────►ProcessedNode───────►BlockRunner │Owns │ │ │ │ │ Owns │ │ Owns │ │ StaticModule◄───────────StaticRuntime───────────►BlockRunner────────►MemoryPlanner │ ▼ │ │ │ │ │ ... Owns│ │ │ │ │ ▼ │ │ │ │ BlockInfo◄├───────────────────────────────────────────┘──────────────────┘ │ │ │ Owns│ │ ▼ │ ProcessedFunction ◄─────────────────────────────────────────────────────────────────────────────┘

5.2StaticModule与StaticRuntime

• StaticModule由torch::jit::Module构造，用于生成StaticRuntime实例。每个StaticModule恰好缓存一个StaticRuntime实例，通过runtime()访问时惰性初始化。
• StaticModule::operator()可直接用于推理，底层只是转发给缓存 runtime 的StaticRuntime::operator()。由此带来的一个后果是：StaticModule::operator()不是线程安全的。
• 多线程场景下的正确用法是给每个线程各自一个StaticRuntime实例。新实例可以直接创建（StaticRuntime(static_module)），也可以从已有实例clone()。
• StaticModule接收一组选项（StaticModuleOptions）控制其派生的所有 runtime 实例的行为。
• 内部地，StaticRuntime拥有一个IValue数组，被所有BlockRunner和ProcessedNode引用；运行时产生的所有值都存储在该数组中。

5.3BlockRunner

BlockRunner代表图中的单个子块（sub-block）。每个图至少有一个对应顶层块的BlockRunner，StaticRuntime通过调用(*top_level_block)(args, kwargs)开始推理。每个BlockRunner拥有自己的MemoryPlanner和一组ProcessedNode。带子块的特殊节点（如prim::If）可能拥有BlockRunner，由算子实现负责调用对应子块的BlockRunner。

5.4MemoryPlanner

MemoryPlanner是抽象基类，每个子类实现一种内存规划算法（详见"内存规划"一节）。除了张量内存规划外，它还封装了另两类优化：

• 受管输出张量（Managed output tensors，见 3.2 节）；
• 借用IValue（Borrowed IValues）：只做输入解包的算子（如dict_unpack）可能产生弱引用以避免引用计数 bump，MemoryPlanner需要负责适当地销毁这些借用。

5.5ProcessedNode与ProcessedFunction

• ProcessedNode是单个算子的抽象。每个ProcessedNode存储对StaticRuntime的IValue数组的非拥有引用，并知道如何把输入/输出索引映射到数组下标（processed_node->output(i)返回ivalue_array[some_set_of_indices[i]]的引用）。
• 每个ProcessedNode持有一个ProcessedFunction，表示实际要执行的算子。ProcessedFunction在StaticModule构造时按上文"算子选择流程"初始化。
• 注意：所有ProcessedFunction在全部 runtime 实例之间共享，因此所有ProcessedFunction必须是线程安全的。

5.6ProcessedNodeMetadata

ProcessedNodeMetadata为ProcessedNode保存各种"额外"字段，通常未被使用，但少数算子需要额外机制：

• prim::If：根据条件判断，持有 true/false 两个子块的BlockRunner；
• prim::Loop：持有循环子块执行的BlockRunner；
• prim::fork：持有torch::jit::TaskLauncher（即std::function<void(std::function<void()>)>），负责分叉图的执行。

5.7 异步执行（Asynchronous Execution）

StaticRuntime::runAsync()API 允许把异步操作放到传入的TaskLauncher上执行：

• 父图（parent graph）在调用线程上内联执行（inline execution）；
• 像prim::fork这样的异步操作在传入的 launcher 上执行；
• 若未提供 launcher，则通过at::launch执行，即运行在inter-op 线程池上。

六、Python 端使用与测试佐证

Static Runtime 的 Python 绑定通过 init.cpp 暴露，核心入口是torch::jit::StaticModule类，提供：

• __call__：直接推理，内部把 py::args/kwargs 转成IValue后调用self(arg_ivalues, kwarg_ivalues)；
• benchmark(args, kwargs, warmup_runs, main_runs)：整体基准测试；
• benchmark_individual_ops(args, kwargs, warmup_runs, main_runs)：逐算子基准测试，返回IndividualMetrics，包含setup_time、memory_alloc_time、memory_dealloc_time、output_dealloc_time、first_iter_time、total_time、out_nodes_count、total_nodes_count、time_per_node、time_per_node_type、percent_per_node_type、instances_per_node_type、out_nodes等统计字段；
• runAsync(args, kwargs)：异步执行入口。

Python 侧的用户包装类位于 torch/testing/_internal/static_module.py，其StaticModule类对 C++ 绑定做了轻量封装（__call__、benchmark、runAsync、benchmark_individual_ops）。

仓库测试 test/test_static_runtime.py 是理解实际用法的绝佳参考，覆盖了：

• fork/wait 系列：test_fork_wait_1至test_fork_wait_4及对应_async变体，验证StaticModule(torch_graph)后调用static_runtime_module.runAsync((inp1, inp2), {})的异步执行路径，包括嵌套 fork/wait、多 fork/wait、异常传播；
• 真实模型：test_multihead_attention_layer直接把attention模块传入StaticModule进行推理；test_multihead_attention_layer_benchmark展示benchmark([src, src, src, src_mask], {}, 2, 2)与benchmark_individual_ops的基准用法；test_mlp复用了./bench/dlrm_s_benchmark.sh的参数构造多层感知机验证。

七、使用建议与限制总结

基于官方文档与源码，使用 Static Runtime 时请留意以下要点：

1. 确认模型适合：控制流极少、张量分配耗时占比高、推理性能要求苛刻的 CPU 模型收益最大；输入 shape 尽量静态化，避免触发动态分配回退。
2. 先冻结再构造：虽然构造StaticModule时会自动freeze_module（若传入的模块未冻结，见 impl.cpp），但建议在模型侧预先eval()与冻结，使图优化更可控。
3. 多线程推理遵循"一模型一StaticModule、一线程一StaticRuntime"，并用 lockfree 栈缓存 runtime 实例；切勿在StaticModule析构后继续使用其StaticRuntime。
4. 按需开启选项：默认enable_out_variant=true、optimize_memory=true；若模型输出需在 run 之后长期存活，可考虑manage_output_tensors=true；如需 TensorExpr 算子融合，开启enable_tensorexpr_fusion并提供sample_inputs，但需知悉它会抑制use_copy_variants/use_maybe_copy_variants对应的替换 pass。
5. 善用基准工具：通过 Python 绑定的benchmark与benchmark_individual_ops（配合IndividualMetrics）可以量化 setup、内存分配/释放、逐节点耗时，快速定位收益点与瓶颈。

Static Runtime 作为实验性功能，其适用范围与前提假设（仅推理、仅 CPU、静态 dtype/shape 优先）直接决定了收益上限。理解其内存规划与算子注册机制后，你可以在自己的推理服务中精准判断它是否适用，并针对性地配置StaticModuleOptions以获得最佳的 CPU 推理性能。

【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch