handwired_rs/
batch.rs

1//! A CPU executor for many records at once.
2//!
3//! Values are stored per wire as a block of [`LANES`] records, so every row of
4//! a layer is a loop over records that the compiler can vectorize. Each lane
5//! still computes `bias + w0 × x0 + w1 × x1 + ...` in the stored order with
6//! FP32 rounding after every step (Rust never fuses or regroups float
7//! arithmetic), so every value is bit for bit the one [`TensorExecutor`]
8//! computes. [`BatchExecutor::run_parallel`] splits the records across threads.
9//!
10//! [`TensorExecutor`]: crate::TensorExecutor
11
12use crate::circuit::{OutOfRange, resolve_index_position};
13use crate::numeric::{evaluate_boolean, relu32};
14use crate::tensor::{GatherCompute, GatherRead, TensorProgram, resolve_gather};
15
16/// Records evaluated together.
17pub const LANES: usize = 64;
18
19type Block = [f32; LANES];
20
21struct Layer {
22    projection_size: usize,
23    output_size: usize,
24    bias: Vec<f32>,
25    row_offsets: Vec<u32>,
26    /// Packed value positions, read directly instead of through a gathered input.
27    positions: Vec<u32>,
28    weights: Vec<f32>,
29    gate_rows: Vec<u32>,
30    boolean_rows: Vec<u32>,
31    gather_computes: Vec<GatherCompute>,
32    gather_reads: Vec<GatherRead>,
33    index_reads: Vec<IndexRead>,
34}
35
36struct IndexRead {
37    row: usize,
38    index_row: usize,
39    below: OutOfRange,
40    above: OutOfRange,
41    candidate_positions: Vec<u32>,
42}
43
44/// Reusable buffers for running one program on batches of records.
45pub struct BatchExecutor<'program> {
46    program: &'program TensorProgram,
47    layers: Vec<Layer>,
48    root_positions: Vec<usize>,
49    values: Vec<Block>,
50    projection: Vec<Block>,
51    gather_indices: Vec<[i32; LANES]>,
52    gather_masses: Vec<Block>,
53}
54
55impl<'program> BatchExecutor<'program> {
56    pub fn new(program: &'program TensorProgram) -> BatchExecutor<'program> {
57        let layers: Vec<Layer> = program
58            .layers
59            .iter()
60            .map(|layer| Layer {
61                projection_size: layer.bias.len(),
62                output_size: layer.output_size,
63                bias: layer.bias.clone(),
64                row_offsets: layer.row_offsets.clone(),
65                positions: layer.columns.iter().map(|&column| layer.input_indices[column as usize]).collect(),
66                weights: layer.weights.clone(),
67                gate_rows: layer.gate_rows.clone(),
68                boolean_rows: layer.boolean_rows.clone(),
69                gather_computes: layer.gather_computes.clone(),
70                gather_reads: layer.gather_reads.clone(),
71                index_reads: layer
72                    .index_reads
73                    .iter()
74                    .map(|read| IndexRead {
75                        row: read.row as usize,
76                        index_row: read.index_row as usize,
77                        below: if read.below_zero { OutOfRange::Zero } else { OutOfRange::Clamp },
78                        above: if read.above_zero { OutOfRange::Zero } else { OutOfRange::Clamp },
79                        candidate_positions: read.candidate_positions.clone(),
80                    })
81                    .collect(),
82            })
83            .collect();
84        let max_projection = layers.iter().map(|layer| layer.projection_size).max().unwrap_or(0);
85        let root_positions = program
86            .roots
87            .iter()
88            .map(|root| program.wire_order.iter().position(|wire| wire == root).expect("root was removed"))
89            .collect();
90        BatchExecutor {
91            program,
92            layers,
93            root_positions,
94            values: vec![[0.0; LANES]; program.value_count],
95            projection: vec![[0.0; LANES]; max_projection],
96            gather_indices: vec![[0; LANES]; program.gather_slot_count],
97            gather_masses: vec![[0.0; LANES]; program.gather_slot_count],
98        }
99    }
100
101    /// Evaluate the program on consecutive source tables, `program.source_size`
102    /// values each, and return the roots' values, `program.roots.len()` per
103    /// record. Sources aren't range-checked; see [`Circuit::validate_source_ranges`].
104    ///
105    /// [`Circuit::validate_source_ranges`]: crate::Circuit::validate_source_ranges
106    pub fn run(&mut self, sources: &[f32]) -> Vec<f32> {
107        let mut outputs = Vec::new();
108        self.run_into(sources, &mut outputs);
109        outputs
110    }
111
112    /// [`run`](Self::run), appending to `outputs`.
113    pub fn run_into(&mut self, sources: &[f32], outputs: &mut Vec<f32>) {
114        let source_size = self.program.source_size;
115        let root_count = self.program.roots.len();
116        let records = sources.len().checked_div(source_size).unwrap_or(0);
117        assert_eq!(records * source_size, sources.len(), "sources aren't whole source tables");
118        outputs.reserve(records * root_count);
119        for first in (0..records).step_by(LANES) {
120            let count = (records - first).min(LANES);
121            for cell in 0..source_size {
122                let block = &mut self.values[cell];
123                for (lane, value) in block.iter_mut().enumerate() {
124                    // A partial batch repeats its last record in the unused lanes.
125                    *value = sources[(first + lane.min(count - 1)) * source_size + cell];
126                }
127            }
128            self.evaluate();
129            for lane in 0..count {
130                outputs.extend(self.root_positions.iter().map(|&position| self.values[position][lane]));
131            }
132        }
133    }
134
135    fn evaluate(&mut self) {
136        #[cfg(target_arch = "x86_64")]
137        if std::arch::is_x86_feature_detected!("avx2") {
138            // SAFETY: the CPU supports AVX2.
139            return unsafe { self.evaluate_avx2() };
140        }
141        self.evaluate_portable();
142    }
143
144    /// The same code with 256-bit vectors. Only AVX2 is enabled, not FMA: Rust never contracts `a + w * x` into a
145    /// fused multiply-add anyway, but this keeps it out of reach.
146    #[cfg(target_arch = "x86_64")]
147    #[target_feature(enable = "avx2")]
148    unsafe fn evaluate_avx2(&mut self) {
149        self.evaluate_portable();
150    }
151
152    // Lane loops index several blocks at once.
153    #[allow(clippy::needless_range_loop)]
154    #[inline(always)]
155    fn evaluate_portable(&mut self) {
156        let values = &mut self.values;
157        let mut available = self.program.source_size;
158        for layer in &self.layers {
159            let projection = &mut self.projection[..layer.projection_size];
160            for (row, block) in projection.iter_mut().enumerate() {
161                *block = [layer.bias[row]; LANES];
162                let entries = layer.row_offsets[row] as usize..layer.row_offsets[row + 1] as usize;
163                for (&position, &weight) in layer.positions[entries.clone()].iter().zip(&layer.weights[entries]) {
164                    let input = &values[position as usize];
165                    for lane in 0..LANES {
166                        block[lane] += weight * input[lane];
167                    }
168                }
169            }
170            for compute in &layer.gather_computes {
171                let slot = compute.slot as usize;
172                for lane in 0..LANES {
173                    let bits = compute.bit_positions.iter().map(|&bit| values[bit as usize][lane]);
174                    (self.gather_indices[slot][lane], self.gather_masses[slot][lane]) = resolve_gather(bits);
175                }
176            }
177            for read in &layer.gather_reads {
178                let slot = read.slot as usize;
179                let limit = read.candidate_positions.len() as i32 - 1;
180                for lane in 0..LANES {
181                    // As in `TensorExecutor`: the clamp keeps a violated declaration inside the table.
182                    let position = read.candidate_positions[self.gather_indices[slot][lane].clamp(0, limit) as usize];
183                    let candidate = values[position as usize][lane];
184                    let contribution = if read.masked { self.gather_masses[slot][lane] * candidate } else { candidate };
185                    projection[read.row as usize][lane] += contribution;
186                }
187            }
188            for read in &layer.index_reads {
189                for lane in 0..LANES {
190                    let index = projection[read.index_row][lane] as f64;
191                    let count = read.candidate_positions.len();
192                    projection[read.row][lane] = match resolve_index_position(index, count, read.below, read.above) {
193                        -2 => f32::NAN,
194                        -1 => 0.0,
195                        position => values[read.candidate_positions[position as usize] as usize][lane],
196                    };
197                }
198            }
199            for (gate, &row) in layer.gate_rows.iter().enumerate() {
200                let gate_block = projection[layer.output_size + gate];
201                let block = &mut projection[row as usize];
202                for lane in 0..LANES {
203                    block[lane] *= relu32(gate_block[lane]);
204                }
205            }
206            for &row in &layer.boolean_rows {
207                for value in projection[row as usize].iter_mut() {
208                    *value = evaluate_boolean(*value);
209                }
210            }
211            values[available..available + layer.output_size].copy_from_slice(&projection[..layer.output_size]);
212            available += layer.output_size;
213        }
214    }
215
216    /// [`BatchExecutor::run`] split across `threads` threads (all available cores
217    /// when `None`), with the outputs in record order.
218    pub fn run_parallel(program: &TensorProgram, sources: &[f32], threads: Option<usize>) -> Vec<f32> {
219        let source_size = program.source_size.max(1);
220        let records = sources.len() / source_size;
221        let threads = threads
222            .unwrap_or_else(|| std::thread::available_parallelism().map_or(1, |count| count.get()))
223            .clamp(1, records.div_ceil(LANES).max(1));
224        // Whole batches per thread, so only the last batch can be partial.
225        let records_per_thread = records.div_ceil(threads).div_ceil(LANES) * LANES;
226        if threads == 1 || records_per_thread == 0 {
227            return BatchExecutor::new(program).run(sources);
228        }
229        let chunks: Vec<Vec<f32>> = std::thread::scope(|scope| {
230            let handles: Vec<_> = sources
231                .chunks(records_per_thread * source_size)
232                .map(|chunk| scope.spawn(move || BatchExecutor::new(program).run(chunk)))
233                .collect();
234            handles.into_iter().map(|handle| handle.join().expect("a batch thread panicked")).collect()
235        });
236        chunks.concat()
237    }
238}
handwired_rs/batch.rs

handwired_rs/
batch.rs
