use std::sync::Arc;
use rayon::{prelude::*, ThreadPool};
use serde_json::json;
use silvic_core::{
backend::{BagSpec, BinSplit, DeviceTree, GradientStats, LeafRange, ObjectiveSetup},
objective,
quantize::{self, QuantSpec},
Backend, BinnedDataset, Result, SilvicError,
};
use crate::runtime::{Buffer, DType, Device, Execution, TensorSpec};
const SOURCE: &str = include_str!("../kernels/training.py");
const TILE: usize = 128;
const HIST_BINS: usize = 128;
const CHUNKS: usize = 16;
const LIMBS: usize = 8;
const NODE_FIELDS: usize = 11;
fn error(message: impl Into<Arc<str>>) -> SilvicError {
SilvicError::device(message)
}
fn install<R: Send>(pool: &Option<Arc<ThreadPool>>, f: impl FnOnce() -> R + Send) -> R {
match pool {
Some(pool) => pool.install(f),
None => f(),
}
}
fn spec(name: &str, dtype: DType, shape: &[usize]) -> TensorSpec {
let mut shape = shape.to_vec();
shape.push(1);
TensorSpec::new(name, dtype, shape)
}
fn padded(rows: usize) -> Result<usize> {
rows.max(1)
.checked_add(TILE - 1)
.map(|n| n / TILE * TILE)
.ok_or_else(|| error("Neuron row capacity overflow"))
}
fn zeros(device: &Device, words: usize) -> Result<Buffer> {
let bytes = words
.max(1)
.checked_mul(4)
.ok_or_else(|| error("Neuron buffer size overflow"))?;
device.allocate(bytes)
}
struct Dataset {
rows: usize,
capacity: usize,
bins: Buffer,
scores: Buffer,
next_scores: Buffer,
columns: usize,
group_columns: Vec<usize>,
}
struct Work {
partition: Buffer,
values: Buffer,
limbs: Buffer,
meta: Buffer,
split: Buffer,
counts: Buffer,
offsets: Buffer,
partials: Buffer,
leaf_partials: Buffer,
histogram: Execution,
leaf_sum: Execution,
partition_flags: Execution,
partition_scatter: Execution,
copy_partition: Execution,
host_partials: Vec<f32>,
host_leaf_partials: Vec<f32>,
host_counts: Vec<u32>,
host_offsets: Vec<u32>,
}
struct TreeWork {
capacity: usize,
nodes: Buffer,
leaves: Buffer,
meta: Buffer,
executions: Vec<(Execution, Execution)>,
}
struct Loaded {
datasets: Vec<Dataset>,
train: BinnedDataset,
objective: ObjectiveSetup,
grad: Vec<f32>,
hess: Vec<f32>,
host_values: Vec<f32>,
quant_values: Vec<f32>,
host_limbs: Vec<i32>,
host_scores: Vec<f32>,
bag: Option<(BagSpec, Vec<u32>)>,
partition_len: usize,
quantized: bool,
quant_ready: bool,
work: Work,
tree: Option<TreeWork>,
}
pub struct NeuronBackend {
device: Device,
pool: Option<Arc<ThreadPool>>,
loaded: Option<Loaded>,
}
impl NeuronBackend {
pub fn new(ordinal: Option<usize>) -> Result<Self> {
Ok(Self {
device: Device::new(ordinal)?,
pool: None,
loaded: None,
})
}
fn state(&mut self) -> Result<&mut Loaded> {
self.loaded
.as_mut()
.ok_or_else(|| error("Neuron backend used before load"))
}
}
fn upload_dataset(device: &Device, data: &BinnedDataset) -> Result<Dataset> {
if data.num_rows > u32::MAX as usize {
return Err(error("Neuron supports at most 2^32 - 1 rows per rank"));
}
let capacity = padded(data.num_rows)?;
let mut group_columns = Vec::with_capacity(data.groups.len());
let mut columns = 0usize;
for group in &data.groups {
group_columns.push(columns);
columns = columns
.checked_add(group.row_stride as usize)
.ok_or_else(|| error("Neuron column overflow"))?;
if group.data.len()
!= data
.num_rows
.checked_mul(group.row_stride as usize)
.ok_or_else(|| error("group size overflow"))?
{
return Err(SilvicError::data("feature group does not match its row stride"));
}
}
let words = columns
.max(1)
.checked_mul(capacity)
.ok_or_else(|| error("Neuron dataset size overflow"))?;
if words > i32::MAX as usize {
return Err(error("Neuron binned tensor exceeds signed 32-bit gather addressing"));
}
let mut bins = vec![0u32; words];
for (g, group) in data.groups.iter().enumerate() {
for row in 0..data.num_rows {
for col in 0..group.row_stride as usize {
bins[(group_columns[g] + col) * capacity + row] =
u32::from(group.data[row * group.row_stride as usize + col]);
}
}
}
Ok(Dataset {
rows: data.num_rows,
capacity,
bins: device.upload(&bins)?,
scores: device.upload(&vec![0f32; capacity])?,
next_scores: zeros(device, capacity)?,
columns: columns.max(1),
group_columns,
})
}
fn bind(
device: &Device,
entry: &str,
constants: serde_json::Value,
inputs: Vec<TensorSpec>,
outputs: Vec<TensorSpec>,
ib: &[Buffer],
ob: &[Buffer],
) -> Result<Execution> {
device
.compile(SOURCE, entry, &constants, &inputs, &outputs)?
.bind(ib, ob)
}
fn work(device: &Device, data: &Dataset) -> Result<Work> {
let n = data.capacity;
let blocks = n / TILE;
let partition = zeros(device, n)?;
let next_partition = zeros(device, n)?;
let values = zeros(device, 2 * n)?;
let limbs = zeros(device, LIMBS * n)?;
let meta = device.upload(&[0u32; 8])?;
let split = device.upload(&[0u32; 8])?;
let flags = zeros(device, n)?;
let counts = zeros(device, blocks)?;
let offsets = zeros(device, 2 * blocks)?;
let partials = zeros(device, CHUNKS * 2 * HIST_BINS)?;
let leaf_partials = zeros(device, CHUNKS * LIMBS)?;
let constants = json!({"ROWS": n, "COLUMNS": data.columns, "TILE": TILE, "CHUNKS": CHUNKS, "BINS": HIST_BINS});
let histogram = bind(
device,
"histogram",
constants,
vec![
spec("bins", DType::I32, &[data.columns * n]),
spec("partition", DType::I32, &[n]),
spec("values", DType::F32, &[2, n]),
spec("meta", DType::I32, &[8]),
],
vec![spec("partials", DType::F32, &[CHUNKS, 2, HIST_BINS])],
&[data.bins.clone(), partition.clone(), values.clone(), meta.clone()],
std::slice::from_ref(&partials),
)?;
let leaf_sum = bind(
device,
"leaf_sums",
json!({"ROWS":n,"TILE":TILE,"CHUNKS":CHUNKS,"LIMBS":LIMBS}),
vec![
spec("partition", DType::I32, &[n]),
spec("limbs", DType::I32, &[LIMBS, n]),
spec("meta", DType::I32, &[8]),
],
vec![spec("partials", DType::F32, &[CHUNKS, LIMBS])],
&[partition.clone(), limbs.clone(), meta.clone()],
std::slice::from_ref(&leaf_partials),
)?;
let partition_flags = bind(
device,
"partition_flags",
json!({"ROWS":n,"TILE":TILE}),
vec![
spec("bins", DType::I32, &[data.columns * n]),
spec("partition", DType::I32, &[n]),
spec("meta", DType::I32, &[8]),
spec("split", DType::I32, &[8]),
],
vec![spec("flags", DType::I32, &[n]), spec("counts", DType::I32, &[blocks])],
&[data.bins.clone(), partition.clone(), meta.clone(), split.clone()],
&[flags.clone(), counts.clone()],
)?;
let partition_scatter = bind(
device,
"partition_scatter",
json!({"ROWS":n,"TILE":TILE}),
vec![
spec("partition", DType::I32, &[n]),
spec("flags", DType::I32, &[n]),
spec("offsets", DType::I32, &[2 * blocks]),
spec("meta", DType::I32, &[8]),
],
vec![spec("result", DType::I32, &[n])],
&[partition.clone(), flags.clone(), offsets.clone(), meta.clone()],
std::slice::from_ref(&next_partition),
)?;
let copy_partition = bind(
device,
"copy_u32",
json!({"ROWS":n,"TILE":TILE}),
vec![spec("source", DType::I32, &[n])],
vec![spec("result", DType::I32, &[n])],
std::slice::from_ref(&next_partition),
std::slice::from_ref(&partition),
)?;
Ok(Work {
partition,
values,
limbs,
meta,
split,
counts,
offsets,
partials,
leaf_partials,
histogram,
leaf_sum,
partition_flags,
partition_scatter,
copy_partition,
host_partials: vec![0.0; CHUNKS * 2 * HIST_BINS],
host_leaf_partials: vec![0.0; CHUNKS * LIMBS],
host_counts: vec![0; blocks],
host_offsets: vec![0; 2 * blocks],
})
}
fn check_range(state: &Loaded, leaf: LeafRange) -> Result<()> {
if u64::from(leaf.begin) + u64::from(leaf.count) > state.partition_len as u64 {
return Err(error("leaf lies outside the active partition"));
}
Ok(())
}
fn encode_fixed(value: i64, out: &mut [i32; 4]) {
for (k, digit) in out.iter_mut().enumerate() {
*digit = if k == 3 {
(value >> 48) as i32
} else {
((value as u64 >> (16 * k)) & 65535) as i32
};
}
}
fn fold_fixed(digits: &[f32]) -> Result<i128> {
let mut value = 0i128;
for (k, &digit) in digits.iter().enumerate() {
if !digit.is_finite() || digit.fract() != 0.0 || digit.abs() > (TILE * 65535) as f32 {
return Err(error("invalid Neuron fixed-point limb reduction"));
}
value += (digit as i128) << (16 * k);
}
Ok(value)
}
impl Backend for NeuronBackend {
fn name(&self) -> String {
format!("{} (neuron; host objectives)", self.device.name())
}
fn set_thread_pool(&mut self, pool: Arc<ThreadPool>) {
self.pool = Some(pool);
}
fn load(&mut self, train: &BinnedDataset, valid: &[&BinnedDataset], objective: &ObjectiveSetup) -> Result<()> {
if train.label.len() != train.num_rows || train.weight.as_ref().is_some_and(|w| w.len() != train.num_rows) {
return Err(SilvicError::data("training label/weight length differs from row count"));
}
if let ObjectiveSetup::LambdaRank { inverse_max_dcg, .. } = objective {
let q = train
.query_boundaries
.as_ref()
.ok_or_else(|| SilvicError::data("ranking requires query boundaries"))?;
if q.first() != Some(&0)
|| q.last().copied().map(u64::from) != Some(train.num_rows as u64)
|| q.windows(2).any(|w| w[0] > w[1])
|| inverse_max_dcg.len() + 1 != q.len()
{
return Err(SilvicError::data("invalid ranking query boundaries"));
}
}
for (f, meta) in train.features.iter().enumerate().filter(|(_, m)| m.num_bins != 0) {
let group = train
.groups
.get(meta.group as usize)
.ok_or_else(|| SilvicError::data("feature references missing group"))?;
if meta.column >= group.row_stride
|| !group.features.contains(&f)
|| meta.num_bins > 256
|| meta
.hist_offset
.checked_add(meta.num_bins)
.is_none_or(|e| e as usize > train.total_bins)
{
return Err(SilvicError::data("invalid feature histogram layout"));
}
for row in 0..train.num_rows {
if group
.data
.get(row * group.row_stride as usize + meta.column as usize)
.is_none_or(|&bin| u32::from(bin) >= meta.num_bins)
{
return Err(SilvicError::data("feature bin outside histogram layout"));
}
}
}
let mut datasets = Vec::with_capacity(valid.len() + 1);
for data in std::iter::once(train).chain(valid.iter().copied()) {
if data.features != train.features
|| data.groups.len() != train.groups.len()
|| data.groups.iter().zip(&train.groups).any(|(a, b)| {
a.row_stride != b.row_stride
|| a.features != b.features
|| a.hist_offset != b.hist_offset
|| a.total_bins != b.total_bins
})
{
return Err(SilvicError::data("validation set has a different binned layout"));
}
datasets.push(upload_dataset(&self.device, data)?);
}
let work = work(&self.device, &datasets[0])?;
let n = datasets[0].capacity;
let host_train = BinnedDataset {
num_rows: train.num_rows,
row_offset: train.row_offset,
features: train.features.clone(),
groups: train
.groups
.iter()
.map(|group| silvic_core::data::FeatureGroup {
features: group.features.clone(),
row_stride: group.row_stride,
hist_offset: group.hist_offset,
total_bins: group.total_bins,
data: Vec::new(),
})
.collect(),
total_bins: train.total_bins,
label: train.label.clone(),
weight: train.weight.clone(),
query_boundaries: train.query_boundaries.clone(),
};
self.loaded = Some(Loaded {
datasets,
train: host_train,
objective: objective.clone(),
grad: vec![0.0; train.num_rows],
hess: vec![0.0; train.num_rows],
host_values: vec![0.0; 2 * n],
quant_values: vec![0.0; 2 * n],
host_limbs: vec![0; LIMBS * n],
host_scores: vec![0.0; train.num_rows],
bag: None,
partition_len: 0,
quantized: false,
quant_ready: false,
work,
tree: None,
});
Ok(())
}
fn init_scores(&mut self, dataset: usize, value: f64) -> Result<()> {
let data = self
.state()?
.datasets
.get_mut(dataset)
.ok_or_else(|| error("missing score dataset"))?;
let values = vec![value as f32; data.capacity];
data.scores.write(&values)
}
fn compute_gradients(&mut self) -> Result<GradientStats> {
let pool = self.pool.clone();
let state = self.state()?;
state.datasets[0].scores.read(&mut state.host_scores)?;
install(&pool, || {
objective::gradients(
&state.objective,
&state.host_scores,
&state.train.label,
state.train.weight.as_deref(),
state.train.query_boundaries.as_deref(),
&mut state.grad,
&mut state.hess,
)
});
let n = state.datasets[0].capacity;
state.host_values[..state.grad.len()].copy_from_slice(&state.grad);
state.host_values[n..n + state.hess.len()].copy_from_slice(&state.hess);
state.work.values.write(&state.host_values)?;
state.quantized = false;
state.quant_ready = false;
Ok(install(&pool, || GradientStats {
max_abs_grad: state
.grad
.par_iter()
.map(|v| f64::from(v.abs()))
.reduce(|| 0.0, f64::max),
max_hess: state
.hess
.par_iter()
.map(|v| f64::from(v.abs()))
.reduce(|| 0.0, f64::max),
}))
}
fn quantize(&mut self, spec: &QuantSpec) -> Result<()> {
if spec.grad_max_q < 0 || spec.grad_max_q > 32767 || spec.hess_max_q < 0 || spec.hess_max_q > 32767 {
return Err(error(
"quantized gradients/hessians must fit signed 16-bit histogram units",
));
}
let pool = self.pool.clone();
let state = self.state()?;
let n = state.datasets[0].capacity;
let (g, h) = state.quant_values.split_at_mut(n);
install(&pool, || {
g[..state.grad.len()]
.par_iter_mut()
.zip(h[..state.hess.len()].par_iter_mut())
.enumerate()
.for_each(|(row, (g, h))| {
let (qg, qh) = quantize::unpack(quantize::quantize_row(
spec,
state.grad[row],
state.hess[row],
state.train.row_offset + row as u64,
));
*g = qg as f32;
*h = qh as f32;
})
});
state.work.values.write(&state.quant_values)?;
state.quantized = true;
state.quant_ready = true;
Ok(())
}
fn reset_partition(&mut self, bag: &BagSpec) -> Result<u32> {
let pool = self.pool.clone();
let state = self.state()?;
let rows = state.train.num_rows as u32;
if state.bag.as_ref().is_none_or(|(cached, _)| cached != bag) {
let selected = install(&pool, || {
(0..rows)
.into_par_iter()
.filter(|&row| match *bag {
BagSpec::All => true,
BagSpec::Bernoulli { seed, round, fraction } => {
quantize::in_bag(seed, round, state.train.row_offset + u64::from(row), fraction)
}
})
.collect::<Vec<_>>()
});
state.bag = Some((*bag, selected));
}
let selected = &state.bag.as_ref().expect("bag initialized").1;
state.work.partition.write(selected)?;
state.partition_len = selected.len();
Ok(selected.len() as u32)
}
fn histogram(&mut self, leaf: LeafRange, quantized: bool, out: &mut [f64]) -> Result<()> {
let state = self.state()?;
check_range(state, leaf)?;
if out.len()
!= state
.train
.total_bins
.checked_mul(2)
.ok_or_else(|| error("histogram size overflow"))?
{
return Err(error("histogram output has wrong length"));
}
if quantized && !state.quant_ready {
return Err(error("quantized histogram requested before quantize"));
}
if quantized != state.quantized {
state.work.values.write(if quantized {
&state.quant_values
} else {
&state.host_values
})?;
state.quantized = quantized;
}
out.fill(0.0);
let n = state.datasets[0].capacity;
for feature in state.train.features.iter().filter(|f| f.num_bins != 0) {
let column = state.datasets[0].group_columns[feature.group as usize] + feature.column as usize;
for bin in (0..feature.num_bins as usize).step_by(HIST_BINS) {
for start in (0..leaf.count as usize).step_by(CHUNKS * TILE) {
let count = (leaf.count as usize - start).min(CHUNKS * TILE);
state.work.meta.write(&[
(leaf.begin as usize + start) as u32,
count as u32,
(column * n) as u32,
bin as u32,
0,
0,
0,
0,
])?;
state.work.histogram.execute()?;
state.work.partials.read(&mut state.work.host_partials)?;
for chunk in 0..CHUNKS {
for lane in 0..HIST_BINS.min(feature.num_bins as usize - bin) {
for component in 0..2 {
let value = state.work.host_partials[(chunk * 2 + component) * HIST_BINS + lane];
if quantized
&& (!value.is_finite()
|| value.fract() != 0.0
|| value.abs() > (TILE * 32767) as f32)
{
return Err(error("invalid exact Neuron quantized histogram partial"));
}
out[2 * (feature.hist_offset as usize + bin + lane) + component] += f64::from(value);
}
}
}
}
}
}
Ok(())
}
fn partition(&mut self, leaf: LeafRange, split: &BinSplit) -> Result<u32> {
let state = self.state()?;
check_range(state, leaf)?;
let group = state
.train
.groups
.get(split.group as usize)
.ok_or_else(|| error("split references missing group"))?;
if split.column >= group.row_stride {
return Err(error("split column outside group"));
}
if leaf.count == 0 {
return Ok(0);
}
let n = state.datasets[0].capacity;
let column = state.datasets[0].group_columns[split.group as usize] + split.column as usize;
state.work.meta.write(&[
leaf.begin,
leaf.count,
(column * n) as u32,
state.partition_len as u32,
0,
0,
0,
0,
])?;
state.work.split.write(&split.left_bins)?;
state.work.partition_flags.execute()?;
state.work.counts.read(&mut state.work.host_counts)?;
let left: u32 = state.work.host_counts.iter().try_fold(0u32, |s, &c| {
s.checked_add(c).ok_or_else(|| error("partition count overflow"))
})?;
if left > leaf.count {
return Err(error("invalid Neuron partition count"));
}
let blocks = n / TILE;
let mut l = leaf.begin;
let mut r = leaf.begin + left;
for block in 0..blocks {
state.work.host_offsets[block] = l;
state.work.host_offsets[blocks + block] = r;
let begin = (block * TILE).max(leaf.begin as usize);
let end = ((block + 1) * TILE).min(leaf.begin as usize + leaf.count as usize);
let count = end.saturating_sub(begin) as u32;
let lc = state.work.host_counts[block];
if lc > count {
return Err(error("invalid Neuron tile partition count"));
}
l += lc;
r += count - lc;
}
state.work.offsets.write(&state.work.host_offsets)?;
state.work.partition_scatter.execute()?;
state.work.copy_partition.execute()?;
Ok(left)
}
fn leaf_sums(&mut self, leaves: &[LeafRange], grad_scale: f32, hess_scale: f32) -> Result<Vec<[i64; 2]>> {
let pool = self.pool.clone();
let state = self.state()?;
for &leaf in leaves {
check_range(state, leaf)?;
}
let n = state.datasets[0].capacity;
install(&pool, || {
state
.host_limbs
.par_chunks_mut(4 * n)
.enumerate()
.for_each(|(component, out)| {
let (values, scale) = if component == 0 {
(&state.grad, grad_scale)
} else {
(&state.hess, hess_scale)
};
let (lo, rest) = out.split_at_mut(n);
let (mid, rest) = rest.split_at_mut(n);
let (hi, sign) = rest.split_at_mut(n);
lo.par_iter_mut()
.zip(mid.par_iter_mut())
.zip(hi.par_iter_mut())
.zip(sign.par_iter_mut())
.zip(values.par_iter())
.for_each(|((((lo, mid), hi), sign), &value)| {
let mut digits = [0; 4];
encode_fixed(quantize::to_fixed(value, scale), &mut digits);
[*lo, *mid, *hi, *sign] = digits;
});
})
});
state.work.limbs.write(&state.host_limbs)?;
let mut result = Vec::with_capacity(leaves.len());
for leaf in leaves {
let mut sums = [0i128; 2];
for start in (0..leaf.count as usize).step_by(CHUNKS * TILE) {
state.work.meta.write(&[
(leaf.begin as usize + start) as u32,
(leaf.count as usize - start).min(CHUNKS * TILE) as u32,
0,
0,
0,
0,
0,
0,
])?;
state.work.leaf_sum.execute()?;
state.work.leaf_partials.read(&mut state.work.host_leaf_partials)?;
for chunk in 0..CHUNKS {
for (component, sum) in sums.iter_mut().enumerate() {
let offset = chunk * LIMBS + component * 4;
*sum += fold_fixed(&state.work.host_leaf_partials[offset..offset + 4])?;
}
}
}
result.push([
i64::try_from(sums[0]).map_err(|_| error("fixed-point gradient sum exceeds i64"))?,
i64::try_from(sums[1]).map_err(|_| error("fixed-point hessian sum exceeds i64"))?,
]);
}
Ok(result)
}
fn add_tree(&mut self, tree: &DeviceTree) -> Result<()> {
let device = self.device.clone();
let state = self.state()?;
if tree.leaf_values.is_empty() || tree.nodes.len().checked_add(1) != Some(tree.leaf_values.len()) {
return Err(error("malformed device tree"));
}
let count = tree
.nodes
.len()
.checked_add(tree.leaf_values.len())
.ok_or_else(|| error("Neuron tree size overflow"))?;
let capacity = count
.checked_next_power_of_two()
.ok_or_else(|| error("Neuron tree capacity overflow"))?;
let words = capacity
.checked_mul(NODE_FIELDS)
.filter(|&words| words <= i32::MAX as usize)
.ok_or_else(|| error("Neuron tree exceeds signed 32-bit gather addressing"))?;
let mut nodes = vec![0u32; words];
let mut values = vec![0f32; capacity];
let internal = tree.nodes.len();
let mut depths = vec![0usize; internal];
let mut depth = 0usize;
for (i, node) in tree.nodes.iter().enumerate() {
let group = state
.train
.groups
.get(node.split.group as usize)
.ok_or_else(|| error("tree references missing group"))?;
if node.split.column >= group.row_stride {
return Err(error("tree split column outside group"));
}
nodes[i] = (state.datasets[0].group_columns[node.split.group as usize] + node.split.column as usize) as u32;
for (field, child) in [(1, node.left), (2, node.right)] {
let child = if child >= 0 {
let child = child as usize;
if child <= i || child >= internal {
return Err(error("tree children must follow their parent"));
}
depths[child] = depths[i] + 1;
child
} else {
let leaf = (!child) as usize;
if leaf >= tree.leaf_values.len() {
return Err(error("invalid tree leaf index"));
}
depth = depth.max(depths[i] + 1);
internal + leaf
};
nodes[field * capacity + i] = child as u32;
}
for word in 0..8 {
nodes[(3 + word) * capacity + i] = node.split.left_bins[word];
}
}
for (leaf, &value) in tree.leaf_values.iter().enumerate() {
let i = internal + leaf;
nodes[capacity + i] = i as u32;
nodes[2 * capacity + i] = i as u32;
values[i] = value;
}
if state.tree.as_ref().is_none_or(|work| work.capacity != capacity) {
let node_buffer = device.upload(&nodes)?;
let leaf_buffer = device.upload(&values)?;
let meta = device.upload(&[0u32; 8])?;
let mut executions = Vec::with_capacity(state.datasets.len());
for data in &state.datasets {
let add = bind(
&device,
"add_tree",
json!({"ROWS":data.capacity,"COUNT":data.rows,"TILE":TILE,"NODES":capacity}),
vec![
spec("bins", DType::I32, &[data.columns * data.capacity]),
spec("scores", DType::F32, &[data.capacity]),
spec("nodes", DType::I32, &[NODE_FIELDS, capacity]),
spec("leaves", DType::F32, &[capacity]),
spec("meta", DType::I32, &[8]),
],
vec![spec("result", DType::F32, &[data.capacity])],
&[
data.bins.clone(),
data.scores.clone(),
node_buffer.clone(),
leaf_buffer.clone(),
meta.clone(),
],
std::slice::from_ref(&data.next_scores),
)?;
let copy = bind(
&device,
"copy_f32",
json!({"ROWS":data.capacity,"TILE":TILE}),
vec![spec("source", DType::F32, &[data.capacity])],
vec![spec("result", DType::F32, &[data.capacity])],
std::slice::from_ref(&data.next_scores),
std::slice::from_ref(&data.scores),
)?;
executions.push((add, copy));
}
state.tree = Some(TreeWork {
capacity,
nodes: node_buffer,
leaves: leaf_buffer,
meta,
executions,
});
}
let work = state.tree.as_mut().expect("tree workspace initialized");
work.nodes.write(&nodes)?;
work.leaves.write(&values)?;
work.meta.write(&[depth as u32, 0, 0, 0, 0, 0, 0, 0])?;
for (add, copy) in &mut work.executions {
add.execute()?;
copy.execute()?;
}
Ok(())
}
fn scores(&mut self, dataset: usize) -> Result<Vec<f32>> {
let data = self
.state()?
.datasets
.get(dataset)
.ok_or_else(|| error("missing score dataset"))?;
let mut scores = vec![0f32; data.rows];
data.scores.read(&mut scores)?;
Ok(scores)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn signed_limbs_cover_i64_edges() {
for value in [i64::MIN, i64::MIN + 1, -65537, -65536, -1, 0, 1, 65535, 65536, i64::MAX] {
let mut digits = [0; 4];
encode_fixed(value, &mut digits);
assert_eq!(fold_fixed(&digits.map(|x| x as f32)).unwrap(), i128::from(value));
}
}
#[test]
fn worst_case_tile_limb_reductions_are_exact() {
for value in [i64::MIN, i64::MAX, -1, 1, (1i64 << 30) - 1, -(1i64 << 30)] {
let mut digits = [0; 4];
encode_fixed(value, &mut digits);
let sums = digits.map(|digit| {
let mut sum = 0f32;
for _ in 0..TILE {
sum += digit as f32;
}
sum
});
assert_eq!(fold_fixed(&sums).unwrap(), i128::from(value) * TILE as i128);
}
}
#[test]
fn mixed_signed_limb_fold_matches_wide_sum() {
let values = [i64::MAX, i64::MIN, 65535, -65536, 123456789012345, -123456789012345];
let mut sums = [0f32; 4];
let mut exact = 0i128;
for value in values {
let mut digits = [0; 4];
encode_fixed(value, &mut digits);
for k in 0..4 {
sums[k] += digits[k] as f32;
}
exact += i128::from(value);
}
assert_eq!(fold_fixed(&sums).unwrap(), exact);
}
fn dataset(rows: usize) -> BinnedDataset {
use silvic_core::data::{BinMapper, FeatureGroup, FeatureMeta};
let numerical = BinMapper::Numerical {
upper_bounds: (0..254)
.map(|i| i as f64)
.chain(std::iter::once(f64::INFINITY))
.collect(),
nan_bin: true,
min: 0.0,
max: 254.0,
};
BinnedDataset {
num_rows: rows,
row_offset: (1u64 << 32) + 113,
features: vec![
FeatureMeta {
name: "category".into(),
bins: BinMapper::Categorical {
categories: (1..256).collect(),
},
num_bins: 256,
group: 0,
column: 0,
hist_offset: 0,
},
FeatureMeta {
name: "numeric_nan".into(),
bins: numerical,
num_bins: 256,
group: 0,
column: 1,
hist_offset: 256,
},
FeatureMeta {
name: "other_group".into(),
bins: BinMapper::Categorical {
categories: vec![1, 2, 3, 4],
},
num_bins: 5,
group: 1,
column: 0,
hist_offset: 512,
},
FeatureMeta {
name: "trivial".into(),
bins: BinMapper::Trivial,
num_bins: 0,
group: 0,
column: 0,
hist_offset: 517,
},
],
groups: vec![
FeatureGroup {
features: vec![0, 1],
row_stride: 4,
hist_offset: 0,
total_bins: 512,
data: (0..rows)
.flat_map(|r| [(r % 256) as u8, ((r * 13 + 7) % 256) as u8, 0, 0])
.collect(),
},
FeatureGroup {
features: vec![2],
row_stride: 4,
hist_offset: 512,
total_bins: 5,
data: (0..rows).flat_map(|r| [(r % 5) as u8, 0, 0, 0]).collect(),
},
],
total_bins: 517,
label: (0..rows).map(|r| (r % 3) as f32).collect(),
weight: Some((0..rows).map(|r| 0.25 + (r % 4) as f32 * 0.25).collect()),
query_boundaries: Some(
(0..rows as u32)
.step_by(31)
.chain(std::iter::once(rows as u32))
.collect(),
),
}
}
#[test]
fn device_training_matches_core_across_objectives_and_modes() {
if std::env::var("SILVIC_TEST_NEURON").as_deref() != Ok("1") {
eprintln!("skipped Neuron device training regression; set SILVIC_TEST_NEURON=1 on real hardware");
return;
}
let data = dataset(4099);
let valid = dataset(259);
let objectives = [
ObjectiveSetup::Regression,
ObjectiveSetup::Binary { sigmoid: 1.3 },
ObjectiveSetup::LambdaRank {
sigmoid: 1.0,
norm: true,
truncation_level: 10,
label_gain: vec![0.0, 1.0, 3.0],
inverse_max_dcg: vec![1.0; data.query_boundaries.as_ref().unwrap().len() - 1],
},
];
for objective in objectives {
let mut data = data.clone();
if matches!(&objective, ObjectiveSetup::Regression) {
data.weight = None;
}
let mut neuron = NeuronBackend::new(None).expect("real Neuron device");
let mut cpu = silvic_core::CpuBackend::new();
neuron.load(&data, &[&valid], &objective).unwrap();
cpu.load(&data, &[&valid], &objective).unwrap();
for backend in [&mut neuron as &mut dyn Backend, &mut cpu as &mut dyn Backend] {
backend.init_scores(0, 0.375).unwrap();
backend.init_scores(1, -0.125).unwrap();
}
let stats = neuron.compute_gradients().unwrap();
let expected_stats = cpu.compute_gradients().unwrap();
assert_eq!(stats.max_abs_grad, expected_stats.max_abs_grad);
assert_eq!(stats.max_hess, expected_stats.max_hess);
let constant_hessian = objective::constant_hessian(&objective, data.weight.as_deref());
let quant = QuantSpec::new(stats.max_abs_grad, stats.max_hess, constant_hessian, 255, true, 71, 4);
neuron.quantize(&quant).unwrap();
cpu.quantize(&quant).unwrap();
let bag = BagSpec::Bernoulli {
seed: 19,
round: 3,
fraction: 0.73,
};
let count = neuron.reset_partition(&bag).unwrap();
assert_eq!(count, cpu.reset_partition(&bag).unwrap());
let leaf = LeafRange {
begin: 117,
count: count - 233,
};
let mut actual = vec![0.0; 2 * data.total_bins];
let mut expected = actual.clone();
for quantized in [true, false, true] {
neuron.histogram(leaf, quantized, &mut actual).unwrap();
cpu.histogram(leaf, quantized, &mut expected).unwrap();
for (i, (&a, &e)) in actual.iter().zip(&expected).enumerate() {
if quantized {
assert_eq!(a, e, "quantized bin {i}");
} else {
assert!((a - e).abs() <= 2e-5 * (1.0 + e.abs()), "float bin {i}: {a} vs {e}");
}
}
}
let mut left_bins = [0u32; 8];
left_bins[0] = 1;
left_bins[7] = 0x8000_0001;
let split = BinSplit {
group: 0,
column: 1,
left_bins,
};
let left = neuron.partition(leaf, &split).unwrap();
assert_eq!(left, cpu.partition(leaf, &split).unwrap());
let selected: Vec<u32> = (0..data.num_rows as u32)
.filter(|&r| quantize::in_bag(19, 3, data.row_offset + u64::from(r), 0.73))
.collect();
let range = leaf.begin as usize..(leaf.begin + leaf.count) as usize;
let mut stable = selected.clone();
let middle = &selected[range.clone()];
let reordered: Vec<u32> = middle
.iter()
.copied()
.filter(|&r| split.goes_left(data.bin(1, r as usize)))
.chain(
middle
.iter()
.copied()
.filter(|&r| !split.goes_left(data.bin(1, r as usize))),
)
.collect();
stable[range].copy_from_slice(&reordered);
let mut actual_rows = vec![0u32; count as usize];
neuron
.loaded
.as_ref()
.unwrap()
.work
.partition
.read(&mut actual_rows)
.unwrap();
assert_eq!(
actual_rows, stable,
"partition must be stable and preserve rows outside the leaf"
);
let leaves = [
LeafRange {
begin: leaf.begin,
count: left,
},
LeafRange {
begin: leaf.begin + left,
count: leaf.count - left,
},
LeafRange { begin: count, count: 0 },
];
let gs = quantize::fixed_point_scale(stats.max_abs_grad);
let hs = quantize::fixed_point_scale(stats.max_hess);
assert_eq!(
neuron.leaf_sums(&leaves, gs, hs).unwrap(),
cpu.leaf_sums(&leaves, gs, hs).unwrap()
);
let tree = DeviceTree {
nodes: vec![silvic_core::backend::DeviceNode {
split,
left: !0,
right: !1,
}],
leaf_values: vec![0.25, -0.125],
};
for backend in [&mut neuron as &mut dyn Backend, &mut cpu as &mut dyn Backend] {
backend.add_tree(&tree).unwrap();
}
for dataset in 0..2 {
assert_eq!(neuron.scores(dataset).unwrap(), cpu.scores(dataset).unwrap());
}
let constant = DeviceTree {
nodes: vec![],
leaf_values: vec![0.0625],
};
for backend in [&mut neuron as &mut dyn Backend, &mut cpu as &mut dyn Backend] {
backend.add_tree(&constant).unwrap();
}
for dataset in 0..2 {
assert_eq!(neuron.scores(dataset).unwrap(), cpu.scores(dataset).unwrap());
}
assert!(neuron
.histogram(LeafRange { begin: count, count: 1 }, true, &mut actual)
.is_err());
assert!(neuron
.partition(
leaf,
&BinSplit {
group: 99,
column: 0,
left_bins: [0; 8]
}
)
.is_err());
}
}
}