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

Question

简介：如何避免因不同线程的不同工作负载而导致的性能损失？ （每个线程上带有while循环的内核）

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

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

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

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

串行伪代码：

// 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;
    ...
}