c++ 把数和运算写成本征Tensor

ykejflvf  于 2023-01-18  发布在  其他
关注(0)|答案(1)|浏览(127)

bounty将在7天后过期。回答此问题可获得+500声望奖励。Niteya Shah希望引起更多人对此问题的关注:一个工作实现,可以生成由GCC或clang自动矢量化的最佳代码。

我想写下面的numpy einsum作为特征Tensor运算

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循环不允许特征正确地向量化操作,因为我有复杂的标量类型。

goucqfw6

goucqfw61#

最小工作示例

下面是一个工作示例,请访问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会打印出非常详细的信息,说明哪些矢量化被遗漏了。如果你真的想知道编译器为什么没有按照你希望的方式进行优化,可以使用optview2this great talk from CPPCON by Ofek Shilon这样的工具。希望这能有所帮助。

相关问题