logicaffeine_compile/codegen_dram/
kernels.rs

1//! DRAM application kernels — the workloads in-memory compute is lightning-fast at: bitwise + COUNT
2//! over massive data, at zero data movement. Popcount is the linchpin — the reduction that turns the
3//! bitwise/XNOR machinery into counts, distances, and dot products:
4//!
5//!   • [`popcount`] — per-lane population count, folded bit-serially into an in-DRAM counter;
6//!   • [`hamming_distance`] = popcount(a XOR b) — vector similarity search (ANN over binary codes),
7//!     binary-net inference, bioinformatics distance;
8//!   • [`agreement_count`] = popcount(a XNOR b) — the 1-bit / binary neural-net DOT PRODUCT (for ±1
9//!     codes, dot = 2·agreements − K). This is how a BitNet-style matmul runs with NO multiply.
10//!
11//! All are composed from the `{And, Or, Not}` basis over a bit-serial (bitplane) layout: bitplane j
12//! is a row whose bit `i` is bit `j` of lane `i`'s value, so one kernel processes all lanes in
13//! parallel — exactly the data-parallel, bandwidth-bound shape the DRAM plane wins on.
14
15use super::dataflow::{bit_serial_sub, DramDataflow};
16
17/// Add a single-bit plane `bit` (0/1 per lane) into a per-lane counter `counter` (W planes, LSB
18/// first) — a half-adder ripple. Overflow beyond W is dropped, so size W ≥ ⌈log2(max_count+1)⌉.
19fn increment_by_bit(df: &mut DramDataflow, counter: &[usize], bit: usize) -> Vec<usize> {
20    let mut carry = bit;
21    let mut out = Vec::with_capacity(counter.len());
22    for &c in counter {
23        out.push(df.xor(c, carry)); // sum bit
24        carry = df.and(c, carry); // carry out
25    }
26    out
27}
28
29/// Per-lane POPCOUNT of `planes`: for each lane, the number of set bits across the planes, as an
30/// `out_width`-plane per-lane integer (LSB first). A reduction realized ENTIRELY in the DRAM plane by
31/// folding each plane into a bit-serial counter. `zero` is a zero row (the initial counter value).
32pub fn popcount(
33    df: &mut DramDataflow,
34    planes: &[usize],
35    out_width: usize,
36    zero: usize,
37) -> Vec<usize> {
38    let mut counter: Vec<usize> = vec![zero; out_width];
39    for &p in planes {
40        counter = increment_by_bit(df, &counter, p);
41    }
42    counter
43}
44
45/// Per-lane HAMMING DISTANCE between two bit-serial codes `a`/`b` (K planes each): popcount(a XOR b).
46/// The core kernel of vector similarity search (ANN over binary codes) and binary-net inference.
47pub fn hamming_distance(
48    df: &mut DramDataflow,
49    a: &[usize],
50    b: &[usize],
51    out_width: usize,
52    zero: usize,
53) -> Vec<usize> {
54    let xored: Vec<usize> = a.iter().zip(b).map(|(&ai, &bi)| df.xor(ai, bi)).collect();
55    popcount(df, &xored, out_width, zero)
56}
57
58/// Per-lane AGREEMENT COUNT popcount(a XNOR b) — the binary / 1-bit neural-net DOT PRODUCT. For ±1
59/// codes packed as bits, the true dot product is `2·agreements − K`; the agreement count is the
60/// in-DRAM half (XNOR + popcount, no multiply), the final affine map is a scalar host step.
61pub fn agreement_count(
62    df: &mut DramDataflow,
63    a: &[usize],
64    b: &[usize],
65    out_width: usize,
66    zero: usize,
67) -> Vec<usize> {
68    let xnored: Vec<usize> = a.iter().zip(b).map(|(&ai, &bi)| df.xnor(ai, bi)).collect();
69    popcount(df, &xnored, out_width, zero)
70}
71
72/// BitNet-style TERNARY·BINARY dot product — one matmul row, entirely in-DRAM, NO multiply. Ternary
73/// weights `w[k]` ∈ {−1,0,+1} are carried as two disjoint bitplane sets: `wp[k]`=1 iff w=+1, `wn[k]`=1
74/// iff w=−1. Binary activations `x[k]` ∈ {0,1}. Then, per lane (per output neuron),
75///   `dot = Σ w[k]·x[k] = popcount(Wp AND X) − popcount(Wn AND X)`
76/// as a two's-complement `out_width`-plane integer. `zero` is a zero row, `ones` an all-ones carry
77/// row (for the subtract). This is how a 1.58-bit LLM layer runs where the weights live: AND +
78/// popcount + subtract, no multiplier.
79pub fn ternary_binary_dot(
80    df: &mut DramDataflow,
81    wp: &[usize],
82    wn: &[usize],
83    x: &[usize],
84    out_width: usize,
85    zero: usize,
86    ones: usize,
87) -> Vec<usize> {
88    let wpx: Vec<usize> = wp.iter().zip(x).map(|(&w, &xi)| df.and(w, xi)).collect();
89    let wnx: Vec<usize> = wn.iter().zip(x).map(|(&w, &xi)| df.and(w, xi)).collect();
90    let pos = popcount(df, &wpx, out_width, zero);
91    let neg = popcount(df, &wnx, out_width, zero);
92    bit_serial_sub(df, &pos, &neg, ones)
93}
94
95#[cfg(test)]
96#[path = "../../../logicaffeine_tests/unit/logicaffeine_compile/codegen_dram/kernels/tests.rs"]
97mod tests;
logicaffeine_compile/codegen_dram/kernels.rs

logicaffeine_compile/codegen_dram/
kernels.rs
