「窥孔优化与超优化」

// 窥孔优化最经典的几个例子
// 1) 冗余加载：mov [rbp-8],rax ; mov rax,[rbp-8]  ->  mov [rbp-8],rax
// 2) 强度削减：imul rax, 8                        ->  shl rax, 3
// 3) 常量折叠：mov rax,4 ; add rax,6              ->  mov rax, 10
// 4) 空跳转：  jmp .Lnext ; .Lnext: ...           ->  .Lnext: ...

# 一个极简的窥孔优化器骨架
RULES = [
    # (模式, 替换)：None 为通配符
    ((("mov", "mem", "reg"), ("mov", "reg", "mem")), (("mov", "mem", "reg"),)),
    ((("imul", "reg", ("const", 8)),), (("shl", "reg", 3),)),
    ((("jmp", "L"), ("label", "L")), ()),                  # 跳转到下一条
]

def match(pattern, window):
    """长度一致且逐项匹配（含嵌套元组）则返回 True"""
    if len(pattern) != len(window):
        return False
    for p, w in zip(pattern, window):
        if p is None:                       # None 通配任意一条
            continue
        if isinstance(p, tuple):
            if not isinstance(w, tuple) or len(p) != len(w):
                return False
            if any(pi is not None and pi != wi for pi, wi in zip(p, w)):
                return False
        elif p != w:
            return False
    return True

def peephole(instrs, max_window=3):
    changed = True
    while changed:                          # 迭代到不动点
        changed = False
        for i in range(len(instrs)):
            for size in range(max_window, 1, -1):   # 先试长窗口
                win = instrs[i:i + size]
                if len(win) < size:
                    continue
                for pat, rep in RULES:
                    if match(pat, win):
                        instrs[i:i + size] = list(rep)
                        changed = True
                        break
                if changed:
                    break
    return instrs

// 强度削减与指令合并的具体效果
// 改前（x86-64，-O1）
//   lea  rax, [rdi + rdi*8]     ; rax = 9 * x
//   mov  ecx, [rsi]             ; 载入两个 32 位
//   mov  edx, [rsi+4]
// 改后（窥孔 + 合并）
//   lea  rax, [rdi + rdi*8]
//   mov  rcx, [rsi]             ; 一条 64 位载入代替两条 32 位
long f(long x, const int *p) {
    long a = x * 9;
    long b = p[0];
    long c = p[1];
    return a + b + c;
}

# 一个简化但可用的代价模型
# 每条指令 -> (latency, reciprocal_throughput, size_bytes)
COST_TABLE = {
    "mov_rr":   (0, 0.25, 3),    # 寄存器间移动，被重命名消除
    "mov_rm":   (4, 0.5,  4),    # 从内存加载（L1 命中）
    "mov_mr":   (3, 1.0,  4),    # 存到内存
    "add_rr":   (1, 0.25, 3),
    "imul_rr":  (3, 1.0,  4),
    "shl_ri":   (1, 0.5,  4),
    "div_rr":   (20, 6.0, 3),    # 除法很贵
    "jmp":      (1, 1.0,  2),
    "jcc":      (1, 0.5,  2),
    "call":     (2, 1.0,  5),
}

def cost(seq, mode="speed", size_weight=0.05):
    lat, thr, sz = 0.0, 0.0, 0
    for op in seq:
        l, t, s = COST_TABLE[op]
        lat += l; thr += t; sz += s
    if mode == "speed":
        return thr + size_weight * sz   # 吞吐 + 体积惩罚
    return sz                           # 体积优先：只关心字节数

# 用 llvm-mca 分析一段机器码的代价
cat > hot.s <<'EOF'
    movq   (%rdi), %rax
    imulq  $9, %rax, %rax
    addq   %rsi, %rax
    movq   %rax, (%rdx)
EOF
llvm-mca -mcpu=skylake -iterations=100 hot.s
# 输出里会给出：Block RThroughput、每次迭代的周期数、各端口压力

// 符号函数：返回 x 的符号（-1、0、1）
int sign_naive(int x) {              // 教科书实现：分支 + 3 条指令
    if (x > 0) return 1;
    if (x < 0) return -1;
    return 0;
}

// 超级优化器找到的无分支实现（x86-64，无跳转）
//   mov edx,1 ; mov eax,0 ; test edi,edi ; cmovg eax,edx
// 消除分支预测失败，随机输入下比教科书版本快数倍
int sign_branchless(int x) {
    return (x > 0) - (x < 0);   // 编译器通常也能生成无分支版本
}

# E-graph 的核心操作：把等价形式并入同一个类
class EGraph:
    def __init__(self):
        self.classes = {}       # eclass_id -> set of enodes
        self.uf = {}            # 并查集，维护类之间的合并

    def add(self, op, children):
        """添加节点，返回它所属的等价类 id"""
        key = (op, tuple(children))
        cid = self.uf.setdefault(key, len(self.classes))
        self.classes.setdefault(cid, set()).add(key)
        return cid

    def merge(self, a, b):
        """合并两个等价类，这是所有改写规则的统一形式"""
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        self.uf[rb] = ra
        self.classes[ra] |= self.classes[rb]
        del self.classes[rb]
        return True

    def rebuild(self):
        """合并后重建：把子节点重定向到已合并的类（需 worklist 收敛）"""
        pass

# 用 E-graph 表示 (x * 2) + (x * 2)，同时应用 a*2 => a<<1 与 a+a => a*2
# 结果：移位形式、乘 2 形式、4*x 形式共存于同一等价类，最后按代价挑选

import random

def differential_test(orig, cand, n=10000, bits=32):
    """差分测试：随机输入下比较两个序列的输出"""
    mask = (1 << bits) - 1
    for _ in range(n):
        args = [random.getrandbits(bits) for _ in range(arity(orig))]
        if (eval_seq(orig, args) & mask) != (eval_seq(cand, args) & mask):
            return False, args       # 找到反例，立即拒绝
    return True, None                # 未找到反例，进入下一道筛子

# 用 Z3 证明两个位向量表达式等价
from z3 import BitVec, Solver

x = BitVec("x", 32)
s = Solver()
s.add((x * 2) + (x * 2) != x * 4)    # 断言两者不等
print(s.check())                     # unsat -> 等价成立
s.reset()
s.add((x << 1) != x * 2)
print(s.check())                     # unsat -> 移位与乘 2 等价

# STOKE 的典型用法：MCMC 搜索 + 随机测试验证，是超级优化的工程化代表
stoke optimize --backend sandbox --init asm_orig.s --out asm_opt.s

// IR: t1 = x + y ; t2 = t1 * 4 ; t3 = t2 + z
// 方案 A（三条指令）：add t1,x,y ; shl t2,t1,2 ; add t3,t2,z
// 方案 B（两条指令，利用 lea 的 base + index*scale + disp 能力）：
//   add t1, x, y
//   lea t3, [z + t1*4]
// 方案 B 少一条指令，且不占用乘法端口
int sel(int x, int y, int z) {
    return (x + y) * 4 + z;
}

// Go 编译器窥孔规则的形式（简化示意）
// (Add64 x (Const64 [c])) && is32Bit(c) -> (ADDQconst [c] x)
// (Mul64 x (Const64 [c])) && isPowerOfTwo(c) -> (SHLQconst [log2(c)] x)
// (Eq64 x (Const64 [0])) -> (TESTQ x x)

# Souper 的工作流：为 IR 片段合成更优实现
souper-check -infer-rhs -souper-external-cache ir.ll
# 输出形如：infer %x
#           %r = mul %x, 8
#         replace %r with shl %x, 3
# 人工审查后可作为 peephole 规则加入后端

// 一个真实的踩坑例子：看似无害的代数简化
// 规则：x - x => 0
// 但如果 x 是浮点数，x - x 在 x = NaN 或 x = Inf 时是 NaN，不是 0
// 规则：x * 0 => 0
// 同样在浮点下失效，且会丢失 x 求值可能触发的异常
double f(double x) {
    return x - x;      // 对 NaN 返回 NaN，不是 0.0
}

// 正确的规则必须区分域：
//   整数域：x - x => 0 合法
//   浮点域：需要 -ffast-math 或 -ffinite-math-only 才允许

# 观察窥孔规则是否生效
cc -O2 -fdump-rtl-peephole2 -c file.c      # GCC：peephole2 改写记录
llc -debug-only=isel,dagcombine file.ll    # LLVM：DAG combine 每一步