基于Thrust的GPU复数计算优化：寻求高效迁移方案

刚接触Thrust，正尝试将以下CPU代码迁移至GPU，但过程存在困难。

原始CPU代码 #

#include <iostream>
#include <complex>
#include <random>
#include <omp.h>
#include <chrono>

using Complex = std::complex<double>;


void fill_random_complex(Complex *array, int n) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<double> dis(0.0, 1.0);

    for (int i = 0; i < n; ++i) {
        double real_part = dis(gen);
        double imag_part = dis(gen);
        array[i] = Complex(real_part, imag_part);
    }
}

void process_loop(
    Complex *ag,
    Complex *bg,
    double *c_distr,
    int *nbnd_loc,
    int *ngk,
    int l1_e, int l2_s, int l2_e,
    int kpt_pool_nloc, int gstart,
    int nimage, int idx, int pert_nglob,
    int ag_dim1, int ag_dim2, int ag_dim3,
    int bg_dim1, int bg_dim2, int bg_dim3,
    int c_distr_dim1) 
{
    auto start_time = std::chrono::high_resolution_clock::now();

    for (int il1 = 0; il1 < l1_e; ++il1) {
        for (int il2 = l2_s - 1; il2 < l2_e; ++il2) {
            int ig1 = nimage * il1 + idx;
            double reduce = 0.0;

            for (int iks = 0; iks < kpt_pool_nloc; ++iks) {
                int nbndval = nbnd_loc[iks];
                int npw = ngk[iks];

                for (int lbnd = 0; lbnd < nbndval; ++lbnd) {
                    for (int il3 = 0; il3 < npw; ++il3) {
                        int ag_index = il3 + nbndval * (lbnd + ag_dim3 * (iks + ag_dim4 * il1));
                        int bg_index = il3 + nbndval * (lbnd + bg_dim3 * (iks + bg_dim4 * il2));

                        reduce += 2.0 * std::real(ag[ag_index]) * std::real(bg[bg_index])
                                + 2.0 * std::imag(ag[ag_index]) * std::imag(bg[bg_index]);
                    }
                }

                if (gstart == 2) {
                    for (int lbnd = 0; lbnd < nbndval; ++lbnd) {
                        int ag_index = 0 + nbndval * (lbnd + ag_dim3 * (iks + ag_dim4 * il1));
                        int bg_index = 0 + nbndval * (lbnd + bg_dim3 * (iks + bg_dim4 * il2));

                        reduce -= std::real(ag[ag_index]) * std::real(bg[bg_index]);
                    }
                }
            }

            int c_distr_index = ig1 * c_distr_dim1 + il2;
            c_distr[c_distr_index] = reduce;
        }
    }

    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed_time = end_time - start_time;
    std::cout << "Time for process_loop: " << elapsed_time.count() << " seconds" << std::endl;
}

int main() {
    const int nimg = 1;
    const int idx = 0;
    const int kpt_pool_nloc = 1;
    const int l1_e = 200;
    const int l2_s = 1;
    const int l2_e = 100;
    const int pert_nglob = 200;
    const int gstart = 2;

    const int ag_dim1 = 16984;
    const int ag_dim2 = 128;
    const int ag_dim3 = kpt_pool_nloc;
    const int ag_dim4 = l1_e;

    const int bg_dim1 = 16984;
    const int bg_dim2 = 128;
    const int bg_dim3 = kpt_pool_nloc;
    const int bg_dim4 = l2_e;

    const int c_distr_dim1 = pert_nglob;

    int *nbnd_loc = new int[kpt_pool_nloc];
    int *ngk      = new int[kpt_pool_nloc];

    nbnd_loc[0] = 128;  // Number of bands
    ngk[0] = 16984;     // Number of grid points

    int ag_size = ag_dim1 * ag_dim2 * ag_dim3 * ag_dim4;
    int bg_size = bg_dim1 * bg_dim2 * bg_dim3 * bg_dim4;
    int c_distr_size = c_distr_dim1 * pert_nglob;

    Complex *ag = new Complex[ag_size];
    Complex *bg = new Complex[bg_size];
    double *c_distr   = new double[c_distr_size];

    fill_random_complex(ag, ag_size);
    fill_random_complex(bg, bg_size);

    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<double> dis(0.0, 1.0);

    for (int i = 0; i < c_distr_size; ++i) {
        c_distr[i] = dis(gen);
    }

    int nimage = nimg;

    process_loop(ag, bg, c_distr, nbnd_loc, ngk, l1_e, l2_s, l2_e, kpt_pool_nloc, gstart, nimage, idx, pert_nglob, ag_dim1, ag_dim2, ag_dim3, bg_dim1, bg_dim2, bg_dim3, c_distr_dim1);

    delete[] ag;
    delete[] bg;
    delete[] c_distr;
    delete[] nbnd_loc;
    delete[] ngk;

    return 0;
}

当前尝试的Thrust实现 #

已尝试对循环进行融合，并为两类求和操作定义了如下functor：

struct calculate_ag_bg {
    Complex* ag; Complex* bg;
    int npw;

    calculate_ag_bg(Complex* _ag,  Complex* _bg, int _npw) 
    : ag(_ag), bg(_bg), npw(_npw) {}

    __host__ __device__
    double operator()(int idx) const {
        int lbnd = idx / npw, il3 = idx % npw;
        double a_real = ag[il3 + lbnd * npw].real() * bg[il3 + lbnd * npw].real();
        double b_imag = ag[il3 + lbnd * npw].imag() * bg[il3 + lbnd * npw].imag();
        double a = 2.0 * a_real;
        double b = 2.0 * b_imag;
        return a + b;
    }
};

struct correct_ag_bg_prod_kernel {
    Complex* ag; Complex* bg; int npw;

    correct_ag_bg_prod_kernel(Complex* _ag,  Complex* _bg, int _npw) 
    : ag(_ag), bg(_bg), npw(_npw) {}

    __host__ __device__
    double operator()(int idx) const {
        return ag[idx * npw].real() * bg[idx * npw].real();
    }
};

随后使用thrust::transform_reduce执行归约，但性能表现较差。此外，kpt_pool_nloc可能大于1，且nbnd_loc和ngk的值不保证均匀，希望了解是否有更优的实现方式，比如使用transform_iterator。

针对你的场景，以下几种方式可以显著提升Thrust实现的性能，同时适配kpt_pool_nloc可变、nbnd_loc/ngk非均匀的情况：

1. 重构内存访问模式，提升缓存命中率 #

原始索引计算中，ag_index和bg_index的跨步较大，GPU缓存无法有效利用。可以预先计算每个iks+il1/il2对应的ag/bg子块起始地址，让内存访问尽可能连续：

// 预计算每个iks+il1对应的ag起始偏移
size_t ag_base = ag_dim3 * ag_dim2 * ag_dim1 * il1 + ag_dim2 * ag_dim1 * iks;
size_t bg_base = bg_dim3 * bg_dim2 * bg_dim1 * il2 + bg_dim2 * bg_dim1 * iks;

后续在functor中通过ag[ag_base + lbnd*ag_dim1 + il3]的方式访问，能大幅提升缓存效率。

2. 使用thrust::transform_iterator批量处理非均匀数据 #

当nbnd_loc和ngk值不均匀时，transform_iterator可将多个独立归约任务合并为一个整体Thrust操作，减少内核启动开销：

• 构建任务列表，每个任务包含ag/bg起始地址、nbndval、npw、是否需要校正等信息
• 用transform_iterator将每个任务映射为归约结果
• 最后将结果写入c_distr对应位置

示例框架：

// 定义任务结构体
struct ReductionTask {
    Complex* ag_ptr;
    Complex* bg_ptr;
    int nbnd;
    int npw;
    bool apply_correction;
};

// 任务处理functor
struct ComputeTaskReduction {
    __host__ __device__
    double operator()(const ReductionTask& task) const {
        double sum = 0.0;
        // 主求和逻辑
        for (int lbnd = 0; lbnd < task.nbnd; ++lbnd) {
            for (int il3 = 0; il3 < task.npw; ++il3) {
                size_t ag_idx = lbnd * task.npw + il3;
                size_t bg_idx = lbnd * task.npw + il3;
                sum += 2.0 * (task.ag_ptr[ag_idx].real() * task.bg_ptr[bg_idx].real() + task.ag_ptr[ag_idx].imag() * task.bg_ptr[bg_idx].imag());
            }
        }
        // 校正逻辑
        if (task.apply_correction) {
            for (int lbnd = 0; lbnd < task.nbnd; ++lbnd) {
                size_t ag_idx = lbnd * task.npw;
                size_t bg_idx = lbnd * task.npw;
                sum -= task.ag_ptr[ag_idx].real() * task.bg_ptr[bg_idx].real();
            }
        }
        return sum;
    }
};

// 构建任务列表并拷贝到设备
std::vector<ReductionTask> tasks;
for (int il1 = 0; il1 < l1_e; ++il1) {
    for (int il2 = l2_s - 1; il2 < l2_e; ++il2) {
        for (int iks = 0; iks < kpt_pool_nloc; ++iks) {
            int nbndval = nbnd_loc[iks];
            int npw = ngk[iks];
            size_t ag_base = ag_dim3 * ag_dim2 * ag_dim1 * il1 + ag_dim2 * ag_dim1 * iks;
            size_t bg_base = bg_dim3 * bg_dim2 * bg_dim1 * il2 + bg_dim2 * bg_dim1 * iks;
            tasks.push_back({ag + ag_base, bg + bg_base, nbndval, npw, (gstart == 2)});
        }
    }
}

thrust::device_vector<ReductionTask> d_tasks(tasks.begin(), tasks.end());
thrust::device_vector<double> d_results(tasks.size());

// 批量处理所有任务
thrust::transform(d_tasks.begin(), d_tasks.end(), d_results.begin(), ComputeTaskReduction());

// 将结果映射到c_distr对应位置（可通过另一个transform完成索引映射）

3. 合并循环层级，减少内核启动次数 #

原始CPU代码中il1和il2的双重循环会导致多次调用transform_reduce，每次启动GPU内核都会产生开销。将所有il1×il2×iks的归约任务合并为一个整体Thrust操作，一次性完成计算，避免多次内核启动的额外消耗。

4. 显式指定执行策略 #

调用Thrust函数时，显式指定thrust::cuda::par执行策略，确保代码运行在GPU上；还可配合thrust::cuda::stream实现异步执行，进一步提升效率。

5. 共享内存优化（可选） #

若性能瓶颈仍明显，可手动编写CUDA内核，将线程块内的ag/bg数据加载到共享内存中，减少全局内存访问次数。该方式实现复杂度较高，但能最大化内存访问效率。

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