tract_linalg/
lib.rs

1#![allow(clippy::missing_safety_doc)]
2#![allow(clippy::redundant_closure_call)]
3#![allow(clippy::len_zero)]
4#![allow(clippy::excessive_precision)]
5#![allow(clippy::approx_constant)]
6#![allow(clippy::manual_is_multiple_of)]
7#![allow(unexpected_cfgs)]
8#![allow(unused_macros)]
9#[macro_use]
10extern crate derive_new;
11extern crate lazy_static;
12extern crate log;
13extern crate num_traits;
14#[macro_use]
15extern crate pastey;
16#[cfg(test)]
17extern crate proptest;
18
19include!(concat!(env!("OUT_DIR"), "/extern_kernel_macro.rs"));
20
21/// Stands in for a function whose body only compiles in builds carrying the leading arch's
22/// instructions — an asm block, an intrinsic, a CPUID probe — taking the argument types of
23/// the real item and bailing when called. `wasm32` means wasm32 *with* `simd128`, the two
24/// conditions the wasm kernels need. Needed only where something names the function on every
25/// arch: a codegen macro, or a descriptor the arch tree declares everywhere while only the
26/// native host ever calls what it names. A plain `#[cfg]` covers the rest.
27macro_rules! bail_stub {
28    (arm; $($rest:tt)*) => { bail_stub!(@ target_arch = "arm"; $($rest)*); };
29    (aarch64; $($rest:tt)*) => { bail_stub!(@ target_arch = "aarch64"; $($rest)*); };
30    (x86_64; $($rest:tt)*) => { bail_stub!(@ target_arch = "x86_64"; $($rest)*); };
31    (riscv64; $($rest:tt)*) => { bail_stub!(@ target_arch = "riscv64"; $($rest)*); };
32    (wasm32; $($rest:tt)*) => {
33        bail_stub!(@ all(target_arch = "wasm32", target_feature = "simd128"); $($rest)*);
34    };
35
36    (@ $built:meta; $vis:vis unsafe fn $name:ident($($ty:ty),* $(,)?) $(-> $ret:ty)?) => {
37        #[cfg(not($built))]
38        $vis unsafe fn $name($(_: $ty),*) $(-> $ret)? {
39            panic!(concat!(stringify!($name), ": not built for this target"))
40        }
41    };
42
43    (@ $built:meta; $vis:vis fn $name:ident($($ty:ty),* $(,)?) $(-> $ret:ty)?) => {
44        #[cfg(not($built))]
45        $vis fn $name($(_: $ty),*) $(-> $ret)? {
46            panic!(concat!(stringify!($name), ": not built for this target"))
47        }
48    };
49}
50
51#[macro_use]
52mod frame;
53#[macro_use]
54pub mod routines;
55pub mod cache;
56pub mod generic;
57pub mod knobs;
58pub mod multithread;
59pub use frame::weights::WeightType;
60pub use generic::{ScaleShiftAndRound, Scaler};
61use tract_data::internal::TensorView;
62// An arch tree compiles when this build can run its kernels, and — for enumeration only —
63// when `foreign-inventory` asks for the others as well.
64#[cfg(any(target_arch = "x86_64", feature = "foreign-inventory"))]
65pub mod x86_64;
66
67pub mod hwbench;
68
69#[cfg(any(target_arch = "aarch64", feature = "foreign-inventory"))]
70pub mod arm64;
71
72#[cfg(any(target_arch = "aarch64", feature = "foreign-inventory"))]
73pub use arm64::has_fp16;
74
75/// True when the running CPU implements FEAT_FP16. No arm64 tree in this build, hence no
76/// kernel that could use it.
77#[cfg(not(any(target_arch = "aarch64", feature = "foreign-inventory")))]
78pub fn has_fp16() -> bool {
79    false
80}
81
82#[cfg(any(target_arch = "arm", feature = "foreign-inventory"))]
83pub mod arm32;
84
85#[cfg(any(target_arch = "riscv64", feature = "foreign-inventory"))]
86pub mod riscv64;
87
88#[cfg(any(all(target_arch = "wasm32", target_feature = "simd128"), feature = "foreign-inventory"))]
89pub mod wasm;
90
91pub mod isa;
92pub mod mmm_routines;
93pub mod mmm_tiers;
94
95pub use self::frame::mmm::MmmDispatch;
96pub use self::frame::*;
97pub use self::routines::Func;
98
99use tract_data::prelude::*;
100
101#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
102pub enum BinOp {
103    Min,
104    Max,
105    Add,
106    Mul,
107    Sub,
108    SubF,
109}
110
111impl BinOp {
112    pub fn flip(&self) -> BinOp {
113        use BinOp::*;
114        match self {
115            Sub => SubF,
116            SubF => Sub,
117            sym => *sym,
118        }
119    }
120}
121
122/// A binary operation over two tensor views, writing its result over the left one. What the two
123/// binary layouts erase to -- [`by_scalar::ByScalarKer::bin`] broadcasts a one-element right
124/// operand, [`unicast::UnicastKer::bin`] walks a right operand of the same length -- so a caller
125/// holding one needs to know neither which layout nor which kernel answered.
126pub type BinFn = dyn Fn(&mut TensorView, &TensorView) -> TractResult<()> + Send + Sync;
127use num_traits::*;
128use std::fmt::Debug;
129use std::ops::*;
130
131pub trait LADatum:
132    Sized
133    + std::fmt::Display
134    + Debug
135    + Copy
136    + Clone
137    + Zero
138    + One
139    + 'static
140    + Add<Output = Self>
141    + Sub<Output = Self>
142    + Mul
143    + AddAssign
144    + PartialOrd
145    + Bounded
146    + tract_data::prelude::Datum
147{
148    #[cfg(test)]
149    fn strat() -> proptest::prelude::BoxedStrategy<Self>;
150}
151
152#[cfg(test)]
153use proptest::prelude::*;
154
155impl LADatum for f16 {
156    #[cfg(test)]
157    fn strat() -> BoxedStrategy<Self> {
158        f32::strat().prop_map(|f| f.as_()).boxed()
159    }
160}
161
162impl LADatum for f32 {
163    #[cfg(test)]
164    fn strat() -> BoxedStrategy<Self> {
165        (-1000isize..1000).prop_map(|i| i as f32 / 1000.0).boxed()
166    }
167}
168
169impl LADatum for f64 {
170    #[cfg(test)]
171    fn strat() -> BoxedStrategy<Self> {
172        (-1000isize..1000).prop_map(|i| i as f64 / 1000.0).boxed()
173    }
174}
175
176impl LADatum for u8 {
177    #[cfg(test)]
178    fn strat() -> BoxedStrategy<Self> {
179        any::<u8>().boxed()
180    }
181}
182
183impl LADatum for i8 {
184    #[cfg(test)]
185    fn strat() -> BoxedStrategy<Self> {
186        any::<i8>().boxed()
187    }
188}
189
190impl LADatum for i32 {
191    #[cfg(test)]
192    fn strat() -> BoxedStrategy<Self> {
193        any::<i32>().boxed()
194    }
195}
196
197#[cfg(test)]
198#[allow(dead_code)]
199fn setup_test_logger() {
200    let _ = env_logger::Builder::from_env("TRACT_LOG").try_init();
201}
tract_linalg/lib.rs

tract_linalg/
lib.rs
