tropical_gemm/types/
bitwise.rs

1use super::scalar::TropicalScalar;
2use super::traits::{SimdTropical, TropicalSemiring};
3use std::fmt;
4use std::ops::{Add, BitAnd, BitOr, Mul};
5
6mod sealed {
7    pub trait Sealed {}
8    impl Sealed for u32 {}
9    impl Sealed for u64 {}
10}
11
12/// Unsigned-integer element types valid as a `TropicalBitwise` lane container.
13///
14/// Sealed: only `u32` (32 lanes) and `u64` (64 lanes) are permitted.
15pub trait BitwiseScalar:
16    TropicalScalar + sealed::Sealed + BitOr<Output = Self> + BitAnd<Output = Self>
17{
18    /// All lanes false (tropical zero).
19    const ZERO: Self;
20    /// All lanes true (tropical one), i.e. `!0`.
21    const ONES: Self;
22}
23
24impl BitwiseScalar for u32 {
25    const ZERO: u32 = 0;
26    const ONES: u32 = u32::MAX;
27}
28
29impl BitwiseScalar for u64 {
30    const ZERO: u64 = 0;
31    const ONES: u64 = u64::MAX;
32}
33
34/// TropicalBitwise semiring: `(uint, |, &, 0, ~0)` — bit-packed boolean.
35///
36/// Each bit-lane of the wrapped word is an **independent** boolean problem
37/// (bit-slicing): one GEMM computes 32 (`u32`) or 64 (`u64`) boolean matmuls at
38/// once. `⊕ = |`, `⊗ = &`, zero = `0`, one = `!0`.
39///
40/// CPU dispatch uses AVX2 or NEON where available, with a portable fallback.
41/// There is no single argmax index for a packed word: each bit lane can have a
42/// different winner. Per-lane argmax is not implemented, so this type does not
43/// implement `TropicalWithArgmax`.
44///
45/// This is for **many independent dense boolean problems**. For a single large
46/// (sparse) boolean graph, use a sparse GraphBLAS tool (GraphBLAST / cuBool /
47/// Bit-GraphBLAS) — that is out of scope for this dense library.
48#[derive(Copy, Clone, PartialEq, Eq)]
49#[repr(transparent)]
50pub struct TropicalBitwise<T: BitwiseScalar>(pub T);
51
52impl<T: BitwiseScalar> TropicalBitwise<T> {
53    /// Create a new TropicalBitwise value from a packed word.
54    #[inline(always)]
55    pub fn new(value: T) -> Self {
56        Self(value)
57    }
58}
59
60impl<T: BitwiseScalar> TropicalSemiring for TropicalBitwise<T> {
61    type Scalar = T;
62
63    fn scalar_slice(values: &[Self]) -> Option<&[Self::Scalar]> {
64        // SAFETY: this type is repr(transparent) over its scalar field.
65        Some(unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) })
66    }
67
68    #[inline(always)]
69    fn tropical_zero() -> Self {
70        Self(T::ZERO)
71    }
72
73    #[inline(always)]
74    fn tropical_one() -> Self {
75        Self(T::ONES)
76    }
77
78    #[inline(always)]
79    fn tropical_add(self, rhs: Self) -> Self {
80        Self(self.0 | rhs.0)
81    }
82
83    #[inline(always)]
84    fn tropical_mul(self, rhs: Self) -> Self {
85        Self(self.0 & rhs.0)
86    }
87
88    #[inline(always)]
89    fn value(&self) -> T {
90        self.0
91    }
92
93    #[inline(always)]
94    fn from_scalar(s: T) -> Self {
95        Self(s)
96    }
97}
98
99impl<T: BitwiseScalar> SimdTropical for TropicalBitwise<T> {
100    const SIMD_AVAILABLE: bool = cfg!(any(target_arch = "x86_64", target_arch = "aarch64"));
101    const SIMD_WIDTH: usize = if cfg!(target_arch = "x86_64") {
102        32 / std::mem::size_of::<T>()
103    } else if cfg!(target_arch = "aarch64") {
104        16 / std::mem::size_of::<T>()
105    } else {
106        0
107    };
108}
109
110impl<T: BitwiseScalar> Add for TropicalBitwise<T> {
111    type Output = Self;
112
113    #[inline(always)]
114    fn add(self, rhs: Self) -> Self::Output {
115        self.tropical_add(rhs)
116    }
117}
118
119impl<T: BitwiseScalar> Mul for TropicalBitwise<T> {
120    type Output = Self;
121
122    #[inline(always)]
123    fn mul(self, rhs: Self) -> Self::Output {
124        self.tropical_mul(rhs)
125    }
126}
127
128impl<T: BitwiseScalar> Default for TropicalBitwise<T> {
129    #[inline(always)]
130    fn default() -> Self {
131        Self::tropical_zero()
132    }
133}
134
135impl<T: BitwiseScalar> fmt::Debug for TropicalBitwise<T> {
136    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137        write!(f, "TropicalBitwise({})", self.0)
138    }
139}
140
141impl<T: BitwiseScalar> fmt::Display for TropicalBitwise<T> {
142    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
143        write!(f, "{}", self.0)
144    }
145}
146
147impl<T: BitwiseScalar> From<T> for TropicalBitwise<T> {
148    #[inline(always)]
149    fn from(value: T) -> Self {
150        Self(value)
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn additive_identity() {
160        let a = TropicalBitwise::<u32>(0b1011);
161        assert_eq!(a.tropical_add(TropicalBitwise::tropical_zero()), a);
162    }
163
164    #[test]
165    fn multiplicative_identity() {
166        let a = TropicalBitwise::<u32>(0b1011);
167        assert_eq!(a.tropical_mul(TropicalBitwise::tropical_one()), a);
168    }
169
170    #[test]
171    fn absorbing_zero() {
172        let a = TropicalBitwise::<u32>(0xDEADBEEF);
173        assert_eq!(
174            a.tropical_mul(TropicalBitwise::tropical_zero()),
175            TropicalBitwise::tropical_zero()
176        );
177    }
178
179    #[test]
180    fn ops_are_bitwise() {
181        let a = TropicalBitwise::<u64>(0b1100);
182        let b = TropicalBitwise::<u64>(0b1010);
183        assert_eq!(a.tropical_add(b).0, 0b1110); // OR
184        assert_eq!(a.tropical_mul(b).0, 0b1000); // AND
185    }
186
187    #[test]
188    fn zero_and_one_values() {
189        assert_eq!(TropicalBitwise::<u32>::tropical_zero().0, 0u32);
190        assert_eq!(TropicalBitwise::<u32>::tropical_one().0, u32::MAX);
191        assert_eq!(TropicalBitwise::<u64>::tropical_one().0, u64::MAX);
192    }
193
194    // Bit-lane 0 of a TropicalBitwise<u32> GEMM must equal the same problem run
195    // as TropicalAndOr (one lane is one AndOr problem). Uses the public matmul
196    // API, which requires KernelDispatch (added in this task).
197    #[test]
198    fn lane0_matches_andor() {
199        use crate::tropical_matmul;
200        use crate::types::TropicalAndOr;
201
202        // 2x3 * 3x2 column-major boolean problem (same as the GPU AndOr test).
203        let a_bool = [true, false, false, false, true, false];
204        let b_bool = [true, true, false, false, true, true];
205
206        let a_u32: Vec<u32> = a_bool.iter().map(|&x| x as u32).collect();
207        let b_u32: Vec<u32> = b_bool.iter().map(|&x| x as u32).collect();
208
209        let c_bw = tropical_matmul::<TropicalBitwise<u32>>(&a_u32, 2, 3, &b_u32, 2);
210
211        let mut c_ao = vec![TropicalAndOr(false); 4];
212        // Safety: 2x3 * 3x2, all leading dims match, no aliasing.
213        unsafe {
214            crate::core::tropical_gemm_portable::<TropicalAndOr>(
215                2,
216                2,
217                3,
218                a_bool.as_ptr(),
219                2,
220                crate::Transpose::NoTrans,
221                b_bool.as_ptr(),
222                3,
223                crate::Transpose::NoTrans,
224                c_ao.as_mut_ptr(),
225                2,
226            );
227        }
228
229        for i in 0..4 {
230            let lane0 = (c_bw[i].0 & 1) == 1;
231            assert_eq!(lane0, c_ao[i].0, "cell {i}: bitwise lane0 != andor");
232        }
233    }
234}
tropical_gemm/types/bitwise.rs

tropical_gemm/types/
bitwise.rs
