Skip to main content

prism_tensor/
tensor.rs

1//! `TensorAxis` declaration + parametric square-matmul impl + shape.
2//!
3//! Per [Wiki ADR-031][09-adr-031] the tensor sub-crate exposes
4//! `TensorAxis` as the canonical Layer-3 surface for tensor compute.
5//! The reference impl [`CpuI8MatmulSquare`] is generic over the square
6//! dimension `DIM`, with `i8` inputs and saturating-`i16` outputs —
7//! the integer-arithmetic determinism contract ADR-030 names as the
8//! axis substitution-determinism baseline.
9//!
10//! Variable-rank tensor compute composes through verbs over
11//! `partition_product!`-declared shapes per ADR-033/044; the axis's
12//! role is the fixed-shape atomic primitive.
13//!
14//! # ADR-055 substrate-Term verb body discipline
15//!
16//! Per [Wiki ADR-055](https://github.com/UOR-Foundation/UOR-Framework/wiki/09-Architecture-Decisions)
17//! every `AxisExtension` impl satisfies the substrate-Term verb body
18//! discipline; the hand-written kernel below uses the default empty
19//! `body_arena()` emitted by foundation-sdk 0.4.11's `axis!`
20//! companion macro (the primitive-fast-path-equivalent realization).
21//!
22//! Explicit substrate-Term decomposition of
23//! `CpuI8MatmulSquare<DIM>::matmul` — `fold_n(DIM, ...)` over rows ×
24//! `fold_n(DIM, ...)` over columns × `fold_n(DIM, ...)` over
25//! reductions, with a `sign_extend` sub-verb (matching `Ge(operand,
26//! Literal(0x80, W8))` to select between `Concat(0x00, operand)` and
27//! `Concat(0xff, operand)`) plus W16 `Mul` + W16 `Add` accumulation
28//! plus saturation via `Match` over `Ge(acc, Literal(0x7fff, W16))` /
29//! `Lt(acc, Literal(0x8000, W16))` per ADR-054 § Substrate-Term
30//! realization examples — is **syntactically expressible** in
31//! foundation-sdk 0.4.11's verb-body grammar. ADR-056 admits
32//! `le`/`lt`/`ge`/`gt` and `concat` in verb/axis bodies (only the
33//! route body's syntactic surface retains the ψ-residuals rejection);
34//! foundation-sdk 0.4.11's depth-2 const-generic-leaf partition-product
35//! projection covers the fold-n composition over matrix shapes. The
36//! remaining work is **operational composition**: the architectural
37//! witness verbs in [`crate::verbs`] (saturating-xor + concat-bytes)
38//! demonstrate the per-element primitives; the unfolded
39//! fold-over-rows-and-columns matmul body is a published-roster
40//! follow-on.
41//!
42//! The hand-written `for`-loop kernel below is the operational form;
43//! byte-output equivalence with BLAS reference outputs at integer
44//! precision is checked at `tests/conformance.rs`.
45//!
46//! [09-adr-031]: https://github.com/UOR-Foundation/UOR-Framework/wiki/09-Architecture-Decisions
47//! [09-adr-054]: https://github.com/UOR-Foundation/UOR-Framework/wiki/09-Architecture-Decisions
48
49#![allow(missing_docs)]
50
51use uor_foundation::enforcement::{GroundedShape, ShapeViolation};
52use uor_foundation::pipeline::{
53    AxisExtension, ConstrainedTypeShape, ConstraintRef, IntoBindingValue,
54};
55use uor_foundation_sdk::axis;
56
57axis! {
58    /// Wiki ADR-031 tensor-compute axis.
59    ///
60    /// The reference impl `CpuI8MatmulSquare<DIM>` is parametric in
61    /// `DIM` for square `DIM × DIM` `i8` matrices, emitting a `DIM ×
62    /// DIM` `i16` product (saturating) per ADR-030's bit-determinism
63    /// commitment.
64    pub trait TensorAxis: AxisExtension {
65        const AXIS_ADDRESS: &'static str = "https://uor.foundation/axis/TensorAxis";
66        /// `2 * MAX_TENSOR_DIM * MAX_TENSOR_DIM` = 512 bytes for
67        /// 16×16. Overridden per impl.
68        const MAX_OUTPUT_BYTES: usize = 32;
69        /// Multiply two row-major `DIM × DIM` `i8` matrices into a
70        /// `DIM × DIM` `i16` product (saturating). Input is `A || B`
71        /// (`2 * DIM * DIM` bytes); output is `2 * DIM * DIM` bytes.
72        ///
73        /// # Errors
74        ///
75        /// Returns `ShapeViolation` on input/output byte-length mismatch.
76        fn matmul(input: &[u8], out: &mut [u8]) -> Result<usize, ShapeViolation>;
77    }
78}
79
80/// Maximum square dimension any [`CpuI8MatmulSquare`] instantiation
81/// supports. Cap at 16: output buffer = `2 * 16 * 16` = 512 bytes,
82/// inputs = 256 bytes each, total kernel byte budget bounded.
83pub const MAX_TENSOR_DIM: usize = 16;
84
85fn arity_violation(constraint: &'static str) -> ShapeViolation {
86    ShapeViolation {
87        shape_iri: "https://uor.foundation/axis/TensorAxisShape",
88        constraint_iri: constraint,
89        property_iri: "https://uor.foundation/axis/inputBytes",
90        expected_range: "https://uor.foundation/axis/TensorInputArity",
91        min_count: 0,
92        max_count: 0,
93        kind: uor_foundation::ViolationKind::ValueCheck,
94    }
95}
96
97/// Parametric square `DIM × DIM` `i8` × `i8` → `i16` matmul.
98///
99/// Determinism: per ADR-030's per-axis substitution-determinism note,
100/// the integer-arithmetic CPU impl preserves bit-identity across
101/// targets. `DIM` is the square dimension; for non-square /
102/// non-integer / variable-shape tensor compute the wiki's pattern is
103/// to compose this axis kernel through verbs over `partition_product!`
104/// (per ADR-033/044) — the axis layer fixes the atom shape.
105#[derive(Debug, Clone, Copy)]
106pub struct CpuI8MatmulSquare<const DIM: usize>;
107
108impl<const DIM: usize> Default for CpuI8MatmulSquare<DIM> {
109    fn default() -> Self {
110        Self
111    }
112}
113
114impl<const DIM: usize> CpuI8MatmulSquare<DIM> {
115    const fn idx(row: usize, col: usize) -> usize {
116        row * DIM + col
117    }
118}
119
120impl<const DIM: usize> TensorAxis for CpuI8MatmulSquare<DIM> {
121    const AXIS_ADDRESS: &'static str = "https://uor.foundation/axis/TensorAxis/CpuI8MatmulSquare";
122    const MAX_OUTPUT_BYTES: usize = 2 * DIM * DIM;
123
124    fn matmul(input: &[u8], out: &mut [u8]) -> Result<usize, ShapeViolation> {
125        if DIM == 0 || DIM > MAX_TENSOR_DIM {
126            return Err(arity_violation(
127                "https://uor.foundation/axis/TensorAxisShape/dimInRange",
128            ));
129        }
130        let mat_bytes = DIM * DIM;
131        let input_bytes = 2 * mat_bytes;
132        let output_bytes = 2 * mat_bytes;
133        if input.len() != input_bytes {
134            return Err(arity_violation(
135                "https://uor.foundation/axis/TensorAxisShape/inputByteLength",
136            ));
137        }
138        if out.len() < output_bytes {
139            return Err(arity_violation(
140                "https://uor.foundation/axis/TensorAxisShape/outputByteLength",
141            ));
142        }
143        let (a_bytes, b_bytes) = input.split_at(mat_bytes);
144        for row in 0..DIM {
145            for col in 0..DIM {
146                let mut acc: i32 = 0;
147                for k in 0..DIM {
148                    #[allow(clippy::cast_possible_wrap)]
149                    let a = i32::from(a_bytes[Self::idx(row, k)] as i8);
150                    #[allow(clippy::cast_possible_wrap)]
151                    let b = i32::from(b_bytes[Self::idx(k, col)] as i8);
152                    acc += a * b;
153                }
154                let saturated: i16 = if acc > i32::from(i16::MAX) {
155                    i16::MAX
156                } else if acc < i32::from(i16::MIN) {
157                    i16::MIN
158                } else {
159                    #[allow(clippy::cast_possible_truncation)]
160                    {
161                        acc as i16
162                    }
163                };
164                let cell = Self::idx(row, col);
165                out[2 * cell..2 * cell + 2].copy_from_slice(&saturated.to_be_bytes());
166            }
167        }
168        Ok(output_bytes)
169    }
170}
171
172// ADR-052 generic-form companion.
173axis_extension_impl_for_tensor_axis!(@generic CpuI8MatmulSquare<DIM>, [const DIM: usize]);
174
175/// 4×4 `i8` matmul — the canonical small-tensor reference.
176pub type CpuI8Tensor4x4Matmul = CpuI8MatmulSquare<4>;
177/// 8×8 `i8` matmul.
178pub type CpuI8Tensor8x8Matmul = CpuI8MatmulSquare<8>;
179/// 16×16 `i8` matmul (the `MAX_TENSOR_DIM` ceiling).
180pub type CpuI8Tensor16x16Matmul = CpuI8MatmulSquare<16>;
181
182// ---- MatrixShape: ConstrainedTypeShape carrier ----
183
184/// Parametric ConstrainedTypeShape for a row-major `ROWS × COLS`
185/// matrix of `ELEM_BYTES`-byte elements. Per ADR-031's `Tensor<Element,
186/// Shape>` shape commitment, restricted to matrix rank-2 here; higher
187/// ranks compose through `partition_product!` per ADR-033/044.
188#[derive(Debug, Clone, Copy)]
189pub struct MatrixShape<const ROWS: usize, const COLS: usize, const ELEM_BYTES: usize>;
190
191impl<const ROWS: usize, const COLS: usize, const ELEM_BYTES: usize> Default
192    for MatrixShape<ROWS, COLS, ELEM_BYTES>
193{
194    fn default() -> Self {
195        Self
196    }
197}
198
199impl<const ROWS: usize, const COLS: usize, const ELEM_BYTES: usize> ConstrainedTypeShape
200    for MatrixShape<ROWS, COLS, ELEM_BYTES>
201{
202    const IRI: &'static str = "https://uor.foundation/type/ConstrainedType";
203    const SITE_COUNT: usize = ROWS * COLS * ELEM_BYTES;
204    const CONSTRAINTS: &'static [ConstraintRef] = &[];
205    #[allow(clippy::cast_possible_truncation)]
206    const CYCLE_SIZE: u64 = 256u64.saturating_pow((ROWS * COLS * ELEM_BYTES) as u32);
207}
208
209impl<const ROWS: usize, const COLS: usize, const ELEM_BYTES: usize>
210    uor_foundation::pipeline::__sdk_seal::Sealed for MatrixShape<ROWS, COLS, ELEM_BYTES>
211{
212}
213impl<const ROWS: usize, const COLS: usize, const ELEM_BYTES: usize> GroundedShape
214    for MatrixShape<ROWS, COLS, ELEM_BYTES>
215{
216}
217impl<const ROWS: usize, const COLS: usize, const ELEM_BYTES: usize> IntoBindingValue
218    for MatrixShape<ROWS, COLS, ELEM_BYTES>
219{
220    const MAX_BYTES: usize = ROWS * COLS * ELEM_BYTES;
221
222    fn into_binding_bytes(&self, _out: &mut [u8]) -> Result<usize, ShapeViolation> {
223        Ok(0)
224    }
225}
226
227/// Parametric ConstrainedTypeShape for a length-`N` vector of
228/// `ELEM_BYTES`-byte elements. Per ADR-031's `Tensor<Element, Shape>`
229/// for rank-1.
230#[derive(Debug, Clone, Copy)]
231pub struct VectorShape<const N: usize, const ELEM_BYTES: usize>;
232
233impl<const N: usize, const ELEM_BYTES: usize> Default for VectorShape<N, ELEM_BYTES> {
234    fn default() -> Self {
235        Self
236    }
237}
238
239impl<const N: usize, const ELEM_BYTES: usize> ConstrainedTypeShape for VectorShape<N, ELEM_BYTES> {
240    const IRI: &'static str = "https://uor.foundation/type/ConstrainedType";
241    const SITE_COUNT: usize = N * ELEM_BYTES;
242    const CONSTRAINTS: &'static [ConstraintRef] = &[];
243    #[allow(clippy::cast_possible_truncation)]
244    const CYCLE_SIZE: u64 = 256u64.saturating_pow((N * ELEM_BYTES) as u32);
245}
246
247impl<const N: usize, const ELEM_BYTES: usize> uor_foundation::pipeline::__sdk_seal::Sealed
248    for VectorShape<N, ELEM_BYTES>
249{
250}
251impl<const N: usize, const ELEM_BYTES: usize> GroundedShape for VectorShape<N, ELEM_BYTES> {}
252impl<const N: usize, const ELEM_BYTES: usize> IntoBindingValue for VectorShape<N, ELEM_BYTES> {
253    const MAX_BYTES: usize = N * ELEM_BYTES;
254
255    fn into_binding_bytes(&self, _out: &mut [u8]) -> Result<usize, ShapeViolation> {
256        Ok(0)
257    }
258}