tritium_core/
shape.rs

1//! GEMM problem geometry.
2
3/// Dimensions of a matmul `C[M,N] = A[M,K] · Wᵀ`, where `W` is the `[N, K]`
4/// ternary weight (row-major, output-major) and `A` is the activation `[M, K]`.
5///
6/// - `m` — activation rows (batch × sequence).
7/// - `n` — output features (weight rows).
8/// - `k` — contraction / input features (weight + activation columns).
9#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
10pub struct GemmShape {
11    /// Activation rows (batch × sequence).
12    pub m: usize,
13    /// Output features (weight rows).
14    pub n: usize,
15    /// Contraction / input features (weight + activation columns).
16    pub k: usize,
17}
18
19impl GemmShape {
20    /// Construct a shape from its `m`, `n`, `k` dimensions.
21    #[inline]
22    pub const fn new(m: usize, n: usize, k: usize) -> Self {
23        Self { m, n, k }
24    }
25
26    /// Multiply-accumulate count `M·N·K`. For ternary this is the count of
27    /// add/sub/skip ops, not true FMAs — the metric backends report throughput on.
28    #[inline]
29    pub const fn macs(&self) -> u64 {
30        (self.m as u64) * (self.n as u64) * (self.k as u64)
31    }
32
33    /// Whether the operand/output buffer lengths are internally consistent.
34    #[inline]
35    pub const fn buffers_fit(&self, act_len: usize, weight_len: usize, out_len: usize) -> bool {
36        act_len == self.m * self.k && weight_len == self.n * self.k && out_len == self.m * self.n
37    }
38}
39
40/// Geometry of a ternary 1-D convolution `Y[B, C_out, L_out] = scale ⊙ conv1d(X, W)` — the codec's
41/// conv op (ADR 0030). The weight is packed 2-D `[C_out, (C_in/groups)·K]` (the per-output-channel
42/// ternary reshape), so `k_g()` is the matmul contraction and `n_g()` the per-group output count.
43#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
44pub struct ConvShape {
45    /// Batch size.
46    pub batch: usize,
47    /// Input channels (divisible by `groups`).
48    pub c_in: usize,
49    /// Output channels (divisible by `groups`).
50    pub c_out: usize,
51    /// Input length.
52    pub l_in: usize,
53    /// Kernel size.
54    pub k: usize,
55    /// Stride (≥ 1).
56    pub stride: usize,
57    /// Dilation (≥ 1).
58    pub dilation: usize,
59    /// Left zero-padding.
60    pub pad_left: usize,
61    /// Right zero-padding.
62    pub pad_right: usize,
63    /// Convolution groups (≥ 1).
64    pub groups: usize,
65}
66
67impl ConvShape {
68    /// Output length, or `0` if the dilated kernel is wider than the padded input (or the geometry is
69    /// degenerate).
70    #[inline]
71    pub const fn l_out(&self) -> usize {
72        if self.k == 0 || self.stride == 0 {
73            return 0;
74        }
75        let eff = self.dilation * (self.k - 1) + 1;
76        let padded = self.l_in + self.pad_left + self.pad_right;
77        if padded < eff {
78            return 0;
79        }
80        (padded - eff) / self.stride + 1
81    }
82
83    /// Input channels per group `C_in/groups`.
84    #[inline]
85    pub const fn c_in_pg(&self) -> usize {
86        self.c_in / self.groups
87    }
88
89    /// Output channels per group `N_g = C_out/groups`.
90    #[inline]
91    pub const fn n_g(&self) -> usize {
92        self.c_out / self.groups
93    }
94
95    /// Flattened per-output-channel weight width `K_g = (C_in/groups)·K`.
96    #[inline]
97    pub const fn k_g(&self) -> usize {
98        self.c_in_pg() * self.k
99    }
100
101    /// Whether the geometry is well-formed and the buffers are the right length (`x=B·C_in·L_in`,
102    /// `weights=C_out·K_g`, `scale=C_out`, `out=B·C_out·L_out`).
103    #[inline]
104    pub fn buffers_fit(
105        &self,
106        x_len: usize,
107        weight_len: usize,
108        scale_len: usize,
109        out_len: usize,
110    ) -> bool {
111        self.groups != 0
112            && self.k != 0
113            && self.stride != 0
114            && self.dilation != 0
115            && self.c_in.is_multiple_of(self.groups)
116            && self.c_out.is_multiple_of(self.groups)
117            && self.l_out() > 0
118            && x_len == self.batch * self.c_in * self.l_in
119            && weight_len == self.c_out * self.k_g()
120            && scale_len == self.c_out
121            && out_len == self.batch * self.c_out * self.l_out()
122    }
123}
tritium_core/shape.rs

tritium_core/
shape.rs
