Skip to main content

Module kernels

Module kernels 

Source
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:

  1. Add cubecl to rlevo-evolution’s dependencies (gated on the custom-kernels feature).
  2. Use #[cube(launch_unchecked)] with Burn’s backend::custom API to plug into the Backend trait.
  3. Provide a pure-tensor fallback (which is the current implementation) for backends that don’t support CubeCL.
  4. Expose a toggle at the operator level so benchmarks can A/B the two paths.