在C++中使用NVidia可以使for循环更快吗?

ojsjcaue  于 2023-06-07  发布在  其他
关注(0)|答案(1)|浏览(188)

我想让一个C++函数更快。我在问你可能的方法。
我可以使用多达32个OMP线程。
我可以使用NVIDIA GPU。
函数的MWE为:

#include <iostream>
#include <complex>
#include <cmath>

typedef std::numeric_limits<double> dbl;
#define _USE_MATH_DEFINES

#include <omp.h>

const std::complex<double> I(0.0, 1.0); // imaginary unit, I*I = -1
std::complex<double> zero_imag (0.0, 0.0);

const int N_rs = 1500;
const int l_max = 70;
const int lmax = 70;
const int N_thetas = l_max + 1;
const int N_phis = 2 * l_max + 2;
const int N_ps = 600;
const int nphi = 2 * l_max + 2;
const double sqrt_of_2_over_pi = sqrt( 2.0 / M_PI );

void rtop(std::complex<double> * Psi_outer_spec,
          std::complex<double> * Psi_outer_spec_plm,
          double * BJ,
          double * wrk,
          std::complex<double> * wrk2,
          double * ris_without_ends,
          double * r_primes_without_ends,
          double * weights_Lobatto_without_ends
         )
{

    int l, kk, kkk, m;
    long int idx, idxx, idxxx;

    // #pragma omp parallel for firstprivate (wrk2) private(l, kkk, idx, m, kk, idxx, idxxx) schedule(static)
    // #pragma omp target teams distribute parallel for firstprivate(wrk2) private(l, kkk, idx, m, kk, idxx, idxxx)
    for (int i = 0; i <= (N_ps - 1); i++) { // THIS IS THE BOTTLENECK !!!
       
        std::complex<double> sum1 = std::complex<double> (0.0, 0.0); // each thread creates a sum1 on its own

        for (l = 0; l <= lmax; l++) {

            for (kkk = 0; kkk <= (N_rs-1); kkk++) {
                idx = i * (N_rs*(l_max+1)) + kkk * (l_max+1) + l;
                wrk2[kkk] = pow(-I, l) * BJ[idx] * wrk[kkk];
            }

            for (m = 0; m <= (nphi-1); m++) {

                sum1 = zero_imag;
                for (kk = 0; kk <= (N_rs-1); kk++) {
                    idxx = kk * (N_thetas*N_phis) + l * N_phis + m;
                    sum1 += Psi_outer_spec[idxx] * wrk2[kk];

                }

                idxxx = i * (N_thetas*N_phis) + l * N_phis + m;
                Psi_outer_spec_plm[idxxx] = sum1 * sqrt_of_2_over_pi;
                                       
            }
            // END for m loop
        }
        // END for l loop
    }    
    // END for i loop
}

int main() {

    double * wrk = new double [N_rs];
    std::complex<double> * wrk2 = new std::complex<double> [N_rs];

    double * ris_without_ends = new double [N_rs];
    double * r_primes_without_ends = new double [N_rs];
    double * weights_Lobatto_without_ends = new double [N_rs];

    double * BJ = new double [N_ps * N_rs * (l_max+1)];

    std::complex<double> * Psi_outer_spec = new std::complex<double> [N_rs * N_thetas * N_phis];
    std::complex<double> * Psi_outer_spec_plm = new std::complex<double> [N_ps * N_thetas * N_phis];

    rtop(Psi_outer_spec, Psi_outer_spec_plm, BJ, wrk, wrk2, ris_without_ends, r_primes_without_ends, weights_Lobatto_without_ends);
   
    return 0;
}

关联的CMakeLists.txt为:

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)

set(CMAKE_VERBOSE_MAKEFILE ON)

set(CMAKE_C_COMPILER "gcc")
set(CMAKE_CXX_COMPILER "g++")

project(trial)

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -pedantic -Wall")

find_package(OpenMP)

add_executable(trial trial.cpp)

if(OpenMP_CXX_FOUND)
target_link_libraries(trial PUBLIC OpenMP::OpenMP_CXX)
endif()

set_property(TARGET trial PROPERTY CXX_STANDARD 17)

编译为:$ cmake ..然后$ cmake --build . --config Release
我的输出是:

-- The C compiler identification is GNU 11.3.0
-- The CXX compiler identification is GNU 11.3.0
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Check for working C compiler: /apps20/sw/eb/software/GCCcore/11.3.0/bin/gcc - skipped
-- Detecting C compile features
-- Detecting C compile features - done
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Check for working CXX compiler: /apps20/sw/eb/software/GCCcore/11.3.0/bin/g++ - skipped
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Found OpenMP_C: -fopenmp (found version "4.5")
-- Found OpenMP_CXX: -fopenmp (found version "4.5")
-- Found OpenMP: TRUE (found version "4.5")
-- Configuring done
-- Generating done
-- Build files have been written to: /work4/clf/ouatu/trial_for_SO/build

然后对于构建:

[ 50%] Building CXX object CMakeFiles/trial.dir/trial.cpp.o
[100%] Linking CXX executable trial
[100%] Built target trial

我所尝试的:

  • 使用OpenMP并行for,我确实获得了加速。
  • OpenMP GPU卸载失败(似乎我的编译器标志无法实现卸载)。(这些标志在此MWE显示的CMakeLists.txt中隐藏)
  • 我愿意接受任何其他建议。

例如,rtop作为CUDA内核是否会受益?很难做到吗?
谢谢你!

shyt4zoc

shyt4zoc1#

我建议一个OpenMP版本与一些优化和调整。快速回顾一些变化以及需要注意的事项:
使用wrk2[kkk] = pow(-I, l) * ...的整个业务是双重冗余的。首先,pow(-I, l)是一种优雅但昂贵的方式,只表达4个不同的值。其次,它只在点积中用作因子。你可以把整个过程折叠成最后的乘法sum1 * sqrt_of_2_over_pi。这也允许wrk2是实值的,这也将最里面的循环从复数-复数点积变成复数-实数点积。
idx = i * (N_rs*(l_max+1)) + kkk * (l_max+1) + l这样的多维索引计算应该在Horner method之后进行,以避免冗余乘法。更多的是吹毛求疵,但也更清楚。例如idx = (i * N_rs + kkk) * (l_max+1) + l。当我们这样做的时候,要小心你的索引变量。它们都是int。特别是三维数组可以快速增长到多个GiB的大小,此时您将遇到整数溢出。如果您担心这可能会成为一个问题,请切换到std::ptrdiff_t
BJPsi_outer_spec_plm上的迭代顺序并不理想。如果可能的话,BJ应该交换两个内部维度以获得更好的数据局部性,这也将允许初始化wrk2的循环的向量化。Psi_outer_spec更糟糕,因为在最里面的循环中沿着外部维度进行迭代。然而,我假设这个顺序是选择的,所以它与Psi_outer_spec_plm相同,因此它是好的。在任何情况下,该较高步幅防止向量化。
我不明白为什么你要在使用它们的范围之外声明计数器和索引变量。即使是现代的C标准也允许在for循环中声明它们,更不用说C++了。对于并行化,您希望限制共享或意外共享的变量的数量。
说到共享数据,据我所知,线程可能重叠的唯一共享内存是wrk2数组。这可以简单地为每个线程分配,这将我们带到了最终的实现。

#   pragma omp parallel
    {
        auto wrk2 = std::make_unique<double[]>(N_rs);
#       pragma omp for collapse(2) nowait
        for (int i = 0; i <= (N_ps - 1); i++) {
            for (int l = 0; l <= lmax; l++) {
                for (int kkk = 0; kkk <= (N_rs-1); kkk++) {
                    int idx = (i * N_rs + kkk) * (lmax + 1) + l;
                    wrk2[kkk] = BJ[idx] * wrk[kkk];
                }
                constexpr std::complex<double> I(0., 1.);
                std::complex<double> factor(-sqrt_of_2_over_pi);
                if(l & 1)
                    factor *= I;
                if(l & 2)
                    factor = -factor;
                for (int m = 0; m <= (N_phis-1); m++) {
                    std::complex<double> sum1;
                    for (int kk = 0; kk <= (N_rs-1); kk++) {
                        int idx = (kk * N_thetas + l) * N_phis + m;
                        sum1 += Psi_outer_spec[idx] * wrk2[kk];
                    }
                    int idx = (i * N_thetas + l) * N_phis + m;
                    Psi_outer_spec_plm[idx] = sum1 * factor;
                }
            }
        }
    }

请注意,通常的pragma omp parallel for是如何拆分为omp parallel和单独的omp for以允许分配临时内存的。collapse(2)表示两个外部循环都是并行的。
其他要考虑的事项:

  • 内部点积可以通过加速的BLAS库或类似的东西更快地计算。我认为Eigen在这里应该可以很好地工作,但是可能需要对它进行一点强制,使其能够使用这种内存布局
  • 看起来我们可以将m循环改为矩阵向量乘积,这可能通过BLAS库解决一些向量化/内存访问问题
  • 既然您询问了编译选项,那么-march=native或您想要的任何基线架构都应该值得在这里使用。-mavx2 -mfma可能是一个很好的折衷方案,可以处理所有相对较新的CPU,而不必过多地专门化二进制文件

编辑:矩阵向量积

回到将m上的循环卸载为矩阵向量积的想法,我们必须重新解释我们用作矩阵的Psi_outer_spec的切片。我选择了一个列主矩阵,因为我想在这一步使用Eigen 3。

  • 行数为N_phi(循环计数器m
  • 列数为N_rs(循环计数器kk
  • 从一列到下一列,我们有一个步幅/ a.k.a. N_phi * N_theta的前导尺寸
  • 左上角的偏移量为l * N_phis

假设这是正确的,我们可以将数组Map到特征向量和矩阵,并让它处理转置访问。这会将wrk2初始化下面的所有内容转换为以下代码

using MatrixMap = Eigen::Map<const Eigen::MatrixXcd,
        Eigen::Unaligned, Eigen::OuterStride<>>;
MatrixMap Psi_slice(
        Psi_outer_spec + l * N_phis /*top left corner*/,
        N_phis /*rows*/, N_rs /*cols*/,
        Eigen::OuterStride<>(N_phis * N_thetas));
const auto wrk2_mapped = Eigen::VectorXd::Map(wrk2.get(), N_rs);
auto Psi_plm_mapped = Eigen::VectorXcd::Map(
        Psi_outer_spec_plm + (i * N_thetas + l) * N_phis, N_phis);
Psi_plm_mapped.noalias() = Psi_slice * wrk2_mapped * factor;

现在,这一步显然提出了一个问题,即我们是否可以通过一些预处理或后处理将整个事情转变为矩阵-矩阵乘积,这可能会照顾到整个并行化和潜在的GPU卸载。这就是为什么我要求一个数学描述,而不是通过代码做这种徒劳的追逐

编辑2:矩阵-矩阵积

实际上可以将其重写为矩阵-矩阵乘积。诀窍在于观察到Psi_outer_speci是独立的。因此,如果我们切换两个外部循环,我们可以在一次操作中计算一个l在所有i上的所有值。
在这样做的时候,我切换回wrk2,因为它是复杂的,并且包含了因子。从技术上讲,这需要更多的计算时间和内存,但对于矩阵-矩阵产品,您可能希望直接通过OpenBLASEigen's backends或甚至通过GPU加速(如CuBLAS)调度到BLAS后端。为此,你需要一个复数-复数乘法。

Eigen::MatrixXcd wrk2mat(N_rs, N_ps);
for (int l = 0; l <= lmax; l++) {
    std::complex<double> factor(-sqrt_of_2_over_pi);
    if(l & 1)
        factor *= I;
    if(l & 2)
        factor = -factor;
#   pragma omp parallel for
    for (int i = 0; i <= N_ps - 1; i++) {
        for (int k = 0; k <= N_rs - 1; ++k) {
            int idx = (i * N_rs + k) * (lmax + 1) + l;
            wrk2mat(k, i) = BJ[idx] * wrk[k] * factor;
        }
    }
    using ConstMatrixMap = Eigen::Map<const Eigen::MatrixXcd,
            Eigen::Unaligned, Eigen::OuterStride<>>;
    ConstMatrixMap Psi_slice(
            Psi_outer_spec + l * N_phis /*top left corner*/,
            N_phis /*rows*/, N_rs /*cols*/,
            Eigen::OuterStride<>(N_phis * N_thetas));
    using MatrixMap = Eigen::Map<Eigen::MatrixXcd,
            Eigen::Unaligned, Eigen::OuterStride<>>;
    MatrixMap Psi_plm_mapped(
            Psi_outer_spec_plm + l * N_phis,
            N_phis, N_ps,
            Eigen::OuterStride<>((lmax + 1) * N_phis));
    Psi_plm_mapped.noalias() = Psi_slice * wrk2mat;
}

只要矩阵足够大,矩阵-矩阵乘积就应该在内部并行化。如果情况并非总是如此,您可以将整个事情 Package 到运行时可选的并行块中。大致如下:

bool small_matrices = ...;
#pragma omp parallel if(small_matrices)
{
    Eigen::MatrixXcd wrk2mat(N_rs, N_ps);
#   pragma omp for nowait
    for (int l = 0; l <= lmax; l++) {
        ...
    }
}

由于OpenMP通常会停用嵌套并行化,因此这将自动停用所有内部parallel部分并按顺序运行它们。

相关问题