如何用std::transform现代化改造并行C风格矩阵-向量乘法实现

问题背景 #

现有基于OpenMP并行的C风格矩阵-向量乘法实现（C[N] = A[N,M] * B[M]），需要改用STL的std::transform结合并行执行策略（如std::execution::par）实现现代化版本，核心问题是如何在std::transform的函数体内获取计算每个C元素所需的行索引，同时寻求其他STL替代方案。

基于std::transform的实现方案 #

std::transform默认迭代容器元素，但我们需要的是每个C元素的行索引来定位矩阵A中对应的行与向量B做点积。因此可以构造一个索引序列，通过迭代索引完成计算：

C++17兼容版本 #

由于C++17无范围库，先创建存储索引的临时向量，再执行std::transform：

template<typename T>
bool run_transform_based_solution(std::vector<T> &A, std::vector<T> &B, std::vector<T> &C, size_t N, size_t M)
{
    if (A.size() != N*M || B.size() != M || C.size() != N) return 1;

    // 构造索引序列：0,1,...,N-1
    std::vector<size_t> indices(N);
    std::iota(indices.begin(), indices.end(), 0);

    // 并行transform：每个索引i对应计算C[i] = A[i*M ... i*M+M-1] · B
    std::transform(std::execution::par,
                   indices.begin(), indices.end(),
                   C.begin(),
                   [&](size_t i) {
                       T sum = 0.0;
                       for (size_t j = 0; j < M; ++j) {
                           sum += A[i*M + j] * B[j];
                       }
                       return sum;
                   });

    return 0;
}

C++20优化版本（可选） #

若升级到C++20，利用范围库std::views::iota可省去临时索引向量的开销：

#include <ranges>

template<typename T>
bool run_transform_based_solution(std::vector<T> &A, std::vector<T> &B, std::vector<T> &C, size_t N, size_t M)
{
    if (A.size() != N*M || B.size() != M || C.size() != N) return 1;

    std::transform(std::execution::par,
                   std::views::iota(0u, N).begin(), std::views::iota(0u, N).end(),
                   C.begin(),
                   [&](size_t i) {
                       T sum = 0.0;
                       for (size_t j = 0; j < M; ++j) {
                           sum += A[i*M + j] * B[j];
                       }
                       return sum;
                   });

    return 0;
}

其他STL替代方案 #

方案1：std::for_each并行版本 #

std::for_each同样支持并行执行策略，逻辑更贴近原OpenMP版本：

template<typename T>
bool run_for_each_based_solution(std::vector<T> &A, std::vector<T> &B, std::vector<T> &C, size_t N, size_t M)
{
    if (A.size() != N*M || B.size() != M || C.size() != N) return 1;

    std::vector<size_t> indices(N);
    std::iota(indices.begin(), indices.end(), 0);

    std::for_each(std::execution::par,
                  indices.begin(), indices.end(),
                  [&](size_t i) {
                      T sum = 0.0;
                      for (size_t j = 0; j < M; ++j) {
                          sum += A[i*M + j] * B[j];
                      }
                      C[i] = sum;
                  });

    return 0;
}

方案2：std::generate并行版本 #

若无需保留C初始值（已提前fill为0），可直接用std::generate生成每个C元素：

template<typename T>
bool run_generate_based_solution(std::vector<T> &A, std::vector<T> &B, std::vector<T> &C, size_t N, size_t M)
{
    if (A.size() != N*M || B.size() != M || C.size() != N) return 1;

    size_t i = 0;
    std::generate(std::execution::par,
                  C.begin(), C.end(),
                  [&]() mutable {
                      T sum = 0.0;
                      for (size_t j = 0; j < M; ++j) {
                          sum += A[i*M + j] * B[j];
                      }
                      i++;
                      return sum;
                  });

    return 0;
}

编译说明 #

所有代码需以-std=c++17（或-std=c++20）编译，部分编译器需额外添加-ltbb或-fopenmp以启用并行执行策略。

内容的提问来源于stack exchange，提问作者Andreas Hadjigeorgiou