浅谈指令集优化

makerli · 2025-12-06 23:25:48 · 科技·工程

0. 前言

指令集的作用

SIMD（单指令多数据）指令集能够在常数时间内对连续内存中的一段数据进行“批量”并行操作。对于矩阵乘法、区间加/区间修改、向量点乘等对连续内存多次重复操作的场景，合理使用指令集可以显著降低常数项，从而加速代码，是卡常利器。

例如 AVX2 的寄存器宽度为 256 位，可以一次性并行处理 8 个 int（8 × 32 = 256），或 4 个 long long（4 × 64 = 256）。在合适场景下，可以将部分常数降至原来的 1/8。

1. 准备

1.1 了解可用的指令集

不要尝试在正式OI竞赛中（如 CSP）使用指令集优化!

能使用什么指令集主要看 CPU 型号。以洛谷为例，洛谷的评测机 CPU 为 Intel(R) Xeon(R) Platinum 8369HC，支持的指令集有 SSE4.2, AVX, AVX2, AVX-512。

如果不知道评测系统的 CPU 型号，可以用下面这段代码获取：

#include <stdint.h>
#include <iostream>
#include <cpuid.h>
static void cpuid(uint32_t func, uint32_t sub, uint32_t data[4]) 
{
    __cpuid_count(func, sub, data[0], data[1], data[2], data[3]);
}
int main() 
{
    uint32_t data[4];
    char str[48];
    for(int i = 0; i < 3; ++i) 
    {
        cpuid(0x80000002 + i, 0, data);
        for(int j = 0; j < 4; ++j)
            reinterpret_cast<uint32_t*>(str)[i * 4 + j] = data[j];
    }
    std::cout << str;
}

这个程序会直接输出 CPU 型号。知道型号后，可以去对应产品官网搜索 CPU 详细信息，里面会有支持的指令集。

对于 Windows，用 CPU-Z 就能直接获取到当前 CPU 所支持的指令集了。Linux 下可通过查阅 /proc/cpuinfo 或使用 cpuid 获取。

1.2 编码前准备

在程序前加入这些东西：

#pragma GCC target("sse,sse2,sse3,ssse3,sse4.1,sse4.2,avx,avx2,popcnt,tune=native")
#include <immintrin.h>
#include <emmintrin.h>

#pragma GCC target 类似于指令集的开关，括号里面那一串是各种指令集的名字，后面两个头文件里面有指令集的函数。

2. 使用

Intel 给了一份指令集函数清单，里面有所有的指令集函数：Intel® Intrinsics Guide。

Intel 还贴心地提供一份离线的版本，供机房断网考试无网络时查询：离线版。

阅读这些还是蛮需要英语水平的。

下面的内容主要针对 AVX2 指令集。

2.1 变量类型

指令集的变量叫做向量（与数学中的向量不同）。对应地，下文将普通数组称为标量。

向量的命名如下：

__<位长度><变量类型>

如果是 avx2，位长度有 128、256 两种，如果用 avx-512，还有 512 位的，一般使用256位的。

向量可以存三种类型：整形（以 i 结尾，包括 int 和 long long，short 之类的也可以）、单精度浮点数（末尾不用加任何东西），和双精度浮点数（以 d 结尾）。

举个例子，__m256i 能存 8 个 int 或 4 个 long long，__m256d 能存 4 个 double。

2.2 指令函数

指令集的指令名字大部分遵循以下规则：

_mm<位长度>_<操作>_<类型>

这里的位长度与变量中的位长度类似，要注意的是，当位长度为 128 时，位长度不用体现在函数名中。

操作有很多种，比如 add、mul、load 和 store 等，详细内容可以去 Intel® Intrinsics Guide 查找。

这里的类型指的是向量中存储的数据的类型，常用的变量类型如下表：

举个例子：对存储了 8 个 int 的 __m256i 进行加法运算，函数名就是_mm256_add_epi32。

2.3 数据存储 & 读取

假设要操作的数据类型都是 int。

可以用这个函数把 8 个 int 类型的数据加载到向量中：__m256i _mm256_set_epi32(int e7, int e6, int e5, int e4, int e3, int e2, int e1, int e0)。 long long 同理，把函数名中的 epi32 改为 epi64 即可，注意 __m256i 只能存 4 个 long long。

__m256i A;
int a[7];
for(int i=0;i<8;i++)
    cin>>a[i];
A=_mm256_set_epi32(a[7],a[6],a[5],a[4],a[3],a[2],a[1],a[0]);    //注意，这里要倒序存入

或者可以使用 _mm256_loadu_si256 直接加载对应长度的内存到向量中。

__m256i A;
int a[7];
for(int i=0;i<8;i++)
    cin>>a[i];
A=_mm256_loadu_si256((m256i*)a[0]);

还有一种方法，用指针和强制类型转换，把向量当普通数组用（可能不太安全）。

__m256i A[100];
int n,*a=(*int)&A;  //用指针访问向量数据
cin>>n;
for(int i=0;i<n;i++)
    cin>>a[i];  //直接当普通数组用，也不用考虑倒序问题
for(int i=0;i<n;i++)
    cout<<a[i]<<' ';

对于向量内数据的读取，可以使用 _mm256_storeu_si256 将向量的数据存入标量数组中。

__m256i A;
/*一些奇妙的操作*/
int t[8];
_mm256_storeu_si256((__m256i*)t,A); //把向量 A 的数据存入标量 t 中

2.4 数据计算

指令集支持加法、乘法等基础运算，还有大小比较等操作。

下面给出一个例子，分别读入长度为 800 的数组 A 和 B，将 A_i \times B_i 存入 C_i 中并输出。

//启用指令集
#include <immintrin.h>
#include <emmintrin.h>
#pragma GCC target("sse,sse2,sse3,ssse3,sse4.1,sse4.2,avx,avx2,tune=native")
#include<bits/stdc++.h>
#define endl '\n'
using namespace std;
__m256i A[100],B[100],C[100];   //指令集变量，256 位
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    int *a=(int*)&A,*b=(int*)&B,*c=(int*)&C;    //用指针，IO 时直接当普通数组用。这里(int*)强制存入 int 类型数据
    for(int i=0;i<800;i++)
        cin>>a[i];
    for(int i=0;i<800;i++)
        cin>>b[i];
    for(int i=0;i<100;i++)  //因为下面的乘法操作一次能计算 8 组数据，所以这里只用循环 800/8=100 次
        C[i]=_mm256_mullo_epi32(A[i],B[i]); //指令集乘法操作
    for(int i=0;i<800;i++)
        cout<<c[i]<<' ';
    return 0;
}

3. 举个例子

以 B2105 矩阵乘法为例，源代码如下：

#include<bits/stdc++.h>
#define endl '\n'
using namespace std;
int a[101][101],b[101][101],ans[101][101];
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    int n,m,l;
    cin>>n>>m>>l;
    for(int i=1;i<=n;i++)
        for(int j=1;j<=m;j++)
            cin>>a[i][j];
    for(int i=1;i<=m;i++)
        for(int j=1;j<=l;j++)
            cin>>b[i][j];
    for(int i=1;i<=n;i++)
        for(int j=1;j<=l;j++)
            for(int k=1;k<=m;k++)
                ans[i][j]+=a[i][k]*b[k][j];
    for(int i=1;i<=n;i++)
    {
        for(int j=1;j<=l;j++)
            cout<<ans[i][j]<<' ';
        cout<<endl;
    }
    return 0;
}

用指令集改写代码的计算部分。

先把目光放在计算部分：

for(int i=1;i<=n;i++)
    for(int j=1;j<=l;j++)
        for(int k=1;k<=m;k++)
            ans[i][j]+=a[i][k]*b[k][j];

第四行的 b[k][j] 在循环计算时会按照 k 的顺序加载，内存不连续，不太方便加载指令集。如果变成 b[j][k]，计算时读取的内存就连续了，这也为后续的操作提供了空间。

因此，在读入矩阵 b 时，可以把行和列调换（也就是矩阵转置）。

    for(int i=1;i<=m;i++)
        for(int j=1;j<=l;j++)
            cin>>b[j][i];

转置完矩阵 b 后，原来的计算公式就变为 ans[i][j]+=a[i][k]*b[j][k]。在这个公式的基础上，就可以请出指令集了！

先用 _mm256_loadu_si256 加载数据到向量中，再用 _mm256_mullo_epi32 相乘，用 _mm256_add_epi32 相加，最后把向量内的每一个数加到 ans 中。

__m256i vsum=_mm256_setzero_si256()/*初始化为全零向量*/,va,vb,prod;
for(int k;k<=m;k+=8)
{
    //加载 8 位 int 到 __m256i 中
    va=_mm256_loadu_si256((__m256i*)(&a[i][k]));
    vb=_mm256_loadu_si256((__m256i*)(&b[j][k]));
    //相乘再加和
    prod=_mm256_mullo_epi32(va,vb);
    vsum=_mm256_add_epi32(vsum,prod);
    //看似步骤变多了，但上面每一个函数都对应一条指令，且总循环次数减少了，所以效率会有提升
}
int tmp[8],sum=0;
_mm256_storeu_si256((__m256i*)tmp,vsum);    //将向量内数据存到标量中，便于下方求和
for(int t=0;t<8;t++)
    sum+=tmp[t];
ans[i][j]=sum;

但是上面代码有一个小问题：m 的值并不一定都能被 8 整除，上面的代码遇到不能被 8 整除的部分就会直接摆烂。

对于不能被 8 整除的部分，就只能老老实实算。

int k=1;
__m256i vsum=_mm256_setzero_si256()/*初始化为全零向量*/,va,vb,prod;
//能被 8 整除的部分
for(;k+7<=m;k+=8)
{
    //加载 8 位 int 到 __m256i 中
    va=_mm256_loadu_si256((__m256i*)(&a[i][k]));
    vb=_mm256_loadu_si256((__m256i*)(&b[j][k]));
    //相乘再加和
    prod=_mm256_mullo_epi32(va,vb);
    vsum=_mm256_add_epi32(vsum,prod);
    //看似步骤变多了，但上面每一个函数都对应一条指令，且总循环次数减少了，所以效率会有提升
}
int tmp[8],sum=0;
//不能被 8 整除的部分
for(;k<=m;k++)
    sum+=a[i][k]*b[j][k];
_mm256_storeu_si256((__m256i*)tmp,vsum);    //将向量内数据存到标量中，便于下方求和
for(int t=0;t<8;t++)
    sum+=tmp[t];
ans[i][j]=sum;

最终代码如下：

#include<bits/stdc++.h>
#include<immintrin.h>
#pragma GCC target("avx2,tune=native")  //只用到了 AVX2 指令集
#define endl '\n'
using namespace std;
int a[101][101],b[101][101],ans[101][101];
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    int n,m,l;
    cin>>n>>m>>l;
    for(int i=1;i<=n;i++)
        for(int j=1;j<=m;j++)
            cin>>a[i][j];
    for(int i=1;i<=m;i++)
        for(int j=1;j<=l;j++)
            cin>>b[j][i];   //转置输入矩阵
    for (int i=1;i<=n;i++)
        for (int j=1;j<=l;j++)
        {
            int k=1;
            __m256i vsum=_mm256_setzero_si256()/*初始化为全零向量*/,va,vb,prod;
            //能被 8 整除的部分
            for(;k+7<=m;k+=8)
            {
                //加载 8 位 int 到 __m256i 中
                va=_mm256_loadu_si256((__m256i*)(&a[i][k]));
                vb=_mm256_loadu_si256((__m256i*)(&b[j][k]));
                //相乘再加和
                prod=_mm256_mullo_epi32(va,vb);
                vsum=_mm256_add_epi32(vsum,prod);
                //看似步骤变多了，但上面每一个函数都对应一条指令，且总循环次数减少了，所以效率会有提升
            }
            int tmp[8],sum=0;
            //不能被 8 整除的部分
            for(;k<=m;k++)
                sum+=a[i][k]*b[j][k];
            _mm256_storeu_si256((__m256i*)tmp,vsum);    //将向量内数据存到标量中，便于下方求和
            for(int t=0;t<8;t++)
                sum+=tmp[t];
            ans[i][j]=sum;
        }
    for(int i=1;i<=n;i++)
    {
        for(int j=1;j<=l;j++)
            cout<<ans[i][j]<<' ';
        cout<<endl;
    }
    return 0;
}

4. 后记

作者能力有限，如有错误请在评论区指出。

祝各位算法效率爆棚，不用指令集卡常。