GPU Parallelism
The model used in this example represents nbuses electrical generators, each tied to all the others through a network of transmission lines. Each generator bus is described by the swing equation, a second-order ODE describing the dynamics of the generator's rotor angle and frequency. The implementation writes this as a system of first-order ODEs.
We introduce uncertainty into the model by adding an uncertain additional demand at bus 0, drawn from a normal distribution. The goal is to solve many instances of the model and obtain the maximum speed deviation for each ensemble member.
We will solve this ensemble of models in parallel, both with the GPU and CPU, using diffsol's batching features and using the cuda-oxide and nalgebra backends.
The swing equation is a non-linear ODE, and is given by:
$$ \begin{align} \frac{d\delta_i}{dt} &= \omega_i \\ M \frac{d\omega_i}{dt} &= -D\omega_i - B\sum_{j=0}^{n-1} \sin(\delta_i-\delta_j) - d\ \mathbf{1}_{i=0} \end{align} $$
where:
- \(\delta_i\): rotor angle at bus \(i\) in the reference frame, radians
- \(\omega_i\): speed deviation at bus \(i\) from the reference frame, radians per second
- \(B\): pairwise line susceptance (how strongly each transmission line transfers power between two buses)
- \(d\): uncertain additional demand at bus \(0\)
- \(\mathbf{1}_{i=0}\): equals (1) at bus \(0\), otherwise \(0\)
First, we write the model using diffsol's OdeEquations trait. We cannot use the easier builder and closures API because it is not supported for GPU models.
We will define the generator inertia M, damping D and per-line susceptance B:
pub const INERTIA: f64 = 4.0;
pub const DAMPING: f64 = 0.6;
pub const SUSCEPTANCE: f64 = 0.25;Then we define the swing equations over a batched ensemble using a single SwingEqn struct with one demand value per batch lane. We hold the demand values in a vector demand on the struct. The network has nbuses generators, each tied to every other. State i < nbuses is the rotor angle of bus i, and state nbuses + g is the speed deviation of bus g, so each state vector holds 2 * nbuses states.
We define the equations using the for_each_elem method, which is called for each element i of the state vector in each batch lane. Since the rotor angles are in the first half of the state vector and the speed deviations are in the second, we can use the index i to determine which equation to evaluate.
We will not show much of the boilerplate code here, but the full code is available in the examples/performance-gpu-ensemble directory of the diffsol repository.
pub struct SwingEqn<M: Matrix> {
ctx: M::C,
demand: M::V,
nbuses: usize,
inertia: M::T,
damping: M::T,
susceptance: M::T,
}
impl<M: Matrix> SwingEqn<M> {
pub fn new(nbuses: usize, ctx: M::C) -> Self {
assert!(
nbuses >= 3 && nbuses.is_multiple_of(2),
"need an even network of 4+ buses"
);
// one parameter per lane; the builder fills it from `OdeBuilder::p`
let demand = M::V::zeros(1, ctx.clone());
Self {
ctx,
demand,
nbuses,
inertia: M::T::from_f64(INERTIA).unwrap(),
damping: M::T::from_f64(DAMPING).unwrap(),
susceptance: M::T::from_f64(SUSCEPTANCE).unwrap(),
}
}
}
pub struct SwingRhs<'a, M: Matrix> {
eqn: &'a SwingEqn<M>,
}
pub struct SwingInit<'a, M: Matrix> {
eqn: &'a SwingEqn<M>,
}
impl<M: Matrix> NonLinearOp for SwingRhs<'_, M> {
fn call_inplace(&self, x: &M::V, _t: M::T, y: &mut M::V) {
// captured by value so the closure stays `Copy + Send`
// and compiles for the device
let n = self.eqn.nbuses;
let (m, d, b) = (self.eqn.inertia, self.eqn.damping, self.eqn.susceptance);
y.for_each_elem(
[x, &self.eqn.demand],
move |y: &mut M::T, [x, demand]: [&[M::T]; 2], _lane: usize, i: usize| {
if i < n {
// rotor angle: d(delta)/dt = omega
*y = x[n + i];
} else {
let g = i - n;
// every bus is tied to every other, as a dense admittance matrix.
let mut flow = M::T::zero();
for j in 0..n {
flow += (x[g] - x[j]).sin();
}
let flow = b * flow;
// the uncertain load at bus 0 drives the system
let inj = if g == 0 { -demand[0] } else { M::T::zero() };
*y = (inj - d * x[n + g] - flow) / m;
}
},
);
}
}
impl<M: Matrix> ConstantOp for SwingInit<'_, M> {
fn call_inplace(&self, _t: M::T, y: &mut M::V) {
// the nominal grid is balanced, so it starts at rest
y.for_each_elem(
[],
|y: &mut M::T, _: [&[M::T]; 0], _lane: usize, _i: usize| *y = M::T::zero(),
);
}
}Once we have defined our equations, we can create an OdeSolverProblem using the builder. To set up the batches on the GPU, we create a vector of N_SAMPLES demand values drawn from a normal distribution with DEMAND_MEAN = 0.3 and DEMAND_SD = 0.08, then pass them to the builder through the demands argument.
/// One problem holding `demands.len()` independent grids, one per batch lane.
#[allow(clippy::type_complexity)]
pub fn swing_problem<M: Matrix + 'static>(
nbuses: usize,
demands: &[f64],
) -> OdeSolverProblem<impl OdeEquations<M = M, V = M::V, T = M::T, C = M::C>> {
let ctx = M::C::default().clone_with_nbatch(demands.len()).unwrap();
OdeBuilder::<M>::new()
.context(ctx.clone())
// one value per lane, so the parameter vector is `nparams * nbatch` long
.p(demands.iter().copied())
.rtol(1e-6)
.atol([1e-8])
.build_from_eqn(SwingEqn::new(nbuses, ctx))
.unwrap()
}We can then solve the ensemble of ODEs on the GPU using the standard diffsol solve API. Under the hood, diffsol solves nbatch ODEs in lockstep on the GPU using the cuda-oxide backend.
/// Every grid in one batched solve.
fn solve_gpu(nbuses: usize, demands: &[f64], t_final: f64) -> Vec<f64> {
use diffsol::OxideMat;
let problem = swing_problem::<OxideMat>(nbuses, demands);
let mut solver = problem.tsit45().unwrap();
max_deviation_hz(&mut solver, nbuses, t_final)
}We want to display the distribution of maximum speed deviation across the ensemble of ODEs, so we use the lower-level step API to advance the equations in time. At each time step, reduce_elem computes one maximum speed deviation per batch lane, avoiding the need to hold the entire trajectory in memory.
/// Worst speed deviation of any generator over the whole run, one value per grid.
fn max_deviation_hz<'a, Solver, Eqn>(solver: &mut Solver, nbuses: usize, t_final: f64) -> Vec<f64>
where
Solver: OdeSolverMethod<'a, Eqn>,
Eqn: OdeEquations<T = f64> + 'a,
{
let ctx = solver.problem().context().clone();
let mut worst = Eqn::V::zeros(1, ctx.clone());
let mut next = Eqn::V::zeros(1, ctx);
// find the maximum speed deviation in Hz across all generators
fn fold<V: Vector<T = f64>>(y: &V, worst: &mut V, next: &mut V, nbuses: usize) {
V::reduce_elem(
next,
[y, worst],
0.0,
move |[x, w], _lane, i| {
// rad/s in the rotating frame, so Hz is omega / 2*pi
if i >= nbuses {
f64::max(x[i].abs() / (2.0 * PI), w[0])
} else {
w[0]
}
},
f64::max,
);
std::mem::swap(worst, next);
}
solver.set_stop_time(t_final).unwrap();
loop {
fold(solver.state().y, &mut worst, &mut next, nbuses);
match solver.step() {
Ok(OdeSolverStopReason::TstopReached) => break,
Ok(_) => (),
Err(e) => panic!("solver failed: {e}"),
}
}
fold(solver.state().y, &mut worst, &mut next, nbuses);
worst.clone_as_vec()
}We can then plot the histogram of the maximum speed deviation across the ensemble, which is shown below:
GPU versus CPU
It is useful to compare the GPU batch solves against a CPU parallel implementation using the rayon crate. The CPU version uses map_init to create a SwingEqn problem for each worker, then solves one ODE per sample in parallel using the thread pool.
/// One solve per grid, spread over rayon's threads.
fn solve_cpu(nbuses: usize, demands: &[f64], t_final: f64) -> Vec<f64> {
demands
.par_iter()
.map_init(
|| swing_problem::<CpuM>(nbuses, &[0.0]),
|problem, &d| {
let p = NalgebraVec::from_vec(vec![d], *problem.eqn.context());
problem.eqn_mut().set_params(&p);
let mut solver = problem.tsit45().unwrap();
max_deviation_hz(&mut solver, nbuses, t_final)[0]
},
)
.collect()
}To show the relative performance, we loop over different ensemble sizes and, for each size, time the GPU and CPU solves. We take the median of 15 runs to reduce noise. The results are plotted below.
Below an ensemble size of 100, fixed costs such as kernel launches, allocation, and setup dominate, so GPU execution time is approximately flat for this workload. Between 100 and 500 samples, computation becomes a larger part of the runtime. At 1000 samples, the timing is approximately linear. In this case, the GPU appears to reach its throughput limit before all of its theoretical resident threads are occupied, GPU profiling would be required to confirm the limiting factor here. CPU scaling becomes linear at smaller ensemble sizes because the CPU has fewer workers available for parallel work. Once both CPU and GPU are in the linear regime, the GPU is about twice as fast as the CPU. The A40's relatively low FP64 peak rate, approximately 1/64 of its FP32 rate, is one factor that limits its advantage, others include memory access, kernel structure, and launch overhead.