kv_store/
kv_store.rs

1//! A fixed-capacity key-value store whose every operation is an exact circuit:
2//! the store's state goes in, the next state and the answer come out.
3//!
4//! A `get` compares the key with every slot and picks the matching value,
5//! which is hard attention: a one-hot match over keys, multiplied into the
6//! values. A `put` overwrites the matching slot, or fills the first empty one,
7//! as a select per slot. The function below is plain Rust, and the circuit is
8//! checked against it.
9//!
10//! cargo run --release -p handwired-rs-gpu --example kv_store
11
12use std::time::Instant;
13
14use handwired_rs::TensorExecutor;
15use handwired_rs::dsl::Params;
16use handwired_rs::handwired;
17use handwired_rs_gpu::{Gpu, GpuBatch};
18
19const SLOTS: usize = 16;
20const GET: u8 = 0;
21const PUT: u8 = 1;
22const DELETE: u8 = 2;
23
24/// One operation on the store. A key of 0 marks an empty slot. Returns the
25/// next keys and values, whether the key was present, and its value before
26/// the operation.
27///
28/// Keys are declared within 0..15 here so that random tests hit present keys
29/// often; any range within u16 works.
30#[handwired]
31fn kv_step(
32    #[range(0, 15)] keys: [u16; 16],
33    values: [u16; 16],
34    #[range(0, 2)] op: u8,
35    #[range(1, 15)] key: u16,
36    value: u16,
37) -> ([u16; 16], [u16; 16], bool, u16) {
38    let mut keys = keys;
39    let mut values = values;
40    let mut found = false;
41    let mut previous: u16 = 0;
42    for slot in 0..16 {
43        if keys[slot] == key {
44            found = true;
45            previous = values[slot];
46            if op == PUT {
47                values[slot] = value;
48            }
49            if op == DELETE {
50                keys[slot] = 0;
51                values[slot] = 0;
52            }
53            break;
54        }
55    }
56    if op == PUT && !found {
57        // Insert into the first empty slot; a full store drops the write.
58        for slot in 0..16 {
59            if keys[slot] == 0 {
60                keys[slot] = key;
61                values[slot] = value;
62                break;
63            }
64        }
65    }
66    (keys, values, found, previous)
67}
68
69type State = ([u16; SLOTS], [u16; SLOTS]);
70
71struct Rng(u64);
72impl Rng {
73    fn next(&mut self) -> u64 {
74        self.0 ^= self.0 << 13;
75        self.0 ^= self.0 >> 7;
76        self.0 ^= self.0 << 17;
77        self.0
78    }
79    fn below(&mut self, bound: u64) -> u64 {
80        self.next() % bound
81    }
82}
83
84fn main() {
85    let start = Instant::now();
86    let compiled = kv_step_circuit();
87    println!(
88        "kv_step: {} live operations in {} layers, built in {:.1?}",
89        compiled.program.live_operation_count(),
90        compiled.program.layers.len(),
91        start.elapsed()
92    );
93    let report = compiled.verify(20_000).unwrap();
94    println!("verified against the function on {} sampled inputs", report.checked);
95
96    // A session: the circuit holds the store's state from one operation to the next, like a database.
97    let mut executor = TensorExecutor::new(&compiled.program);
98    let mut circuit_state: State = ([0; SLOTS], [0; SLOTS]);
99    let mut native_state: State = ([0; SLOTS], [0; SLOTS]);
100    let script = [(PUT, 3, 300), (PUT, 7, 700), (GET, 3, 0), (PUT, 3, 333), (GET, 3, 0), (DELETE, 7, 0), (GET, 7, 0)];
101    let mut source = Vec::new();
102    for (op, key, value) in script {
103        source.clear();
104        (circuit_state.0, circuit_state.1, op, key, value).encode(&mut source);
105        let outputs = executor.run(&source);
106        let keys: [u16; SLOTS] = std::array::from_fn(|slot| outputs[slot] as u16);
107        let values: [u16; SLOTS] = std::array::from_fn(|slot| outputs[SLOTS + slot] as u16);
108        let (found, previous) = (outputs[2 * SLOTS] == 1.0, outputs[2 * SLOTS + 1] as u16);
109        circuit_state = (keys, values);
110        let native = kv_step(native_state.0, native_state.1, op, key, value);
111        native_state = (native.0, native.1);
112        assert_eq!((circuit_state, found, previous), (native_state, native.2, native.3));
113        let name = ["get", "put", "delete"][op as usize];
114        println!("  {name:6} {key:>2} {value:>4}  ->  found {found:5}, previous {previous}");
115    }
116
117    // 100,000 independent operations on 100,000 different stores, in one GPU dispatch.
118    let mut rng = Rng(0x9E37_79B9_7F4A_7C15);
119    let records: Vec<([u16; SLOTS], [u16; SLOTS], u8, u16, u16)> = (0..100_000)
120        .map(|_| {
121            let keys: [u16; SLOTS] = std::array::from_fn(|_| rng.below(16) as u16);
122            let values: [u16; SLOTS] = std::array::from_fn(|_| rng.below(65_536) as u16);
123            (keys, values, rng.below(3) as u8, 1 + rng.below(15) as u16, rng.below(65_536) as u16)
124        })
125        .collect();
126    let start = Instant::now();
127    let native: Vec<_> =
128        records.iter().map(|&(keys, values, op, key, value)| kv_step(keys, values, op, key, value)).collect();
129    let native_time = start.elapsed();
130    let gpu = Gpu::new().expect("a GPU adapter");
131    let batch = GpuBatch::new(&gpu, &compiled).unwrap();
132    let _ = batch.run(&gpu, &records[..1024]).unwrap();
133    let start = Instant::now();
134    let on_gpu = batch.run(&gpu, &records).unwrap();
135    let gpu_time = start.elapsed();
136    assert_eq!(on_gpu, native, "the GPU disagrees with the function");
137    println!("100,000 operations on separate stores; every result matches the function");
138    println!("plain Rust: {native_time:>9.1?}");
139    println!("GPU ({}, {:?}): {gpu_time:>9.1?}  (upload, run, read back)", gpu.info.name, gpu.info.device_type);
140}
kv_store/kv_store.rs

kv_store/
kv_store.rs
