与 Python 相比，优化 C++ 中的张量收缩

Question

我在 C++ 和 Python 中有 3 个向量（Python 中的 numpy 数组），我希望进行以下张量收缩：

import numpy as np
import time

N_t, N = 400, 400
a = np.random.rand(N_t, 2, 2, N)
b = np.random.rand(2, 2, N)
c = np.random.rand(N_t, 2, 2, N)

start_time = time.time()
d = np.einsum('iacm, cdm, jbdm -> ijabm', a, b, c)
print(time.time() - start_time)

为了简单起见，随机生成 3 个数组。 Python 大约需要 2 秒。

现在，在 C++ 中，无需任何优化，我就可以编写（借助 ChatGPT 以避免敲出一些常见功能的劳动）

#include <iostream>
#include <vector>
#include <chrono>
#include <random>

using namespace std;

// Function to generate a random 4D vector with given shape
std::vector<std::vector<std::vector<std::vector<double> > > > generateRandom4DVector(int dim1, int dim2, int dim3, int dim4) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<double> dis(-1.0, 1.0); // Random numbers between -1 and 1
    
    std::vector<std::vector<std::vector<std::vector<double> > > > vec(dim1,
                                                                   std::vector<std::vector<std::vector<double> > >(dim2,
                                                                                                                 std::vector<std::vector<double> >(dim3,std::vector<double>(dim4))));
    
    // Populate the vector with random values
    for (int i = 0; i < dim1; ++i) {
        for (int j = 0; j < dim2; ++j) {
            for (int k = 0; k < dim3; ++k) {
                for (int l = 0; l < dim4; ++l) {
                    vec[i][j][k][l] = dis(gen); // Generate random number and assign to vector element
                }
            }
        }
    }
    
    return vec;
}

int main() {
    
    int dim1 = 400, dim2 = 2, dim3 = 2, dim4 = 400;
    
    std::vector<std::vector<std::vector<std::vector<double> > > > x = generateRandom4DVector(dim1, dim2, dim3, dim4);
    std::vector<std::vector<std::vector<std::vector<double> > > > y = generateRandom4DVector(dim1, dim2, dim3, dim4);
    std::vector<std::vector<std::vector<std::vector<double> > > > z = generateRandom4DVector(dim1, dim2, dim3, dim4);
    std::vector<std::vector<std::vector<std::vector<std::vector<int> > > > > w(
            dim1, std::vector<std::vector<std::vector<std::vector<int> > > >(
                    dim1, std::vector<std::vector<std::vector<int> > >(
                            2, std::vector<std::vector<int> >(
                                    2, std::vector<int>(dim4)
                            )
                    )
            ));
    
    auto start = std::chrono::high_resolution_clock::now();
    
    for (int i = 0; i < dim1; i++) {
        for (int j = 0; j < dim1; j++) {
            for (int a = 0; a < 2; a++) {
                for (int b = 0; b < 2; b++) {
                    for (int m = 0; m < dim4; m++) {
                        for (int c = 0; c < 2; c++) {
                            for (int d = 0; d < 2; d++) {
                                w[i][j][a][b][m] += x[i][a][c][m] * y[0][c][d][m] * z[j][b][d][m];
                            }
                        }
                    }
                }
            }
        }
    }
    
    
    // Stop measuring time
    auto end = std::chrono::high_resolution_clock::now();
    
    // Calculate duration
    std::chrono::duration<double> duration = end - start;
    
    // Output duration in seconds
    std::cout << "Elapsed time: " << duration.count() << " seconds" << std::endl;
    
    
    return 0;
}

大约需要16秒，非常慢。

一个非常幼稚的解决方案是通过将非常大的总和放入函数中来多重处理这些 for 循环。并且利用

np.einsum

中没有更多空间进行多处理的优势，并且纯 Python 的 for 循环非常慢。如果我有很多 CPU，我总是可以利用这一点让我的 C++ 比 Python 更快，因为在纯 C++ 中，for 循环要快得多。我正在尝试寻找更聪明的例程来解决这个问题。如何明智地使用 C++ 库来加速此类类型的求和？

Answer 1

一样简单

numpy

但是，我看到你会用它来做数学。使用矩阵库，如 Eigen、Armadillo 等。我认为 Boost Ublas 也支持它。他们将使用表达式模板来大量优化子表达式，或识别特殊情况算法的情况，这将为您带来 NumPy 性能。