使用特征值快速计算平方范数和归一化向量

Question

我的一段代码需要运行数百万次。看起来成本最高的部分就是这个计算：

Vector2d directedInverseSquare(Vector2d &pos1, Vector2d &pos2) {

    Vector2d diff = pos2 - pos1;
    Vector2d rHat = diff.normalized();
    double rSquared = diff.squaredNorm();

    return rHat / rSquared;
}

它在

1/r^2

->

pos1

的方向上计算

pos2

，其中

r

是它们之间的距离。 这也相当于

R/r^3

，其中 R 是原始差异向量。但我发现这里编写代码的形式更快。

在我看来，使代码更快的明显方法是以某种方式同时“生成”平方距离和标准化向量。他们的计算过程非常相似，但我目前仍然一个接一个地称呼他们，没有任何这样的优化。

我尝试过自己做，但 Eigen 总是出来得更快。所以也许执行此操作的特征函数是最好的，但我还没有找到这样的函数。

Answer 1

好吧，随着讨论的结束，这是我的测试代码和改进版本。

首先，我的基准。没什么花哨的。

struct AoS
{
    Eigen::Vector2d a, b, c;
};

int main()
{
    const int repetitions = 2000000;
    constexpr int blocks = 340; // ~ 16 kiB
    AoS inbufs[blocks], outbufs[blocks];
    for(AoS& buf: inbufs)
        for(Eigen::Vector2d* vec: {&buf.a, &buf.b, &buf.c})
            *vec = Eigen::Vector2d::Random();
    for(int i = 0; i < repetitions; ++i) {
        for(int j = 0; j < blocks; ++j) {
            const AoS& in = inbufs[j];
            AoS& out = outbufs[j];
            out.a = directedInverseSquare(in.a, in.b);
            out.b = directedInverseSquare(in.b, in.c);
            out.c = directedInverseSquare(in.c, in.a);
        }
        // don't optimize the store away
        asm volatile ("" ::: "memory");
    }
}

编译标志

-DNDEBUG -march=native -fno-math-errno

使用 gcc-12.3 编译。在我的 TigerLake CPU 上，您的实现需要 6.2 秒。

切换到简单的

diff / r**3

给我们 4.5 秒的时间：

Eigen::Vector2d directedInverseSquare2(
        const Eigen::Vector2d &pos1, const Eigen::Vector2d &pos2)
{
    Eigen::Vector2d diff = pos2 - pos1;
    double r = diff.norm();
    return diff / (r * r * r);
}

chtz 的版本速度大致相同，但指令较少。客观地说我认为这是最好的。

Eigen::Vector2d directedInverseSquare3(
        const Eigen::Vector2d &pos1, const Eigen::Vector2d &pos2)
{
    Eigen::Vector2d diff = pos2 - pos1;
    double rSquared = diff.squaredNorm();
    return diff / (rSquared * std::sqrt(rSquared));
}

为了改进这一点，我们应该使用 2x3 矩阵进行矢量化。这使我们能够一步完成

b-a

和

c-b

的情况。

a-c

操作必须单独完成。

using Matrix23d = Eigen::Matrix<double, 2, 3>;

void directedInverseSquare(Matrix23d& out, const Matrix23d& in)
{
    Eigen::Array22d diffs = in.rightCols<2>() - in.leftCols<2>();
    Eigen::Array2d rSquared = diffs.colwise().squaredNorm();
    out.leftCols<2>() = diffs.rowwise() / (rSquared * rSquared.sqrt()).transpose();
    Eigen::Vector2d diff3 = in.col(0) - in.col(2);
    double rSquared3 = diff3.squaredNorm();
    out.col(2) = diff3 / (rSquared3 * std::sqrt(rSquared3));
}
int main()
{
    const int repetitions = 2000000;
    constexpr int blocks = 340; // ~ 16 kiB
    Matrix23d inbufs[blocks], outbufs[blocks];
    for(Matrix23d& buf: inbufs)
        buf = Matrix23d::Random();
    for(int i = 0; i < repetitions; ++i) {
        for(int j = 0; j < blocks; ++j)
            directedInverseSquare(outbufs[j], inbufs[j]);
        asm volatile ("" ::: "memory");
    }
}

不幸的是，在我的测试中这只是稍微快一些； 4.3秒。不过，组装看起来不错。-我还测试了一个转置版本，效果稍差。

但是，如果您有更多向量，这种矩阵样式（尤其是转置版本）效果非常好。矢量化的好处增加，最后一个标量元素的相对成本减少。

void directedInverseSquare(Eigen::MatrixX2d& out, const Eigen::MatrixX2d& in)
{
    const Eigen::Index n = in.rows();
    const auto& diffs = in.bottomRows(n - 1) - in.topRows(n - 1);
    const auto& rsquared = diffs.rowwise().squaredNorm().array();
    out.topRows(n - 1) = diffs.array().colwise() / (rsquared * rsquared.sqrt());
    Eigen::Vector2d diff3 = in.row(0) - in.row(n - 1);
    double rSquared3 = diff3.squaredNorm();
    // note the eval() to force vectorized evaluation
    out.row(n - 1) = (diff3 / (rSquared3 * std::sqrt(rSquared3))).eval();
}
int main()
{
    const int repetitions = 2000000;
    constexpr int blocks = 340; // ~ 16 kiB
    Eigen::MatrixX2d inbufs, outbufs(340 * 3, 2);
    inbufs = Eigen::MatrixX2d::Random(340 * 3, 2);
    for(int i = 0; i < repetitions; ++i) {
        directedInverseSquare(outbufs, inbufs);
        asm volatile ("" ::: "memory");
    }
}

这需要 2.8 秒。

作为奖励，这是我可以使用 Intel 内在函数 和 AVX + FMA 指令为 2x3 矩阵版本想出的最佳版本。这需要 4.1 秒。

#include <immintrin.h>

void directedInverseSquare(Matrix23d& out, const Matrix23d& in)
{
    const double* inptr = in.data();
    // [b0, b1, c0, c1]
    __m256d rightcols = _mm256_loadu_pd(inptr + 2);
    // [a0, a1, b0, b1]
    __m256d leftcols = _mm256_loadu_pd(inptr);
    // [b0 - a0, b1 - a1, c0 - b0, c1 - b1]
    __m256d diffs = _mm256_sub_pd(rightcols, leftcols);
    // [c0, c1]
    __m128d cvec = _mm256_extractf128_pd(rightcols, 1);
    // [a0, a1]
    __m128d avec = _mm256_castpd256_pd128(leftcols);
    // [a0 - c0, a1 - c1]
    __m128d taildiff = _mm_sub_pd(avec, cvec);
    __m256d taildiff256 = _mm256_castpd128_pd256(taildiff);
    // [b0 - a0, a0 - c0, c0 - b0, x]
    __m256d lowvals = _mm256_unpacklo_pd(diffs, taildiff256);
    // [b1 - a1, a1 - c1, c1 - b1, x]
    __m256d highvals = _mm256_unpackhi_pd(diffs, taildiff256);
    // [norm(b-a)**2, norm(a-c)**2, norm(c-b)**2, x]
    __m256d rsquared = _mm256_mul_pd(lowvals, lowvals);
    rsquared = _mm256_fmadd_pd(highvals, highvals, rsquared);
    // [norm(b-a), norm(a-c), norm(c-b), x]
    __m256d r = _mm256_sqrt_pd(rsquared);
    // [norm(b-a)**3, norm(a-c)**3, norm(c-b)**3, x]
    __m256d rcubed = _mm256_mul_pd(r, rsquared);
    // [norm(b-a)**3, norm(b-a)**3, norm(c-b)**3, norm(c-b)**3]
    __m256d bcast1 = _mm256_permute_pd(rcubed, _MM_SHUFFLE(0, 0, 0, 0));
    // [norm(b-a)**3, norm(a-c)**3]
    __m128d rcubed128 = _mm256_castpd256_pd128(rcubed);
    // [norm(a-c)**3, norm(a-c)**3]
    __m128d bcast2 = _mm_permute_pd(rcubed128, _MM_SHUFFLE2(1, 1));
    // [(b0-a0)/norm(b-a)**3, ..., (c1-b1)/norm(c-b)**3]
    __m256d result1 = _mm256_div_pd(diffs, bcast1);
    // [(a0-c0)/norm(a-c)**3, (a1-c1)/norm(a-c)**3]
    __m128d result2 = _mm_div_pd(taildiff, bcast2);
    double* outptr = out.data();
    _mm256_storeu_pd(outptr, result1);
    _mm_store_pd(outptr + 4, result2);
}