09. AI Compilers and Graph Optimization | AI 编译器与计算图优化

难度： Hard | 环境： CPU-first | 标签： 编译系统, AI Compiler, Graph Optimization | 目标人群： 编译与系统入门者

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行：

(国内推荐：魔搭社区免费实例)

---

本节导读

同一段“矩阵乘 → 加偏置 → 激活”，可以逐个调用算子，也可以交给编译器改写。两种方式计算相同的结果，却可能使用不同的中间张量、启动次数和后端路径。

本节从这条计算链出发，比较图变换改变的读写与启动，再检查目标后端是否支持所需算子、精度和工作区。完成后，你应能区分“变换满足条件”和“实际执行更快”；融合细节进入算子融合页，真实执行与计时进入算子优化专题。

关键词： AI 编译器、计算图、中间张量、后端路径

---

前置阅读

先从 1D：执行模型、异步调度与算子映射理解程序如何映射到设备执行单元。本节不要求预先掌握具体编译框架。

Q1：图变换能减少哪些启动与中间张量读写？

把四个串行算子拆成四组执行时，组间结果需要写出并由下一组读取；合并为更少的组，有机会减少这些边界和启动次数。若变换还需要调整布局，就会增加额外读写，应与减少量分别计算。

下面假设算子构成单链、各边界张量等大、每组对应一次启动。字节量是逻辑读写估算，启动耗时来自给定假设；它们不能直接相加，也不能代替端到端计时。

Q1 机制计算器：分别记录启动开销与逻辑读写

固定算子链与张量大小，只改变分组，再加入布局转换字节数作对照。

import math


def graph_cost(num_ops, fused_groups, intermediate_elements, bytes_per_element, layout_bytes=0, launch_overhead_us=5.0):
    """估算 kernel launch 与中间张量读写，不把不同单位相加成任意分数。"""
    if (type(num_ops) is not int or type(fused_groups) is not int
            or num_ops <= 0 or not 1 <= fused_groups <= num_ops):
        raise ValueError('fused_groups must be in [1, num_ops]')
    if (type(intermediate_elements) is not int or intermediate_elements < 0
            or type(bytes_per_element) is not int or bytes_per_element <= 0
            or type(layout_bytes) is not int or layout_bytes < 0
            or not math.isfinite(launch_overhead_us) or launch_overhead_us < 0):
        raise ValueError('元素数和布局字节数须为非负整数，元素字节数须为正整数，开销须为有限非负数')
    materialized_boundaries = fused_groups - 1
    intermediate_bytes = materialized_boundaries * intermediate_elements * bytes_per_element * 2
    return {
        'launches': fused_groups,
        'launch_overhead_us': fused_groups * launch_overhead_us,
        'materialized_boundaries': materialized_boundaries,
        'intermediate_traffic_bytes': intermediate_bytes,
        'layout_traffic_bytes': layout_bytes,
        'total_extra_traffic_bytes': intermediate_bytes + layout_bytes,
    }

naive = graph_cost(4, 4, 1_000_000, 2)
fused = graph_cost(4, 1, 1_000_000, 2, layout_bytes=1_000_000)
assert fused['launches'] < naive['launches']
assert fused['total_extra_traffic_bytes'] < naive['total_extra_traffic_bytes']
print('naive:', naive)
print('fused:', fused)

Q2：融合条件满足后，为什么还要单独记录代价？

共享中间结果提供了融合机会，但还要检查语义、副作用和资源预算。这里用“无副作用”作为简化示例的保守前提，并不表示所有带副作用的操作都无法合法融合。

满足这些前提后，再计算避免的中间读写减去新增布局读写。有净读写减少并不保证更快；没有净读写减少也不意味着不可执行，启动次数、计算和资源占用仍会影响耗时。布尔输入表示已给定的检查结果，函数本身不证明语义等价。

Q2 对照验证：分开记录前提和读写减少量

def fusion_candidate(shared_bytes, layout_rewrite_bytes, no_side_effect,
                     register_bytes, register_budget_bytes):
    """记录示例前提与逻辑读写差额，不把成本增加当作执行失败。"""
    if any(type(value) is not int or value < 0
           for value in (shared_bytes, layout_rewrite_bytes, register_bytes)):
        raise ValueError('字节数必须是非负整数')
    if type(register_budget_bytes) is not int or register_budget_bytes <= 0:
        raise ValueError('寄存器预算必须是正整数字节数')
    if type(no_side_effect) is not bool:
        raise ValueError('副作用检查结果必须是布尔值')
    reasons = []
    if not no_side_effect:
        reasons.append('side_effect_requires_analysis')
    if register_bytes > register_budget_bytes:
        reasons.append('register_budget')
    net = shared_bytes - layout_rewrite_bytes
    return {
        'stated_conditions_met': not reasons,
        'net_logical_traffic_saved_bytes': net,
        'has_net_traffic_saving': net > 0,
        'reasons': reasons,
    }


ideal = fusion_candidate(8_000_000, 0, True, 48_000, 64_000)
layout_cost = fusion_candidate(8_000_000, 10_000_000, True, 48_000, 64_000)
side_effect = fusion_candidate(8_000_000, 0, False, 48_000, 64_000)
assert ideal['stated_conditions_met'] and ideal['has_net_traffic_saving']
assert layout_cost['stated_conditions_met']
assert layout_cost['net_logical_traffic_saved_bytes'] == -2_000_000
assert not side_effect['stated_conditions_met']
print(ideal)
print(layout_cost)
print(side_effect)

Q3：同一张图为什么会匹配不同后端路径？

目标后端支持的算子、数据类型、形状与工作区不同，同一张图因而可能需要不同的变换或回退路径。先核对声明的能力与预算，再比较可行候选的时间；实际部署还需检查设备、版本和执行结果。

下面使用虚构的后端能力和延迟输入练习筛选，不探测本机后端，也不表示 CPU 或 GPU 的真实速度。工作区统一使用 MiB；未满足形状或内存要求的候选，即使假设延迟更低也不能入选。

Q3 约束筛选：先检查给定条件，再比较假设延迟

def choose_backend(requirements, candidates):
    """使用算子、dtype、动态形状和显存硬约束筛选后端。"""
    if (not math.isfinite(requirements['workspace_limit_mib'])
            or requirements['workspace_limit_mib'] < 0
            or type(requirements['dynamic_shape']) is not bool):
        raise ValueError('工作区预算与动态形状条件不合法')
    for candidate in candidates.values():
        if (not math.isfinite(candidate['workspace_mib'])
                or candidate['workspace_mib'] < 0
                or not math.isfinite(candidate['estimated_latency_ms'])
                or candidate['estimated_latency_ms'] <= 0
                or type(candidate['dynamic_shape']) is not bool):
            raise ValueError('候选工作区、假设延迟与动态形状条件不合法')
    eligible = []
    rejected = {}
    for name, candidate in candidates.items():
        reasons = []
        if not requirements['ops'].issubset(candidate['ops']):
            reasons.append('missing_ops')
        if requirements['dtype'] not in candidate['dtypes']:
            reasons.append('unsupported_dtype')
        if requirements['dynamic_shape'] and not candidate['dynamic_shape']:
            reasons.append('dynamic_shape')
        if candidate['workspace_mib'] > requirements['workspace_limit_mib']:
            reasons.append('workspace')
        if reasons:
            rejected[name] = reasons
        else:
            eligible.append(name)
    selected = min(eligible, key=lambda name: candidates[name]['estimated_latency_ms']) if eligible else None
    return {'selected': selected, 'eligible': eligible, 'rejected': rejected}

requirements = {'ops': {'matmul', 'softmax'}, 'dtype': 'bf16', 'dynamic_shape': True, 'workspace_limit_mib': 512}
candidates = {
    'host_example': {'ops': {'matmul', 'softmax'}, 'dtypes': {'fp32', 'bf16'}, 'dynamic_shape': True, 'workspace_mib': 64, 'estimated_latency_ms': 20.0},
    'accelerator_example': {'ops': {'matmul', 'softmax'}, 'dtypes': {'fp16', 'bf16'}, 'dynamic_shape': True, 'workspace_mib': 256, 'estimated_latency_ms': 3.0},
    'static_only': {'ops': {'matmul', 'softmax'}, 'dtypes': {'bf16'}, 'dynamic_shape': False, 'workspace_mib': 128, 'estimated_latency_ms': 1.5},
}
decision = choose_backend(requirements, candidates)
assert decision['selected'] == 'accelerator_example'
assert decision['rejected']['static_only'] == ['dynamic_shape']
print(decision)

相关阅读

导语： 如果想继续把图优化和执行后端的关系补完整，可以接着看这些页。

• 04. Triton 矩阵乘法 (GEMM) 与自动调优 (Autotune)
• 05. Triton 性能调优与基准测试 (Autotune & Profiling)
• 18. CUDA Graph and JIT Compile | CUDA Graph 与 JIT 编译

---