将 numpy einsum 运算编写为特征张量

Nit*_*hah 5 c++ matrix-multiplication eigen numpy-einsum tensor

我想将以下 numpy einsum 编写为特征张量操作

import numpy as np

L = np.random.rand(2, 2, 136)
U = np.random.rand(2, 2, 136)

result = np.einsum('ijl,jkl->ikl', U, L)
Run Code Online (Sandbox Code Playgroud)

我可以像 C++ 中那样用 for 循环编写它

  for (int i = 0; i < 2; i++) {
    for (int j = 0; j < 2; j++) {
      for (int k = 0; k < 2; k++) {
        for (int l = 0; l < 136; l++) {
          result(i, k, l) += U(i, j, l) * L(j, k, l);
        }
      }
    }
  }

Run Code Online (Sandbox Code Playgroud)

如何使用其运算写入特征表示法?使用 for 循环不允许 eigen 正确矢量化操作,因为我有复杂的标量类型。

编辑。

正如所要求的,Jet 是对偶数的扩展,其中每个元素都是一个数字,后面跟着该数字与某些参数的梯度数组。 http://ceres-solver.org/automatic_derivatives.html

一个幼稚的实现可能看起来像

template<typename T, int N>
struct Jet
{
    T a;
    T v[N];
};

Run Code Online (Sandbox Code Playgroud)

如果jet是使用eigen ops编写的,那么想法是使用表达式模板,eigen应该直接向量化所有操作。

Gug*_*ugi 2

在你的情况下,第三维“”没有发生收缩l。因此,从某种意义上说,LU是长度为 136 的 2x2 矩阵数组,并且您将该矩阵U[l]与相乘L[l]。因此,我认为做与 Eigen 类似的事情np.einsum是不可能的;Eigen::Tensor::contract仅支持“真实”收缩。但当然,我们总是可以手动在第三维上进行循环。但如下所示,这表现非常糟糕。

尽管如此,还是有一些方法可以加快速度并对循环进行矢量化,方法是依靠自动矢量化(对我来说效果不佳)或提供额外的编译器提示(通过 OpenMP SIMD)。

下面,我定义cDim12=2第一维和第二维的大小,以及cDim13=136第三维的大小。对于时间安排,所有代码均使用-O3 -mavxgcc 11.2 和 clang 15.0.2 进行编译。我使用谷歌基准测试来获取 Intel Core i7-4770K 的计时(是的,已经有好几年了,抱歉)。使用了 2023 年 1 月 20 日的 Eigen trunk (08c961e83)。

TL;DR:总结如下结果:

  • 使用手动循环(具有更好的迭代顺序)、通过 OpenMP-SIMD 编译指示进行自动矢量化以及使用gcc进行原始访问是最快的(“DirectAccessWithOMP”):比使用 AVX 的直接循环快 2.8 倍。我想这接近“最佳”(参见godbolt)。
  • 我无法让 clang 正确地矢量化循环。既然你提到 gcc 或 clang 都可以,你似乎可以选择,我会坚持使用 gcc。
  • 与最快的 gcc 结果相比,Python 似乎慢了一个数量级或更多。

注意:测量您的现实世界应用程序,因为那里的情况可能完全不同!

来自原始帖子的代码作为基线(“FromOriginalPost”)

原始帖子中的简单代码如下所示,并用作基线。

Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
result.setZero();
for (int i = 0; i < cDim12; i++) {
  for (int j = 0; j < cDim12; j++) {
    for (int k = 0; k < cDim12; k++) {
      for (int l = 0; l < cDim3; l++) {
        result(i, k, l) += U(i, j, l) * L(j, k, l);
      }
    }
  }
}
Run Code Online (Sandbox Code Playgroud)

优化循环顺序(“OptimizedOrder”)

请注意,默认情况下Eigen::Tensor使用列主要顺序(不建议使用行主要顺序)。因此,在诸如 之类的表达式中U(i, j, l)i应该是最快(最内层)循环和l最慢(最外层)循环。尽我所能重新排序:

for (int l = 0; l < cDim3; l++) {
  for (int j = 0; j < cDim12; j++) {
    for (int k = 0; k < cDim12; k++) {
      for (int i = 0; i < cDim12; i++) {
        result(i, k, l) += U(i, j, l) * L(j, k, l);
      }
    }
  }
}
Run Code Online (Sandbox Code Playgroud)

速度提高了 1.3 倍至 1.4 倍。

使用Eigen::Tensor::chipand contract(“EigenChipAndContract”)

尽可能多地使用特征特征,我想出了以下结论:

Eigen::array<Eigen::IndexPair<int>, 1> productDims = {Eigen::IndexPair<int>(1, 0)};
Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
for (int l = 0; l < cDim3; l++) {
  result.chip(l, 2) = U.chip(l, 2).contract(L.chip(l, 2), productDims);
}
Run Code Online (Sandbox Code Playgroud)

其性能非常糟糕:与“FromOriginalPost”相比,它在 gcc 上慢 18 倍,在 clang 上慢 24 倍。

使用Eigen::TensorMapand contract(“EigenMapAndContract”)

“EigenChipAndContract”可能会进行大量复制,因此另一个想法是用于Eigen::TensorMap获取每个必要的数据“切片”的“引用”。对于原始数组访问,请再次注意 Eigen 使用列优先顺序。

Eigen::array<Eigen::IndexPair<int>, 1> productDims = {Eigen::IndexPair<int>(1, 0)};
Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
for (int l = 0; l < cDim3; l++) {
  Eigen::TensorMap<Eigen::Tensor<double, 2>> U_chip(U.data() + l * cDim12 * cDim12, cDim12, cDim12);
  Eigen::TensorMap<Eigen::Tensor<double, 2>> L_chip(L.data() + l * cDim12 * cDim12, cDim12, cDim12);
  Eigen::TensorMap<Eigen::Tensor<double, 2>> result_chip(result.data() + l * cDim12 * cDim12, cDim12, cDim12);
  result_chip = U_chip.contract(L_chip, productDims);
}
Run Code Online (Sandbox Code Playgroud)

这实际上比“EigenChipAndContract”要快一些,但仍然很慢。与“FromOriginalPost”相比,gcc 慢 14 倍,clang 慢 19 倍。

使用 OpenMP 进行矢量化(“EigenAccessWithOMP”)

尽管gccclang都可以进行自动矢量化,但如果没有额外的提示,它们不会产生良好的结果。#pragma omp simd collapse(4)但是,当使用以下命令编译时,两者都支持 OpenMP pragma -fopenmp

#pragma omp simd collapse(4)
for (int l = 0; l < cDim3; l++) {
  for (int j = 0; j < cDim12; j++) {
    for (int k = 0; k < cDim12; k++) {
      for (int i = 0; i < cDim12; i++) {
        result(i, k, l) += U(i, j, l) * L(j, k, l);
      }
    }
  }
}
Run Code Online (Sandbox Code Playgroud)

编译-O3 -mavx -fopenmp结果为

  • 与 gcc 的原始代码(“FromOriginalPost”)相比,运行时间2.5 倍,
  • 但是对于 clang,代码要2.5 倍。搜索 clang + OpenMP-SIMD 问题,显然 clang 有时确实会遇到麻烦(例如在这篇文章中)。在godbolt上检查结果, clang 确实产生了相当长的结果。

使用 OpenMP + 直接原始访问进行矢量化(“DirectAccessWithOMP”)

前面的代码使用了Eigen::Tensor::operator(),它应该内联到原始数组访问。然而,记住列主布局,我们还可以直接访问底层数组并检查这是否有任何改进。它还允许再次向编译器提示数据已正确对齐(尽管 Eigen 已经这样定义了它们)。

double * pR = result.data();
double * pU = U.data();
double * pL = L.data();

#pragma omp simd collapse(4) aligned(pR, pU, pL: 32) // 32: For AVX
for (int l = 0; l < cDim3; l++) {
  for (int j = 0; j < cDim12; j++) {
    for (int k = 0; k < cDim12; k++) {
      for (int i = 0; i < cDim12; i++) {
        pR[i + cDim12*(k + cDim12*l)] += pU[i + cDim12*(j + cDim12*l)] * pL[j + cDim12*(k + cDim12*l)];
      }
    }
  }
}
Run Code Online (Sandbox Code Playgroud)

有点令人惊讶的是,与“EigenAccessWithOMP”相比,gcc 的速度快了 1.1 倍,clang 的速度快了 1.4 倍。与原始的“FromOriginalPost”相比,gcc 快 2.8 倍,clang 慢 2.5 倍。

当在godbolt上查看时,gcc 确实生成了一些非常简洁的汇编。

Python

不确定将调用的绝对执行时间与 C++ 版本进行比较有多远np.einsum,因为 Python 需要进行额外的字符串解析等。不过,这里是代码:

Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
result.setZero();
for (int i = 0; i < cDim12; i++) {
  for (int j = 0; j < cDim12; j++) {
    for (int k = 0; k < cDim12; k++) {
      for (int l = 0; l < cDim3; l++) {
        result(i, k, l) += U(i, j, l) * L(j, k, l);
      }
    }
  }
}
Run Code Online (Sandbox Code Playgroud)

对于 Python 3.9 和 numpy-1.24.1,与 gcc 的“FromOriginalPost”相比,这大约慢 6 倍,与“DirectAccessWithOMP”相比,慢 16 倍。

原始计时

对于海湾合作委员会:

---------------------------------------------------------------
Benchmark                     Time             CPU   Iterations
---------------------------------------------------------------
FromOriginalPost            823 ns          823 ns      3397793
OptimizedOrder              573 ns          573 ns      4895246
DirectAccess               1306 ns         1306 ns      2142826
EigenAccessWithOMP          324 ns          324 ns      8655549
DirectAccessWithOMP         296 ns          296 ns      9418635
EigenChipAndContract      14405 ns        14405 ns       193548
EigenMapAndContract       11390 ns        11390 ns       243122
Run Code Online (Sandbox Code Playgroud)

对于叮当声:

---------------------------------------------------------------
Benchmark                     Time             CPU   Iterations
---------------------------------------------------------------
FromOriginalPost            753 ns          753 ns      3714543
OptimizedOrder              570 ns          570 ns      4921914
DirectAccess                569 ns          569 ns      4929755
EigenAccessWithOMP         2704 ns         2704 ns      1037819
DirectAccessWithOMP        1908 ns         1908 ns      1466390
EigenChipAndContract      17713 ns        17713 ns       157427
EigenMapAndContract       14064 ns        14064 ns       198875
Run Code Online (Sandbox Code Playgroud)

Python:

np.einsum (per iteration): 4873.6035999991145 ns
Run Code Online (Sandbox Code Playgroud)

完整代码

也在godbolt上,但是并不是很有用,因为编译器经常超时。我在本地编译了-O3 -DNDEBUG -std=c++17 -mavx -fopenmp -Wall -Wextra.

#include <iostream>
#include <iomanip>
#include <cmath>

#include <unsupported/Eigen/CXX11/Tensor>
#include <benchmark/benchmark.h>

//====================================================
// Globals
//====================================================

static constexpr int cDim12 = 2;
static constexpr int cDim3 = 136;


Eigen::Tensor<double, 3> CreateRandomTensor()
{
  Eigen::Tensor<double, 3> m(cDim12, cDim12, cDim3);
  m.setRandom();
  return m;
}


Eigen::Tensor<double, 3> const L = CreateRandomTensor();
Eigen::Tensor<double, 3> const U = CreateRandomTensor();


//====================================================
// Helpers
//====================================================

Eigen::Tensor<double, 3> ReferenceResult() 
{
  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  result.setZero();
  for (int i = 0; i < cDim12; i++) {
    for (int j = 0; j < cDim12; j++) {
      for (int k = 0; k < cDim12; k++) {
        for (int l = 0; l < cDim3; l++) {
          result(i, k, l) += U(i, j, l) * L(j, k, l);
        }
      }
    }
  }

  return result;
}



void CheckResult(Eigen::Tensor<double, 3> const & result) 
{
  Eigen::Tensor<double, 3> const ref = ReferenceResult();
  Eigen::Tensor<double, 3> const diff = ref - result;
  Eigen::Tensor<double, 0> const max = diff.maximum();
  Eigen::Tensor<double, 0> const min = diff.minimum();
  double const maxDiff = std::max(std::abs(max(0)), std::abs(min(0)));
  if (maxDiff > 1e-14) {
    std::cerr << "ERROR! Max Diff = " << std::setprecision(17) << maxDiff << std::endl;
  }
}


//====================================================
// Benchmarks
//====================================================


static void FromOriginalPost(benchmark::State& state) 
{
  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    result.setZero();
    for (int i = 0; i < cDim12; i++) {
      for (int j = 0; j < cDim12; j++) {
        for (int k = 0; k < cDim12; k++) {
          for (int l = 0; l < cDim3; l++) {
            result(i, k, l) += U(i, j, l) * L(j, k, l);
          }
        }
      }
    }

    benchmark::DoNotOptimize(result.data());
  }

  CheckResult(result);
}
BENCHMARK(FromOriginalPost);


static void OptimizedOrder(benchmark::State& state) 
{
  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    result.setZero();
    for (int l = 0; l < cDim3; l++) {
      for (int j = 0; j < cDim12; j++) {
        for (int k = 0; k < cDim12; k++) {
          for (int i = 0; i < cDim12; i++) {
            result(i, k, l) += U(i, j, l) * L(j, k, l);
          }
        }
      }
    }

    benchmark::DoNotOptimize(result.data());
  }

  CheckResult(result);
}
BENCHMARK(OptimizedOrder);


static void DirectAccess(benchmark::State& state) 
{
  Eigen::Tensor<double, 3> U = ::U;
  Eigen::Tensor<double, 3> L = ::L;

  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    result.setZero();
    double * pR = result.data();
    double * pU = U.data();
    double * pL = L.data();

    for (int l = 0; l < cDim3; l++) {
      for (int j = 0; j < cDim12; j++) {
        for (int k = 0; k < cDim12; k++) {
          for (int i = 0; i < cDim12; i++) {
            pR[i + cDim12*(k + cDim12*l)] += pU[i + cDim12*(j + cDim12*l)] * pL[j + cDim12*(k + cDim12*l)];
          }
        }
      }
    }

    benchmark::DoNotOptimize(result.data());
  }

  CheckResult(result);
}
BENCHMARK(DirectAccess);


static void EigenAccessWithOMP(benchmark::State& state) 
{
  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    result.setZero();
    #pragma omp simd collapse(4)
    for (int l = 0; l < cDim3; l++) {
      for (int j = 0; j < cDim12; j++) {
        for (int k = 0; k < cDim12; k++) {
          for (int i = 0; i < cDim12; i++) {
            result(i, k, l) += U(i, j, l) * L(j, k, l);
          }
        }
      }
    }

    benchmark::DoNotOptimize(result.data());
  }

  CheckResult(result);
}
BENCHMARK(EigenAccessWithOMP);


static void DirectAccessWithOMP(benchmark::State& state) 
{
  Eigen::Tensor<double, 3> U = ::U;
  Eigen::Tensor<double, 3> L = ::L;

  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    result.setZero();
    double * pR = result.data();
    double * pU = U.data();
    double * pL = L.data();

    #pragma omp simd collapse(4) aligned(pR, pU, pL: 32) // 32: For AVX
    for (int l = 0; l < cDim3; l++) {
      for (int j = 0; j < cDim12; j++) {
        for (int k = 0; k < cDim12; k++) {
          for (int i = 0; i < cDim12; i++) {
            pR[i + cDim12*(k + cDim12*l)] += pU[i + cDim12*(j + cDim12*l)] * pL[j + cDim12*(k + cDim12*l)];
          }
        }
      }
    }

    benchmark::DoNotOptimize(result.data());
  }

  CheckResult(result);
}
BENCHMARK(DirectAccessWithOMP);


static void EigenChipAndContract(benchmark::State& state) 
{
  Eigen::array<Eigen::IndexPair<int>, 1> productDims = {Eigen::IndexPair<int>(1, 0)};
  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    result.setZero();
    for (int l = 0; l < cDim3; l++) {
      result.chip(l, 2) = U.chip(l, 2).contract(L.chip(l, 2), productDims);
    }
    benchmark::DoNotOptimize(result.data());
  }

  CheckResult(result);
}
BENCHMARK(EigenChipAndContract);


static void EigenMapAndContract(benchmark::State& state) 
{
  Eigen::Tensor<double, 3> U = ::U;
  Eigen::Tensor<double, 3> L = ::L;
  Eigen::array<Eigen::IndexPair<int>, 1> productDims = {Eigen::IndexPair<int>(1, 0)};

  Eigen::Tensor<double, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    result.setZero();
    for (int l = 0; l < cDim3; l++) {
      Eigen::TensorMap<Eigen::Tensor<double, 2>> U_chip(U.data() + l * cDim12 * cDim12, cDim12, cDim12);
      Eigen::TensorMap<Eigen::Tensor<double, 2>> L_chip(L.data() + l * cDim12 * cDim12, cDim12, cDim12);
      Eigen::TensorMap<Eigen::Tensor<double, 2>> result_chip(result.data() + l * cDim12 * cDim12, cDim12, cDim12);
      result_chip = U_chip.contract(L_chip, productDims);
    }
    benchmark::DoNotOptimize(result.data());
  }

  CheckResult(result);
}
BENCHMARK(EigenMapAndContract);


BENCHMARK_MAIN();
Run Code Online (Sandbox Code Playgroud)

编辑喷气机

在编辑原始帖子后,使用的算术类型并不是真正的内置类型,而是jets。Eigen 可以扩展以支持自定义类型(如此处简要概述。然而,该Eigen::Tensor::contract()函数仍然不会“神奇地”支持等效的,np.einsum('ijl,jkl->ikl', U, L)因为最后一个维度l并没有真正执行收缩。当然,人们可以写一个,但这似乎远非微不足道。

如果唯一需要的类似收缩的操作是原始帖子中的操作,并且张量没有进一步相乘/相加/等,那么最简单的事情就是手动实现单个循环并使用编译器、编译器设置,编译指示等来找出最佳性能。

喷射类型(改编自此处):

template<int N> struct Jet {
  double a = 0.0;
  Eigen::Matrix<double, 1, N> v = Eigen::Matrix<double, 1, N>::Zero();
};

template<int N> 
EIGEN_STRONG_INLINE Jet<N> operator+(const Jet<N>& f, const Jet<N>& g) {
  return Jet<N>{f.a + g.a, f.v + g.v};
}

template<int N> 
EIGEN_STRONG_INLINE Jet<N> operator*(const Jet<N>& f, const Jet<N>& g) {
  return Jet<N>{f.a * g.a, f.a * g.v + f.v * g.a};
}
Run Code Online (Sandbox Code Playgroud)

例如(列主)

Eigen::Tensor<Jet<N>, 3> L = CreateRandomTensor<Eigen::ColMajor>();
Eigen::Tensor<Jet<N>, 3> U = CreateRandomTensor<Eigen::ColMajor>();

Eigen::Tensor<Jet<N>, 3> result(cDim12, cDim12, cDim3);
SetToZero(result);

for (int l = 0; l < cDim3; l++) {
  for (int j = 0; j < cDim12; j++) {
    for (int k = 0; k < cDim12; k++) {
      for (int i = 0; i < cDim12; i++) {
        Jet<N> & r = result(i, k, l);
        r = r + U(i, j, l) * L(j, k, l);
      }
    }
  }
}
Run Code Online (Sandbox Code Playgroud)

或按行主序:

Eigen::Tensor<Jet<N>, 3, Eigen::RowMajor> L = CreateRandomTensor<Eigen::RowMajor>();
Eigen::Tensor<Jet<N>, 3, Eigen::RowMajor> U = CreateRandomTensor<Eigen::RowMajor>();

Eigen::Tensor<Jet<N>, 3, Eigen::RowMajor> result(cDim12, cDim12, cDim3);
SetToZero(result);

for (int i = 0; i < cDim12; i++) {
  for (int k = 0; k < cDim12; k++) {
    for (int j = 0; j < cDim12; j++) {
      for (int l = 0; l < cDim3; l++) {
        Jet<N> & r = result(i, k, l);
        r = r + U(i, j, l) * L(j, k, l);
      }
    }
  }
}
Run Code Online (Sandbox Code Playgroud)

gcc 和 clang 产生相同的性能。它们自动向量化列主循环,但显然不是行主循环。直接访问底层数据并不能改善情况。此外,在这两种情况下,添加#pragma omp simd collapse(4)都会导致性能更差(clang 还警告循环无法矢量化);Jet::v我猜想原因是 Eigen 内部矩阵乘法中使用的显式 SIMD 。

再次作为附加说明:Eigen 文档说您不应该真正将行主序与Eigen::Tensor

张量库支持 2 种布局:ColMajor(默认)和 RowMajor。目前仅完全支持默认的列主要布局,因此目前不建议尝试使用行主要布局。

完整代码:

#include <iostream>
#include <iomanip>
#include <cmath>

#include <unsupported/Eigen/CXX11/Tensor>
#include <benchmark/benchmark.h>


static constexpr int cDim12 = 2;
static constexpr int cDim3 = 136;


template<int N> struct Jet {
  double a = 0.0;
  Eigen::Matrix<double, 1, N> v = Eigen::Matrix<double, 1, N>::Zero();
};

template<int N> 
EIGEN_STRONG_INLINE Jet<N> operator+(const Jet<N>& f, const Jet<N>& g) {
  return Jet<N>{f.a + g.a, f.v + g.v};
}

template<int N> 
EIGEN_STRONG_INLINE Jet<N> operator-(const Jet<N>& f, const Jet<N>& g) {
  return Jet<N>{f.a - g.a, f.v - g.v};
}

template<int N> 
EIGEN_STRONG_INLINE Jet<N> operator*(const Jet<N>& f, const Jet<N>& g) {
  return Jet<N>{f.a * g.a, f.a * g.v + f.v * g.a};
}

template<int N> 
EIGEN_STRONG_INLINE Jet<N> operator/(const Jet<N>& f, const Jet<N>& g) {
  return Jet<N>{f.a / g.a, f.v / g.a - f.a * g.v / (g.a * g.a)};
}



static constexpr int N = 8;


template <Eigen::StorageOptions storage>
auto CreateRandomTensor()
{
  Eigen::Tensor<Jet<N>, 3, storage> result(cDim12, cDim12, cDim3);
  for (int l = 0; l < cDim3; l++) {
    for (int k = 0; k < cDim12; k++) {
      for (int i = 0; i < cDim12; i++) {
        Jet<N> jet;
        jet.a = (double)rand() / RAND_MAX;
        jet.v.setRandom();
        result(i, k, l) = jet;
      }
    }
  }
  return result;
}


template <class T>
void SetToZero(T & result)
{
  for (int l = 0; l < cDim3; l++) {
    for (int k = 0; k < cDim12; k++) {
      for (int i = 0; i < cDim12; i++) {
        result(i, k, l) = Jet<N>{};
      }
    }
  }
}


static void EigenAccessNoOMP(benchmark::State& state) 
{
  srand(42);
  Eigen::Tensor<Jet<N>, 3> L = CreateRandomTensor<Eigen::ColMajor>();
  Eigen::Tensor<Jet<N>, 3> U = CreateRandomTensor<Eigen::ColMajor>();

  Eigen::Tensor<Jet<N>, 3> result(cDim12, cDim12, cDim3);
  for (auto _ : state) {
    SetToZero(result);

    for (int l = 0; l < cDim3; l++) {
      for (int j = 0; j < cDim12; j++) {
        for (int k = 0; k < cDim12; k++) {
          for (int i = 0; i < cDim12; i