奇妙之旅：SIMD加速矩阵运算

前言

游戏中的矩阵运算往往涉及大量4x4矩阵乘法。最直接的实现方法是通过循环实现4×4×4次乘法及相应的加法运算来得到结果。然而，当计算量达到数百万级别时，传统的实现方式会暴露出性能瓶颈。

在此之前，我通过查阅相关文献，发现了一种名为SIMD（Single Instruction Multiple Data，单指令多数据流）的技术。这种技术允许CPU在执行单个指令时同时处理多个数据，从而显著提升了计算效率。

预备知识

原理

SIMD技术的核心在于指令的多数据处理能力。与传统的SISD（单指令单数据）CPU相比，SIMD CPU能够在同一指令下同时访问和处理多个数据项。这种特性使得SIMD特别适合于处理数据密集型任务，如矩阵运算。

AVX和AVX2指令集

以下是AVX和AVX2指令集中的一些函数原型：

__m256 _mm256_i32gather_ps

概要：从内存中按32位索引提取单精度浮点数数据，并存储到256位寄存器中。

描述：从内存中的base_addr地址开始，每隔32位读取一份数据，通过vindex提供的索引和scale倍数来确定读取的位置。

伪代码：

FOR j := 0 to 7    i := j*32    m := j*32    addr := base_addr + SignExtend64(vindex[m+31:m]) * ZeroExtend64(scale) * 8    dst[i+31:i] := MEM[addr+31:addr]ENDFORdst[MAX:256] := 0

__m256 _mm256_dp_ps

概要：对两个256位寄存器中的单精度浮点数数据进行乘法和求和操作，并将结果存储到目标寄存器中。

描述：根据imm8的高4位决定是否进行乘法操作，并根据低4位决定结果的存储位置。

伪代码：

DEFINE DP(a[127:0], b[127:0], imm8[7:0]) {    FOR j := 0 to 3        i := j*32        IF imm8[(4+j)%8]            temp[i+31:i] := a[i+31:i] * b[i+31:i]        ELSE            temp[i+31:i] := 0        FI    ENDFOR    sum[31:0] := (temp[127:96] + temp[95:64]) + (temp[63:32] + temp[31:0])    FOR j := 0 to 3        i := j*32        IF imm8[j%8]            tmpdst[i+31:i] := sum[31:0]        ELSE            tmpdst[i+31:i] := 0        FI    FI    RETURN tmpdst[127:0]}

通过以上两个函数，我们可以将4x4矩阵乘法中的乘法次数从64次降低到8次。

计算逻辑

由于我的CPU仅支持AVX2指令集，不具备512位的寄存器，因此在实现中使用256位寄存器进行演示。具体来说，通过将矩阵乘法中的乘法和加法操作分解到256位寄存器中，能够显著提高运算效率。

一次性计算两次点积

temp = _mm256_dp_ps(a12, b11, mask);c.data[0] = temp[0];c.data[1] = temp[4];

剩下的类似操作

temp = _mm256_dp_ps(a12, b22, mask);c.data[4] = temp[0];c.data[5] = temp[4];

最终结果

通过上述方法，一次完整的矩阵乘法运算即可完成。

代码实战

为了验证SIMD技术的性能优势，我编写了以下代码进行测试。

测试结果

SIMD-AVX版本：仅需约0.1ms完成10,000,000次循环运算。

SISD TRIVIAL版本：需要约100ms完成相同次数的运算。

代码解析

#include 
   
    #include 
    
     #include 
     
      #include 
      
       class matrix4 {public:    matrix4() : data{0} {}    matrix4(float v) : data{v,0,0,0,0,v,0,0,0,0,v,0,0,0,0,v} {}    public:        union {            float data[16];            float ptr[4][4];        };};__declspec(align(16)) matrix4 A(1.f);__declspec(align(16)) matrix4 B(2.f);__declspec(align(16)) float v[16] = {0, 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15};// 生成索引数组__declspec(align(16)) __m256i gatherA12 = _mm256_set_epi32(13, 9, 5, 1, 12, 8, 4, 0);__declspec(align(16)) __m256i gatherA34 = _mm256_set_epi32(15, 11, 7, 3, 14, 10, 6, 2);__declspec(align(16)) __m256i gatherB11 = _mm256_set_epi32(3, 2, 1, 0, 3, 2, 1, 0);__declspec(align(16)) __m256i gatherB22 = _mm256_set_epi32(7, 6, 5, 4, 7, 6, 5, 4);__declspec(align(16)) __m256i gatherB33 = _mm256_set_epi32(11, 10, 9, 8, 11, 10, 9, 8);__declspec(align(16)) __m256i gatherB44 = _mm256_set_epi32(15, 14, 13, 12, 15, 14, 13, 12);__declspec(align(16)) __m256 temp;__declspec(align(16)) __m256 a12, a34;__declspec(align(16)) __m256 b11, b22, b33, b44;// 加载数据到寄存器a12 = _mm256_i32gather_ps(A.data, gatherA12, sizeof(float));a34 = _mm256_i32gather_ps(A.data, gatherA34, sizeof(float));b11 = _mm256_i32gather_ps(B.data, gatherB11, sizeof(float));b22 = _mm256_i32gather_ps(B.data, gatherB22, sizeof(float));b33 = _mm256_i32gather_ps(B.data, gatherB33, sizeof(float));b44 = _mm256_i32gather_ps(B.data, gatherB44, sizeof(float));// 执行乘法和加法temp = _mm256_dp_ps(a12, b11, 0b11110001);ret.data[0] = temp.m256_f32[0];ret.data[1] = temp.m256_f32[4];temp = _mm256_dp_ps(a34, b11, 0b11110001);ret.data[2] = temp.m256_f32[0];ret.data[3] = temp.m256_f32[4];// 依此类推...

主函数

int main() {    time_begin();    for (int i = 0; i < 10000000; ++i) {        volatile auto res = mm_mul_mat(A, B);    }    std::cout << "SIMD-AVX time: " << time_end() << std::endl;    time_begin();    for (int i = 0; i < 10000000; ++i) {        volatile auto m = mul_1(A, B);    }    std::cout << "SISD TRIVIAL time: " << time_end() << std::endl;    return 0;}

结论

通过上述优化，SIMD技术显著提升了矩阵乘法的性能。这次实践让我深刻体会到技术的魅力，也激发了我进一步探索高性能计算领域的兴趣。

转载地址：http://tewp.baihongyu.com/

你可能感兴趣的文章

Objective-C实现knapsack背包问题算法(附完整源码)