前言

调GEMM算子性能，手写Tiling参数试了21种组合，最好的是tile_m=64, tile_k=64, tile_n=64，吞吐89 tokens/s。用AOE调优引擎自动搜索，搜出tile_m=128, tile_k=64, tile_n=128，吞吐102 tokens/s，比手写最优还高15%。

很多人以为AOE就是"网格搜索"，其实它用贝叶斯优化+强化学习，搜索效率比网格搜索高100倍，能在10分钟内找到全局最优Tiling参数。

AOE 的定位

AOE（Ascend Optimization Engine）是CANN内置的调优引擎，自动搜索算子的最优Tiling参数、编译选项、内存布局。

CANN 架构中的 AOE：
AscendCL（编程接口层）
  ↓
AOL 算子库（ops-math/ops-nn/ops-blas/...）
  ↓
GE（图引擎）
  ↓
AOE 调优引擎 ← 你在这（离线调优，生成最优配置）
  ↓
Runtime（运行时）
  ↓
驱动层

AOE不是在线调优（不拖慢推理），是离线调优（提前搜好最优配置，推理时直接加载）。

工程经验： 不复用AOE手动调Tiling，要试几十种组合，耗时2-3天。用AOE自动调优，10分钟搜完，性能还更高。不是AOE多智能，是它不基于直觉搜，能找到人想不到的组合。

AOE 的核心技术

1. 搜索空间建模

AOE把Tiling参数、编译选项、内存布局统一建模成搜索空间。

Tiling参数搜索空间：

# GEMM 算子的 Tiling 搜索空间
tiling_search_space = {
    # tile_m: L0A 容量约束
    "tile_m": list(range(16, 256+1, 16)),  # [16, 32, 48, ..., 256]
    
    # tile_k: L0B 容量约束
    "tile_k": list(range(16, 256+1, 16)),
    
    # tile_n: L0C 容量约束
    "tile_n": list(range(16, 256+1, 16)),
    
    # 流水线策略
    "pipeline": ["none", "single_buffer", "double_buffer"],
    
    # L1 预取策略
    "l1_prefetch": [True, False],
    
    # 输出对齐策略
    "output_align": [True, False],
}

# 搜索空间大小：
# len(tile_m) × len(tile_k) × len(tile_n) × len(pipeline) × len(l1_prefetch) × len(output_align)
# = 16 × 16 × 16 × 3 × 2 × 2 = 49,152 种组合

编译选项搜索空间：

# 编译选项搜索空间
compile_search_space = {
    # 优化等级
    "opt_level": ["O0", "O1", "O2", "O3"],
    
    # 指令调度策略
    "schedule": ["default", "balanced", "throughput", "latency"],
    
    # 循环展开次数
    "unroll_factor": [0, 2, 4, 8, 16],
    
    # 内存布局
    "layout": ["row_major", "col_major", "z_major"],
}

总搜索空间：49,152 × 240 = 11,796,480 种组合（千万级）。

2. 贝叶斯优化搜索

网格搜索要穷举千万种组合，10年都搜不完。AOE用贝叶斯优化，根据已有结果建代理模型（Gaussian Process），预测哪些组合可能最优，优先搜那些。

# AOE 贝叶斯优化（伪代码）
import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor

class AOE_BayesianOptimization:
    def __init__(self, search_space):
        self.search_space = search_space
        self.X = []  # 已搜的配置
        self.y = []  # 对应的性能（tokens/s）
        self.gp = GaussianProcessRegressor()  # 代理模型
        
    def suggest_next(self):
        # 训练代理模型
        self.gp.fit(self.X, self.y)
        
        # 预测所有未搜配置的期望性能
        X_candidates = self._get_all_candidates()
        mu, sigma = self.gp.predict(X_candidates, return_std=True)
        
        # Upper Confidence Bound (UCB) 采集函数
        # 选"期望性能高"且"不确定性大"的配置（平衡探索和利用）
        ucb = mu + 1.96 * sigma
        best_idx = np.argmax(ucb)
        
        return X_candidates[best_idx]
    
    def evaluate(self, config):
        # 编译算子（用config中的Tiling参数、编译选项）
        kernel = compile_kernel(config)
        
        # 跑benchmark（测吞吐）
        throughput = benchmark(kernel, iterations=100)
        
        return throughput
    
    def run(self, max_iter=100):
        for i in range(max_iter):
            # 建议下一个要搜的配置
            config = self.suggest_next()
            
            # 评估性能
            throughput = self.evaluate(config)
            
            # 记录结果
            self.X.append(config)
            self.y.append(throughput)
            
            # 如果找到够好的配置，提前停止
            if throughput >= target_throughput:
                break
        
        # 返回最优配置
        best_idx = np.argmax(self.y)
        return self.X[best_idx]

贝叶斯优化只需要搜100-200个配置，就能找到全局最优（搜索效率比网格搜索高100倍）。

工程经验： AOE的贝叶斯优化能找到人想不到的配置。比如GEMM算子，直觉是"tile_m/tile_k/tile_n都相等（64/64/64）最优"，但AOE搜出tile_m=128, tile_k=64, tile_n=128，性能高15%。原因是L0A/L0B/L0C容量不是完全对称的。

3. 强化学习增强

对于超大搜索空间（比如Transformer推理要调Attention+FFN+LayerNorm的Tiling，搜索空间亿级），贝叶斯优化也不够快。AOE用强化学习（Reinforcement Learning）进一步增强。

# AOE 强化学习（伪代码）
import torch
import torch.nn as nn
import torch.optim as optim

class AOE_RL:
    def __init__(self, search_space):
        self.search_space = search_space
        
        # 策略网络（输入：算子特征，输出：Tiling参数）
        self.policy_net = nn.Sequential(
            nn.Linear(in_features=64, out_features=128),
            nn.ReLU(),
            nn.Linear(in_features=128, out_features=len(search_space)),
            nn.Softmax(dim=-1),
        )
        
        self.optimizer = optim.Adam(self.policy_net.parameters())
        
    def select_action(self, op_features):
        # 策略网络输出Tiling参数的概率分布
        probs = self.policy_net(op_features)
        
        # 采样一个配置
        action = torch.multinomial(probs, 1)
        
        return action
    
    def evaluate_action(self, action):
        # 编译算子
        config = self._action_to_config(action)
        kernel = compile_kernel(config)
        
        # 跑benchmark
        throughput = benchmark(kernel, iterations=100)
        
        return throughput
    
    def update_policy(self, op_features, action, reward):
        # REINFORCE 算法更新策略网络
        log_prob = torch.log(self.policy_net(op_features)[action])
        loss = -log_prob * reward
        
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
    
    def run(self, op_features, max_episode=1000):
        for episode in range(max_episode):
            # 选一个配置
            action = self.select_action(op_features)
            
            # 评估性能（奖励）
            reward = self.evaluate_action(action)
            
            # 更新策略网络
            self.update_policy(op_features, action, reward)
        
        # 返回最优配置
        best_action = self.select_action(op_features)  # 贪心
        return self._action_to_config(best_action)

强化学习能在1小时内搜完亿级搜索空间（比贝叶斯优化快10倍）。

使用流程

1. 准备待调优算子

// 待调优的 GEMM 算子（Ascend C）
#include "kernel_operator.h"

class GemmKernel {
public:
    __aicore__ void Process(GM_ADDR a, GM_ADDR b, GM_ADDR c,
                           int M, int K, int N) {
        // Tiling 参数（待 AOE 调优）
        int tile_m = 64;  // ← AOE 会改这个
        int tile_k = 64;  // ← AOE 会改这个
        int tile_n = 64;  // ← AOE 会改这个
        
        // ... 算子逻辑
    }
};

2. 启动 AOE 调优

# 1. 准备调优配置文件
cat > aoe_config.json << EOF
{
    "op_name": "GemmKernel",
    "search_space": {
        "tile_m": [16, 32, 64, 128, 256],
        "tile_k": [16, 32, 64, 128, 256],
        "tile_n": [16, 32, 64, 128, 256],
        "pipeline": [0, 1, 2],
        "l1_prefetch": [0, 1],
        "output_align": [0, 1]
    },
    "optimization_objective": "throughput",  // 优化目标：吞吐
    "max_search_time": 600,  // 搜索时间上限：10分钟
    "target_throughput": 100  // 目标吞吐：100 tokens/s
}
EOF

# 2. 启动 AOE 调优
aoe_tune --config aoe_config.json \
         --op-source gemm_kernel.cpp \
         --output optimized_kernel.cpp

# 3. 等待调优完成（10分钟）
# 输出：
# [INFO] AOE tuning started...
# [INFO] Search space size: 49,152
# [INFO] Bayesian optimization: 100 iterations, best throughput: 95 tokens/s
# [INFO] Reinforcement learning: 500 episodes, best throughput: 102 tokens/s
# [INFO] Tuning completed. Optimized kernel saved to optimized_kernel.cpp

3. 使用调优后的算子

# 编译调优后的算子
npu-smi set -t mm -s 0 -d optimized_gemm.o optimized_kernel.cpp

# 链接成动态库
ld -shared optimized_gemm.o -o liboptimized_gemm.so

# 在 ACL 中调用（性能比调优前高 15%）
aclError ret = aclrtLaunchKernel(optimized_gemm_kernel, grid, block, args, 0, stream);

性能对比

AOE调优 vs 手动调优 vs 默认配置（Qwen2.5-7B，910B单卡，FP16）：

策略	吞吐(tokens/s)	调优时间
默认配置	34	0
手动调优（试21种组合）	89	2天
AOE调优（贝叶斯优化）	95	10分钟
AOE调优（强化学习）	102	1小时

AOE强化学习比手动调优高15%，时间从2天降到1小时。

工程经验： AOE调优不是"万能"的。搜索空间太大（>1亿种组合），强化学习也要搜1小时。要限制搜索空间："tile_m": [32, 64, 128]（只搜常见值），别写range(16, 256+1, 16)（搜16个值）。

踩坑实录

坑1：AOE调优后算子编译失败

AOE搜出的Tiling参数导致L0A/L0B/L0C溢出，编译时报error: L0A buffer overflow。

解决：搜索空间加约束。"constraints": ["tile_m * tile_k * 2 < 64*1024", ...]（保证不溢出）。

坑2：AOE调优后性能反而降

AOE在单卡上调优，多卡推理时通信开销变大，性能反而降。

解决：多卡推理要在多卡环境下调优。aoe_tune --num-devices 8（8卡环境调优）。

坑3：AOE调优时间太长（>2小时）

搜索空间太大（比如Transformer推理要调Attention+FFN+LayerNorm，搜索空间亿级），AOE搜不完。

解决：分算子调优（先调Attention，再调FFN，再调LayerNorm），每个算子搜索空间降到千万级。

坑4：AOE调优后精度掉（FP16 → INT8量化）

AOE为了性能，自动开了INT8量化，精度掉2-3%。

解决：搜索空间禁止量化。"quantization": [false]（只搜FP16配置）。

https://atomgit.com/cann/amct

https://atomgit.com/cann/asc-devkit

https://atomgit.com/cann/cann-samples