Schanuel猜想与Triton的FPSan语义

# Schanuel猜想与FPSan的语义 来源：https://cp4space.hatsya.com/2026/05/03/schanuels-conjecture-and-the-semantics-of-fpsan/ 我最近正与Pawel Szczerbuk合作，将一部分时间投入到开发一个名为FPSan的工具上。它被实现为一个Triton编译器pass（https://github.com/triton-lang/triton/blob/main/lib/Dialect/TritonInstrument/Transforms/FpSanitizer.cpp），但它并不具备编译器pass应有的理想特性：具体而言，它不会保持功能不变，会导致运行变慢，而且迄今为止完全没有文档。（关于后一点，Pawel有一个开放的PR（https://github.com/triton-lang/triton/pull/10177）来添加文档。） 其目的是为了更容易验证用Triton编写的、涉及浮点运算的程序之间的代数等价性。关键问题在于，在浮点运算中，像结合律这样的代数定律并不精确成立：通常情况下，\(*a* + *b*) + *c* 未必等于 *a* + \(*b* + *c*\)。因此，如果你重写一个程序以利用这一点，例如将顺序求和循环替换为并行树状归约，那么程序的行为将不再完全相同。 FPSan可以被视为程序空间上的一个幂等函数，它将所有浮点运算替换为（完全不同的！）整数运算，使得如果 f 和 g 是代数等价的程序，那么给定相同的输入，FPSan(f) 和 FPSan(g) 产生完全相同的结果。 更正式地说，*依赖于实数版本的Schanuel猜想*，只要程序 f 和 g 满足以下性质，该结论就成立： - 每个程序在其浮点输入上实现一个算术电路，并且控制流与那些浮点输入无关； - 算术电路仅包含输入、输出、常数集合 \{-1.0, 0.0, +1.0\}、环运算 \{- , + , ×\} 以及指数函数 exp。 这些运算可能看起来有些限制，但它已经涵盖了机器学习中涉及的大多数常见GPU内核：矩阵乘法和[大部分]自注意力机制都在FPSan的保证范围之内。 证明过程推迟到本文末尾，以避免偏离主线。这很可能是唯一一个正确性依赖于一个极其困难的未解决问题（https://en.wikipedia.org/wiki/Schanuel%27s_conjecture）的编译器消毒器，该问题属于超越数论领域。 ### 实现 具体来说，FPSan 构造了一个双射的“嵌入函数” φ，从 IEEE-754 单精度浮点数集合（共有 2^32 个元素）映射到模 2^32 的整数环。函数 φ 的实现如下： - 使用 IEEE-754 编码将浮点数编码为 32 位字； - 最高位（符号位）保持不变； - 对于剩余的 31 位，首先乘以一个奇数常数（模 2^31），然后进行一次 xorshift，再乘以另一个奇数常数（模 2^31），最后（如果符号位已设置）取补码； - 将 32 位字解释为模 2^32 的整数。 这样设计是为了合理地混合比特，同时满足以下性质：对于所有非零的 *x*，有 φ(−*x*) = −φ(*x*)，φ(0.0) = 0，且 φ(1.0) = 1。“负零”浮点数被映射到 2^31，这是模 2^32 整数环中另一个加性自逆元素。 使用这个函数，FPSan 进行了如下替换： - 浮点加法 fadd(*x*, *y*) 替换为 φ⁻¹(φ(*x*) + φ(*y*))； - 浮点减法 fsub(*x*, *y*) 替换为 φ⁻¹(φ(*x*) − φ(*y*))； - 浮点乘法 fmul(*x*, *y*) 替换为 φ⁻¹(φ(*x*) × φ(*y*))； - 浮点指数运算 exp(*x*) 替换为 φ⁻¹(C^φ(*x*))，其中 C 是一个特定的常数，且满足 C ≡ 5 (mod 8)。 最后一个定义利用了模 2^32 整数乘法群的结构。具体来说： - 只有 2^31 个奇数元素（即满足 mod 2 为 1 的元素）是可逆的，因此属于乘法群； - 在这些元素中，满足 mod 4 为 1 的 2^30 个元素在乘法下构成一个循环群； - 在这些元素中，满足 mod 8 为 5 的 2^29 个元素是生成元，或者说具有最大周期，这就是为什么我们选择 C 满足 C ≡ 5 (mod 8)。 映射 *x* → C^*x* 在模 2^32 下是良定义的，因为 C^(2^32) ≡ 1 (mod 2^32)，所以 C^*x* (mod 2^32) 仅依赖于 *x* 模 2^32 的值。这之所以成立，仅仅是因为我们的模数是 2 的幂；对于一般的模数 *n*，乘法群的指数通常不整除 *n*。 经过重写后的 fadd、fsub、fmul 和 exp 显然遵守所有环公理、恒等式 exp(fadd(*x*, *y*)) = fmul(exp(*x*), exp(*y*)) 以及关系 exp(0.0) = 1.0。 ### 混合精度功能 FPSan 为任意浮点数据类型构造了嵌入函数 φ 的类似物，将其映射到具有相同基数大小的整数环。要从 j 位精度向下转换到 k 位精度，我们先将高精度 j 位浮点数嵌入到模 2^j 的整数环中，然后取模 2^k 下的像，最后解除嵌入得到一个低精度 k 位浮点数。向上转换则是相反的过程，我们选择从模 2^k 整数环到模 2^j 整数环的“符号扩展”（https://en.wikipedia.org/wiki/Sign_extension）提升；这特别意味着常数 \{-1, 0, 1\} 在不同精度之间的任意转换中都能幸存。先向上转换再向下转换会产生恒等映射；反之则不成立，因为向下转换必然会丢失信息。 在嵌入函数及其逆函数中构造乘法器需要能够高效地计算模 2^k 下的逆元；我们使用 ceil(log2(k)) 次 2-adic 牛顿法迭代来实现这一点（https://github.com/triton-lang/triton/blob/22573d88adfee7e117273ca592c84df9b6a0509c/lib/Dialect/TritonInstrument/Transforms/FpSanitizer.cpp#L459）。 Pawel 编写了将 Triton 的混合精度矩阵乘法原语 `tl.dot`（https://triton-lang.org/main/python-api/generated/triton.language.dot.html）转换为 FPSan 等效版本的规则，方法是将它展开为标量浮点运算。混合函数 φ 和 φ⁻¹ 仅需应用于每个输入和输出元素，而矩阵乘法的核心仅涉及 int32 乘法和加法。 ### 证明 现在进入有趣的部分：证明 Schanuel 猜想蕴含 FPSan 的期望性质。 假设我们在实数域上有两个算术电路 *f* 和 *g*，每个电路仅由输入、输出、常数集合 \{-1, 0, +1\}、环运算 \{- , + , ×\} 以及指数函数 exp 构成。进一步假设它们是等价的，即实现了从 ℝᵐ 到 ℝⁿ 的相同函数。 假设实数版本的 Schanuel 猜想成立，那么由 {0, −, +, ×, exp} 生成的 ℝ 的子环 X 同构于**无生成元的自由指数环**，正如 Macintyre 1991（https://www.sciencedirect.com/science/article/pii/016800729190017G）中所证明的那样。如果电路在 ℝ 上等价，那么当限制在 X 上时它们必然等价，并且通过该同构，它们也在无生成元的自由指数环中等价。 模 2^32 的整数环连同单变量函数 C^x（其中 C 是一个满足 C ≡ 5 (mod 8) 的特定常数）是无生成元自由指数环的一个商；特别地，我们可以通过设定 θ(exp(x)) = C^θ(x) 来构造一个从无生成元自由指数环到模 2^32 整数环的满同态 θ。 因此，可以得出电路在 FPSan 下仍然等价，因为 FPSan 仅仅是通过嵌入函数 φ 将操作拉回，从而赋予浮点数一个指数环的结构。 ### 正弦和余弦 我们还在 2-adic 整数上实现了正弦和余弦的类似物，方法是在通过附加满足 i² = −1 的形式符号 i 而得到的二次扩张中，取 (－3/5 + 4/5 i)ⁿ 的实部和虚部。这些函数满足三角函数的和差恒等式，以及通常的范数恒等式 sin(x)² + cos(x)² = 1。 结论（即任何在实数域上成立的关于 {0, 1, −, +, ×, exp} 的有效代数恒等式在 FPSan 下也成立，这假设了 Schanuel 猜想）可以被加强：当包含 sin 和 cos 时，该结论仍然成立。 我们将定义以下两个环序列： - 实数子环的上升链 {Y₀, Y₁, Y₂, ...}，其中 Y₀ = ℤ，Yₙ₊₁ 是由 Yₙ 加上所有 x ∈ Yₙ 的 exp(x)、sin(x)、cos(x) 所生成的环； - 复数子环的上升链 {W₀, W₁, W₂, ...}，其中 W₀ = ℤ[i]，Wₙ₊₁ 是由 Wₙ 加上所有 x ∈ Wₙ 的 exp(x) 所生成的环。 我们可以通过对 n 进行归纳证明 Wₙ 恰好等于 Yₙ[i]（特别地，Yₙ 是 Wₙ 的实部）。具体而言，假设该结论对 n−1 成立，则有： - 如果 x ∈ Yₙ₋₁，那么 exp(x)、cos(x) = (exp(ix) + exp(−ix))/2 和 sin(x) = (exp(ix) − exp(−ix))/(2i) 明显都在 Wₙ 中； - 如果 z ∈ Wₙ₋₁，那么其实部和虚部 a 和 b 属于 Yₙ₋₁，而 exp(z) 的实部和虚部分别是 exp(a) cos(b) 和 exp(a) sin(b)，这些都属于 Yₙ。 定义 Y 为所有 Yₙ 的并集，类似地定义 W 为所有 Wₙ 的并集，则有 W = Y[i]。 我们现在可以重复 Macintyre 的证明思路，但针对环序列 Wₙ：假设复数版本的 Schanuel 猜想成立，那么 W 是由一个满足 i² = −1 的 i 生成的自由指数环。任何在 Y 中关于 {0, 1, −, +, ×, exp, cos, sin} 成立的代数关系都可以转化为在 W 中关于 {0, 1, i, −, +, ×, exp} 成立的代数关系，并且必须在任何满足 i² = −1 的指数环中成立。结论由此得证。