c++ - CUDA:每个 GPU 线程上的 Runge-Kutta 轨迹

Question

总结:如何避免不同线程的不同工作负载导致的性能损失？ (内核在每个线程上都有一个 while 循环)

问题:我想在许多不同的初始条件下使用 Runge-Kutta 求解粒子轨迹(由二阶微分方程描述)。轨迹通常具有不同的长度(每个轨迹在粒子撞击某个目标时结束)。此外，为了确保数值稳定性，自适应设置 Runge-Kutta 步长。这会导致两个嵌套的 while 循环，迭代次数未知(请参见下面的系列示例)。

我想使用 CUDA/C++ 实现 Runge-Kutta 例程以在 GPU 上运行。轨迹彼此没有依赖性，因此作为第一种方法，我将对不同的初始条件进行并行处理，以便每个线程都对应于一个唯一的轨迹。当一个线程完成一个粒子轨迹时，我希望它从一个新的开始。

但是，如果我理解正确的话，每个 while 循环(粒子轨迹)的未知长度意味着不同的线程将获得不同的工作量，这可能会导致 GPU 的严重性能损失.

问题:这是否可以(以简单的方式)克服不同线程的不同工作负载导致的性能损失？例如将每个 warp 大小设置为仅为 1，这样每个线程(warp)就可以独立运行？ r 这会导致其他性能损失(例如，没有合并的内存读取)吗？

串口伪代码:

// Solve a particle trajectory for each inital condition
// N_trajectories: much larger than 1e6
for( int t_i = 0; t_i < N_trajectories; ++t_i )
{
    // Set start coordinates
    double x = x_init[p_i];
    double y = y_init[p_i];
    double vx = vx_init[p_i];
    double vy = vy_init[p_i];
    double stepsize = ...;
    double tolerance = ...;
    ...

    // Solve Runge-Kutta trajectory until convergence
    int converged = 0;
    while ( !converged )
    {
        // Do a Runge-Kutta step, if step-size too large then decrease it
        int goodStepSize = 0
        while( !goodStepSize )
        {
            // Update x, y, vx, vy
            double error = doRungeKutta(x, y, vx, vy, stepsize);

            if( error < tolerance )
                goodStepSize = 1;
            else
                stepsize *= 0.5;
        }

        if( (abs(x-x_final) < epsilon) && (abs(y-y_final) < epsilon) )
            converged = 1;
    }
}

我的代码的简短测试表明，在找到令人满意的 Runge-Kutta 步长之前，内部 while 循环在 99% 的所有情况下运行 2-4 次，在 1% 的所有情况下运行 >10 次。

并行伪代码:

int tpb = 64;
int bpg = (N_trajectories + tpb-1) / tpb;
RungeKuttaKernel<<<bpg, tpb>>>( ... );

__global__ void RungeKuttaKernel( ... )
{
    int idx = ...;

    // Set start coordinates
    double x = x_init[idx];
    ...

    while ( !converged )
    {
        ...

        while( !goodStepSize )
        {
            double error = doRungeKutta( ... );
            ...
        }

        ...
    }
}

Answer 1

我会尝试自己回答这个问题，直到有人提出更好的解决方案。

直接移植串行代码的陷阱:两个 while 循环会导致明显的分支发散和性能损失。外环是“完整”轨迹，而内环是一个具有自适应步长校正的 Runge-Kutta 步。 内循环:如果我们试图用太大的步长求解Runge-Kutta，那么近似误差就会太大，我们需要用较小的步长重做这一步，直到误差变小比我们的宽容。这意味着需要很少迭代才能找到合适步长的线程将不得不等待需要更多迭代的线程。 外环:这反射(reflect)了在轨迹完成之前我们需要多少成功的 Runge-Kutta 步。不同的轨迹将以不同的步数到达目标。在我们完全完成之前，我们将始终必须等待迭代次数最多的轨迹。

提议的并行方法:我们注意到每次迭代都包含执行一个 Runge-Kutta 步骤。分支来自这样一个事实，即我们要么需要减少下一次迭代的步长，要么在步长合适的情况下更新 Runge-Kutta 系数(例如位置/速度)。因此，我建议我们用一个 for 循环替换两个 while 循环。 for 循环的第一步是求解 Runge-Kutta，然后是一个 if 语句来检查步长是否足够小或是否更新位置(并检查总收敛)。所有线程现在一次只能解决一个 Runge-Kutta 步骤，我们用低占用率(所有线程等待需要最多尝试找到正确步长的线程)换取单个 if 的分支发散成本-陈述。就我而言，与此 if 语句的评估相比，求解 Runge-Kutta 的代价更高，因此我们进行了改进。现在的问题在于为 for 循环设置适当的限制并标记需要更多工作的线程。此限制将设置已完成线程必须等待其他线程的最长时间的上限。 伪代码:

int N_trajectories = 1e6;
int trajectoryStepsPerKernel = 50;
thrust::device_vector<int> isConverged(N_trajectories, 0); // Set all trajectories to unconverged
int tpb = 64;
int bpg = (N_trajectories + tpb-1) / tpb;

// Run until all trajectories are converged
while ( vectorSum(isConverged) != N_trajectories )
{
    RungeKuttaKernel<<<bpg, tpb>>>( trajectoryStepsPerKernel, isConverged, ... );
    cudaDeviceSynchronize();
}

__global__ void RungeKuttaKernel( ... )
{
    int idx = ...;

    // Set start coordinates
    int converged = 0;
    double x = x_init[idx];
    ...

    for ( int i = 0; i < trajectoryStepsPerKernel; ++i )
    {
        double error = doRungeKutta( x_new, y_new, ... );

        if( error > tolerance )
        {
            stepsize *= 0.5;
        } else {
            converged = checkConvergence( x, x_new, y, y_new, ... );
            x = x_new;
            y = y_new;
            ...
        }
    }

    // Update start positions in case we need to continue on trajectory 
    isConverged[idx] = converged;
    x_init[idx] = x;
    y_init[idx] = y;
    ...
}