Expand description
Custom CubeCL kernels for hot-path operators.
This module is a design placeholder. The current release ships
only the pure-tensor operator baselines in crate::ops::selection,
crate::ops::crossover, and crate::ops::mutation. Those
compose from Burn tensor primitives and run on every backend Burn
supports (flex, wgpu, …) with no extra work. The custom-kernels
Cargo feature exists so downstream crates can pin the future ABI
when kernels land.
§Why kernels aren’t in the current release
Three operator paths were identified where a fused CubeCL kernel
would eliminate multi-launch overhead. Landing real kernels requires
non-trivial CubeCL integration (cubecl 0.9 ships with Burn 0.20.1)
and device-specific validation on wgpu. None of that work blocks the
core strategy machinery, so it was deferred to keep the release
shippable.
The three designs below document the intended interfaces so a future contributor can write them without re-deriving the motivation.
§Tournament selection
Today the pure-tensor path
(super::selection::tournament_select) samples tournament
indices on the host, packs them into a 1-D Int tensor, and does a
single tensor.select(0, indices) gather. Cost at pop_size = N:
N host-side index draws and one device kernel launch.
A fused kernel would take the form
fn tournament_select_cube<F: Float, I: Int>(
fitness: &Tensor<F>, // (N,)
rng_state: &Tensor<I>, // (N, k) pre-sampled index pairs
winners: &mut Tensor<I>, // (N,) output
)performing the index sampling and comparison in a single launch,
eliminating the host trip entirely. Expected speedup at N ≥ 256
on wgpu: order-of-magnitude.
§DE trial-vector construction
Classical DE computes v_i = x_{r1} + F · (x_{r2} − x_{r3}) plus
a binomial-crossover mask per gene. In
crate::algorithms::de this is composed from three selects,
one subtract, one mul_scalar, one mask-build, and one
mask_where — seven kernel launches per generation.
A fused kernel that takes the whole population plus pre-sampled indices and emits the trial vector in one pass:
fn de_trial_cube<F: Float, I: Int>(
pop: &Tensor<F>, // (N, D)
indices: &Tensor<I>, // (N, k) sampled parent indices
f: F, cr: F, // scalars
rng_bits: &Tensor<I>, // crossover mask seeds
variant: u32, // const-generic DeVariant discriminant
trial: &mut Tensor<F>, // (N, D) output
)Expected impact: DE’s inner loop is dominated by these 7 launches;
collapsing to 1 would likely double throughput at pop_size ≥ 256.
§Fitness-proportionate (roulette) selection
Roulette selection is a prefix-sum + inverse-CDF lookup. Burn’s
cumsum + searchsorted would work but materializes two
intermediate tensors. This kernel is lower priority than the two
above — the pure-tensor path is fine for typical population sizes —
but worth writing if profiling later shows roulette on a hot path.
§Kernel infrastructure
When implementing these:
- Add
cubecltorlevo-evolution’s dependencies (gated on thecustom-kernelsfeature). - Use
#[cube(launch_unchecked)]with Burn’sbackend::customAPI to plug into theBackendtrait. - Provide a pure-tensor fallback (which is the current
implementation) for backends that don’t support
CubeCL. - Expose a toggle at the operator level so benchmarks can A/B the two paths.