CUDA：每个 GPU 线程上的 Runge-Kutta 轨迹

Question

总结：你如何避免不同线程的不同工作负载导致的性能损失？ （每个线程都有一个while循环的内核）

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

我想实现 Runge-Kutta 例程以在具有 CUDA/C++ 的 GPU 上运行。轨迹彼此之间没有依赖关系，因此作为第一种方法，我将在不同的初始条件上并行化，这样每个线程将对应一个唯一的轨迹。当一个线程完成粒子轨迹后，我希望它从一个新的开始。

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

问题：这是否可能（以简单的方式）克服不同线程的不同工作负载导致的性能损失？例如将每个经线大小设置为只有 1，这样每个线程（经线）就可以独立运行？ 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，那么逼近误差会太大，我们需要用更小的步长重做一步，直到误差更小超过我们的容忍度。这意味着需要很少迭代来找到合适步长的线程将不得不等待需要更多迭代的线程。 外循环：这反映了在完成轨迹之前我们需要多少成功的 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;
    ...
}