📅 创建时间：2026-06-03 🏷️ 标签：#计算图 #DataflowGraph #动态图 #静态图 #TorchDynamo #FX #Tracing

📚 前置知识：[[/04-ai/01-llm-engineering/07-llm-evolution]]（LLM 发展脉络） 📚 相关知识：[[02-ir-fundamentals]]（IR 基础）[[04-mlir-architecture]]（MLIR 架构）[[05-operation-semantics]]（算子语义）

---

计算图的构建与表示 / Building and Representing Computational Graphs

┌──────────────────────────────────────────────────────────────────────────────┐
│  🎯 场景：torch.compile 的"图简化"困惑                                            │
├──────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│  你写了一个带条件的 PyTorch 模型：                                              │
│                                                                              │
│  ```python                                                                   │
│  def forward(self, x):                                                        │
│      if x.sum() > 0:                                                         │
│          return self.layer1(x)                                               │
│      else:                                                                   │
│          return self.layer2(x)                                               │
│  ```                                                                         │
│                                                                              │
│  用 torch.compile 编译后，发现 torch.compile 捕获的图里没有 if/else 分支，          │
│  只有一个被"简化"过的版本。你开始怀疑：编译器是不是把我的模型改坏了？                 │
│                                                                              │
└──────────────────────────────────────────────────────────────────────────────┘

第1节：计算图基础——什么是节点，什么是边

1.1 图的基本定义

计算图（Computational Graph） 是深度学习框架表示神经网络计算的核心数据结构。在数学上，它是一个有向无环图（DAG），其中：

元素	数学定义	在深度学习中的含义
节点（Node）	图中的顶点	算子（Operation/Tensor Expression）
边（Edge）	连接节点的线	张量（Tensor）流向
入度（In-degree）	指向节点的边数	算子的输入数量
出度（Out-degree）	从节点发出的边数	算子的输出数量

# 用伪代码展示一个简单计算图的构建过程
class Tensor:
    """张量节点"""
    def __init__(self, name, shape, dtype):
        self.name = name          # 节点标识
        self.shape = shape        # 张量形状
        self.dtype = dtype        # 数据类型
        self.producer = None     # 生产这个张量的算子
        self.consumers = []       # 消费这个张量的算子列表

class Operation:
    """算子节点"""
    def __init__(self, name, op_type):
        self.name = name          # 算子名称
        self.op_type = op_type    # 算子类型：matmul, relu, add...
        self.inputs = []          # 输入张量列表
        self.outputs = []         # 输出张量列表

# 构建一个简单的计算图：y = ReLU(Wx + b)
# 对应的计算图：
#   x (输入) ──→ MatMul ──→ Add(b) ──→ ReLU ──→ y (输出)
#   W (权重) ──↗

def build_simple_graph():
    """
    展示如何手动构建一个计算图
    """
    # 创建张量节点
    x = Tensor("x", shape=(32, 128), dtype="float32")      # 输入张量
    W = Tensor("W", shape=(128, 64), dtype="float32")     # 权重矩阵
    b = Tensor("b", shape=(64,), dtype="float32")          # 偏置向量
    
    # 创建算子节点
    matmul = Operation("matmul", "MatMul")
    matmul.inputs = [x, W]
    
    add = Operation("add", "Add")
    add.inputs = [matmul.outputs[0] if matmul.outputs else None, b]  # 需要 matmul 输出
    
    relu = Operation("relu", "ReLU")
    relu.inputs = [add.outputs[0] if add.outputs else None]
    
    # 建立边关系：生产者-消费者
    matmul.outputs = [Tensor("matmul_out", shape=(32, 64), dtype="float32")]
    matmul.outputs[0].producer = matmul
    
    add.outputs = [Tensor("add_out", shape=(32, 64), dtype="float32")]
    add.outputs[0].producer = add
    
    relu.outputs = [Tensor("y", shape=(32, 64), dtype="float32")]
    relu.outputs[0].producer = relu
    
    # 建立消费者关系
    matmul.outputs[0].consumers.append(add)
    add.outputs[0].consumers.append(relu)
    
    print(f"计算图拓扑：{x.name} → {matmul.name} → {add.name} → {relu.name} → {relu.outputs[0].name}")
    # 输出：计算图拓扑：x → matmul → add → relu → y
    
    return [matmul, add, relu]

build_simple_graph()

1.2 数据流 vs 控制流

计算图有两种基本的执行模式，理解它们的区别对于理解编译器的行为至关重要：

数据流图（Dataflow Graph） 是深度学习框架最常用的表示方式。在数据流图中：

• 边的存在表示数据依赖
• 节点执行顺序由数据依赖决定
• 没有"控制流"——所有条件分支都通过数据选择实现

"""
数据流图的思维实验：
原始代码（带控制流）：
    if condition:
        y = layer1(x)
    else:
        y = layer2(x)

数据流图等价表示（无控制流）：
    # 用 Mask 选择分支输出
    mask = condition.cast(dtype)              # True → 1.0, False → 0.0
    y = mask * layer1(x) + (1 - mask) * layer2(x)
"""

# 实际 PyTorch 中的条件分支编译结果（简化示意）
def dataflow_equivalence():
    """
    展示条件分支如何被编译成数据流图
    
    原始 Python 代码：
        if x.sum() > 0:
            return self.layer1(x)
        else:
            return self.layer2(x)
    
    编译器捕获的可能等价表示：
    """
    # 1. 计算两个分支（两路都执行）
    out1 = "layer1_output"   # 分支1的结果
    out2 = "layer2_output"   # 分支2的结果
    
    # 2. 计算条件
    condition = "x.sum() > 0"  # 布尔条件
    
    # 3. 用选择操作代替 if/else
    # out = tf.where(condition, out1, out2)  ← 典型的数据流表示
    
    print("条件分支的数据流等价表示：")
    print("    mask = (x.sum() > 0) ? 1.0 : 0.0")
    print("    y = mask * layer1(x) + (1-mask) * layer2(x)")
    print()
    print("⚠️ 代价：两个分支都会被执行，只是结果被 mask 选择了")

dataflow_equivalence()

第2节：动态图 vs 静态图——Eager 模式的优势和代价

2.1 动态图（Dynamic Graph / Eager Execution）

动态图模式下，Python 代码被逐行解释执行。每行代码立即执行并返回结果，不需要事先构建计算图。

# PyTorch 动态图示例
import torch

class DynamicGraphDemo(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(128, 64)
    
    def forward(self, x):
        # 每次调用 forward，图都是新构建的
        # 输入形状可以任意变化
        
        # 第1次调用：x.shape = (32, 128)
        # 第2次调用：x.shape = (16, 128)  ← 不同的 batch size
        
        h = self.linear(x)
        h = torch.relu(h)
        
        # 可以在运行时添加任意 Python 控制流
        if h.sum() > 0:
            h = h * 0.5  # 动态分支
        
        return h

# 每次调用都重新构建图
model = DynamicGraphDemo()
output1 = model(torch.randn(32, 128))  # 32 个样本
output2 = model(torch.randn(16, 128))  # 16 个样本，batch size 不同
output3 = model(torch.randn(8, 128))    # 8 个样本

print(f"动态图优势：支持任意形状输入，Python 控制流完全保留")
print(f"output1.shape = {output1.shape}")  # torch.Size([32, 64])
print(f"output2.shape = {output2.shape}")  # torch.Size([16, 64])
print(f"output3.shape = {output3.shape}")  # torch.Size([8, 64])

动态图的优势：

优势	说明
调试友好	可以在任意位置 print、检查变量、和 IDE 集成
灵活性	支持任意 Python 控制流、数据结构
无 Shape 限制	每次输入可以是不同形状
快速迭代	无需编译阶段，适合研究和实验

动态图的代价：

代价	说明
运行开销	每次前向传播都要解释 Python 代码
优化受限	编译器无法看到完整的计算图
部署困难	需要运行时环境，难以序列化

2.2 静态图（Static Graph / Graph Execution）

静态图模式下，Python 代码首先被"捕获"成一个计算图，然后图被编译和优化，最后编译后的图被执行。

"""
TensorFlow 1.x 静态图示例（伪代码）

# 阶段1：定义计算图（不执行任何操作）
with tf.Graph().as_default():
    # 这里的代码只是构建图结构，不会实际执行
    x = tf.placeholder(tf.float32, shape=[None, 128])
    W = tf.Variable(tf.random_normal([128, 64]))
    b = tf.Variable(tf.zeros([64]))
    
    # 矩阵乘法操作被添加到图中
    h = tf.matmul(x, W) + b
    y = tf.nn.relu(h)
    
    # 定义优化操作
    loss = tf.reduce_mean((y - y_target) ** 2)
    train_op = tf.train.AdamOptimizer(0.001).minimize(loss)
    
    # 阶段2：编译和优化（这里才真正优化图）
    # 编译过程会做：
    #   - 常量折叠
    #   - 算子融合
    #   - 内存优化
    #   - 设备放置
    
    # 阶段3：执行（session.run）
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for _ in range(1000):
            sess.run(train_op, feed_dict={x: batch_x, y_target: batch_y})
"""

print("静态图优势：编译优化、执行效率高、易于部署")
print()
print("静态图劣势：")
print("  - 调试困难（错误发生在图执行时，不在 Python 代码处）")
print("  - 灵活性差（不支持动态 Shape、控制流受限）")
print("  - 编译时间长（图越大，编译越慢）")

2.3 两种模式的对比

特性	动态图（PyTorch）	静态图（TF1.x）
执行模式	解释执行	先编译后执行
调试	逐行调试，IDE 完全支持	需要 tfdbg，困难
动态控制流	✅ 完全支持	❌ 有限支持
动态 Shape	✅ 完全支持	❌ 固定 Shape
执行效率	较低（解释开销）	较高（优化后）
部署难度	较难（需要 TorchScript）	较易（已编译）
编译时间	无需编译	较长
内存使用	动态管理	可优化

第3节：Graph Capture 策略——编译器如何捕获计算图

3.1 Tracing（追踪）策略

Tracing 是最直观的图捕获方法：执行一次函数，记录所有操作。

import torch
import torch.nn.functional as F

# ============================================================
# PyTorch FX (torch.fx) Tracing 示例
# ============================================================

class SimpleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(128, 64)
        self.linear2 = torch.nn.Linear(64, 32)
    
    def forward(self, x):
        h = self.linear1(x)
        h = F.relu(h)
        h = self.linear2(h)
        return h

model = SimpleModel()

# 用 torch.fx 进行 symbolic tracing
# 关键点：函数会被"执行一次"，执行过程中的操作被记录下来
traced = torch.fx.symbolic_trace(model)

print("=" * 60)
print("Tracing 捕获的计算图：")
print("=" * 60)
traced.graph.print_tabular()
# 输出示例：
# opcode         name    target                     args        kwargs
# -------------  ------  -------------------------  ----------  --------
# placeholder    x       x                          ()          {}
# call_module    linear1  linear1                    (x,)        {}
# call_function  relu    <built-in method relu...>  (linear1,)  {}
# call_module    linear2  linear2                    (relu,)     {}
# output         output  output                     (linear2,)  {}

print()
print("生成的 Python 代码：")
print(traced.code)
# 输出示例：
# def forward(self, x):
#     linear1 = self.linear1(x)
#     relu = torch.relu(linear1)
#     linear2 = self.linear2(relu)
#     return linear2

Tracing 的核心问题：无法捕获控制流

# ============================================================
# Tracing 无法捕获控制流的例子
# ============================================================

class ModelWithControlFlow(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(128, 64)
        self.layer2 = torch.nn.Linear(128, 64)
    
    def forward(self, x):
        # 动态控制流：取决于输入值
        if x.sum() > 0:
            return self.layer1(x)
        else:
            return self.layer2(x)

model_with_control = ModelWithControlFlow()

# tracing 只会捕获执行时走的分支
# 如果 x.sum() > 0，则只捕获 layer1 分支
traced_control = torch.fx.symbolic_trace(model_with_control)

print("控制流模型的 Tracing 结果：")
print(traced_control.code)
# 只会得到其中一个分支！另一个分支在图中完全消失了！

print()
print("⚠️ 警告：Tracing 的局限性")
print("=" * 60)
print("原始代码有两个分支：")
print("    if x.sum() > 0:")
print("        return self.layer1(x)  ← 分支 A")
print("    else:")
print("        return self.layer2(x)  ← 分支 B")
print()
print("Tracing 结果（假设 x.sum() > 0）：")
print("    return self.layer1(x)  ← 只有分支 A，B 完全丢失！")
print()
print("解决方案：")
print("  1. 使用 make_fx 或 random kwargs 强制走不同分支")
print("  2. 使用 TorchScript 解析 AST")
print("  3. 使用 wrapper 函数包装分支选择")

3.2 AST Parsing（抽象语法树解析）策略

AST Parsing 直接解析 Python 源代码，不执行函数，因此可以保留控制流结构。

# ============================================================
# TorchScript AST 解析示例
# ============================================================

"""
TorchScript 的工作原理：
1. 分析 Python 函数的 AST
2. 将 Python 语法转换为 TorchScript IR
3. TorchScript IR 支持控制流原语（if, while, for）

原始 Python：
    def forward(self, x):
        if x.sum() > 0:
            return self.layer1(x)
        else:
            return self.layer2(x)

TorchScript IR（简化表示）：
    graph(%self, %x):
        %cond = aten::sum(%x)
        %cmp = aten::gt(%cond, 0)
        %result: Tensor = prim::If(%cmp):
            block0():
                %out1 = aten::linear(%x, %self.layer1.weight, %self.layer1.bias)
                -> (%out1)
            block1():
                %out2 = aten::linear(%x, %self.layer2.weight, %self.layer2.bias)
                -> (%out2)
        return (%result)
"""

# TorchScript 支持 if/while/for 等控制流
class TorchScriptControlFlow(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(128, 64)
        self.layer2 = torch.nn.Linear(128, 64)
    
    def forward(self, x):
        # @torch.jit.script 会保留这个 if 结构
        if x.sum() > 0:
            return self.layer1(x)
        else:
            return self.layer2(x)

# 使用 TorchScript 编译
scripted_model = torch.jit.script(TorchScriptControlFlow())

print("TorchScript 保留控制流的 IR：")
print(scripted_model.graph)
# 输出包含 prim::If 节点，带有两个 block

3.3 Symbolic（符号化）策略

Symbolic 方法使用抽象解释（Abstract Interpretation）：不执行具体操作，而是跟踪值的"形状"和"类型"。

# ============================================================
# JAX Symbolic / Jaxpr 示例
# ============================================================

import jax
import jax.numpy as jnp

def simple_function(x, W, b):
    """JAX 函数"""
    h = jnp.matmul(x, W) + b
    return jnp.relu(h)

# 用 jax.make_jaxpr 生成符号化的中间表示
x = jnp.ones((32, 128))
W = jnp.ones((128, 64))
b = jnp.ones((64,))

# make_jaxpr 返回 Jaxpr（ JAX 程序表示）
jaxpr = jax.make_jaxpr(simple_function)(x, W, b)

print("JAX Symbolic 中间表示（Jaxpr）：")
print(jaxpr)
# 输出：
# { lambda ; a:i32[32,128] b:i32[128,64] c:i32[64].
#   let d:i32[32,64] = broadcast_in_dim[
#       broadcast_dimensions=(1, 2)
#       shape=(32, 64)
#     ] c
#       in let e:i32[32,64] = matmul_general[dimension_numbers=((1,),(0,),(0, 1),)]
#           a b
#           f:i32[32,64] = add e d
#           g:i32[32,64] = relu f
#       in { a }

print()
print("Symbolic 方法的特点：")
print("  - 保留数据类型和形状信息（i32 = int32）")
print("  - 操作被标记为抽象的（broadcast_in_dim, matmul_general）")
print("  - 可以进行符号化优化（常量折叠、代数简化）")

3.4 三种策略对比

特性	Tracing	AST Parsing	Symbolic
工作原理	执行并记录	解析源码	抽象解释
控制流	❌ 无法保留	✅ 完整保留	⚠️ 部分保留
副作用	❌ 会被执行	❌ 会执行	❌ 不会执行
执行速度	快（只执行一次）	快（不执行）	快（抽象执行）
动态 Shape	⚠️ 受限	✅ 支持	✅ 支持
实现复杂度	低	高	高
适用场景	大部分情况	复杂控制流	数学表达式
代表框架	torch.compile, FX	TorchScript	JAX, Relay (TVM)

第4节：PyTorch FX 实战

4.1 FX 基本用法

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.fx as fx

# ============================================================
# 完整的 FX Workflow
# ============================================================

class FeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        return x

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128 * 8 * 8, 256)
        self.fc2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

class FullModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = FeatureExtractor()
        self.classifier = Classifier()
    
    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

# 步骤1：创建模型
model = FullModel()
model.eval()  # FX tracing 需要 eval 模式

# 步骤2：Symbolic Tracing
# 创建假的输入来指导 tracing
example_input = torch.randn(1, 3, 32, 32)
traced_model = fx.symbolic_trace(model)

print("=" * 60)
print("FX 捕获的计算图：")
print("=" * 60)
traced_model.graph.print_tabular()

4.2 手动图操作

# ============================================================
# FX 图操作：手动插入、删除、替换节点
# ============================================================

class ConvBatchNormFusion(nn.Module):
    """演示如何手动融合 Conv + BatchNorm"""
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3, padding=1)
        self.bn = nn.BatchNorm2d(64)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x

model = ConvBatchNormFusion()
traced = fx.symbolic_trace(model)

# 图操作需要获取 Graph 对象
graph = traced.graph

print("原始图结构：")
for node in graph.nodes:
    print(f"  {node.op}: {node.name} = {node.target}")

# 手动遍历图找到 Conv 和 BatchNorm 节点
def fuse_conv_bn(graph_module):
    """
    融合 Conv + BatchNorm 的简化实现
    真实实现要复杂得多，这里演示核心思路
    """
    graph = graph_module.graph
    
    # 遍历所有节点
    for node in list(graph.nodes):
        if node.op == 'call_module' and isinstance(graph_module.get_submodule(node.target), nn.BatchNorm2d):
            # 找到 BatchNorm 节点，检查它的输入
            bn_node = node
            conv_node = None
            
            # 向上遍历找 Conv
            for inp in bn_node.all_input_nodes:
                if inp.op == 'call_module':
                    conv_node = inp
                    break
            
            if conv_node is not None:
                print(f"  发现 Fusion 候选: {conv_node.name} -> {bn_node.name}")
                # 实际融合：修改权重、删除 BN 节点
                # 这里简化处理，真实实现需要：
                # 1. 计算融合后的 BatchNorm 参数
                # 2. 更新 Conv 的权重和偏置
                # 3. 将 BN 的后续节点指向 Conv

fuse_conv_bn(traced)

4.3 使用 Proxy 进行调试

# ============================================================
# FX Proxy 机制：允许在 tracing 时插入打印
# ============================================================

def register_hooks(module, prefix=''):
    """为模块的每个子模块注册前向钩子"""
    for name, child in module.named_children():
        full_name = f"{prefix}.{name}" if prefix else name
        def forward_hook(mod, input, output, name=full_name):
            print(f"  [HOOK] {name}: input shape = {input[0].shape}, output shape = {output.shape}")
        child.register_forward_hook(forward_hook)
        register_hooks(child, full_name)

class DebugModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(128, 64)
        self.linear2 = nn.Linear(64, 32)
    
    def forward(self, x):
        h = self.linear1(x)
        print(f"  [DEBUG] After linear1: shape={h.shape}, mean={h.mean().item():.4f}")
        h = F.relu(h)
        h = self.linear2(h)
        print(f"  [DEBUG] After linear2: shape={h.shape}, mean={h.mean().item():.4f}")
        return h

# 注册钩子
register_hooks(DebugModel())

# tracing 时钩子不会生效（因为是符号执行）
# 但我们可以手动遍历
model = DebugModel()
traced = fx.symbolic_trace(model)

# 打印所有节点
print("调试追踪：手动遍历节点")
for node in traced.graph.nodes:
    if node.op == 'call_module':
        print(f"  节点: {node.name}, 目标: {node.target}")
        print(f"    输入: {[inp.name for inp in node.all_input_nodes]}")

第5节：动态图的特殊挑战

5.1 Python 特定语法的处理

# ============================================================
# 动态图挑战1：print 语句
# ============================================================

class PrintDebugModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(128, 64)
    
    def forward(self, x):
        # print 在 Python 中是语句，不是表达式
        print(f"Input shape: {x.shape}")
        return self.linear(x)

# FX 无法处理 print（不在 forward 的返回值中）
try:
    traced = fx.symbolic_trace(PrintDebugModel())
except Exception as e:
    print(f"⚠️ FX 无法捕获 print：{e}")

# 解决方案1：使用 assert
class PrintFixModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(128, 64)
    
    def forward(self, x):
        # 用 assert 代替 print（assert 也是语句）
        assert x.shape[-1] == 128, f"Expected last dim 128, got {x.shape[-1]}"
        return self.linear(x)

# 解决方案2：使用 wrapper
def wrap_with_print(forward_fn):
    """包装函数，添加打印逻辑"""
    def wrapped(*args, **kwargs):
        result = forward_fn(*args, **kwargs)
        return result
    return wrapped

# ============================================================
# 动态图挑战2：字典和集合
# ============================================================

class DictModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = {'a': nn.Linear(128, 64), 'b': nn.Linear(128, 64)}
    
    def forward(self, x, key):
        # 字典索引是动态的
        return self.weights[key](x)

# FX 无法捕获字典索引（需要静态分析 key 的所有可能值）
try:
    traced = fx.symbolic_trace(DictModel())
except Exception as e:
    print(f"⚠️ FX 无法捕获字典索引：{e}")

# 解决方案：使用 nn.ModuleList 或 nn.ModuleDict
class FixedDictModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.ModuleDict({
            'a': nn.Linear(128, 64),
            'b': nn.Linear(128, 64)
        })
    
    def forward(self, x, key):
        return self.weights[key](x)

# 这样 FX 可以捕获（ModuleDict 有确定性索引）
traced = fx.symbolic_trace(FixedDictModel())
print("✓ ModuleDict 可以被 FX 捕获")

5.2 数据依赖的条件执行

# ============================================================
# 动态图挑战3：数据依赖的条件
# ============================================================

class DataDependentCondition(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 32)
        )
    
    def forward(self, x):
        # 条件依赖于 x 的值！这使得静态分析困难
        threshold = x.mean()  # 动态计算的阈值
        if threshold > 0:
            return self.net(x)
        else:
            return x

# FX tracing 会执行一次函数
# 假设 x.mean() > 0，则只捕获 then 分支
traced = fx.symbolic_trace(DataDependentCondition())

print("⚠️ Data Dependent Condition 的 Tracing 结果：")
print(traced.code)
# 只有一个分支！另一个分支完全消失！

print()
print("⚠️ 运行时问题：")
print("  当 threshold <= 0 时，模型会执行未追踪的分支 B")
print("  可能导致：")
print("    1. 运行时错误（形状不匹配）")
print("    2. 性能问题（重新走 Python 解释器）")
print("    3. 数值问题（梯度无法正确计算）")

5.3 副作用的处理

# ============================================================
# 动态图挑战4：副作用
# ============================================================

class SideEffectModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.counter = 0  # 模型内部状态
    
    def forward(self, x):
        self.counter += 1  # 修改内部状态！
        print(f"Forward pass #{self.counter}")
        return x * self.counter

model = SideEffectModel()

# Tracing 时 counter 被记录为常量 0
traced = fx.symbolic_trace(model)

print("副作用问题演示：")
print("  原始模型：每次调用 counter 增加")
print("  Traced 模型：counter 被固定为 0（静态值）")
print()
print("  原始模型调用 3 次：")
for i in range(3):
    result = model(torch.ones(1, 8))
    print(f"    Pass {i+1}: counter = {model.counter}")

print()
print("  Traced 模型调用 3 次：")
for i in range(3):
    result = traced(torch.ones(1, 8))
    print(f"    Pass {i+1}: 无副作用")

第6节：图的等价性——编译器捕获的图和原始代码语义等价吗

6.1 等价性的定义

计算图的语义等价性有几种不同的定义：

类型	定义	检查方法
结构等价	节点和边完全相同	简单图同构
计算等价	相同输入产生相同输出	随机测试
渐近等价	输出在数值误差内相同	浮点比较
语义等价	包括副作用的等价	需要规范定义

6.2 常见的不等价情况

# ============================================================
# 图捕获可能引入的不等价情况
# ============================================================

print("=" * 60)
print("图捕获可能导致不等价的几种情况：")
print("=" * 60)

print()
print("1️⃣ 控制流丢失")
print("   原始：if x.sum() > 0: return A else: return B")
print("   Traced：return A  （假设 x.sum() > 0 为 True）")
print()
print("2️⃣ 数值误差累积")
print("   静态图可能做常数折叠：(W1 @ W2) @ x = W1 @ (W2 @ x)")
print("   浮点运算顺序不同可能导致数值差异")
print()
print("3️⃣ 副作用丢失")
print("   原始：x = layer(x); print(x.sum())")
print("   Traced：x = layer(x)  （print 丢失）")
print()
print("4️⃣ 动态 Shape/dtype 丢失")
print("   原始：处理任意 batch size")
print("   Traced：固定 batch size = 1")
print()
print("5️⃣ 随机操作")
print("   原始：torch.randn()")
print("   Traced：固定随机种子")

第7节：子图编译——哪些部分值得编译

7.1 编译收益分析

# ============================================================
# 子图编译策略
# ============================================================

def analyze_subgraph_benefits():
    """
    分析哪些子图值得编译
    """
    print("=" * 60)
    print("子图编译收益分析")
    print("=" * 60)
    
    scenarios = [
        ("小模型 (<1M 参数)", "低", "编译开销 > 执行收益"),
        ("大模型 (100M+ 参数)", "高", "执行时间长，编译一次收益多次"),
        ("计算密集型 (CNN/Transformer)", "高", "矩阵乘法可大量优化"),
        ("内存密集型 (大量小操作)", "中", "算子融合可减少内存访问"),
        ("IO 密集型 (数据加载瓶颈)", "低", "计算优化无意义"),
        ("动态控制流多", "低", "图捕获困难，收益不确定"),
    ]
    
    print()
    print("| 场景 | 编译收益 | 原因 |")
    print("|------|----------|------|")
    for scene, benefit, reason in scenarios:
        print(f"| {scene} | {benefit} | {reason} |")

analyze_subgraph_benefits()

print()
print("=" * 60)
print("子图编译策略")
print("=" * 60)

print("""
策略1：整图编译
    torch.compile(model)
    优点：全局优化
    缺点：编译慢，无法处理动态部分

策略2：部分编译（torch.compile 的 backend='aot'）
    # 只编译部分模块
    compiled_layer = torch.compile(model.layer)
    
策略3：分离编译
    # 动态部分用解释器，静态部分用编译器
    def forward(self, x):
        if x.shape[0] < 32:  # 动态分支
            return self.fast_path(x)
        else:  # 静态分支
            return self.compiled_path(x)

策略4：JIT 编译（延迟编译）
    # 第一次执行时编译
    @torch.compile
    def my_function(x):
        return x @ x.T
""")

7.2 TorchDynamo 的智能图捕获

# ============================================================
# TorchDynamo 的自适应编译
# ============================================================

"""
TorchDynamo（torch.compile 的核心）是基于 Python Bytecode  interception 的图捕获系统：

工作原理：
1. 拦截 Python Bytecode（使用 eval_frame API）
2. 将 Python 代码翻译成 FX Graph
3. 对图进行优化
4. 生成优化后的 Python Bytecode

优势：
- 可以处理任意的 Python 代码
- 支持动态控制流（通过 guards）
- 支持动态 Shape
- 增量编译（只编译变化的图）

Guards 机制：
DYNAMO 会在编译时添加"守卫"（guards），用于检查运行时的条件：
- 输入 Shape
- 输入 dtype
- 模型参数是否改变
- 全局变量是否改变

如果 guards 失败，DYNAMO 会重新编译
"""

print("TorchDynamo 的图捕获流程：")
print("""
  Python 代码
      ↓
  Bytecode interception (eval_frame)
      ↓
  FX Graph construction (AC (Ahead-of-time) or AOT)
      ↓
  Optimization passes
      ↓
  Guard generation
      ↓
  Compiled Python Bytecode
      ↓
  Guard failure? → Re-compile
""")

---

升华

┌──────────────────────────────────────────────────────────────────────────────┐
│  📚 计算图表示：核心原则                                                       │
├──────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│  1. **图是程序的中间表示**：计算图是连接高级语言和硬件指令的桥梁                 │
│                                                                              │
│  2. **动态 vs 静态是权衡**：动态图灵活但慢，静态图快但不灵活                    │
│                                                                              │
│  3. **图捕获策略决定能力边界**：Tracing 简单但不支持控制流                      │
│                                                                              │
│  4. **等价性需要验证**：编译器捕获的图不一定和原始代码语义等价                   │
│                                                                              │
└──────────────────────────────────────────────────────────────────────────────┘

"AI 可查 vs 必须理解"清单

必须理解（不理解就等于不会）：

• 🔴 计算图的基本概念：节点=算子，边=张量，这是所有深度学习框架的基础
• 🔴 动态图 vs 静态图的权衡：为什么 PyTorch 选择动态图，TensorFlow 为什么转向 eager 模式
• 🔴 Tracing 的局限性：为什么 Tracing 无法捕获控制流，替代方案是什么
• 🔴 图捕获的三种策略：Tracing、AST Parsing、Symbolic 的适用场景
• 🔴 FX 的基本用法：symbolic_trace 能做什么，不能做什么

AI 可查（知道去哪查就行）：

• ✅ 具体框架的 API 细节：torch.fx 的 API 变化很快，可以查文档
• ✅ 特定模型的图结构：用 model.graph.print_tabular() 查看
• ✅ TorchDynamo 的 guards 机制：内部实现复杂，用到时查文档
• ✅ 特定编译优化的效果：benchmark 数据，不同硬件差异很大
• ✅ 子图编译的边界情况：特定模型可能有特殊问题

学习状态：🟡 开始学习