CUDA: Runge-Kutta trajectory on each GPU thread - c++

Question

Summary: How do you avoid performance loss caused by different work loads for different threads? (Kernel with a while loop on each thread)
Problem:
I want to solve particle trajectories (described by a 2nd order differential equation) using Runge-Kutta for many different initial conditions. The trajectories will generally have different lengths (each trajectory ends when a particle hits some target). Furthermore, to ensure numerical stability, the Runge-Kutta stepsize is set adaptively. This leads to two nested while-loops, with unknown number of iterations (see serial example below).
I want to implement the Runge-Kutta routine to run on a GPU with CUDA/C++. The trajectories have no dependency of each other, so as a first approach, I will just parallelize over the different initial conditions such that each thread will correspond to a unique trajectory. When a thread is done with a particle trajectory, I want it to start with a new one.
If I understand it correctly, however, the unknown length of each while loop (particle trajectory) means that different threads will get different amounts of work, which might lead to a severe performance loss on GPU.
Question: Is this possible to overcome (in a simple way) the performance losses caused by different work load for different threads? For example setting each warp size to be only 1, such that each thread(warp) can then run independently? r will this lead to other performance losses (e.g. no coalesced memory reads)?
Serial pseudo-code:
// 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;
}
}
A short test of my code shows that the inner while-loop runs 2-4 times in 99% of all cases and >10 times in 1% of all cases, before a satisfactory Runge-Kutta step-size was found.
Parallel pseudo-code:
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

I will attempt to answer the question myself, until someone comes up with a better solution.
Pitfalls with directly porting the serial code:
The two while loops will lead to significant branch divergence and performance loss. The outer loop is the "full" trajectory, while the inner loop is one Runge-Kutta step with adaptive step size correction. Inner loop: If we attempt to solve Runge-Kutta with a too large step size then the approximation error will be too large, and we need to redo the step with a smaller step size until the error is smaller than our tolerance. This means that threads that need very few iterations to find an appropriate step size will have to wait for threads that need more iterations. Outer loop: this reflects how many successful Runge-Kutta steps we need before the trajectory is completed. Different trajectories will reach their target in different amount of steps. We will always have to wait for the trajectory with the most iterations before we are completely done.
Proposed parallel approach:
We notice that every iteration consists of doing one Runge-Kutta step. The branching comes from the fact that we either need to reduce the step size for the next iteration, or update the Runge-Kutta coefficients (e.g. positon/velocity) if the step size was OK. I therefore propose that we replace the two while-loops with one for-loop. The first step of the for-loop is to solve Runge-Kutta, followed by an if-statement to check if the step size was small enough or if updating the positions (and checking for total convergence). All threads will now solve only one Runge-Kutta step at a time, and we trade away low occupancy (all threads wait for the thread that need the most attempts to find the correct step size) for the cost of branch divergence of a single if-statement. In my case, solving Runge-Kutta is expensive compared with the evaluations of this if-statement, so we have made an improvement. The issue now lies in setting an appropriate limit on the for-loop and flagging the threads that need more work. This limit will set an upper bound on the longest time a finished thread has to wait for others. Pseudo-code:
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;
...
}