使用
使用
我想将以下 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)
我可以像这样用 for 循环来写它C++
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);
}
}
}
}
如何使用其操作编写特征表示法？使用 for 循环不允许特征正确矢量化操作，因为我有复杂的标量类型。
编辑。
正如要求的那样，Jet 是双数的扩展，其中每个元素都是一个数字，后跟该数字的梯度数组以及一些参数。 http://ceres-solver.org/automatic_derivatives.html
天真的修饰可能看起来像
template<typename T, int N>
struct Jet
{
T a;
T v[N];
};
如果 jet 是使用特征运算编写的，则想法是使用表达式模板，特征应直接矢量化所有操作。
在你的情况下，在三维"l"中没有收缩发生。因此，从某种意义上说，L和U是长度为 136 的 2x2 矩阵的数组，您将矩阵U[l]乘以L[l].因此，我认为不可能对 Eigen 进行类似于np.einsum的事情;Eigen::Tensor::contract只支持"真正的"收缩。但是，人们当然可以始终手动在3rd维度上进行循环。但如下所示，这表现非常糟糕。
尽管如此，还是有一些方法可以加快速度并矢量化循环，方法是依靠自动矢量化(对我来说效果不佳)或提供额外的编译器提示(通过 OpenMP SIMD)。
在下文中，我将cDim12=2定义为第一和第二维度的大小，cDim13=136定义为第三维度。 对于计时，所有代码都是使用 gcc 11.2 和 clang 15.0.2 的-O3 -mavx编译的。我使用谷歌基准测试来获取英特尔酷睿 i7-4770K 的时间(是的，对不起，已经有好几年的历史了)。使用了 2023 年 1 月 20 日的特征干线 (08c961e83)。
TL;DR：总结以下结果：
- 使用手动循环(具有更好的迭代顺序)，通过 OpenMP-SIMD 编译指示进行自动矢量化以及使用gcc进行原始访问是最快的("DirectAccessWithOMP")：比使用 AVX 的简单循环快 2.8 倍。我想这接近于"最佳"(参见Godbolt)。
- 我无法让叮当正确地矢量化循环。既然你提到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);
}
}
}
}
循环顺序("优化顺序")
请注意，默认情况下，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);
}
}
}
}
这快了 1.3-1.4 倍。
使用Eigen::Tensor::chip和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);
}
这表现非常糟糕：与"FromOriginalPost"相比，它在gcc上慢18倍，在clang上慢24倍。
使用Eigen::TensorMap和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);
}
这实际上比"EigenChipAndContract"快一些，但仍然很慢。与"FromOriginalPost"相比，gcc 慢 14 倍，clang 慢 19 倍。
Vectorization with OpenMP ("EigenAccessWithOMP")
尽管 gcc 和 clang 都可以进行自动矢量化，但如果没有额外的提示，它们不会产生良好的结果。但是，当使用-fopenmp编译时，两者都支持OpenMP编译指示#pragma omp simd collapse(4)：
#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);
}
}
}
}
编译与-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)];
}
}
}
}
有点令人惊讶的是，与"EigenAccessWithOMP"相比，gcc快1.1倍，clang快1.4倍。 与原始的"FromOriginalPost"相比，gcc快2.8倍，clang慢2.5倍。
在 godbolt 上查看时，gcc 确实会产生一些相当简洁的组装。
蟒
不确定将调用np.einsum的绝对执行时间与C++版本进行比较有多牵强，因为 Python 需要执行额外的字符串解析等。不过，这是代码：
import numpy as np
import timeit
L = np.random.rand(2, 2, 136)
U = np.random.rand(2, 2, 136)
numIterations = 1000000
timing = timeit.timeit(lambda: np.einsum('ijl,jkl->ikl', U, L), number=numIterations)
print(f"np.einsum (per iteration): {timing.real/(numIterations*1e-9)}ns")
对于Python 3.9和numpy-1.24.1，这比"FromOriginalPost"慢6倍，比gcc的"DirectAccessWithOMP"慢16倍。
原始计时
对于 gcc：
---------------------------------------------------------------
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
对于叮当声：
---------------------------------------------------------------
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
蟒：
np.einsum (per iteration): 4873.6035999991145 ns
完整代码
同样在 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();
>编辑喷气式飞机原始帖子被编辑后，使用的算术类型并不是真正的内置类型，而是喷气式飞机。特征可以扩展以支持自定义类型(如此处简要概述)。然而，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};
}
例如(列主)
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);
}
}
}
}
或按主排顺序：
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);
}
}
}
}
GCC 和 Clang 产生相同的性能。它们自动矢量化列主循环，但显然不是行主循环。直接访问基础数据并不能改善情况。此外，在这两种情况下，添加#pragma omp simd collapse(4)都会导致性能变差(clang 还警告说，循环无法矢量化);我想在内部将Jet::v矩阵乘以特征时使用的显式 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++) {
Jet<N> & r = result(i, k, l);
r = r + U(i, j, l) * L(j, k, l);
}
}
}
}
benchmark::DoNotOptimize(result.data());
}
}
BENCHMARK(EigenAccessNoOMP);
static void EigenAccessNoOMPRowMajor(benchmark::State& state)
{
srand(42);
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);
for (auto _ : state) {
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);
}
}
}
}
benchmark::DoNotOptimize(result.data());
}
}
BENCHMARK(EigenAccessNoOMPRowMajor);
static void DirectAccessNoOMP(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);
Jet<N> * pR = result.data();
Jet<N> * pU = U.data();
Jet<N> * 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++) {
Jet<N> & r = pR[i + cDim12*(k + cDim12*l)];
r = r + pU[i + cDim12*(j + cDim12*l)] * pL[j + cDim12*(k + cDim12*l)];
}
}
}
}
benchmark::DoNotOptimize(result.data());
}
}
BENCHMARK(DirectAccessNoOMP);
static void EigenAccessWithOMP(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);
#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++) {
Jet<N> & r = result(i, k, l);
r = r + U(i, j, l) * L(j, k, l);
}
}
}
}
benchmark::DoNotOptimize(result.data());
}
}
BENCHMARK(EigenAccessWithOMP);
static void DirectAccessWithOMP(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);
Jet<N> * pR = result.data();
Jet<N> * pU = U.data();
Jet<N> * pL = L.data();
#pragma omp simd collapse(4) aligned(pR, pU, pL: 32)
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 = pR[i + cDim12*(k + cDim12*l)];
r = r + pU[i + cDim12*(j + cDim12*l)] * pL[j + cDim12*(k + cDim12*l)];
}
}
}
}
benchmark::DoNotOptimize(result.data());
}
}
BENCHMARK(DirectAccessWithOMP);
BENCHMARK_MAIN();
最小工作示例
这是一个工作示例。请参阅 godbolt.org 以运行代码。
#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>
int main() {
// Setup tensors
Eigen::Tensor<double, 3> U(2, 2, 136);
Eigen::Tensor<double, 3> L(2, 2, 136);
// Fill with random vars
U.setRandom();
L.setRandom();
// Create a vector of dimension pairs you want to contract over
// Since j is the second index in the second tensor (U) we specify index 1, since j is
// the first index for the second tensor (L), we specify index 0.
Eigen::array<Eigen::IndexPair<int>, 1> contraction_dims = {Eigen::IndexPair<int>(1,0)};
// Perform contraction and save result
Eigen::Tensor<double, 3> result = U.contract(L, contraction_dims);
}
矢 量化
矢量化是一件棘手的事情。您可能希望使用-O3 -fopt-info-vec-missed编译代码，-fopt-info-vec-missed将打印出有关遗漏的矢量化的非常详细的信息。如果你真的想进一步了解为什么你的编译器没有按照你希望的方式优化事物，请查看像optview2这样的工具，以及Ofek Shilon的CPPCON的精彩演讲。希望这有帮助。