PyTorch 算子融合详解：如何利用 NNC 编译器减少显存读写以提升推理效率

在 AI 模型推理阶段，提升速度的关键往往不在于计算本身的复杂度，而在于数据在存储介质（如GPU HBM）和计算单元（CUDA Core）之间传输的效率。算子融合（Operator Fusion）正是解决这一问题的核心技术，它通过将多个计算核（Kernel）合并成一个，极大地减少了中间结果的显存读写次数，从而提高计算密度和推理速度。

PyTorch 内部集成了强大的优化器，如基于 TorchScript 的 NNC（Neural Network Compiler），它能够自动识别并融合一系列连续的、且符合融合条件的算子，例如卷积(Conv)、批量归一化(BN)、ReLU等。

本文将详细演示如何使用 torch.jit.script 触发 PyTorch 的编译器进行算子融合，并对比融合前后模型的推理性能。

算子融合的核心原理

想象一个序列操作：A = Conv(X) -> B = ReLU(A)。在未融合的情况下，GPU需要执行两个独立的 Kernel：
1. 执行 Conv Kernel，结果 A 写入显存。
2. 执行 ReLU Kernel，从显存读取 A，计算 B，结果 B 写入显存。

如果将这两个操作融合，GPU只执行一个融合后的 Kernel：
1. 执行 Fused Kernel (Conv + ReLU)，计算过程中数据 A 保留在寄存器或 L1/L2 缓存中，直接传递给 ReLU 部分计算 B，最终 B 写入显存。

显而易见，融合后的版本消除了中间数据 A 的显存读写开销，这在端侧或高并发推理场景中，能带来显著的性能提升。

实战：使用 PyTorch JIT 脚本化实现融合

我们以一个简单的 Conv2d 后接 ReLU 的模型为例进行演示。请注意，为了获得准确的性能对比，以下代码必须在支持 CUDA 的环境中运行。

1. 定义模型和基线测试

首先，我们定义一个标准的 PyTorch 模型，并使用 Eager Mode 进行基准性能测试。

import torch
import torch.nn as nn
import time

# 1. 定义一个适合融合的简单模型 (Conv -> ReLU)
class SimpleFusionModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 关键的融合目标：Conv 和 ReLU
        self.conv = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x) # 算子融合主要发生在此处
        x = self.pool(x)
        return x

# 设置设备和输入数据
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if not torch.cuda.is_available():
    print("Warning: CUDA not available. Results may not show typical fusion gains.")

model = SimpleFusionModel().to(device).eval()
# 使用较大的 Batch Size 模拟实际推理场景
input_data = torch.randn(64, 32, 56, 56).to(device)

# 预热
for _ in range(10):
    _ = model(input_data)

# 定义基准测试函数
def benchmark(model, input_data, iterations=100):
    torch.cuda.synchronize()
    start_time = time.time()
    with torch.no_grad():
        for _ in range(iterations):
            _ = model(input_data)
    torch.cuda.synchronize()
    end_time = time.time()
    return (end_time - start_time) * 1000 / iterations

# --- 基准测试：Eager Mode (未融合) ---
eager_latency = benchmark(model, input_data)
print(f"\n1. Eager Mode (Baseline) Latency: {eager_latency:.3f} ms")

2. 使用 JIT 脚本化并测试融合性能

使用 torch.jit.script 将模型转换为 TorchScript 格式。在这一过程中，PyTorch 的编译器后端（如 NNC）会自动分析计算图，识别可融合的模式，并生成优化的融合 Kernel。

# --- 算子融合：JIT Scripting (利用 NNC) ---
# 使用 torch.jit.script 编译模型，触发 PyTorch 后端进行优化和融合
try:
    # 必须保证模型在评估模式且在GPU上
    scripted_model = torch.jit.script(model)
except Exception as e:
    print(f"Scripting failed: {e}")
    exit()

# 预热 JIT 模型
for _ in range(10):
    _ = scripted_model(input_data)

# 基准测试：Scripted Mode (融合后)
scripted_latency = benchmark(scripted_model, input_data)
print(f"2. Scripted (Fused) Mode Latency: {scripted_latency:.3f} ms")

# 观察优化效果
if eager_latency > scripted_latency:
    speedup = ((eager_latency - scripted_latency) / eager_latency) * 100
    print(f"\n== 性能总结 ==")
    print(f"JIT 融合带来的性能提升: {speedup:.2f}%")
    print("融合成功，有效减少了显存读写开销。")
else:
    print("性能提升不明显或未发生融合。请检查模型结构或运行环境。")

3. 结果分析

运行上述代码，你会观察到 Scripted (Fused) Mode 的延迟通常低于 Eager Mode。这是因为 NNC 成功地将 Conv2d 和紧随其后的 ReLU 算子融合为一个统一的 CUDA Kernel。数据在 Conv 阶段计算完成后，无需离开 L1/L2 缓存，直接传递给 ReLU 部分完成计算，从而避免了昂贵的 HBM 显存访问。

关键提醒：
1. 适用性： 算子融合主要对连续的 Element-wise 运算（如加法、激活函数、归一化）以及某些简单的逐点操作（如 Conv + Bias）效果显著。
2. 限制： JIT 脚本化并非总是完美的，复杂的控制流（如 if/else）或 Python 原生容器可能导致脚本化失败或优化受限。在这种情况下，推荐使用 torch.jit.trace，尽管 Tracing 牺牲了动态特性但通常能更好地进行优化。