sim_lib_numbers_tensor_bit/
bit_tensor.rs

1//! Bit-packed boolean tensor storage, its `SpecTensor` backend, and the
2//! library that registers it as the `bool` element-type backend.
3
4use sim_kernel::{
5    AbiVersion, DefaultFactory, Dependency, Export, Factory, Lib, LibManifest, LibTarget, Linker,
6    Result, Symbol, Value, Version,
7};
8use sim_lib_numbers_tensor::{
9    SpecTensor, SpecTensorDescriptor, Tensor, domains, element_count, spec_tensor_descriptor_value,
10    spec_tensor_symbol,
11};
12
13/// A boolean tensor stored as bit-packed `u64` words.
14///
15/// Each element occupies a single bit, so an `n`-element tensor is held in
16/// `ceil(n / 64)` words. The logical [`shape`](Self::shape) and element count
17/// drive layout; bitwise operations work directly on the packed words.
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct BitTensor {
20    shape: Vec<usize>,
21    len: usize,
22    words: Vec<u64>,
23}
24
25impl BitTensor {
26    /// Packs a slice of booleans into a bit tensor of the given shape.
27    ///
28    /// Returns `None` when `bits.len()` does not match the element count
29    /// implied by `shape`.
30    pub fn from_bools(shape: Vec<usize>, bits: &[bool]) -> Option<Self> {
31        let len = element_count(&shape);
32        if len != bits.len() {
33            return None;
34        }
35        let mut words = vec![0u64; len.div_ceil(64)];
36        for (index, bit) in bits.iter().enumerate() {
37            if *bit {
38                words[index / 64] |= 1u64 << (index % 64);
39            }
40        }
41        Some(Self { shape, len, words })
42    }
43
44    /// Unpacks the tensor back into one boolean per element, in flat order.
45    pub fn to_bools(&self) -> Vec<bool> {
46        (0..self.len)
47            .map(|index| ((self.words[index / 64] >> (index % 64)) & 1) == 1)
48            .collect()
49    }
50
51    /// Element-wise bitwise OR with another bit tensor of the same shape.
52    ///
53    /// Returns `None` when the shapes differ.
54    pub fn bit_or(&self, other: &Self) -> Option<Self> {
55        map_words(self, other, |left, right| left | right)
56    }
57
58    /// Element-wise bitwise XOR with another bit tensor of the same shape.
59    ///
60    /// Returns `None` when the shapes differ.
61    pub fn bit_xor(&self, other: &Self) -> Option<Self> {
62        map_words(self, other, |left, right| left ^ right)
63    }
64
65    /// Element-wise bitwise AND with another bit tensor of the same shape.
66    ///
67    /// Returns `None` when the shapes differ.
68    pub fn bit_and(&self, other: &Self) -> Option<Self> {
69        map_words(self, other, |left, right| left & right)
70    }
71}
72
73impl SpecTensor for BitTensor {
74    fn shape(&self) -> &[usize] {
75        &self.shape
76    }
77
78    fn dtype(&self) -> Symbol {
79        domains::bool()
80    }
81
82    fn to_uniform(&self) -> Tensor {
83        Tensor {
84            shape: self.shape.clone(),
85            dtype: self.dtype(),
86            data: self
87                .to_bools()
88                .into_iter()
89                .map(bool_value)
90                .collect::<Option<Vec<_>>>()
91                .expect("bool tensor values should always encode"),
92        }
93    }
94
95    fn from_uniform(tensor: &Tensor) -> Option<Self> {
96        let bits = tensor
97            .data
98            .iter()
99            .map(parse_bool_cell)
100            .collect::<Option<Vec<_>>>()?;
101        Self::from_bools(tensor.shape.clone(), &bits)
102    }
103}
104
105/// Registered library that installs the bit-packed boolean tensor backend.
106///
107/// Loading this [`Lib`] registers a [`SpecTensor`] descriptor binding the
108/// `bool` element type to the [`BitTensor`] storage, so the base tensor domain
109/// can construct and round-trip boolean tensors through packed `u64` words.
110pub struct BitTensorLib;
111
112impl BitTensorLib {
113    /// Creates the bit-tensor library. The value is stateless; the spec-tensor
114    /// descriptor is installed when it is loaded into a
115    /// [`Cx`](sim_kernel::Cx).
116    pub fn new() -> Self {
117        Self
118    }
119}
120
121impl Default for BitTensorLib {
122    fn default() -> Self {
123        Self::new()
124    }
125}
126
127impl Lib for BitTensorLib {
128    fn manifest(&self) -> LibManifest {
129        LibManifest {
130            id: tensor_lib_symbol(),
131            version: Version(env!("CARGO_PKG_VERSION").to_owned()),
132            abi: AbiVersion { major: 0, minor: 1 },
133            target: LibTarget::HostRegistered,
134            requires: Vec::<Dependency>::new(),
135            capabilities: Vec::new(),
136            exports: vec![Export::Value {
137                symbol: tensor_spec_symbol(),
138            }],
139        }
140    }
141
142    fn load(&self, _cx: &mut sim_kernel::LoadCx, linker: &mut Linker<'_>) -> Result<()> {
143        linker.value(
144            tensor_spec_symbol(),
145            spec_tensor_descriptor_value(
146                &DefaultFactory,
147                SpecTensorDescriptor {
148                    symbol: tensor_spec_symbol(),
149                    dtype: domains::bool(),
150                    implementation: "BitTensor",
151                    storage: "bit-packed u64 words",
152                },
153            )?,
154        )
155    }
156}
157
158/// The manifest id symbol for this library (`numbers/tensor-bit`).
159pub fn tensor_lib_symbol() -> Symbol {
160    domains::domain("tensor-bit")
161}
162
163/// The symbol under which the bit-tensor [`SpecTensor`] descriptor is exported.
164pub fn tensor_spec_symbol() -> Symbol {
165    spec_tensor_symbol("bit")
166}
167
168fn bool_value(value: bool) -> Option<Value> {
169    DefaultFactory
170        .number_literal(domains::bool(), value.to_string())
171        .ok()
172}
173
174fn parse_bool_cell(value: &Value) -> Option<bool> {
175    let mut cx = sim_kernel::Cx::new(
176        std::sync::Arc::new(sim_kernel::NoopEvalPolicy),
177        std::sync::Arc::new(DefaultFactory),
178    );
179    let literal = value
180        .object()
181        .as_number_value()?
182        .number_literal(&mut cx)
183        .ok()??;
184    (literal.domain == domains::bool())
185        .then(|| literal.canonical.parse::<bool>().ok())
186        .flatten()
187}
188
189fn map_words(
190    left: &BitTensor,
191    right: &BitTensor,
192    f: impl Fn(u64, u64) -> u64,
193) -> Option<BitTensor> {
194    (left.shape == right.shape).then(|| BitTensor {
195        shape: left.shape.clone(),
196        len: left.len,
197        words: left
198            .words
199            .iter()
200            .zip(right.words.iter())
201            .map(|(left, right)| f(*left, *right))
202            .collect(),
203    })
204}
205
206#[cfg(test)]
207mod tests {
208    use sim_kernel::Lib;
209
210    use super::{BitTensor, BitTensorLib, SpecTensor, tensor_spec_symbol};
211
212    #[test]
213    fn bit_tensor_and_matches_bool_and() {
214        let left = BitTensor::from_bools(vec![4], &[true, false, true, true]).unwrap();
215        let right = BitTensor::from_bools(vec![4], &[true, true, false, true]).unwrap();
216        let out = left.bit_and(&right).unwrap();
217        assert_eq!(out.to_bools(), vec![true, false, false, true]);
218        let uniform = out.to_uniform();
219        assert_eq!(uniform.shape, vec![4]);
220    }
221
222    #[test]
223    fn lib_exports_spec_tensor_descriptor() {
224        assert_eq!(
225            BitTensorLib::new().manifest().exports[0].symbol(),
226            &tensor_spec_symbol()
227        );
228    }
229}
sim_lib_numbers_tensor_bit/bit_tensor.rs

sim_lib_numbers_tensor_bit/
bit_tensor.rs
