rustc无法为闭区间循环生成SIMD指令的排查与解决

问题现象 #

• 编写逻辑完全等价的三重循环（统计2000范围内勾股三元组数量），分别用rustc、g++/clang++开启-O3优化编译，Rust版本运行耗时1.38s，C++版本仅耗时0.42s，性能差3倍，排查确认差距来自rustc未自动生成AVX向量指令。
• 尝试给rustc添加-C target-feature=+avx编译参数，未实现预期的向量化优化。
• 已知rustc基于LLVM后端，理论支持AVX指令生成，核心疑问为：Rust是否对AVX使用有安全限制，或是开启AVX的配置方式有误。

复现环境与代码 #

编译器版本 #

g++ 9.4.0
clang++ 10.0.0
rustc 1.64.0-nightly

测试CPU参数 #

Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   39 bits physical, 48 bits virtual
CPU(s):                          8
On-line CPU(s) list:             0-7
Thread(s) per core:              2
Core(s) per socket:              4
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           60
Model name:                      Intel(R) Core(TM) i7-4710HQ CPU @ 2.50GHz
Stepping:                        3
CPU MHz:                         800.000
CPU max MHz:                     3500,0000
CPU min MHz:                     800,0000
BogoMIPS:                        4988.85
Virtualization:                  VT-x
L1d cache:                       128 KiB
L1i cache:                       128 KiB
L2 cache:                        1 MiB
L3 cache:                        6 MiB
NUMA node0 CPU(s):               0-7
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon pebs
                                  bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_
                                 deadline_timer aes xsave avx f16c rdrand lahf_lm abm cpuid_fault epb invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 
                                 avx2 smep bmi2 erms invpcid xsaveopt dtherm ida arat pln pts md_clear flush_l1d

Rust测试版本 #

编译命令：

rustc -C opt-level=3

源码：

use std::time;

const N: i32 = 2000;

fn main() {
    println!("n: {}", N);
    let start = time::Instant::now();

    let mut count: i32 = 0;
    for a in 1..=N {
        for b in 1..a {
            for c in 1..=b {
                if a*a == b*b + c*c {
                    count += 1;
                }
            }
        }
    }

    let duration = start.elapsed();
    println!("found: {}", count);
    println!("elapsed: {:?}", duration);
}

运行结果：

n: 2000
found: 1981
elapsed: 1.382231558s

C++测试版本 #

初始编译命令：

g++ -O3

clang++ -O3

源码：

#include <iostream>
#include <chrono>

constexpr int N = 2000;
constexpr double MS_TO_SEC = 1e-3;

int main(int argc, char **argv) {
    std::cout << "n: " << N << std::endl;
    auto start = std::chrono::high_resolution_clock::now();

    int count = 0;
    for (int a = 1; a <= N; ++a) {
        for (int b = 1; b < a; ++b) {
            for (int c = 1; c <= b; ++c) {
                if (a*a == b*b + c*c) {
                    count += 1;
                }
            }
        }
    }

    auto elapsed = std::chrono::high_resolution_clock::now() - start;
    std::cout << "found: " << count << std::endl;
    std::cout << "elapsed: " 
              << double(std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count())*MS_TO_SEC 
              << "s" << std::endl;
}

运行结果：

n: 2000
found: 1981
elapsed: 0.419s

补充测试 #

给C编译参数添加-mavx2启用AVX2指令集后，C版本性能进一步提升，此时Rust与C++的性能差距拉大到9倍。

原因与解决方案 #

Rust不存在AVX等SIMD指令集的安全层面使用限制，未触发自动向量化是两个问题共同导致的，调整后即可达到和同参数C++一致的性能：

1. 指令集参数配置错误：仅添加-C target-feature=+avx不足以触发对应循环的向量化优化，需要匹配CPU支持的指令集版本，测试CPU支持AVX2，需添加-C target-feature=+avx2参数；本地测试也可以直接使用-C target-cpu=native参数，让rustc自动识别当前CPU支持的所有指令集并启用。
   • 注意：如果编译需要分发的通用二进制，不要使用target-cpu=native，否则生成的程序在不支持对应指令集的CPU上会触发非法指令崩溃，需要按目标用户的最低硬件配置指定对应target-feature。
2. 闭区间Range写法阻碍优化：1.64版本的Rust中，1..=N这类闭区间范围（RangeInclusive）的实现会干扰LLVM的循环分析，导致无法自动做向量化变换，把所有闭区间循环改成半开区间写法即可解决，比如将for a in 1..=N改为for a in 1..(N+1)，for c in 1..=b改为for c in 1..(b+1)。

内容的提问来源于stack exchange，提问作者Fedor Lebed