对比Python，优化C++中张量收缩计算的方法探究

问题背景 #

需要实现张量收缩计算 iacm, cdm, jbdm -> ijabm，Python中用NumPy einsum耗时约2秒，但未优化的C嵌套vector实现耗时约16秒，希望找到更高效的C实现方案。

Python参考实现 #

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)

未优化的C++实现（耗时约16秒） #

#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;
}

优化方案 #

1. 内存布局优化：使用连续内存替代嵌套vector #

嵌套std::vector的内存分散，缓存命中率极低。改用连续一维数组模拟多维结构，手动计算索引访问元素，大幅提升缓存利用率。

#include <vector>
#include <random>

// 4D数组索引计算：i * dim2*dim3*dim4 + a * dim3*dim4 + c * dim4 + m
inline int idx4d(int i, int a, int c, int m, int dim2, int dim3, int dim4) {
    return i * dim2*dim3*dim4 + a * dim3*dim4 + c * dim4 + m;
}

// 3D数组索引计算：c * dim2*dim4 + d * dim4 + m
inline int idx3d(int c, int d, int m, int dim2, int dim4) {
    return c * dim2*dim4 + d * dim4 + m;
}

// 5D结果数组索引计算：i * dim1*2*2*dim4 + j * 2*2*dim4 + a * 2*dim4 + b * dim4 + m
inline int idx5d(int i, int j, int a, int b, int m, int dim1, int dim4) {
    return i * dim1*2*2*dim4 + j * 2*2*dim4 + a * 2*dim4 + b * dim4 + m;
}

std::vector<double> generateRandom4D(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);
    
    std::vector<double> vec(dim1 * dim2 * dim3 * dim4);
    for (size_t k = 0; k < vec.size(); ++k) {
        vec[k] = dis(gen);
    }
    return vec;
}

std::vector<double> generateRandom3D(int dim1, int dim2, int dim3) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<double> dis(-1.0, 1.0);
    
    std::vector<double> vec(dim1 * dim2 * dim3);
    for (size_t k = 0; k < vec.size(); ++k) {
        vec[k] = dis(gen);
    }
    return vec;
}

2. 循环重排与循环展开 #

原循环顺序的内存访问模式不连续，小维度（c、d均为2）的循环开销占比高。通过重排循环顺序提升缓存局部性，同时手动展开小维度循环消除循环控制开销。

int main() {
    const int N_t = 400, N = 400;
    const int dim2 = 2, dim3 = 2;

    std::vector<double> a = generateRandom4D(N_t, dim2, dim3, N);
    std::vector<double> b = generateRandom3D(dim3, dim2, N); // 对应原b的(2,2,N)
    std::vector<double> c = generateRandom4D(N_t, dim2, dim3, N);
    std::vector<double> w(N_t * N_t * dim2 * dim2 * N, 0.0); // 用double替代int，避免类型溢出

    auto start = std::chrono::high_resolution_clock::now();

    // 重排循环：将m放到外层，让同一m的所有计算集中，提升缓存命中率
    for (int m = 0; m < N; ++m) {
        // 手动展开c、d的循环（因为维度只有2）
        for (int i = 0; i < N_t; ++i) {
            for (int j = 0; j < N_t; ++j) {
                // 计算所有a,b组合
                for (int ai = 0; ai < 2; ++ai) {
                    for (int bi = 0; bi < 2; ++bi) {
                        double sum = 0.0;
                        // 展开c=0,1; d=0,1
                        sum += a[idx4d(i, ai, 0, m, 2, 2, N)] * b[idx3d(0, 0, m, 2, N)] * c[idx4d(j, bi, 0, m, 2, 2, N)];
                        sum += a[idx4d(i, ai, 0, m, 2, 2, N)] * b[idx3d(0, 1, m, 2, N)] * c[idx4d(j, bi, 1, m, 2, 2, N)];
                        sum += a[idx4d(i, ai, 1, m, 2, 2, N)] * b[idx3d(1, 0, m, 2, N)] * c[idx4d(j, bi, 0, m, 2, 2, N)];
                        sum += a[idx4d(i, ai, 1, m, 2, 2, N)] * b[idx3d(1, 1, m, 2, N)] * c[idx4d(j, bi, 1, m, 2, 2, N)];
                        w[idx5d(i, j, ai, bi, m, N_t, N)] = sum;
                    }
                }
            }
        }
    }

    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration = end - start;
    std::cout << "Elapsed time: " << duration.count() << " seconds" << std::endl;

    return 0;
}

3. 多线程并行：基于OpenMP的粗粒度并行 #

i和j的计算相互独立，可将外层的i-j循环并行化，利用多CPU核心提升效率。编译时需添加OpenMP选项（如GCC的-fopenmp）。

// 在循环前添加并行指令
#pragma omp parallel for collapse(2)
for (int i = 0; i < N_t; ++i) {
    for (int j = 0; j < N_t; ++j) {
        for (int m = 0; m < N; ++m) {
            for (int ai = 0; ai < 2; ++ai) {
                for (int bi = 0; bi < 2; ++bi) {
                    double sum = 0.0;
                    sum += a[idx4d(i, ai, 0, m, 2, 2, N)] * b[idx3d(0, 0, m, 2, N)] * c[idx4d(j, bi, 0, m, 2, 2, N)];
                    sum += a[idx4d(i, ai, 0, m, 2, 2, N)] * b[idx3d(0, 1, m, 2, N)] * c[idx4d(j, bi, 1, m, 2, 2, N)];
                    sum += a[idx4d(i, ai, 1, m, 2, 2, N)] * b[idx3d(1, 0, m, 2, N)] * c[idx4d(j, bi, 0, m, 2, 2, N)];
                    sum += a[idx4d(i, ai, 1, m, 2, 2, N)] * b[idx3d(1, 1, m, 2, N)] * c[idx4d(j, bi, 1, m, 2, 2, N)];
                    w[idx5d(i, j, ai, bi, m, N_t, N)] = sum;
                }
            }
        }
    }
}

4. 向量化优化：利用SIMD指令 #

现代CPU支持SIMD指令，可一次计算多个浮点数。编译时添加-O3 -mavx2（针对支持AVX2的CPU），让编译器自动生成向量化代码；也可使用Intel Intrinsics手动实现SIMD计算，进一步挖掘硬件性能。

5. 使用专业数值计算库 #

借助Eigen、Blaze或Intel MKL等高度优化的数值库，将张量收缩转化为矩阵乘法组合，利用库的优化实现高效计算。

以Eigen为例：

#include <Eigen/Dense>
#include <vector>
#include <chrono>
#include <random>

int main() {
    const int N_t = 400, N = 400;
    using Mat2d = Eigen::Matrix<double, 2, 2>;

    // 存储每个m对应的矩阵
    std::vector<std::vector<Mat2d>> a(N_t, std::vector<Mat2d>(N));
    std::vector<Mat2d> b(N);
    std::vector<std::vector<Mat2d>> c(N_t, std::vector<Mat2d>(N));
    std::vector<std::vector<std::vector<std::vector<double>>>> w(N_t, std::vector<std::vector<std::vector<double>>>(N_t, std::vector<std::vector<double>>(2, std::vector<double>(2, 0.0))));

    // 初始化数据
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<double> dis(-1.0, 1.0);
    for (int i = 0; i < N_t; ++i) {
        for (int m = 0; m < N; ++m) {
            a[i][m] = Mat2d::Random();
            c[i][m] = Mat2d::Random();
        }
    }
    for (int m = 0; m < N; ++m) {
        b[m] = Mat2d::Random();
    }

    auto start = std::chrono::high_resolution_clock::now();

    #pragma omp parallel for collapse(2)
    for (int i = 0; i < N_t; ++i) {
        for (int j = 0; j < N_t; ++j) {
            Eigen::Matrix<double, 2, 2> sum_mat = Eigen::Matrix<double, 2, 2>::Zero();
            for (int m = 0; m < N; ++m) {
                // 计算：a[i][m] * b[m] * c[j][m].transpose()，然后累加
                sum_mat += a[i][m] * b[m] * c[j][m].transpose();
            }
            // 将结果赋值到w
            for (int ai = 0; ai < 2; ++ai) {
                for (int bi = 0; bi < 2; ++bi) {
                    w[i][j][ai][bi] = sum_mat(ai, bi);
                }
            }
        }
    }

    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration = end - start;
    std::cout << "Elapsed time: " << duration.count() << " seconds" << std::endl;

    return 0;
}

总结 #

通过连续内存布局、循环优化、多线程并行、SIMD向量化或专业数值库的组合优化，C++实现的性能可显著超越NumPy einsum。其中，使用Eigen等库的方案开发效率最高，同时能获得接近硬件极限的性能。

内容的提问来源于stack exchange，提问作者John