使用 Eigen 和 MEX 快速评估三角函数的性能瓶颈

Question

在使用 Matlab 的 C++ MEX API 的项目中，我必须计算超过 100,000 个 x 值的值 exp(j * 2pi * x)，其中 x 始终为正双精度值。我编写了一些辅助函数，它们使用欧拉公式将计算分解为 sin/cos。然后我应用范围缩减的方法将我的值减少到它们在域 [0,T/4] 中的对应点，其中 T 是我正在计算的指数的周期。我跟踪 [0, T] 中的哪个象限，原始值稍后会落入。然后，我可以使用 horner 形式的泰勒级数多项式计算三角函数，并根据原始值所在的象限应用适当的移位。有关此技术中某些概念的更多信息，请查看 this answer。下面是这个函数的代码：

Eigen::VectorXcd calcRot2(const Eigen::Ref<const Eigen::VectorXd>& idxt) {

    Eigen::VectorXd vidxt = idxt.array() - idxt.array().floor();
    Eigen::VectorXd quadrant = (vidxt.array()*2+0.5).floor();
    vidxt.array() -= (quadrant.array()*0.5);
    vidxt.array() *= 2*3.14159265358979;
    const Eigen::VectorXd sq = vidxt.array()*vidxt.array();

    Eigen::VectorXcd M(vidxt.size());
    M.real() = fastCos2(sq);
    M.imag() = fastSin2(vidxt,sq);
    M = (quadrant.array() == 1).select(-M,M);
    return M;
}

我分析了使用 std::chrono 调用此函数的代码段，并对函数的平均调用次数超过 500 次（其中对 mex 函数的每次调用通过在循环中调用 calcRot2 来处理所有 100,000 多个值。每次迭代都通过大约 200 个值到 calcRot2)。我发现以下平均运行时间：

runtime with calcRot2:                   75.4694 ms
runtime with fastSin/Cos commented out:  50.2409 ms
runtime with calcRot2 commented out:     30.2547 ms

看看这两种极端情况的区别，似乎 calcRot 对运行时的贡献很大。但是，其中只有一部分来自 sin/cos 计算。我会假设 Eigen 的隐式矢量化和编译器将使函数中其他操作的运行时有效地忽略不计。 (floor shouldn't be a problem!) 性能瓶颈到底在哪里？

这是我正在执行的编译命令（它使用我认为与 gcc 相同的 MinGW64）：

mex(ipath,'CFLAGS="$CFLAGS -O3 -fno-math-errno -ffast-math -fopenmp -mavx2"','LDFLAGS="$LDFLAGS -fopenmp"','DAS.cpp','DAShelper.cpp')

参考代码

作为参考，这里是调用定时器的主mex函数中的代码段，以及调用calcRot2()的辅助函数：

MEX 函数调用：

chk1 = std::chrono::steady_clock::now();
// Calculate beamformed signal at each point
Eigen::MatrixXcd bfVec(p.nPoints,1);
#pragma omp parallel for
for (int i = 0; i < p.nPoints; i++) {
    calcPoint(idxt.col(i),SIG,p,bfVec(i));
}
chk2 = std::chrono::steady_clock::now();
auto diff3 = chk2 - chk1;

计算点：

void calcPoint(const Eigen::Ref<const Eigen::VectorXd>& idxt,
               const Eigen::Ref<const Eigen::MatrixXcd>& SIG,
               Parameters& p, std::complex<double>& bfVal) {
    
Eigen::VectorXcd pRot = calcRot2(idxt*p.fc/p.fs);

int j = 0;
for (auto x : idxt) {
    if(x >= 0) {
        int vIDX = static_cast<int>(x);
        bfVal += (SIG(vIDX,j)*(vIDX + 1 - x) + SIG(vIDX+1,j)*(x - vIDX))*pRot(j);
    }
    j++;
}
}

澄清

为了澄清，这条线

(vidxt.array()*2+0.5).floor()

意味着屈服：

0 if vidxt is between [0,0.25] 
1 if vidxt is between [0.25,0.75]
2 if vidxt is between [0.75,1]

这里的想法是，当 vidxt 处于第二个区间时（对于周期为 2pi 的函数，单位圆上的象限 2 和 3），那么该值需要映射到它的负值。否则，范围缩小会将值映射到正确的值。

Answer 1

Eigen 向量化的好处超过了，因为您将表达式评估为临时向量。分配、解除分配、填充和读取这些向量的成本似乎很高。尤其如此，因为表达式本身相对简单（只是几个标量操作）。

表达式对象

这里通常有帮助的是聚合成更少的表达式。例如，第 3 行和第 4 行可以合并为一个： vidxt.array() = 2*3.14159265358979 * (vidxt.array() - quadrant.array()*0.5); （顺便说一句：请注意，math.h 包含一个常量 M_PI，其中 pi 为双精度）。

除此之外，Eigen 表达式可以组合和重用。像这样的：

    auto vidxt0 = idxt.array() - idxt.array().floor();
    auto quadrant = (vidxt0*2+0.5).floor();
    auto vidxt = 2*3.14159265358979 * (vidxt0 - quadrant.array()*0.5);
    auto sq = vidxt.array().square();

    Eigen::VectorXcd M(vidxt.size());
    M.real() = fastCos2(sq);
    M.imag() = fastSin2(vidxt,sq);
    M = (quadrant.array() == 1).select(-M,M);

请注意，auto 值都不是向量。它们是表达式对象，行为类似于数组，可以计算为向量或数组。

您可以通过将它们声明为模板来将它们传递给您的 fastCos2 和 fastSin2 函数。典型的 Eigen 模式类似于

template<Derived>
void fastCos2(const Eigen::ArrayBase<Derived>& sq);

这里的想法是，最终，所有内容都会编译成一个巨大的循环，当您将表达式计算为向量或数组时，就会执行该循环。如果多次引用同一个子表达式，编译器或许能够消除冗余计算。

不幸的是，我无法从这段特定的代码中获得更好的性能，所以这里并没有真正的帮助，但在这种情况下仍然值得探索。

fastSin/Cos 返回值

说到临时向量：您没有包含 fastSin/Cos 函数的代码，但看起来很像您返回一个临时向量，然后将其复制到实部和虚部或实际返回值中。这是您可能要避免的另一个临时情况。像这样的：

template<class Derived1, class Derived2>
void fastCos2(const Eigen::MatrixBase<Derived1>& M, const Eigen::MatrixBase<Derived2>& sq)
{
    Eigen::MatrixBase<Derived1>& M_mut = const_cast<Eigen::MatrixBase<Derived1>&>(M);
    M_mut = sq...;
}

fastCos2(M.real(), sq);

关于函数参数的话题请参考Eigen's documentation。

在这种特殊情况下，这种方法的缺点是现在输出不连续（实部和虚部交错）。这可能会对矢量化产生负面影响。您可以通过将 sin 和 cos 函数组合到一个表达式中来解决这个问题。需要进行基准测试。

使用普通循环

正如其他人指出的那样，在这种特殊情况下使用循环可能更容易。你注意到这比较慢。我有一个理论原因：您没有在编译选项中指定-DNDEBUG。如果不这样做，则特征向量中的所有数组索引都使用断言进行范围检查。这些会耗费时间并阻止矢量化。如果你包含这个编译标志，我发现我的代码比使用 Eigen 表达式要快得多。

或者，您可以使用原始 C 指针指向输入和输出向量。像这样的：

std::ptrdiff_t n = idxt.size();
Eigen::VectorXcd M(n);
const double* iidxt = idxt.data();
std::complex<double>* iM = M.data();
for(std::ptrdiff_t j = 0; j < n; ++j) {
    double ival = iidxt[j];
    double vidxt = ival - std::floor(ival);
    double quadrant = std::floor(vidxt * 2. + 0.5);
    vidxt = (vidxt - quadrant * 0.5) * (2. * 3.14159265358979);
    double sq = vidxt * vidxt;
    // stand-in for sincos
    std::complex<double> jval(sq, vidxt + sq);
    iM[j] = quadrant == 1. ? -jval : jval;
}

固定大小的数组

为了避免内存分配的成本并使编译器更容易首先避免内存操作，它可以帮助在固定大小的块上运行计算。像这样的：

std::ptrdiff_t n = idxt.size();
Eigen::VectorXcd M(n);
std::ptrdiff_t i;
for(i = 0; i + 4 <= n; i += 4) {
    Eigen::Array4d idxt_i = idxt.segment<4>(i);
    ...
    M.segment<4>(i) = ...;
}
if(i + 2 <= n) {
    Eigen::Array2D idxt_i = idxt.segment<2>(i);
    ...
    M.segment<2>(i) = ...;
    i += 2;
}
if(i < n) {
    // last index scalar
}

这类东西需要仔细调整，以确保生成矢量化代码并且堆栈上没有不必要的临时值。如果你会读汇编，Godbolt 很有帮助。

其他说明

• Eigen 包括 sin 和 cos 的矢量化版本。您是否将您的代码与这些进行了比较，而不是例如Eigen 的复 exp 函数？

• 根据您的数学库，还有一个显式的 sincos 函数可以在一个函数中计算正弦和余弦。它不是矢量化的，但仍然可以节省范围缩小的时间。您可以（通常）通过std::polar 访问它。试试这个：

Eigen::VectorXd scale = ...;
Eigen::VectorXd phase = ...;
// M = scale * exp(-2 pi j phase)
Eigen::VectorXd M = scale.binaryExpr(-2. * M_PI * phase,
      [](double s, double p) noexcept -> std::complex<double> {
          return std::polar(s, p);
});

• 如果您的目标是近似值而不是精确结果，那么您的第一步不应该是转换为单精度吗？也许在缩小范围后以避免丢失太多小数位。至少它会使每个时钟周期完成的工作增加一倍。此外，常规的正弦和余弦实现在浮点上花费的时间更少。

编辑

• 我不得不在演员表中将自己更正为 int64 而不是 int。在 AVX512 之前没有向量化到 int64_t 的转换

• (vidxt.array()*2+0.5).floor() 这行让我有点不爽。这意味着将 [0, 0.5) 向下舍入为负无穷大，将 [0.5, 1) 向上舍入为正无穷大，对吗？ vidxt 永远不会是负数。因此这一行应该等同于(vidxt.array()*2).round()。使用 AVX2 和 -ffast-math 可以节省一条指令。在 SSE2 中，这些都没有真正矢量化，如 Godbolt

  所示