向量化實現矩陣運算最佳化(一)

ChebyshevTST發表於2023-09-28
矩陣

xsimd簡介

  xsimd是C++的一個開源simd庫，實現了對常見simd指令的封裝，從而使得simd的操作更為簡單。接下來先從兩個簡單的例子來入門xsimd。

void average(const std::vector<double>& v1, const std::vector<double>& v2, std::vector<double>& v) {
    int n = v.size();
    int size = xsimd::batch<double, xsimd::avx>::size;
    int loop = n - n % size;

    for (int i = 0; i < loop; i += size) {
        auto a = xsimd::batch<double>::load_unaligned(&v1[i]);
        auto b = xsimd::batch<double>::load_unaligned(&v2[i]);
        auto res = a + b; 
        res.store_unaligned(&v[i]);
    }
    for (int i = loop; i < n; ++i) 
        v[i] = v1[i] + v2[i];
}

  上述demo實現了兩個向量相加的操作，由於每次都能從vector當中載入size個資料，因此對剩餘的不能進行vectorize的資料進行了分別處理。比如說，有一百個資料，每次處理8個資料，到最後剩下4個數不能湊到8，所以用樸素的迭代方式進行求和。這個demo是非對齊記憶體的處理方式。

using vector_type = std::vector<double, xsimd::default_allocator<double>>;
std::vector<double> v1(1000000), v2(1000000), v(1000000);
vector_type s1(1000000), s2(1000000), s(1000000);

void average_aligned(const vector_type& s1, const vector_type& s2, vector_type& s) {
    int n = s.size();
    int size = xsimd::batch<double>::size;
    int loop = n - n % size;

    for (int i = 0; i < loop; i += size) {
        auto a = xsimd::batch<double>::load_aligned(&s1[i]);
        auto b = xsimd::batch<double>::load_aligned(&s2[i]);
        auto res = a + b;
        res.store_aligned(&s[i]);
    }

    for (int i = loop; i < n; ++i) 
        s[i] = s1[i] + s2[i];
}

  要實現對齊記憶體的操作方式，我們必須對vector指定特定的分配器，不然最後執行出來的程式碼會出現segment fault。

  總之，要記住常用的api，load_aligned, store_aligned, load_unaligned, store_unaligned，它們分別對應了記憶體對齊與否的處理方式。接下來我們再講解另外一個demo，並且提供與openmp的效能對比。

auto sum(const std::vector<double>&v) {
    int n = v.size();
    int size = xsimd::batch<int>::size;
    int loop = n - n % size;

    double res{};
    for (int i = 0; i < loop; ++i) {
        auto tmp = xsimd::batch<int>::load_unaligned(&v[i]);
        res += xsimd::hadd(tmp);
    }

    for (int i = loop; i < n; ++i) {
        res += v[i];
    }

    return res;
}

auto aligned_sum(const std::vector<double, xsimd::default_allocator<double>>& v) {
    int n = v.size();
    int size = xsimd::batch<int>::size;
    int loop = n - n % size;

    double res{};
    for (int i = 0; i < loop; ++i) {
        auto tmp = xsimd::batch<int>::load_aligned(&v[i]);
        res += xsimd::hadd(tmp);
    }

    for (int i = loop; i < n; ++i) {
        res += v[i];
    }
    
    return res;
}

  這個例子實現了對向量求和的功能。總體與前面基本一樣，這裡hadd是一個對向量求和的函式。

  對於openmp的向量化實現，則較為簡單，只需要在for迴圈上面加上特定指令即可。不過需要注意的是，openmp支援C語法，有一些C++的新特性可能並不支援，而且需要把花括號放到下一行，我們來看具體操作。

auto parallel_sum(const std::vector<double>& v) {
    double res{};

    int n = v.size();
    #pragma omp simd
    for (int i = 0; i < n; ++i)
        res += v[i];

    return res;
}

  不要忘記加上編譯選項-fopenmp和-march=native，為了效能測試，我開啟了O2最佳化，以下是簡單的測試結果，資料規模是一千萬。

  一般情況下進行了記憶體對齊都會比沒有對齊的要快一些，同時可以看到openmp與xsimd也差了一個量級。當然不同平臺的結果可能會有差異，需要用更專業的工具進行測量比較。

相關文章