1
# MLGO 的核心思路：用强化学习决定内联/不内联
2
# 状态：函数调用图的局部特征（调用深度、指令数、循环数）
3
# 动作：inline / no-inline
4
# 奖励：二进制大小减少量
5
import numpy as np
6

7
class InlineEnv:
8
    def __init__(self, call_graph):
9
        self.call_graph = call_graph
10
        self.state = self._extract_features(call_graph)
11

12
    def step(self, action):
13
        # action: 0=no-inline, 1=inline
14
        if action == 1:
15
            self._inline_current_call()
16
        reward = self._compute_size_reduction()
17
        self.state = self._extract_features(self.call_graph)
18
        return self.state, reward, self._is_done()

1
# 使用 MLGO 训练内联策略
2
# 1. 采集训练数据
3
clang -O2 -fpass-plugin=MLInlineAdvisor \
4
  -mllvm -enable-ml-inline-advisor=training \
5
  -mllvm -ml-inline-advisor-model=unused \
6
  -c training_corpus/*.c
7

8
# 2. 训练模型
9
python3 train_inline_model.py --data=training_log/
10

11
# 3. 使用训练好的模型编译
12
clang -O2 -fpass-plugin=MLInlineAdvisor \
13
  -mllvm -enable-ml-inline-advisor=production \
14
  -mllvm -ml-inline-advisor-model=trained_model/ \
15
  -c my_app.cpp

1
传统启发式：
2
  if (callee_size < threshold && caller_size < max_size)
3
    → inline
4

5
MLGO 方法：
6
  提取调用点特征 → 模型预测 → inline / no-inline
7

8
  特征包括：
9
  - 调用者/被调用者的大小
10
  - 调用频率
11
  - 调用深度
12
  - 函数类型（冷/热）
13
  - 历史内联效果

1
# MLGO 训练流程（简化）
2
class InliningEnv:
3
    def __init__(self, llvm_module):
4
        self.module = llvm_module
5
        self.call_sites = extract_call_sites(llvm_module)
6

7
    def step(self, action):
8
        """执行内联决策，返回奖励"""
9
        if action == 'inline':
10
            inline_function(self.call_sites[self.current_idx])
11

12
        # 编译并测量执行时间
13
        binary = compile_module(self.module)
14
        runtime = measure_runtime(binary)
15

16
        # 奖励：执行时间减少
17
        reward = self.baseline_runtime - runtime
18

19
        self.current_idx += 1
20
        next_state = self.get_state()
21
        done = self.current_idx >= len(self.call_sites)
22

23
        return next_state, reward, done
24

25
    def get_state(self):
26
        """提取调用点特征"""
27
        site = self.call_sites[self.current_idx]
28
        return [
29
            site.caller_size,
30
            site.callee_size,
31
            site.call_depth,
32
            site.is_hot,
33
            site.caller_instruction_count,
34
            site.callee_instruction_count,
35
        ]
36

37
# 训练模型
38
model = PolicyNetwork(input_dim=6, hidden_dim=128, output_dim=2)
39
agent = PPOAgent(model)
40
agent.train(InliningEnv(train_modules), episodes=10000)

1
搜索空间：
2
  tile_size_x: [1, 2, 4, 8, 16, 32, 64, 128]
3
  tile_size_y: [1, 2, 4, 8, 16, 32, 64, 128]
4
  tile_size_k: [1, 2, 4, 8, 16, 32, 64, 128]
5
  unroll_factor: [1, 2, 4, 8]
6
  vectorize: [true, false]
7
  parallel: [true, false]
8

9
  总搜索空间：8 × 8 × 8 × 4 × 2 × 2 = 8192 种配置
10

11
  AutoTVM 目标：在有限时间内找到最优配置

1
# AutoTVM 代价模型
2
class CostModel:
3
    def __init__(self):
4
        self.model = xgboost.XGBRegressor()
5
        self.features = ScheduleFeatureExtractor()
6

7
    def predict(self, schedule):
8
        """预测给定调度的执行时间"""
9
        features = self.features.extract(schedule)
10
        return self.model.predict([features])
11

12
    def update(self, schedule, measured_time):
13
        """用实际测量结果更新模型"""
14
        features = self.features.extract(schedule)
15
        self.model.fit(
16
            np.array([features]),
17
            np.array([measured_time]),
18
            increment=True
19
        )
20

21
# 特征提取
22
class ScheduleFeatureExtractor:
23
    def extract(self, schedule):
24
        return [
25
            schedule.tile_size_x,
26
            schedule.tile_size_y,
27
            schedule.tile_size_k,
28
            schedule.unroll_factor,
29
            schedule.vectorize,
30
            schedule.parallel,
31
            # 计算特征
32
            schedule.arithmetic_intensity,
33
            schedule.memory_access_count,
34
            schedule.compute_ops,
35
        ]

1
# AutoTVM 搜索流程
2
def auto_tune(task, n_trials=1000):
3
    model = CostModel()
4
    best_config = None
5
    best_time = float('inf')
6

7
    for i in range(n_trials):
8
        # 1. 模型推荐候选配置
9
        candidates = model.recommend_top_k(task, k=10)
10

11
        # 2. 实际测量候选配置
12
        for config in candidates:
13
            time = measure_on_device(task, config)
14

15
            # 3. 更新模型
16
            model.update(config, time)
17

18
            # 4. 记录最优
19
            if time < best_time:
20
                best_time = time
21
                best_config = config
22

23
    return best_config, best_time

1
IR2Vec：将 LLVM IR 转换为向量表示
2

3
方法 1：Flow-based（基于数据流）
4
  - 每条指令映射为一个向量
5
  - 通过数据流图传播信息
6
  - 最终得到函数级向量
7

8
方法 2：Symmetric（基于对称）
9
  - 每条指令的向量 = 操作码向量 + 操作数向量
10
  - 函数向量 = 所有指令向量的聚合
11

12
应用：
13
  - 内联决策预测
14
  - 函数相似性检测
15
  - 优化 Pass 选择

1
ProGraML：图神经网络用于程序分析
2

3
输入：程序的控制流图 + 数据流图 + 调用图
4
处理：GNN 在图上传播信息
5
输出：节点级/图级向量表示
6

7
优势：
8
  - 跨语言（任何有 IR 的语言）
9
  - 跨任务（同一个模型可用于多种优化决策）
10
  - 可迁移（在一个程序集训练，迁移到另一个）

1
# TVM：ML-oriented 编译框架
2
import tvm
3
from tvm import relay
4

5
# 定义模型
6
model = relay.nn.conv2d(data, weight, strides=(1,1), padding=(1,1))
7

8
# AutoTVM 调优
9
with tvm.target.Target("cuda"):
10
    best_config = auto_tvm.tune(model, n_trials=1000)
11

12
# 编译优化后的模型
13
lib = tvm.compile(model, target="cuda", config=best_config)

1
MLIR：多级中间表示框架
2
  - 统一的 IR 框架，支持多种方言
3
  - Affine 方言：循环优化
4
  - Tensor 方言：张量计算
5
  - GPU 方言：GPU 编程
6
  - Linalg 方言：线性代数
7

8
  优势：
9
  - 不同优化级别可以混合使用
10
  - ML 方法可以在任何方言级别介入
11
  - 可扩展：添加新方言不需要修改框架

1
1. LLM for 编译器
2
   - 大语言模型理解代码语义
3
   - 自动生成优化规则
4
   - 代码翻译（语言A → 语言B）
5

6
2. 自改进编译器
7
   - 编译器从每次编译中学习
8
   - 持续优化决策策略
9
   - 适应硬件演进
10

11
3. 端到端编译
12
   - 从源码直接生成优化后的机器码
13
   - 跳过中间 IR 层级
14
   - Transformer-based 代码生成

1
# MLGO 的调用点特征（部分）
2
def extract_call_site_features(caller, callee, call_site):
3
    return {
4
        # 调用者特征
5
        "caller_basic_block_count": caller.num_bbs,
6
        "caller_instruction_count": caller.num_insts,
7
        "caller_users_count": caller.num_users,
8

9
        # 被调用者特征
10
        "callee_basic_block_count": callee.num_bbs,
11
        "callee_instruction_count": callee.num_insts,
12
        "callee_callsite_count": callee.num_callsites,
13
        "callee_is_local": callee.is_local,
14

15
        # 调用点特征
16
        "callsite_cost": call_site.cost,
17
        "callsite_is_recursive": call_site.is_recursive,
18
        "callsite_depth": call_site.depth,
19
        "callsite_is_hot": call_site.is_hot,
20

21
        # 历史特征
22
        "prev_inline_count": call_site.prev_inlines,
23
        "prev_inline_size_increase": call_site.size_increase,
24
    }

1
import tvm
2
from tvm import autotvm
3

4
# 定义矩阵乘法的调度模板
5
@autotvm.template("matmul")
6
def matmul(M, N, K):
7
    A = tvm.placeholder((M, K), name='A')
8
    B = tvm.placeholder((K, N), name='B')
9
    C = tvm.compute((M, N), lambda i, j: tvm.sum(A[i, k] * B[k, j], axis=k))
10

11
    s = tvm.create_schedule(C.op)
12

13
    # 可调参数
14
    cfg = autotvm.get_config()
15
    cfg.define_knob("tile_y", [1, 8, 16, 32, 64])   # y 方向 tile
16
    cfg.define_knob("tile_x", [1, 8, 16, 32, 64])   # x 方向 tile
17
    cfg.define_knob("tile_k", [1, 8, 16, 32])        # k 方向 tile
18
    cfg.define_knob("unroll", [0, 1, 2, 4])          # 展开因子
19

20
    # 应用调度
21
    y, x = s[C].op.axis
22
    yo, yi = s[C].split(y, factor=cfg["tile_y"].val)
23
    xo, xi = s[C].split(x, factor=cfg["tile_x"].val)
24
    s[C].reorder(yo, xo, yi, xi)
25

26
    if cfg["unroll"].val > 0:
27
        s[C].unroll(yi)
28
        s[C].unroll(xi)
29

30
    return s, [A, B, C]