Минимальный рабочий пример

Вот рабочий пример. См. 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 и это отличное выступление от CPPCON Офека Шилона. Надеюсь это поможет.

В вашем случае не происходит сжатия в 3-м измерении «l». Итак, в каком-то смысле L и U — это массивы матриц 2x2 длиной 136, и вы умножаете матрицу U[l] на L[l]. Я думаю, что сделать что-то подобное np.einsum с Eigen невозможно; Eigen::Tensor::contract поддерживает только «настоящие» сокращения. Но, конечно, всегда можно выполнить цикл по 3-му измерению вручную. Но, как показано ниже, это работает очень плохо.

Тем не менее, есть способы ускорить работу и векторизовать циклы, либо полагаясь на автоматическую векторизацию (у меня это не сработало), либо давая дополнительные подсказки компилятору (через OpenMP SIMD).

Далее я определяю cDim12=2 как размер первого и второго измерения, а cDim13=136 как третьего измерения. Что касается таймингов, весь код был скомпилирован с помощью -O3 -mavx с gcc 11.2 и clang 15.0.2. Я использовал бенчмарк Google, чтобы узнать тайминги на Intel Core i7-4770K (да, ему уже несколько лет, извините). Был использован ствол Eigen (08c961e83) от 20 января 2023 года.

TL;DR: подведем итоги ниже:

• Использование ручного цикла (с лучшим порядком итерации), автоматической векторизации с помощью прагмы OpenMP-SIMD, а также прямого доступа с помощью gcc было самым быстрым («DirectAccessWithOMP»): в 2,8 раза быстрее, чем ваш простой цикл с AVX. Я думаю, что это близко к «оптимальному» (ср. Godbolt).
• Я не мог заставить clang правильно векторизовать цикл. Поскольку вы упомянули, что подходит либо gcc, либо clang, у вас, похоже, есть выбор, и я бы остановился на gcc.
• Python кажется на порядок или более медленнее по сравнению с самым быстрым результатом gcc.

Примечание. Измерьте свое приложение в реальном мире, так как там все может вести себя совершенно по-другому!

Код из оригинального поста в качестве основы ("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);
      }
    }
  }
}

Оптимизированный порядок циклов ("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);
      }
    }
  }
}

Это в 1,3-1,4 раза быстрее.

Использование Eigen::Tensor::chip и contract ("EigenChipAndContract")

Максимально используя функции 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++) {
  result.chip(l, 2) = U.chip(l, 2).contract(L.chip(l, 2), productDims);
}

Это работает очень плохо: оно в 18 раз медленнее на gcc и в 24 раза медленнее на clang по сравнению с «FromOriginalPost».

Использование 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, это в 14 раз медленнее для gcc и в 19 раз медленнее для clang.

Векторизация с OpenMP ("EigenAccessWithOMP")

Хотя и gcc , и clang умеют делать автоматическую векторизацию, без дополнительных подсказок они не дают хороших результатов. Однако оба поддерживают прагму OpenMP #pragma omp simd collap(4) при компиляции с -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);
      }
    }
  }
}

Компиляция с -O3 -mavx -fopenmp приводит к

• время выполнения в 2,5 раза быстрее по сравнению с исходным кодом («FromOriginalPost») для gcc,
• но для 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)];
      }
    }
  }
}

Несколько удивительно, что это в 1,1 раза быстрее для gcc и в 1,4 раза быстрее для clang по сравнению с «EigenAccessWithOMP». По сравнению с оригинальным «FromOriginalPost» он в 2,8 раза быстрее для gcc и в 2,5 раза медленнее для clang.

При просмотре на 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 это примерно в 6 раз медленнее по сравнению с «FromOriginalPost» и в 16 раз медленнее по сравнению с «DirectAccessWithOMP» для gcc.

Сырые тайминги

Для 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 можно расширить для поддержки пользовательских типов (как кратко описано здесь). Однако функция 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 также предупреждает, что циклы не могут быть векторизованы); Я предполагаю, что явные SIMD, используемые в матричном умножении Jet::v внутри Eigen, являются причиной.

Еще раз в качестве дополнительного примечания: в документации 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();