Add native self-hosted instance connection to fluxer_desktop

Trimmed monorepo checkout (fluxer_desktop + packages/voice_engine_v2 +
tools/ci) with a "Connect to a Different Server" menu item and popout
that lets the desktop app switch to any self-hosted Fluxer instance,
plus fixes for well-known discovery on single-domain self-hosted
deployments and a false-positive ERR_ABORTED on same-origin client
redirects during the switch. Defaults to chat.fluxr.chat and uses an
isolated userData directory from the official build.
This commit is contained in:
2026-07-01 18:22:43 -04:00
commit 682afacd30
1763 changed files with 613720 additions and 0 deletions
@@ -0,0 +1,105 @@
use std::{env, fs};
pub mod armv7neon;
mod armvfpv2;
mod cortex_a7;
mod cortex_a9;
use armv7neon::*;
use crate::frame::element_wise::ElementWiseKer;
use crate::frame::mmm::kernel::MatMatMulKer;
use crate::Ops;
fn has_neon_cpuinfo() -> std::io::Result<bool> {
let cpu_info = fs::read_to_string("/proc/cpuinfo")?;
let neon = cpu_info.split("\n").any(|line| {
line.starts_with("Features") && (line.contains("neon") || line.contains("asimd"))
});
Ok(neon)
}
fn cpu_part() -> Option<usize> {
fs::read_to_string("/proc/cpuinfo")
.ok()
.and_then(|cpuinfo| {
cpuinfo
.lines()
.find(|line| line.starts_with("CPU part"))
.and_then(|s| s.trim().split_whitespace().last())
.and_then(|s| s.strip_prefix("0x"))
.and_then(|s| usize::from_str_radix(s, 16).ok())
})
}
fn has_neon() -> bool {
if let Ok(v) = env::var("TRACT_CPU_ARM32_NEON") {
return v == "true" || v == "1";
}
has_neon_cpuinfo().unwrap_or(false)
}
pub fn plug(ops: &mut Ops) {
let impls = vec![
armv7neon_mmm_f32_8x4_cortexa7::mmm(),
armv7neon_mmm_f32_8x6_cortexa7::mmm(),
armv7neon_mmm_f32_8x4_cortexa9::mmm(),
armv7neon_mmm_f32_8x6_cortexa9::mmm(),
armv7neon_mmm_f32_8x4_generic::mmm(),
armv7neon_mmm_f32_8x6_generic::mmm(),
crate::generic::mmm::generic_f32_4x4::mmm(),
];
ops.mmm_f32_impls = impls.clone();
if has_neon() {
log::info!("armv7neon activated (smmm, ssigmoid), stanh)");
let cpu = cpu_part().unwrap_or(0);
fn prefer_8x4(_m: Option<usize>, _k: Option<usize>, n: Option<usize>) -> bool {
n.map(|n| n % 4 == 0 && n % 6 != 0 && n <= 12)
.unwrap_or(false)
}
ops.mmv_f32 = match cpu {
0xc07 => Box::new(|_, _| armv7neon::armv7neon_mmm_f32_32x1_cortexa7::mmm()),
0xc09 => Box::new(|_, _| armv7neon::armv7neon_mmm_f32_32x1_cortexa9::mmm()),
_ => Box::new(|_, _| armv7neon::armv7neon_mmm_f32_32x1_generic::mmm()),
};
ops.mmm_f32 = match cpu {
0xc07 => {
let model = cortex_a7::model();
Box::new(move |m, k, n| model.pick(&impls, m, k, n))
}
0xc09 => {
let model = cortex_a9::model();
Box::new(move |m, k, n| model.pick(&impls, m, k, n))
}
_ => Box::new(|m, k, n| {
if prefer_8x4(m, k, n) {
armv7neon::armv7neon_mmm_f32_8x4_generic::mmm()
} else {
armv7neon::armv7neon_mmm_f32_8x6_generic::mmm()
}
}),
};
ops.qmmm_i32 = Box::new(|_, _, _| armv7neon::armv7neon_mmm_i32_8x4::mmm());
ops.qmmv_i32 = Box::new(|_, _| armv7neon::armv7neon_mmm_i32_32x1::mmm());
ops.sigmoid_f32 = Box::new(|| armv7neon_sigmoid_f32_4n::ew());
ops.tanh_f32 = Box::new(|| armv7neon_tanh_f32_4n::ew());
} else {
log::info!("armvfpv2 activated for smmm");
ops.mmm_f32 = Box::new(|_, _, _| armvfpv2::armvfpv2_mmm_f32_4x4::mmm());
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn may_have_neon() {
println!("Has neon ? {:?}", has_neon());
if let Ok(neon) = env::var("TRACT_CPU_EXPECT_ARM32_NEON") {
assert_eq!(neon == "true", has_neon());
}
}
}
@@ -0,0 +1,30 @@
use crate::frame::element_wise::*;
use crate::frame::mmm::*;
extern_kernel!(fn armv7neon_prefetch(start: *const u8, end: *const u8) -> ());
#[inline(always)]
pub fn prefetch(start: *const u8, len: usize) {
unsafe { armv7neon_prefetch(start, start.offset(len as isize)) }
}
MMMKernel!(i32, armv7neon_mmm_i32_8x4; 8, 4; 32, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(i32, armv7neon_mmm_i32_32x1; 32,1 ; 32, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_8x4_cortexa7; 8, 4; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_8x4_cortexa9; 8, 4; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_8x4_generic; 8, 4; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_8x6_cortexa7; 8, 6; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_8x6_cortexa9; 8, 6; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_8x6_generic; 8, 6; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_32x1_cortexa7; 32, 1; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_32x1_cortexa9; 32, 1; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
MMMKernel!(f32, armv7neon_mmm_f32_32x1_generic; 32, 1; 4, 4; 0, 0; prefetch, crate::arm32::has_neon());
sigmoid_impl!(
f32,
armv7neon_sigmoid_f32_4n,
4,
4,
crate::arm32::has_neon()
);
tanh_impl!(f32, armv7neon_tanh_f32_4n, 4, 4, crate::arm32::has_neon());
@@ -0,0 +1,3 @@
use crate::frame::mmm::*;
MMMKernel!(f32, armvfpv2_mmm_f32_4x4; 4, 4; 4, 4; 0, 0; no_prefetch, true);
@@ -0,0 +1,861 @@
use crate::frame::mmm::CostModel;
pub fn model() -> CostModel<'static> {
CostModel {
big_product_mkn_threshold: 4193728.0,
big_product_kernel_choice: "armv7neon_mmm_f32_8x6_cortexa7",
kernels: &[
"armv7neon_mmm_f32_8x4_cortexa7",
"armv7neon_mmm_f32_8x4_cortexa9",
"armv7neon_mmm_f32_8x4_generic",
"armv7neon_mmm_f32_8x6_cortexa7",
"armv7neon_mmm_f32_8x6_cortexa9",
"armv7neon_mmm_f32_8x6_generic",
"generic_f32_4x4",
],
mrs: &[4, 8],
nrs: &[4, 6],
feat_norm_mean: &[
4.589878771602424,
4.5739692460187005,
4.598167981532298,
13.762015999153403,
1.5038983903420524,
0.749874245472837,
3.465165995975855,
0.8777665995975855,
1.5022635814889336,
0.7570422535211268,
2.482142857142857,
0.8333752515090543,
],
feat_norm_stddev: &[
1.2587312982588519,
1.2603116830524392,
1.2581181647300588,
1.3169322340874257,
1.1192637768418767,
0.43308528195884044,
2.2762097127791114,
0.32755518043295856,
1.1069539235554247,
0.42886977033219037,
1.7067987601825914,
0.37264049924995035,
],
w1: &[
0.06765510141849518,
0.024555781856179237,
-0.8821254968643188,
-0.004870870150625706,
-0.10525479167699814,
0.1827959418296814,
0.1633400171995163,
-0.2377464473247528,
-0.17880690097808838,
0.19097138941287994,
0.04676022008061409,
-0.11329511553049088,
0.4089120030403137,
-0.3100685477256775,
-0.1652061492204666,
-0.19124962389469147,
-0.03810987249016762,
-0.00785011239349842,
0.09714752435684204,
-0.11142419278621674,
0.19261880218982697,
-0.2893339991569519,
-0.19540216028690338,
0.39759594202041626,
-0.00619965186342597,
-0.8473111391067505,
0.343344122171402,
-0.12575943768024445,
0.029266485944390297,
-0.02900734543800354,
-0.019343264400959015,
0.08306540548801422,
-0.1927606761455536,
0.23312175273895264,
0.2576882541179657,
-0.35881471633911133,
-0.27300119400024414,
-0.2995607852935791,
-0.7934547662734985,
-0.9349930286407471,
-0.011614155024290085,
-0.12521372735500336,
0.011371670290827751,
0.05779163911938667,
0.17875070869922638,
-0.23169392347335815,
-0.09749509394168854,
0.07436174154281616,
0.24035069346427917,
-0.1262669861316681,
0.3874961733818054,
-0.11149000376462936,
0.03639678284525871,
0.17740628123283386,
0.03768332302570343,
-0.20480288565158844,
-0.1955408751964569,
0.44144806265830994,
0.3628064692020416,
-0.2537013292312622,
0.019405143335461617,
0.06186319515109062,
0.5196826457977295,
0.3010406494140625,
0.04013144597411156,
0.03517461195588112,
-0.037290964275598526,
0.009919736534357071,
-0.3135205805301666,
0.4654330909252167,
0.46720823645591736,
0.29665476083755493,
0.09099660068750381,
-0.7376689314842224,
-0.07840575277805328,
-0.5192644000053406,
0.019796665757894516,
-0.021734869107604027,
0.13953897356987,
-0.04154204577207565,
0.10942933708429337,
-0.13621817529201508,
-0.04218055680394173,
0.09188657253980637,
-0.16021296381950378,
-0.19393481314182281,
0.3737955689430237,
0.08288388699293137,
-0.08280416578054428,
-0.13087297976016998,
-0.09470323473215103,
0.2779513895511627,
0.03663017228245735,
0.36601993441581726,
0.8102841377258301,
0.6883901953697205,
-0.33066609501838684,
-0.34960171580314636,
0.923985481262207,
0.5853908061981201,
0.07039576023817062,
-0.11843020468950272,
-0.06797836720943451,
0.0974433571100235,
-0.4707315266132355,
0.37827417254447937,
0.15521520376205444,
-0.7403592467308044,
-0.25005313754081726,
0.596679151058197,
-0.7277861833572388,
-0.6915309429168701,
-0.0050544412806630135,
-0.12311484664678574,
0.04149714484810829,
0.05289606750011444,
0.2448417991399765,
-0.47261708974838257,
-0.3535511791706085,
0.4614925682544708,
0.9230178594589233,
-0.5351396799087524,
0.8224894404411316,
0.37244901061058044,
-0.08826857805252075,
-0.0452042818069458,
0.0035054143518209457,
0.09203510731458664,
0.08918709307909012,
-0.0694250762462616,
-0.053435735404491425,
0.1012222170829773,
0.3401939570903778,
-0.38458573818206787,
0.3040490746498108,
0.7614821791648865,
-0.17064380645751953,
0.22403603792190552,
0.08646601438522339,
-0.08289062976837158,
-0.20126193761825562,
0.2795524299144745,
0.13253425061702728,
-0.07332615554332733,
0.2151418924331665,
0.16798575222492218,
0.003749655559659004,
0.2437056005001068,
-0.09098415076732635,
0.18923071026802063,
0.07854695618152618,
-0.25417080521583557,
0.15693743526935577,
-0.30657434463500977,
-0.19041943550109863,
0.26519766449928284,
0.24278832972049713,
-0.18357035517692566,
-0.015992645174264908,
0.43973660469055176,
0.02785446122288704,
0.3032245934009552,
-0.021606506779789925,
-0.2682349383831024,
-0.10395143181085587,
0.050348248332738876,
0.12892353534698486,
-0.10498340427875519,
-0.027477847412228584,
0.09730125963687897,
-0.16150422394275665,
-0.21831916272640228,
0.10376061499118805,
-0.25544440746307373,
0.031593386083841324,
0.11986788362264633,
0.22690074145793915,
-0.3509098291397095,
-0.1881190538406372,
-0.04210145026445389,
0.6883101463317871,
-0.07829979062080383,
0.4657376706600189,
0.9263871908187866,
0.08322961628437042,
0.04429711028933525,
-0.08905605971813202,
-0.06788893789052963,
-0.056182388216257095,
-0.04881853610277176,
-0.04854113608598709,
0.15449045598506927,
0.32911357283592224,
-0.5772383809089661,
-0.00027374469209462404,
-0.2995521128177643,
-0.027322502806782722,
0.5023694038391113,
0.045783523470163345,
-0.4035968780517578,
0.053967904299497604,
0.00014662329340353608,
0.021607715636491776,
-0.028252260759472847,
-0.05918470770120621,
-0.1273883581161499,
0.0679078996181488,
0.25051605701446533,
-0.0745333656668663,
0.18680104613304138,
-0.12048312276601791,
0.013110226020216942,
-0.07659415900707245,
0.2906968295574188,
0.3136366307735443,
-0.47699007391929626,
0.02583535574376583,
-0.15701107680797577,
0.045304182916879654,
0.23456838726997375,
-0.06186807528138161,
0.3926846981048584,
-0.13252438604831696,
-0.16362214088439941,
0.013557562604546547,
-0.09991434961557388,
0.09150815010070801,
-0.006477471441030502,
0.2915862202644348,
0.5867642164230347,
-0.37984445691108704,
0.033169880509376526,
0.024414243176579475,
-0.0384003147482872,
-0.06395144015550613,
0.07380940765142441,
-0.025898484513163567,
0.03951931372284889,
-0.2343142330646515,
0.27318838238716125,
0.1105947494506836,
0.290696382522583,
-0.17851489782333374,
-0.17699271440505981,
-0.210996612906456,
-0.10575137287378311,
0.15886521339416504,
0.10631759464740753,
0.22946283221244812,
-0.3170112073421478,
-0.49773311614990234,
-0.10753292590379715,
-0.1114523783326149,
-0.10953730344772339,
0.4754663109779358,
0.20793643593788147,
0.021392812952399254,
-0.0691467821598053,
0.03368104621767998,
-0.017844771966338158,
0.1657843142747879,
-0.5556477904319763,
-1.108074426651001,
-0.822117805480957,
-0.06053074076771736,
-0.4072379469871521,
0.09109722077846527,
-0.5544739961624146,
-0.13978064060211182,
-0.36262163519859314,
0.20034632086753845,
0.050625383853912354,
0.1497042030096054,
-0.18745489418506622,
0.0894727036356926,
0.00417149206623435,
0.2228451371192932,
0.00852279644459486,
-0.028313757851719856,
0.04104698821902275,
-0.0874263271689415,
0.19788521528244019,
-0.019343160092830658,
-0.03962515667080879,
0.2092486023902893,
-0.44425246119499207,
-0.48542261123657227,
-0.04222029820084572,
0.7616084218025208,
0.512810468673706,
-0.17871123552322388,
0.5459727644920349,
-0.13069608807563782,
0.09155352413654327,
0.11548610031604767,
-0.15368784964084625,
0.038799818605184555,
-0.049028217792510986,
-0.03215758875012398,
-0.050522346049547195,
0.1663637012243271,
-0.15482299029827118,
-0.9425870180130005,
-0.7017998695373535,
0.04315050691366196,
-0.019968662410974503,
0.03749818727374077,
-0.07611791789531708,
0.32011789083480835,
-0.6925904750823975,
-0.49334919452667236,
0.23214411735534668,
1.1447347402572632,
-0.6757001876831055,
0.7940422296524048,
0.40169182419776917,
-0.018513813614845276,
0.048821814358234406,
-0.016693273559212685,
0.008068449795246124,
0.04566117003560066,
-0.09829569607973099,
-0.026971371844410896,
0.05381541699171066,
-0.3659301698207855,
0.3473235070705414,
0.14521746337413788,
0.11228122562170029,
-0.041056130081415176,
-0.11228874325752258,
0.006667478010058403,
0.15931302309036255,
-0.30010080337524414,
0.3464723229408264,
0.4476386308670044,
-0.3498152494430542,
0.2616507112979889,
-0.19995814561843872,
0.10946320742368698,
0.4034257233142853,
-0.08651446551084518,
0.018647747114300728,
0.11572548002004623,
-0.100877545773983,
-0.16341210901737213,
0.2377898246049881,
0.3417612910270691,
-0.49084869027137756,
-0.02805873565375805,
-0.09811390936374664,
0.17161016166210175,
0.3627470135688782,
-0.08954513072967529,
0.06629404425621033,
0.012786897830665112,
0.01578289456665516,
-0.32630467414855957,
0.4854920506477356,
0.12709765136241913,
-0.4909423291683197,
-0.3745254874229431,
-0.6513142585754395,
-0.040075208991765976,
-0.569782018661499,
-0.009953420609235764,
0.04735071584582329,
0.0230120699852705,
-0.07381311058998108,
-0.06293600797653198,
0.20196016132831573,
0.26551517844200134,
-0.42071688175201416,
0.28809165954589844,
0.19747501611709595,
-0.5686206221580505,
-0.5285986661911011,
0.02009684592485428,
0.11322621256113052,
-0.1082596555352211,
-0.0856761634349823,
-0.04493662342429161,
-0.6179490089416504,
-0.1442672610282898,
0.028762176632881165,
0.12426868081092834,
-0.5771384835243225,
0.1608373522758484,
0.004147801548242569,
-0.047590240836143494,
0.10347189754247665,
0.11780986934900284,
-0.08490656316280365,
-0.0746934711933136,
0.15699702501296997,
0.1298881322145462,
-0.14411042630672455,
-0.08601037412881851,
0.2997709810733795,
-0.05418943241238594,
-0.1772651970386505,
0.04576871916651726,
-0.13510753214359283,
-0.057203926146030426,
0.18647770583629608,
0.0055348677560687065,
-0.12238732725381851,
-0.11199415475130081,
0.43077343702316284,
0.1349855363368988,
0.21327465772628784,
0.05924845486879349,
0.12549948692321777,
-0.060076650232076645,
0.23921678960323334,
0.02152605727314949,
-0.1352948695421219,
0.09325127303600311,
-0.14411674439907074,
0.010495728813111782,
0.11577513813972473,
-0.07580242305994034,
0.42641204595565796,
-0.5557231903076172,
-0.12044595927000046,
0.024152765050530434,
-0.14175696671009064,
0.024960221722722054,
0.10017693042755127,
-0.07402117550373077,
0.09156208485364914,
0.455565482378006,
0.424320250749588,
-0.07668061554431915,
0.10318724811077118,
-0.32521969079971313,
-0.2653461694717407,
-0.03919212520122528,
0.12909358739852905,
-0.17091549932956696,
0.07353391498327255,
0.11510979384183884,
-0.23758216202259064,
-0.3059186339378357,
-0.046047650277614594,
0.17527209222316742,
0.19020265340805054,
-0.20766229927539825,
-0.23476286232471466,
-0.14011070132255554,
0.1085173636674881,
-0.020777594298124313,
0.014691418968141079,
0.21648286283016205,
-0.21576255559921265,
0.28203028440475464,
0.6320008635520935,
-0.23609709739685059,
0.16072526574134827,
0.30149686336517334,
-0.05675647035241127,
-0.018186205998063087,
-0.1844293773174286,
0.13510139286518097,
0.05780869722366333,
0.07202577590942383,
0.07459436357021332,
0.18700383603572845,
-0.09449177235364914,
0.057188909500837326,
0.21453143656253815,
-0.30002379417419434,
-0.12217795103788376,
0.03723505884408951,
-0.18360234797000885,
-0.029992947354912758,
0.10999765247106552,
0.09575961530208588,
-0.36028456687927246,
-0.4311397075653076,
0.5812231302261353,
],
b1: &[
0.3801889419555664,
-0.5001883506774902,
0.19484910368919373,
0.6488791704177856,
0.38620173931121826,
0.8780303597450256,
-0.1126403734087944,
0.021730314940214157,
-0.7806469202041626,
-0.04312174394726753,
0.3102167546749115,
0.9241658449172974,
0.8900863528251648,
-0.2938256561756134,
-0.5012822151184082,
-0.00329477502964437,
0.5169500708580017,
0.4563848376274109,
-0.4903448224067688,
0.27919942140579224,
-0.4288303554058075,
-0.1836952418088913,
-0.09118890762329102,
0.5528226494789124,
-0.19896377623081207,
0.33588215708732605,
0.07895006239414215,
0.07812929153442383,
0.6203332543373108,
0.8427650332450867,
-0.684628427028656,
0.5408275723457336,
-0.5548633933067322,
-0.49557214975357056,
0.7953769564628601,
-0.4109633266925812,
-0.6270897388458252,
-0.43285393714904785,
-0.7562689781188965,
-0.7167727947235107,
],
w2: &[
0.15592391788959503,
0.25119924545288086,
-0.499594122171402,
-0.5441639423370361,
-0.11186911165714264,
-0.6334478855133057,
0.28880706429481506,
-0.592946469783783,
0.7188563942909241,
-0.49322614073753357,
-0.1398385912179947,
-0.1868145614862442,
0.9288992881774902,
-0.07525540888309479,
0.2288437783718109,
0.09932874143123627,
0.2782813012599945,
-0.12644614279270172,
-0.14151062071323395,
0.38845404982566833,
0.2691279947757721,
-0.9148958921432495,
0.19230225682258606,
0.6098687052726746,
-0.24782557785511017,
-0.6989489197731018,
-0.30721813440322876,
-0.4890380799770355,
-0.43724432587623596,
-0.38428765535354614,
-0.6491377353668213,
-0.28134995698928833,
-0.36228886246681213,
-0.05963568389415741,
0.5086851119995117,
0.4664144814014435,
0.3797634541988373,
0.5596290826797485,
-0.1977449357509613,
0.6540879607200623,
-0.24533972144126892,
0.6865915656089783,
-0.18364377319812775,
0.0013501447392627597,
-0.4037604331970215,
-0.287411093711853,
-0.43570032715797424,
-0.4085054099559784,
0.7341827750205994,
-0.29973891377449036,
-0.18240050971508026,
-0.23446109890937805,
0.7225431799888611,
0.008502814918756485,
0.04582007974386215,
0.03352205455303192,
0.12457727640867233,
-0.2019437849521637,
-0.1299249827861786,
-0.09946829080581665,
0.40665051341056824,
-0.6841736435890198,
-0.523845911026001,
0.21656402945518494,
0.6046024560928345,
-0.6393186450004578,
-0.3965637981891632,
-0.7872777581214905,
-0.13687947392463684,
-0.19312888383865356,
-0.5453231930732727,
-0.21912647783756256,
0.011589044705033302,
0.2665385603904724,
0.3249806761741638,
0.293254017829895,
0.1047254130244255,
0.4246895909309387,
-0.0033608688972890377,
0.4066942632198334,
0.06138676777482033,
0.382074236869812,
0.0787188857793808,
-0.28631800413131714,
-0.3500039279460907,
-0.1490340679883957,
-0.14991725981235504,
-0.180477574467659,
0.15140952169895172,
-0.35168370604515076,
0.38904908299446106,
-0.11262823641300201,
-0.18404939770698547,
0.5045862197875977,
0.23344825208187103,
0.6740546226501465,
-0.054060351103544235,
-0.47260594367980957,
0.287933886051178,
0.28975099325180054,
0.2366262525320053,
-0.1751112937927246,
-0.15358465909957886,
-0.062381260097026825,
0.45881521701812744,
-0.12647950649261475,
0.45258036255836487,
-0.21084383130073547,
-0.15994171798229218,
-0.4229416847229004,
-0.18642400205135345,
-0.2506699860095978,
0.20604389905929565,
0.16662882268428802,
-0.23073841631412506,
0.045810505747795105,
0.33520498871803284,
0.37685254216194153,
0.11563336104154587,
0.22259201109409332,
-0.010484708473086357,
-0.45855188369750977,
0.24794596433639526,
0.33667632937431335,
0.20378778874874115,
0.4198003113269806,
0.23384596407413483,
0.23601709306240082,
-0.509751558303833,
0.5694931149482727,
-0.08933047205209732,
0.037133198231458664,
0.20635388791561127,
-0.2857131361961365,
-0.4278101921081543,
-0.26602792739868164,
0.1998632550239563,
0.4324374794960022,
-0.13389578461647034,
0.11837134510278702,
-0.17028754949569702,
0.37928706407546997,
0.10062910616397858,
-0.04736608266830444,
-0.04692180082201958,
0.6633663773536682,
-0.3517492711544037,
0.2055688351392746,
0.44142597913742065,
0.42460545897483826,
0.4567111134529114,
0.3061029016971588,
-0.16390416026115417,
-0.3541538417339325,
0.2544074058532715,
-0.18162837624549866,
-0.21904821693897247,
-0.2520917057991028,
-0.07266020774841309,
-0.23432950675487518,
-0.1989256739616394,
0.09460597485303879,
-0.24563294649124146,
0.9719013571739197,
0.2578149139881134,
0.26680076122283936,
-0.39480605721473694,
0.22382304072380066,
-0.4284250736236572,
0.4294125437736511,
-0.04923247918486595,
0.5011574625968933,
0.1887599676847458,
-0.02984841726720333,
-0.16428305208683014,
-0.33957910537719727,
-0.16184143722057343,
0.37313663959503174,
-0.11775537580251694,
-0.34507161378860474,
-0.24848994612693787,
0.3492432236671448,
-0.2122095823287964,
-0.022055158391594887,
0.07298140972852707,
0.36230477690696716,
-0.2514148950576782,
0.11675992608070374,
0.4010731875896454,
0.31790846586227417,
0.0585796944797039,
0.30878275632858276,
0.5536429286003113,
-0.061644136905670166,
-0.06381722539663315,
-0.1873038411140442,
-0.24746698141098022,
-0.3139619529247284,
-0.19278131425380707,
-0.48264867067337036,
0.5122742056846619,
0.09536745399236679,
0.17870695888996124,
0.18145892024040222,
0.2471739798784256,
-0.16399677097797394,
-0.18874068558216095,
0.21305255591869354,
-0.6930050253868103,
-0.4031701982021332,
0.5250658392906189,
0.4295860230922699,
-0.464653879404068,
-0.026941847056150436,
-0.08213993161916733,
0.34638163447380066,
-0.15401627123355865,
0.021148433908820152,
0.19726167619228363,
-0.25100240111351013,
3.085673233726993e-05,
0.16563303768634796,
-0.008333534933626652,
-0.02890022285282612,
-0.284770667552948,
0.3429299592971802,
0.6073935627937317,
-0.10915102809667587,
0.3420248329639435,
0.07347360253334045,
0.18400518596172333,
0.2084905058145523,
0.3218590021133423,
0.16883575916290283,
-0.6880696415901184,
-0.37455135583877563,
0.04792584478855133,
-0.04572531208395958,
-0.17001567780971527,
-0.12369263172149658,
-0.3716808259487152,
-0.04167286679148674,
0.04307235777378082,
-0.1655367612838745,
-0.47902533411979675,
-0.21886907517910004,
0.4065888226032257,
0.30626556277275085,
0.25965678691864014,
0.07168732583522797,
-0.17138782143592834,
-0.6293558478355408,
-0.6350710988044739,
0.25923609733581543,
0.5668261647224426,
-0.030662082135677338,
-0.7059182524681091,
-0.25901535153388977,
0.25449642539024353,
-0.3232290744781494,
0.42758384346961975,
0.7120643258094788,
0.023215001448988914,
-0.40807682275772095,
0.1332295536994934,
-0.33705568313598633,
0.1038941740989685,
0.39904412627220154,
-0.567590057849884,
-0.26575762033462524,
0.7635160088539124,
-0.38967835903167725,
-0.08988548815250397,
0.4150312840938568,
-0.540441632270813,
0.33467426896095276,
-0.03507159277796745,
0.00720902718603611,
0.6702240109443665,
0.2707512676715851,
],
b2: &[
0.3580038547515869,
0.06861710548400879,
-0.04651366174221039,
0.24638813734054565,
0.1557426154613495,
-0.40271297097206116,
-0.405432790517807,
],
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,861 @@
use crate::frame::mmm::CostModel;
pub fn model() -> CostModel<'static> {
CostModel {
big_product_mkn_threshold: 4194036.0,
big_product_kernel_choice: "armv7neon_mmm_f32_8x6_cortexa9",
kernels: &[
"armv7neon_mmm_f32_8x4_cortexa7",
"armv7neon_mmm_f32_8x4_cortexa9",
"armv7neon_mmm_f32_8x4_generic",
"armv7neon_mmm_f32_8x6_cortexa7",
"armv7neon_mmm_f32_8x6_cortexa9",
"armv7neon_mmm_f32_8x6_generic",
"generic_f32_4x4",
],
mrs: &[4, 8],
nrs: &[4, 6],
feat_norm_mean: &[
4.582296677813486,
4.595402322442016,
4.571260231028445,
13.748959231283994,
1.5179177668804225,
0.7575757575757576,
3.5337608449641644,
0.8831887338111405,
1.5048409405255878,
0.7526719476926946,
2.489123601156796,
0.8326417704011065,
],
feat_norm_stddev: &[
1.2635817489024164,
1.2723436827339079,
1.2620157548883217,
1.3497763942449361,
1.1141159992246472,
0.42854956435545316,
2.2880460409304937,
0.32119525880720723,
1.1154901716833412,
0.43145902105435263,
1.7051378780434328,
0.37329539587896904,
],
w1: &[
0.5391961336135864,
-0.32089367508888245,
0.203999862074852,
-0.10011337697505951,
0.09040801972150803,
-0.14198464155197144,
0.031854499131441116,
0.12334256619215012,
0.15339604020118713,
-0.20091375708580017,
-0.014548280276358128,
0.12154694646596909,
0.31225234270095825,
0.10782113671302795,
0.44618168473243713,
0.8267014026641846,
-0.1204405128955841,
-0.08261110633611679,
-0.052502430975437164,
0.3066086769104004,
0.1493932157754898,
-0.14119412004947662,
-0.1985343098640442,
0.19361039996147156,
-0.4636686146259308,
0.08120443671941757,
0.03210291638970375,
0.17303235828876495,
0.16502155363559723,
-0.19771894812583923,
-0.11060577630996704,
0.08698348701000214,
-0.07793445140123367,
0.32749465107917786,
0.3663202226161957,
-0.4629170894622803,
-0.1586134433746338,
0.4272242486476898,
-0.12016090005636215,
-0.17830348014831543,
-0.05493386462330818,
-0.036517318338155746,
0.01293050218373537,
0.016577009111642838,
0.10738552361726761,
-0.3662779629230499,
-0.2917434275150299,
0.5752639770507812,
0.11406347155570984,
0.8622727394104004,
0.07158719748258591,
0.29530274868011475,
-0.11287810653448105,
0.12262264639139175,
0.02478562481701374,
0.17749948799610138,
-0.036227867007255554,
0.10140471905469894,
-0.011896232143044472,
-0.021761735901236534,
0.06046223267912865,
0.5727048516273499,
-0.007826486602425575,
0.3863913118839264,
-0.04224887117743492,
0.056023009121418,
-0.02467598207294941,
0.0385640449821949,
0.0219524335116148,
-0.03437826409935951,
-0.2060588151216507,
0.2895224988460541,
0.10751669108867645,
0.00845037866383791,
-0.1836385875940323,
-0.24757762253284454,
-0.09606243669986725,
0.03918633610010147,
0.07913251221179962,
0.06499160826206207,
-0.08156774938106537,
0.08835449814796448,
0.13896305859088898,
-0.16936920583248138,
0.010146846994757652,
-0.42553824186325073,
0.39916151762008667,
-0.004584060981869698,
-0.10256388038396835,
0.041573416441679,
0.05155385658144951,
0.015019520185887814,
0.09554271399974823,
-0.20487457513809204,
-0.4146610200405121,
-0.773110032081604,
0.3662724494934082,
-0.23762361705303192,
0.6974321603775024,
0.8990052938461304,
0.02772649936378002,
0.042197681963443756,
-0.0022736566606909037,
-0.028843341395258904,
-0.4559306204319,
0.6326258778572083,
0.4568879008293152,
-0.4892531633377075,
-0.032289132475852966,
0.04378330707550049,
-0.4118069112300873,
0.2493579089641571,
-0.021955665200948715,
-0.01538186427205801,
-0.21400974690914154,
-0.09971866756677628,
0.02185226045548916,
-0.18125569820404053,
-0.13828244805335999,
-0.20846466720104218,
-0.10373540222644806,
0.4842098653316498,
-0.06586655229330063,
0.03369470313191414,
0.013142148964107037,
0.017437899485230446,
0.15891534090042114,
0.5269678831100464,
0.02546108327805996,
-0.004250233061611652,
-5.8676625485531986e-05,
0.06777831166982651,
-0.14051207900047302,
0.6876491904258728,
-0.3455996811389923,
0.0378129817545414,
0.15291574597358704,
-0.03829087316989899,
-0.05761529877781868,
-0.05344394966959953,
0.1421334147453308,
-0.3614322543144226,
-0.21606910228729248,
0.1558765172958374,
0.14480257034301758,
-0.1799984872341156,
0.4238421618938446,
-0.08961529284715652,
-0.04010967165231705,
0.14250615239143372,
-0.0038367861416190863,
-0.044531334191560745,
-0.08958051353693008,
-0.1577986180782318,
-0.5795103907585144,
-1.1048516035079956,
0.16444185376167297,
-0.09989812225103378,
-0.26304998993873596,
0.040687527507543564,
0.065303735435009,
-0.06267901510000229,
0.08742637187242508,
0.02480895072221756,
0.23719966411590576,
-0.09509539604187012,
0.39278310537338257,
0.18978112936019897,
0.11301649361848831,
-0.16268616914749146,
-0.14119602739810944,
-0.04518252611160278,
0.10456270724534988,
0.008367948234081268,
0.004280170891433954,
0.01894286274909973,
-0.1547478288412094,
0.197267547249794,
0.20271208882331848,
-0.28377917408943176,
-0.26751258969306946,
0.15954937040805817,
0.33988064527511597,
0.16848208010196686,
0.11668887734413147,
-0.057433612644672394,
-0.049777109175920486,
0.00744214653968811,
-0.012330793775618076,
-0.08413149416446686,
-0.2053118497133255,
0.09235486388206482,
-0.1354941576719284,
0.41610953211784363,
0.8428494334220886,
0.880882740020752,
0.024029193446040154,
-0.08453702926635742,
0.00771496444940567,
-0.013013732619583607,
-0.23804998397827148,
0.4110376536846161,
0.23720477521419525,
-0.13951541483402252,
-0.1747516244649887,
-0.34215790033340454,
0.014357345178723335,
0.34224632382392883,
0.03783192113041878,
0.01125166192650795,
-0.08253959566354752,
0.015717405825853348,
-0.22759634256362915,
0.3980898857116699,
0.2427154779434204,
-0.3319437801837921,
0.11146843433380127,
-0.9666317105293274,
-0.12227121740579605,
-0.1948898285627365,
-0.030186548829078674,
0.0011711223050951958,
-0.040062546730041504,
-0.16316139698028564,
-0.14714862406253815,
0.13224393129348755,
-0.0019320327555760741,
-0.09674090147018433,
0.3630145490169525,
-0.019513679668307304,
-0.07729464769363403,
-0.34592965245246887,
0.15215164422988892,
0.046678490936756134,
0.06675180792808533,
-0.08943335711956024,
0.006386714521795511,
0.10086977481842041,
-0.07409387081861496,
-0.19604018330574036,
-0.042700666934251785,
0.12124726921319962,
0.5694677233695984,
0.25033196806907654,
0.01862989366054535,
0.0053687929175794125,
-0.0017405126709491014,
-0.01638556271791458,
-0.32222822308540344,
0.5348804593086243,
0.5546748042106628,
1.2770946025848389,
0.11648745834827423,
-0.058405984193086624,
-0.2997635006904602,
-0.2040756195783615,
0.15525077283382416,
-0.12436354905366898,
-0.089121975004673,
0.06441225856542587,
0.2444663643836975,
-0.3495825529098511,
-0.05243751034140587,
0.08752834796905518,
0.08800745010375977,
-0.09807545691728592,
-0.3823537230491638,
-0.13047000765800476,
0.029333092272281647,
0.11618250608444214,
-0.0638590008020401,
-0.09598273783922195,
-0.07390140742063522,
0.09151650220155716,
-0.1700282245874405,
0.23608872294425964,
0.24879834055900574,
-0.15922772884368896,
-0.33795130252838135,
-0.053850702941417694,
0.1014639139175415,
-0.05480973795056343,
-0.06753639131784439,
0.04606246575713158,
-0.07082260400056839,
0.07848796248435974,
0.05011916160583496,
-0.05570689216256142,
-0.14584510028362274,
-0.8908579349517822,
-0.5959509611129761,
-0.8982105255126953,
0.0788002535700798,
-0.03575791418552399,
0.052424680441617966,
-0.08019822835922241,
0.10848221182823181,
0.0957408994436264,
0.1457311511039734,
-0.1956494003534317,
-0.21669772267341614,
0.9854136109352112,
-0.23215851187705994,
0.16359730064868927,
0.02025810070335865,
-0.08975380659103394,
-0.013868067413568497,
-0.22188447415828705,
0.020666224882006645,
-0.22304703295230865,
0.06407633423805237,
0.19804184138774872,
-0.05285267159342766,
-0.5510660409927368,
-0.8522927761077881,
-0.6061599850654602,
0.08484024554491043,
-0.08973539620637894,
0.013228937052190304,
-0.07834818214178085,
0.02858446165919304,
-0.3826225996017456,
0.059726644307374954,
0.1139102503657341,
-0.19311848282814026,
0.05770142376422882,
0.22584261000156403,
0.34312352538108826,
-0.15085645020008087,
0.34372228384017944,
0.08070214092731476,
0.5744000673294067,
-0.08693907409906387,
-0.003695777617394924,
-0.1334235966205597,
0.06418291479349136,
0.02848576195538044,
-0.34958112239837646,
-0.3419312834739685,
-0.09599799662828445,
0.015022341161966324,
0.03255023807287216,
0.09713662415742874,
-0.1730588674545288,
0.1904430240392685,
-0.32815566658973694,
-0.16749203205108643,
0.35736411809921265,
-0.503787100315094,
0.5057004690170288,
-0.47198373079299927,
0.11386436969041824,
-0.0722493901848793,
0.03358639404177666,
0.005928087048232555,
-0.05637047439813614,
0.06552420556545258,
-0.07283362001180649,
-0.09314802289009094,
0.13586974143981934,
-0.5054865479469299,
-0.18127793073654175,
0.08853171765804291,
-0.13333705067634583,
-0.2623322308063507,
0.17757390439510345,
0.04408252611756325,
-0.0277855321764946,
-0.05175777152180672,
0.40444689989089966,
-0.03518976643681526,
-0.36402902007102966,
-0.019589770585298538,
-0.05277400091290474,
-0.27273234724998474,
-0.07373850792646408,
-0.058221735060214996,
0.14292845129966736,
-0.005004828795790672,
-0.05554938316345215,
0.20361287891864777,
-0.30462127923965454,
-0.1140812486410141,
0.16081976890563965,
-0.07133162021636963,
-0.20463652908802032,
0.34733739495277405,
0.17099761962890625,
0.025868643075227737,
-0.02960631065070629,
-0.02717636525630951,
0.02027258090674877,
-0.13165302574634552,
0.36201152205467224,
0.5002728700637817,
0.39691421389579773,
-0.04605599492788315,
0.28801581263542175,
-1.0140656232833862,
-0.5481916666030884,
0.0896061584353447,
-0.049390073865652084,
0.08813252300024033,
-0.1784677952528,
0.34480658173561096,
-0.36402803659439087,
0.16948284208774567,
0.45740315318107605,
-0.23747704923152924,
0.580975353717804,
-0.24338461458683014,
-0.11410018056631088,
0.06431885808706284,
-0.0317281149327755,
-0.024683356285095215,
-0.10083278268575668,
0.024547407403588295,
-0.16270779073238373,
-0.07757837325334549,
0.19732129573822021,
0.03790999948978424,
-0.18804220855236053,
0.8675169348716736,
0.5377629399299622,
-0.0036910742055624723,
-0.0016441351035609841,
-0.030448857694864273,
0.07757671177387238,
-0.1475408971309662,
0.613543689250946,
0.30266445875167847,
0.12106148898601532,
0.05485830456018448,
-0.04748840630054474,
-0.23233623802661896,
-0.1949906051158905,
0.05692804977297783,
0.07474583387374878,
-0.11879625171422958,
0.07200933247804642,
-0.012743310071527958,
-0.02546215057373047,
-0.3765566349029541,
0.28637346625328064,
-0.18051809072494507,
0.5034835934638977,
-0.34970414638519287,
-0.2386687994003296,
-0.03804561868309975,
-0.03649319335818291,
-0.10303670912981033,
0.1299818456172943,
0.24685724079608917,
-0.34168556332588196,
-0.086674265563488,
0.32085898518562317,
0.48488491773605347,
-0.522548258304596,
0.309568852186203,
0.167385995388031,
0.11308691650629044,
0.14733079075813293,
-0.22416195273399353,
0.14763982594013214,
-0.07242503017187119,
0.07601745426654816,
-0.10375087708234787,
-0.03409396857023239,
-0.35759225487709045,
0.18936687707901,
0.28248289227485657,
0.26482364535331726,
0.061123836785554886,
-0.021603189408779144,
-0.13469825685024261,
0.07248867303133011,
-0.03464066982269287,
0.06557167321443558,
0.16093865036964417,
-0.1718607246875763,
],
b1: &[
-0.3893989324569702,
-0.2791002690792084,
0.07853052020072937,
-0.4629746377468109,
-0.7148261070251465,
0.8680436015129089,
-0.46459102630615234,
0.0404132716357708,
-0.44012945890426636,
0.08434166759252548,
0.32190972566604614,
-0.20194832980632782,
-0.3781348764896393,
-0.23968002200126648,
-0.581799328327179,
0.6500483155250549,
-0.6192854046821594,
0.5922245383262634,
0.44006091356277466,
0.2982949912548065,
0.6136102676391602,
-0.597486138343811,
-0.3697699308395386,
-0.45241132378578186,
0.60771644115448,
-0.3373708128929138,
0.5697194337844849,
0.4784911870956421,
-0.49601855874061584,
0.5023709535598755,
0.21592296659946442,
-0.45412343740463257,
0.5104787945747375,
0.558862566947937,
0.4729066491127014,
-0.5520593523979187,
-0.5120576620101929,
-0.7157037258148193,
0.12596718966960907,
0.4773174524307251,
],
w2: &[
0.1379607617855072,
0.09308824688196182,
-0.2596932649612427,
0.4461972713470459,
0.3480601906776428,
0.036684323102235794,
0.4057384729385376,
-0.3081648051738739,
0.4561280608177185,
0.2749394178390503,
-0.1400817334651947,
0.3145979046821594,
-0.16919250786304474,
0.7247185707092285,
0.3479674756526947,
-0.7546817064285278,
0.38135531544685364,
-0.3939172029495239,
-0.038021210581064224,
0.026914050802588463,
-0.5281358361244202,
0.39009571075439453,
0.4090450406074524,
0.5053343772888184,
-0.23938016593456268,
0.488080233335495,
-0.38536468148231506,
-0.23763014376163483,
0.2661689519882202,
-0.14746293425559998,
-0.7541974186897278,
0.27726081013679504,
-0.4072169065475464,
-0.8030230402946472,
-0.386343389749527,
0.6674754619598389,
0.06677238643169403,
0.5055669546127319,
-0.44330647587776184,
-0.3423362970352173,
-0.10948927700519562,
0.11290912330150604,
-0.2759379744529724,
0.5522158741950989,
-0.5766478776931763,
0.7288797497749329,
-0.4967955946922302,
-0.5466133952140808,
0.7254890203475952,
0.1274457424879074,
0.3098924458026886,
0.2524661719799042,
-0.7162019610404968,
0.19503603875637054,
-0.5212412476539612,
0.0968603864312172,
0.4835629463195801,
-0.5865079164505005,
0.27647316455841064,
0.1975109577178955,
-0.845225989818573,
0.4172143042087555,
-0.014424118213355541,
-0.24702520668506622,
-0.16123531758785248,
-0.047759659588336945,
-0.09985388815402985,
0.10430619865655899,
0.53556889295578,
0.2595883011817932,
0.11729882657527924,
0.36996161937713623,
-0.41997936367988586,
-0.3332042694091797,
0.2527308464050293,
0.6039140820503235,
0.35183605551719666,
0.42042237520217896,
-0.2265913337469101,
-0.06852111965417862,
0.3749903440475464,
0.3698897361755371,
-0.43096107244491577,
0.1275794953107834,
0.27926334738731384,
-0.3282606303691864,
0.290679931640625,
-0.14467079937458038,
0.3357028663158417,
-0.0683436468243599,
-0.35492125153541565,
-0.14275093376636505,
-0.1504347324371338,
0.1782987266778946,
0.07464402168989182,
-0.2788643538951874,
0.5896115303039551,
-0.314520001411438,
-0.3235827684402466,
-0.2899278700351715,
-0.21264874935150146,
0.41862159967422485,
0.3237628936767578,
0.2948566973209381,
-0.6101413369178772,
-0.025511808693408966,
-0.4238346517086029,
-0.28283095359802246,
0.32077667117118835,
-0.34138476848602295,
-0.5257527232170105,
0.24129967391490936,
-0.38175472617149353,
-0.20559589564800262,
-0.11267697811126709,
0.32475054264068604,
0.29545050859451294,
0.0010625360300764441,
0.4097916781902313,
-0.3120468556880951,
0.3134985566139221,
0.33620578050613403,
-0.27408266067504883,
-0.0118736382573843,
0.21356475353240967,
-0.6716119647026062,
0.14166241884231567,
0.020748334005475044,
0.27158322930336,
-0.27066248655319214,
-0.5078546404838562,
0.39642488956451416,
0.4044502079486847,
0.1363500952720642,
0.38089585304260254,
-0.18438327312469482,
-0.08652642369270325,
0.05718545988202095,
-0.5758764743804932,
0.0948563665151596,
0.298057496547699,
-0.07299521565437317,
-0.24248233437538147,
0.29135069251060486,
-0.44556060433387756,
0.6689074039459229,
-0.12930674850940704,
-0.12669484317302704,
0.1074564978480339,
-0.20472179353237152,
0.14787982404232025,
-0.13180267810821533,
0.3045596182346344,
-0.3345180153846741,
-0.3405822217464447,
0.22327540814876556,
0.02809770777821541,
0.17404714226722717,
0.22873322665691376,
-0.3915692865848541,
-0.39005470275878906,
-0.4675980806350708,
0.44798821210861206,
-0.31790846586227417,
-0.21734853088855743,
0.2172199934720993,
-0.3485357165336609,
0.1241735890507698,
-0.6933310031890869,
-0.09649480134248734,
0.24731965363025665,
-0.20421941578388214,
0.13033808767795563,
-0.4282769560813904,
-0.22173112630844116,
0.08912057429552078,
-0.3927532434463501,
0.3523387908935547,
0.36073970794677734,
-0.036902282387018204,
0.5880261063575745,
-0.29945725202560425,
-0.40845751762390137,
-0.3265145421028137,
0.370391309261322,
-0.3553546965122223,
0.5133077502250671,
0.1800842434167862,
-0.34683868288993835,
0.28811708092689514,
0.3033837080001831,
-0.4140017628669739,
0.4362258017063141,
0.3689269423484802,
0.3121638596057892,
-0.3287503123283386,
-0.15226924419403076,
-0.17191028594970703,
-0.10683685541152954,
0.34219542145729065,
0.34955963492393494,
0.22892920672893524,
-0.20123478770256042,
-0.3934169411659241,
0.25449705123901367,
-0.541163444519043,
0.21640898287296295,
0.19343338906764984,
-0.14020974934101105,
0.010480044409632683,
-0.24229897558689117,
-0.4682120084762573,
0.02336042746901512,
0.039344485849142075,
0.42446646094322205,
-0.3173693120479584,
0.23609045147895813,
0.20335273444652557,
-0.19347436726093292,
-0.05698636546730995,
0.17990583181381226,
0.30915674567222595,
0.3115670382976532,
0.4147215485572815,
-0.38558056950569153,
-0.12379863113164902,
0.025996098294854164,
-0.3010733425617218,
0.03275908902287483,
-0.6039671897888184,
0.06267470866441727,
-0.012677585706114769,
0.3484704792499542,
0.24301587045192719,
-0.40881243348121643,
-0.16732162237167358,
0.190901979804039,
-0.5619192719459534,
0.30009278655052185,
-0.43359509110450745,
0.26643550395965576,
0.5083268880844116,
0.3491555452346802,
0.4731655716896057,
0.6301924586296082,
-0.8111121654510498,
0.6473397016525269,
-0.001451796037144959,
0.3649038076400757,
-0.6002859473228455,
-0.41925248503685,
0.05584913119673729,
0.7823511362075806,
0.421135276556015,
0.5779385566711426,
-0.49475061893463135,
0.5293950438499451,
-0.45432502031326294,
-0.680946946144104,
-0.3506624102592468,
-0.21028658747673035,
0.4775547385215759,
0.25049126148223877,
0.2707470655441284,
-0.3469635546207428,
0.5959001779556274,
-0.5623777508735657,
-0.6334168910980225,
0.4096938669681549,
-0.3921370208263397,
-0.27649807929992676,
0.4424516260623932,
-0.28308066725730896,
-0.22009265422821045,
-0.386872798204422,
0.5130718350410461,
0.5702601075172424,
0.7469420433044434,
-0.09606175124645233,
-0.4271978437900543,
],
b2: &[
-0.07522959262132645,
0.3644154667854309,
-0.25166040658950806,
-0.12973527610301971,
0.25026997923851013,
-0.2794199585914612,
-0.17614373564720154,
],
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,182 @@
#![allow(clippy::excessive_precision)]
mod arm64simd;
pub mod cortex_a53;
mod cortex_a55;
//mod cortex_a72;
//mod cortex_a73;
pub use arm64simd::*;
use crate::Ops;
use crate::frame::element_wise::ElementWiseKer;
use crate::frame::mmm::kernel::MatMatMulKer;
lazy_static::lazy_static! {
static ref KIND: Kind = Kind::choose();
}
// https://en.wikipedia.org/wiki/Comparison_of_ARMv8-A_cores
const PART_A53: &str = "0xd03";
const PART_A55: &str = "0xd05";
#[allow(dead_code)]
const PART_A72: &str = "0xd08";
#[allow(dead_code)]
const PART_A73: &str = "0xd09";
#[allow(dead_code)]
const PART_A75: &str = "0xd0a";
fn max_cpuid() -> std::io::Result<String> {
let cpu_info = std::fs::read_to_string("/proc/cpuinfo")?;
let max = cpu_info
.lines()
.filter(|line| line.starts_with("CPU part"))
.map(|line| line.split_whitespace().last().unwrap_or(""))
.max();
Ok(max.unwrap_or("").to_string())
}
#[inline]
pub fn has_fp16() -> bool {
cfg!(feature_cpu = "fp16")
}
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
enum Kind {
Generic,
AppleM,
CortexA53,
CortexA55,
CortexA72,
CortexA73,
CortexA75,
}
impl Kind {
fn choose() -> Kind {
let kind = if let Ok(kind) = std::env::var("TRACT_CPU_AARCH64_KIND") {
log::info!("CPU kind forced with TRACT_CPU_AARCH64_KIND: {}", kind);
let kind = kind.to_lowercase();
if kind.contains("a53") {
Kind::CortexA53
} else if kind.contains("a55") {
Kind::CortexA55
} else if kind.contains("a72") {
Kind::CortexA72
} else if kind.contains("a73") {
Kind::CortexA73
} else if kind.contains("a75") {
Kind::CortexA75
} else if kind.contains("applem") {
Kind::AppleM
} else {
Kind::Generic
}
} else if cfg!(target_os = "macos") {
Kind::AppleM
} else {
let part = if let Ok(part) = std::env::var("TRACT_CPU_AARCH64_OVERRIDE_CPU_PART") {
log::info!(
"CPU part forced with TRACT_CPU_AARCH64_OVERRIDE_CPU_PART: {}",
part
);
part
} else if cfg!(target_os = "linux") {
let part = max_cpuid().unwrap_or_else(|_| "0x00".to_string());
log::info!("CPU part auto detected: {}", part);
part
} else {
log::info!("Unknown CPU part");
"0x00".to_string()
};
match &*part {
PART_A53 => Kind::CortexA53,
PART_A55 => Kind::CortexA55,
PART_A72 => Kind::CortexA72,
PART_A73 => Kind::CortexA73,
PART_A75 => Kind::CortexA75,
_ => Kind::Generic,
}
};
log::info!("CPU optimisation: {:?}", kind);
kind
}
}
pub fn plug(ops: &mut Ops) {
let impls = vec![
arm64simd_mmm_f32_12x8_gen::mmm(),
arm64simd_mmm_f32_12x8_a53::mmm(),
arm64simd_mmm_f32_12x8_a55::mmm(),
arm64simd_mmm_f32_8x8_gen::mmm(),
arm64simd_mmm_f32_8x8_a53::mmm(),
arm64simd_mmm_f32_8x8_a55::mmm(),
arm64simd_mmm_f32_16x4_gen::mmm(),
arm64simd_mmm_f32_16x4_a53::mmm(),
arm64simd_mmm_f32_16x4_a55::mmm(),
arm64simd_mmm_f32_24x4_gen::mmm(),
arm64simd_mmm_f32_24x4_a53::mmm(),
arm64simd_mmm_f32_24x4_a55::mmm(),
crate::generic::mmm::generic_f32_4x4::mmm(),
];
ops.mmm_f32_impls = impls.clone();
ops.qmmm_i32 = Box::new(|_, _, _| arm64simd_mmm_i32_8x8::mmm());
ops.qmmv_i32 = Box::new(|_, _| arm64simd_mmm_i32_64x1::mmm());
ops.mmv_f32 = match *KIND {
Kind::CortexA53 => Box::new(|_, _| arm64simd_mmm_f32_64x1_a53::mmm()),
Kind::CortexA55 => Box::new(|_, _| arm64simd_mmm_f32_64x1_a55::mmm()),
_ => Box::new(|_, _| arm64simd_mmm_f32_64x1_gen::mmm()),
};
let model = match *KIND {
Kind::CortexA53 => Some(cortex_a53::model()),
Kind::CortexA55 => Some(cortex_a55::model()),
_ => None,
};
ops.mmm_f32 = if let Some(model) = model {
Box::new(move |m, k, n| model.pick(&impls, m, k, n))
} else {
Box::new(move |_, _, n| {
if n.unwrap_or(8) < 8 {
arm64simd_mmm_f32_16x4_gen::mmm()
} else {
arm64simd_mmm_f32_8x8_gen::mmm()
}
})
};
#[cfg(feature = "no_fp16")]
if has_fp16() {
log::warn!(
"This is a build with fp16 disabled, while your platform CPU seems to support it."
);
}
#[cfg(not(feature = "no_fp16"))]
if has_fp16() {
if *KIND == Kind::CortexA55 {
ops.mmm_f16 = Box::new(|_, _, n| {
use tract_data::internal::DimLike;
if n.unwrap_or(1024).divceil(4) * 4 < n.unwrap_or(1024).divceil(8) * 8 {
arm64fp16_mmm_f16_32x4_a55::mmm()
} else {
arm64fp16_mmm_f16_16x8_a55::mmm()
}
});
ops.mmv_f16 = Box::new(|_, _| arm64fp16_mmm_f16_128x1_a55::mmm());
} else {
ops.mmm_f16 = Box::new(|_, _, n| {
use tract_data::internal::DimLike;
if n.unwrap_or(1024).divceil(4) * 4 < n.unwrap_or(1024).divceil(8) * 8 {
arm64fp16_mmm_f16_32x4_gen::mmm()
} else {
arm64fp16_mmm_f16_16x8_gen::mmm()
}
});
ops.mmv_f16 = Box::new(|_, _| arm64fp16_mmm_f16_128x1_gen::mmm());
}
}
ops.sigmoid_f32 = Box::new(|| arm64simd_sigmoid_f32_4n::ew());
ops.tanh_f32 = Box::new(|| arm64simd_tanh_f32_4n::ew());
#[cfg(not(feature = "no_fp16"))]
if has_fp16() {
ops.tanh_f16 = Box::new(|| arm64fp16_tanh_f16_8n::ew());
ops.sigmoid_f16 = Box::new(|| arm64fp16_sigmoid_f16_8n::ew());
}
}
@@ -0,0 +1,52 @@
use crate::frame::element_wise::ElementWiseKer;
use crate::frame::mmm::*;
#[cfg(not(feature = "no_fp16"))]
use tract_data::half::f16;
MMMKernel!(f32, arm64simd_mmm_f32_8x8_a55; 8, 8; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_12x8_a55; 12, 8; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_16x4_a55; 16, 4; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_24x4_a55; 24, 4; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_64x1_a55; 64, 1; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_16x4_a53; 16, 4; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_24x4_a53; 24, 4; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_8x8_a53; 8, 8; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_12x8_a53; 12, 8; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_64x1_a53; 64, 1; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_16x4_gen; 16, 4; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_24x4_gen; 24, 4; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_8x8_gen; 8, 8; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_12x8_gen; 12, 8; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(f32, arm64simd_mmm_f32_64x1_gen; 64, 1; 16, 16; 1, 1; no_prefetch, true);
MMMKernel!(i32, arm64simd_mmm_i32_8x8; 8, 8; 16, 16; 0,0; no_prefetch, true);
MMMKernel!(i32, arm64simd_mmm_i32_64x1; 64, 1; 16, 1; 0,0; no_prefetch, true);
#[cfg(not(feature = "no_fp16"))]
MMMKernel!(f16, arm64fp16_mmm_f16_16x8_gen; 16, 8; 16, 16; 1, 1; no_prefetch, crate::arm64::has_fp16());
#[cfg(not(feature = "no_fp16"))]
MMMKernel!(f16, arm64fp16_mmm_f16_16x8_a55; 16, 8; 16, 16; 1, 1; no_prefetch, crate::arm64::has_fp16());
#[cfg(not(feature = "no_fp16"))]
MMMKernel!(f16, arm64fp16_mmm_f16_32x4_gen; 32, 4; 16, 16; 1, 1; no_prefetch, crate::arm64::has_fp16());
#[cfg(not(feature = "no_fp16"))]
MMMKernel!(f16, arm64fp16_mmm_f16_32x4_a55; 32, 4; 16, 16; 1, 1; no_prefetch, crate::arm64::has_fp16());
#[cfg(not(feature = "no_fp16"))]
MMMKernel!(f16, arm64fp16_mmm_f16_128x1_gen; 128, 1; 16, 16; 1, 1; no_prefetch, crate::arm64::has_fp16());
#[cfg(not(feature = "no_fp16"))]
MMMKernel!(f16, arm64fp16_mmm_f16_128x1_a55; 128, 1; 16, 16; 1, 1; no_prefetch, crate::arm64::has_fp16());
tanh_impl!(f32, arm64simd_tanh_f32_4n, 4, 4, true);
sigmoid_impl!(f32, arm64simd_sigmoid_f32_4n, 4, 4, true);
#[cfg(not(feature = "no_fp16"))]
tanh_impl!(f16, arm64fp16_tanh_f16_8n, 8, 8, crate::arm64::has_fp16());
#[cfg(not(feature = "no_fp16"))]
sigmoid_impl!(
f16,
arm64fp16_sigmoid_f16_8n,
8,
8,
crate::arm64::has_fp16()
);
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,4 @@
use crate::frame::mmm::cost_model::CostModel;
pub fn models() -> Vec<(&'static str, CostModel<'static>)> {
vec![]
}
@@ -0,0 +1,4 @@
use crate::frame::mmm::cost_model::CostModel;
pub fn models() -> Vec<(&'static str, CostModel<'static>)> {
vec![]
}
@@ -0,0 +1,17 @@
#[macro_use]
pub mod element_wise;
#[macro_use]
pub mod lut;
#[macro_use]
pub mod mmm;
pub mod pack;
#[macro_use]
pub mod sigmoid;
#[macro_use]
pub mod tanh;
pub use pack::Packer;
pub use pack::PackingWriter;
pub use self::element_wise::{ElementWise, ElementWiseImpl};
pub use self::mmm::{MatMatMul, MatMatMulImpl};
@@ -0,0 +1,186 @@
use std::alloc::*;
use std::fmt::Debug;
use std::marker::PhantomData;
use tract_data::anyhow;
use crate::LADatum;
macro_rules! ew_impl {
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr) => {
paste! {
mod [<sys_ $func>] {
#[allow(unused_imports)]
use tract_data::prelude::f16;
extern_kernel!(fn $func(ptr: *mut $ti, count: usize) -> ());
}
#[derive(Copy, Clone, Debug)]
#[allow(non_camel_case_types)]
pub struct $func;
impl ElementWiseKer<$ti> for $func {
#[inline(always)]
fn name() -> &'static str {
stringify!($func)
}
#[inline(always)]
fn nr() -> usize {
$nr
}
#[inline(always)]
fn alignment_items() -> usize {
$alignment_items
}
#[inline(always)]
fn alignment_bytes() -> usize {
$alignment_items * std::mem::size_of::<$ti>()
}
#[inline(never)]
fn run(buf: &mut [$ti]) {
unsafe { [<sys_ $func>]::$func(buf.as_mut_ptr(), buf.len()) }
}
}
}
};
}
struct TempBuffer {
layout: Layout,
buffer: *mut u8,
}
impl Default for TempBuffer {
fn default() -> Self {
TempBuffer {
layout: Layout::new::<()>(),
buffer: std::ptr::null_mut(),
}
}
}
impl TempBuffer {
fn ensure(&mut self, size: usize, alignment: usize) {
unsafe {
if size > self.layout.size() || alignment > self.layout.align() {
let size = size.max(self.layout.size());
let alignment = alignment.max(self.layout.align());
if !self.buffer.is_null() {
std::alloc::dealloc(self.buffer, self.layout);
}
self.layout = Layout::from_size_align_unchecked(size, alignment);
self.buffer = std::alloc::alloc(self.layout);
assert!(!self.buffer.is_null());
}
}
}
}
impl Drop for TempBuffer {
fn drop(&mut self) {
unsafe {
if !self.buffer.is_null() {
std::alloc::dealloc(self.buffer, self.layout);
}
}
}
}
std::thread_local! {
static TMP: std::cell::RefCell<TempBuffer> = std::cell::RefCell::new(TempBuffer::default());
}
pub trait ElementWise<T>: Send + Sync + Debug + dyn_clone::DynClone
where
T: Copy + Debug + PartialEq + Send + Sync,
{
fn run(&self, vec: &mut [T]) -> anyhow::Result<()>;
}
dyn_clone::clone_trait_object!(<T> ElementWise<T> where T: Copy);
#[derive(Debug, Clone, new)]
pub struct ElementWiseImpl<K, T>
where
T: LADatum,
K: ElementWiseKer<T> + Clone,
{
phantom: PhantomData<(K, T)>,
}
impl<K, T> ElementWise<T> for ElementWiseImpl<K, T>
where
T: LADatum,
K: ElementWiseKer<T> + Clone,
{
fn run(&self, vec: &mut [T]) -> anyhow::Result<()> {
if vec.is_empty() {
return Ok(());
}
unsafe {
TMP.with(|buffer| {
let mut buffer = buffer.borrow_mut();
buffer.ensure(K::nr() * T::datum_type().size_of(), K::alignment_bytes());
let tmp = std::slice::from_raw_parts_mut(buffer.buffer as *mut T, K::nr());
let mut compute_via_temp_buffer = |slice: &mut [T]| {
tmp[..slice.len()].copy_from_slice(slice);
K::run(tmp);
slice.copy_from_slice(&tmp[..slice.len()])
};
let prefix_len = vec
.as_ptr()
.align_offset(K::alignment_bytes())
.min(vec.len());
if prefix_len > 0 {
compute_via_temp_buffer(&mut vec[..prefix_len]);
}
let aligned_len = (vec.len() - prefix_len) / K::nr() * K::nr();
if aligned_len > 0 {
K::run(&mut vec[prefix_len..][..aligned_len]);
}
if prefix_len + aligned_len < vec.len() {
compute_via_temp_buffer(&mut vec[prefix_len + aligned_len..]);
}
})
}
Ok(())
}
}
pub trait ElementWiseKer<T>: Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
where
T: LADatum,
{
fn name() -> &'static str;
fn alignment_bytes() -> usize;
fn alignment_items() -> usize;
fn nr() -> usize;
fn run(vec: &mut [T]);
fn ew() -> Box<dyn ElementWise<T>> {
Box::new(ElementWiseImpl::<Self, T>::new())
}
}
#[cfg(test)]
pub mod test {
use crate::{LADatum, frame::element_wise::*};
use proptest::test_runner::{TestCaseError, TestCaseResult};
use tract_data::internal::*;
pub fn test_element_wise<K: ElementWiseKer<T>, T: LADatum, F: Fn(T) -> T>(
values: &[T],
reference: F,
) -> TestCaseResult {
let op = ElementWiseImpl::<K, T>::new();
let mut values = values.to_vec();
while values.len() < K::nr() {
values.push(T::zero());
}
let expected = values.iter().copied().map(reference).collect::<Vec<_>>();
let mut found = values;
op.run(&mut found).unwrap();
tensor1(&found)
.close_enough(&tensor1(&expected), true)
.map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
Ok(())
}
}
@@ -0,0 +1,158 @@
use std::fmt;
use std::hash::Hash;
use std::marker::PhantomData;
use tract_data::internal::*;
pub trait Lut: fmt::Debug + dyn_clone::DynClone + Send + Sync + DynHash {
fn table(&self) -> &[u8];
fn run(&self, buf: &mut [u8]);
}
dyn_clone::clone_trait_object!(Lut);
impl std::hash::Hash for Box<dyn Lut> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
use std::any::Any;
std::hash::Hash::hash(&self.type_id(), state);
self.dyn_hash(state)
}
}
#[derive(Debug, Clone, Hash)]
pub struct LutImpl<K: LutKer> {
table: Tensor,
_boo: PhantomData<K>,
}
impl<K: LutKer> DynHash for LutImpl<K> {
fn dyn_hash(&self, state: &mut dyn std::hash::Hasher) {
tract_data::hash::dyn_hash(self, state)
}
}
impl<K: LutKer> LutImpl<K> {
pub fn new(table: &[u8]) -> LutImpl<K> {
unsafe {
LutImpl {
table: Tensor::from_raw_aligned::<u8>(
&[table.len()],
table,
K::table_alignment_bytes(),
)
.unwrap(),
_boo: PhantomData,
}
}
}
}
impl<K: LutKer> Lut for LutImpl<K> {
fn table(&self) -> &[u8] {
self.table.as_slice().unwrap()
}
fn run(&self, buf: &mut [u8]) {
unsafe {
let table: *const u8 = self.table.as_ptr_unchecked();
let align = K::input_alignment_bytes();
let aligned_start = (buf.as_ptr() as usize + align - 1) / align * align;
let prefix = (aligned_start - buf.as_ptr() as usize).min(buf.len());
for i in 0..(prefix as isize) {
let ptr = buf.as_mut_ptr().offset(i);
*ptr = *table.offset(*ptr as isize);
}
let remaining = buf.len() - prefix;
if remaining == 0 {
return;
}
let n = K::n();
let aligned_len = remaining / n * n;
if aligned_len > 0 {
K::run(buf.as_mut_ptr().add(prefix), aligned_len, table);
}
let remaining = buf.len() - aligned_len - prefix;
for i in 0..remaining {
let ptr = buf.as_mut_ptr().add(i + prefix + aligned_len);
*ptr = *table.offset(*ptr as isize);
}
}
}
}
pub trait LutKer: Clone + fmt::Debug + Send + Sync + Hash {
fn name() -> &'static str;
fn n() -> usize;
fn input_alignment_bytes() -> usize;
fn table_alignment_bytes() -> usize;
unsafe fn run(buf: *mut u8, len: usize, table: *const u8);
}
#[cfg(test)]
#[macro_use]
pub mod test {
use super::*;
use proptest::prelude::*;
#[derive(Debug)]
pub struct LutProblem {
pub table: Vec<u8>,
pub data: Vec<u8>,
}
impl Arbitrary for LutProblem {
type Parameters = ();
type Strategy = BoxedStrategy<Self>;
fn arbitrary_with(_p: ()) -> Self::Strategy {
proptest::collection::vec(any::<u8>(), 1..256)
.prop_flat_map(|table| {
let data = proptest::collection::vec(0..table.len() as u8, 0..100);
(Just(table), data)
})
.prop_map(|(table, data)| LutProblem { table, data })
.boxed()
}
}
impl LutProblem {
pub fn reference(&self) -> Vec<u8> {
self.data.iter().map(|x| self.table[*x as usize]).collect()
}
pub fn test<K: LutKer>(&self) -> Vec<u8> {
let lut = LutImpl::<K>::new(&self.table);
let mut data = self.data.clone();
lut.run(&mut data);
data
}
}
#[macro_export]
macro_rules! lut_frame_tests {
($cond:expr, $ker:ty) => {
mod lut {
use proptest::prelude::*;
#[allow(unused_imports)]
use $crate::frame::lut::test::*;
proptest::proptest! {
#[test]
fn lut_prop(pb in any::<LutProblem>()) {
if $cond {
prop_assert_eq!(pb.test::<$ker>(), pb.reference())
}
}
}
#[test]
fn test_empty() {
let pb = LutProblem {
table: vec![0],
data: vec![],
};
assert_eq!(pb.test::<$ker>(), pb.reference())
}
}
};
}
}
@@ -0,0 +1,95 @@
pub mod cost_model;
#[macro_use]
pub(crate) mod fuse;
#[macro_use]
pub(crate) mod kernel;
pub(crate) mod input_store;
#[macro_use]
#[allow(clippy::module_inception)]
pub(crate) mod mmm;
mod scratch;
mod storage;
#[cfg(test)]
#[macro_use]
pub mod tests;
pub use cost_model::*;
pub use fuse::*;
pub use input_store::*;
pub use kernel::*;
pub use mmm::*;
pub use scratch::*;
pub use storage::*;
pub fn no_prefetch(_ptr: *const u8, _len: usize) {}
macro_rules! MMMKernel {
($ti:ident, $func:ident; $mr: expr, $nr: expr; $alignment_bytes_packed_a: expr, $alignment_bytes_packed_b: expr; $end_padding_packed_a: expr, $end_padding_packed_b: expr ; $prefetch: ident, $cond: expr) => {
paste! {
mod [<sys_ $func>] {
use crate::frame::mmm::*;
#[allow(unused_imports)]
use tract_data::prelude::f16;
extern_kernel!(fn $func(op: *const FusedKerSpec<$ti>) -> isize);
}
#[allow(non_camel_case_types)]
#[derive(Copy, Clone, Debug, new)]
pub struct $func;
impl MatMatMulKer<$ti> for $func {
#[inline(always)]
fn name() -> &'static str {
stringify!($func)
}
#[inline(always)]
fn mr() -> usize {
$mr
}
#[inline(always)]
fn nr() -> usize {
$nr
}
#[inline(always)]
fn alignment_bytes_packed_a() -> usize {
$alignment_bytes_packed_a
}
#[inline(always)]
fn alignment_bytes_packed_b() -> usize {
$alignment_bytes_packed_b
}
#[inline(always)]
fn end_padding_packed_a() -> usize {
$end_padding_packed_a
}
#[inline(always)]
fn end_padding_packed_b() -> usize {
$end_padding_packed_b
}
#[inline(always)]
fn kernel(spec: &[FusedKerSpec<$ti>]) -> isize {
debug_assert!(spec.len() > 0);
debug_assert!(matches!(spec[spec.len() - 1], FusedKerSpec::Done));
unsafe { [<sys_ $func>]::$func(spec.as_ptr()) }
}
#[inline(always)]
fn prefetch(ptr: *const u8, len: usize) {
($prefetch)(ptr, len)
}
}
}
test_mmm_kernel!($ti, $func, $cond);
};
}
macro_rules! test_mmm_kernel {
(f16, $func:ident, $cond: expr) => {
test_mmm_kernel_f16!($func, $cond);
};
(f32, $func:ident, $cond: expr) => {
test_mmm_kernel_f32!($func, $cond);
};
(i32, $func:ident, $cond: expr) => {
test_mmm_kernel_i32!($func, $cond);
};
}
@@ -0,0 +1,94 @@
use tract_data::internal::*;
use tract_data::itertools::{Itertools, izip};
use super::MatMatMul;
fn order_f<F: tract_num_traits::Float>(&a: &F, &b: &F) -> std::cmp::Ordering {
if a < b {
std::cmp::Ordering::Less
} else {
std::cmp::Ordering::Greater
}
}
#[derive(Debug)]
pub struct CostModel<'a> {
pub big_product_mkn_threshold: f32,
pub big_product_kernel_choice: &'a str,
pub kernels: &'a [&'a str],
pub mrs: &'a [u32],
pub nrs: &'a [u32],
pub feat_norm_mean: &'a [f32],
pub feat_norm_stddev: &'a [f32],
pub w1: &'a [f32],
pub b1: &'a [f32],
pub w2: &'a [f32],
pub b2: &'a [f32],
}
impl<'a> CostModel<'a> {
pub fn features(&self, m: usize, k: usize, n: usize) -> Vec<f32> {
let mut feat = vec![
(m as f32).ln(),
(k as f32).ln(),
(n as f32).ln(),
(n as f32 * m as f32 * k as f32).ln(),
];
for &mr in self.mrs {
let mr = mr as usize;
feat.push((m % mr) as f32);
feat.push((m % mr != 0) as usize as f32);
}
for &nr in self.nrs {
let nr = nr as usize;
feat.push((n % nr) as f32);
feat.push((n % nr != 0) as usize as f32);
}
feat
}
fn normalize(&self, feat: &mut [f32]) {
izip!(feat, self.feat_norm_mean, self.feat_norm_stddev)
.for_each(|(x, m, s)| *x = (*x - m) / s)
}
fn dnn(x: &[f32], w: &[f32], b: &[f32]) -> Vec<f32> {
let x = tract_ndarray::Array1::from_vec(x.to_vec());
let w = tract_ndarray::Array2::from_shape_vec([b.len(), x.len()], w.to_vec()).unwrap();
let b = tract_ndarray::Array1::from_vec(b.to_vec());
(w.dot(&x) + b).to_vec()
}
pub fn predict(&self, m: usize, k: usize, n: usize) -> &str {
let mut x = self.features(m, k, n);
self.normalize(&mut x);
let mut hidden = Self::dnn(&x, self.w1, self.b1);
(crate::generic().tanh_f32)().run(&mut hidden).unwrap();
let output = Self::dnn(&hidden, self.w2, self.b2);
let ix = output.iter().copied().position_max_by(order_f).unwrap();
self.kernels[ix]
}
pub fn pick(
&self,
impls: &[Box<dyn MatMatMul>],
m: Option<usize>,
k: Option<usize>,
n: Option<usize>,
) -> Box<dyn MatMatMul> {
if let (Some(m), Some(k), Some(n)) = (m, k, n) {
let choice = self.predict(m, k, n);
impls
.iter()
.find(|k| k.kernel_name() == choice)
.unwrap()
.clone()
} else {
impls
.iter()
.find(|k| k.kernel_name() == self.big_product_kernel_choice)
.unwrap()
.clone()
}
}
}
@@ -0,0 +1,819 @@
use std::fmt::Debug;
use super::{InputStore, OutputStore, OutputStoreKer, PackedStore};
use tract_data::internal::*;
#[repr(usize)]
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum RoundingPolicy {
Native,
Zero,
Away,
MinusInf,
PlusInf,
Even,
Odd,
}
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BinOp {
Min,
Max,
Add,
Mul,
Sub,
SubF,
}
impl BinOp {
pub fn flip(&self) -> BinOp {
use BinOp::*;
match self {
Sub => SubF,
SubF => Sub,
sym => *sym,
}
}
}
#[derive(Clone, Debug)]
pub enum FusedSpec<'t> {
BinScalar(&'t Tensor, BinOp),
BinPerRow(&'t Tensor, BinOp),
BinPerCol(&'t Tensor, BinOp),
AddRowColProducts(&'t Tensor, &'t Tensor),
AddUnicast(OutputStore),
QScale(isize, RoundingPolicy, i32),
RoundingShiftRight(usize, RoundingPolicy),
ShiftLeft(usize),
Store(OutputStore),
AddMatMul {
k: usize,
a: PackedStore,
b: InputStore,
},
}
impl<'t> FusedSpec<'t> {
pub fn prefer_col_outer(&self) -> bool {
if let FusedSpec::AddMatMul { b, .. } = self {
match b {
InputStore::Packed { .. } => false,
InputStore::VirtualPacking { .. } => true,
InputStore::LatePacking { .. } => true,
}
} else {
false
}
}
}
// Careful here, the jump_to comments are used by the build script.
#[repr(C, usize)]
#[derive(PartialEq, Eq, Copy, Clone, Debug)]
#[rustfmt::skip]
pub enum FusedKerSpec<TI: Copy> {
Done, // jump_to:done
Clear, // jump_to:clear
ScalarMin(TI), // jump_to:scalar_min
ScalarMax(TI), // jump_to:scalar_max
ScalarAdd(TI), // jump_to:scalar_add
ScalarMul(TI), // jump_to:scalar_mul
ScalarSub(TI), // jump_to:scalar_sub
ScalarSubF(TI), // jump_to:scalar_sub_flipped
PerRowMin(*const TI), // jump_to:per_row_min
PerRowMax(*const TI), // jump_to:per_row_max
PerRowAdd(*const TI), // jump_to:per_row_add
PerRowMul(*const TI), // jump_to:per_row_mul
PerRowSub(*const TI), // jump_to:per_row_sub
PerRowSubF(*const TI), // jump_to:per_row_sub_flipped
PerColMin(*const TI), // jump_to:per_col_min
PerColMax(*const TI), // jump_to:per_col_max
PerColAdd(*const TI), // jump_to:per_col_add
PerColMul(*const TI), // jump_to:per_col_mul
PerColSub(*const TI), // jump_to:per_col_sub
PerColSubF(*const TI), // jump_to:per_col_sub_flipped
QScale(isize, RoundingPolicy, i32), // jump_to:q_scale
RoundingShiftRight(usize, RoundingPolicy), // jump_to:q_shr
ShiftLeft(usize), // jump_to:q_shl
AddUnicast(OutputStoreKer), // jump_to:add_unicast
AddRowColProducts(*const TI, *const TI), // jump_to:add_row_col_products
Store(OutputStoreKer), // jump_to:store
// jump_to:add_mat_mul
AddMatMul { k: usize, pa: *const u8, pb: *const u8, cpu_variant: usize },
}
#[cfg(test)]
#[macro_use]
pub mod test {
use crate::frame::mmm::storage::*;
use crate::frame::mmm::*;
use crate::generic::{ScaleShiftAndRound, Scaler};
use num_traits::{AsPrimitive, Bounded};
use proptest::prelude::*;
use tract_data::internal::*;
#[test]
fn check_non_linear_enum_size() {
assert_eq!(
std::mem::size_of::<RoundingPolicy>(),
std::mem::size_of::<usize>()
);
assert_eq!(
std::mem::size_of::<super::FusedKerSpec<f32>>(),
std::mem::size_of::<usize>() + std::mem::size_of::<OutputStoreKer>()
);
assert_eq!(
std::mem::size_of::<super::FusedKerSpec<f32>>(),
5 * std::mem::size_of::<usize>()
);
}
#[macro_export]
macro_rules! mmm_kernel_fuse_tests {
($cond:expr, $ker:ident, $tc:ty, $ti: ty) => {
mod fuse {
use super::super::$ker;
#[allow(unused_imports)]
use tract_data::prelude::f16;
#[allow(unused_imports)]
use $crate::frame::mmm::fuse::test;
use $crate::frame::mmm::fuse::test::tile;
#[test]
fn return_zeros() {
if $cond {
test::return_zeros::<$ker, $tc, $ti>()
}
}
proptest::proptest! {
#[test]
fn return_c_prop(c in tile::<$ker, $tc, $ti>()) {
if $cond {
test::return_c::<$ker, $tc, $ti>(&c)
}
}
}
#[test]
fn return_c_min_row() {
if $cond {
test::return_c_min_row::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_max_row() {
if $cond {
test::return_c_max_row::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_add_row() {
if $cond {
test::return_c_add_row::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_mul_row() {
if $cond {
test::return_c_mul_row::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_sub_row() {
if $cond {
test::return_c_sub_row::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_subf_row() {
if $cond {
test::return_c_subf_row::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_mul_col() {
if $cond {
test::return_c_mul_col::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_add_col() {
if $cond {
test::return_c_add_col::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_add_row_col_product() {
if $cond {
test::return_c_add_row_col_product::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_scalar_max() {
if $cond {
test::return_c_scalar_max::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_scalar_min() {
if $cond {
test::return_c_scalar_min::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_scalar_add() {
if $cond {
test::return_c_scalar_add::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_scalar_mul() {
if $cond {
test::return_c_scalar_mul::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_scalar_sub() {
if $cond {
test::return_c_scalar_sub::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_scalar_subf() {
if $cond {
test::return_c_scalar_subf::<$ker, $tc, $ti>()
}
}
#[test]
fn return_c_plus_d() {
if $cond {
test::return_c_plus_d::<$ker, $tc, $ti>()
}
}
}
};
}
#[macro_export]
macro_rules! qmmm_kernel_fuse_tests {
($cond:expr, $ker:ident, $ta:ty, $tb:ty, $tc:ty, $ti: ty) => {
mod fuseq {
use $crate::frame::mmm::fuse::RoundingPolicy;
#[allow(unused_imports)]
use $crate::frame::mmm::fuse::test;
use $crate::frame::mmm::fuse::test::QScaleProblem;
use $crate::frame::mmm::kernel::MatMatMulKer;
use $crate::generic::Scaler;
use proptest::prelude::*;
use super::super::$ker;
// FIXME: Scaler should be arbitrary
macro_rules! test_q_scale {
($policy: ident) => {
paste! {
#[test]
fn [<return_q_scale_halfpos_ $policy:lower>]() {
if $cond {
let len = (<$ker>::mr() * <$ker>::nr()) as i64;
let v = (0..len).map(|i| (i - len / 2) as $tc).collect();
QScaleProblem::<$ker, $tc, $ti>::new(v, Scaler::new(0.5f32, RoundingPolicy::$policy)).run()
}
}
#[test]
fn [<return_q_scale_halfneg_ $policy:lower>]() {
if $cond {
let len = (<$ker>::mr() * <$ker>::nr()) as i64;
let v = (0..len).map(|i| (i - len / 2) as $tc).collect();
QScaleProblem::<$ker, $tc, $ti>::new(v, Scaler::new(-0.5f32, RoundingPolicy::$policy)).run()
}
}
#[test]
fn [<return_q_scale_pot_ $policy:lower>]() {
if $cond {
let len = (<$ker>::mr() * <$ker>::nr()) as i64;
let v = (0..len).map(|i| (i - len / 2) as $tc).collect();
QScaleProblem::<$ker, $tc, $ti>::new(v, Scaler::new(0.25f32, RoundingPolicy::$policy)).run()
}
}
#[test]
fn [<return_q_scale_nonpot_ $policy:lower>]() {
if $cond {
let len = (<$ker>::mr() * <$ker>::nr()) as i64;
let v = (0..len).map(|i| (i - len / 2) as $tc).collect();
QScaleProblem::<$ker, $tc, $ti>::new(v, Scaler::new(1f32 / 5., RoundingPolicy::$policy)).run()
}
}
#[test]
fn [<return_q_scale_bigpot_ $policy:lower>]() {
if $cond {
let len = (<$ker>::mr() * <$ker>::nr()) as i64;
let v = (0..len).map(|i| (i - len / 2) as $tc).collect();
QScaleProblem::<$ker, $tc, $ti>::new(v, Scaler::new(4f32, RoundingPolicy::$policy)).run()
}
}
#[test]
fn [<return_q_scale_bignonpot_ $policy:lower>]() {
if $cond {
let len = (<$ker>::mr() * <$ker>::nr()) as i64;
let v = (0..len).map(|i| (i - len / 2) as $tc).collect();
QScaleProblem::<$ker, $tc, $ti>::new(v, Scaler::new(14., RoundingPolicy::$policy)).run()
}
}
}
}
}
test_q_scale!(Zero);
test_q_scale!(Away);
test_q_scale!(MinusInf);
test_q_scale!(PlusInf);
test_q_scale!(Even);
test_q_scale!(Odd);
proptest::proptest! {
#[test]
fn return_q_scale_prop(pb in any::<QScaleProblem<$ker, $tc, $ti>>()) {
if $cond {
pb.run()
}
}
}
#[test]
fn return_c_scale_bigpot() {
if $cond {
test::return_c_scale_bigpot::<$ker, $tc, $ti>()
}
}
}
};
}
pub fn mmm_stride_storage<T: Copy>(v: &[T], rsc: usize) -> OutputStoreKer {
OutputStoreKer {
ptr: v.as_ptr() as _,
row_byte_stride: (std::mem::size_of::<T>() * rsc) as isize,
col_byte_stride: std::mem::size_of::<T>() as isize,
item_size: std::mem::size_of::<T>(),
}
}
use crate::LADatum;
pub fn return_zeros<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum,
TI: LADatum + Bounded + PartialEq,
{
let v = vec![TC::max_value(); K::mr() * K::nr()];
let c = mmm_stride_storage(&v, K::nr());
let non_linear = tvec![
FusedKerSpec::Clear,
FusedKerSpec::Store(c),
FusedKerSpec::Done
];
let err = K::kernel(&non_linear);
assert_eq!(err, 0);
let expected = vec![TC::zero(); v.len()];
assert_eq!(v, expected);
}
pub fn fused_ops<K, TC, TI, E>(c: &[TC], ops: &[FusedKerSpec<TI>], expect: E)
where
K: MatMatMulKer<TI>,
TC: Datum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
E: Fn(usize, usize, TI) -> TI,
{
assert!(c.len() == K::mr() * K::nr());
let v = c.to_vec();
let c = mmm_stride_storage(&v, K::nr());
let mut ops = ops.to_vec();
ops.insert(0, FusedKerSpec::AddUnicast(c));
ops.insert(0, FusedKerSpec::Clear);
ops.push(FusedKerSpec::Store(c));
ops.push(FusedKerSpec::Done);
let expected = (0..v.len())
.map(|ix| expect(ix / K::nr(), ix % K::nr(), v[ix].as_()).as_())
.collect::<Vec<TC>>();
let err = K::kernel(&ops);
assert_eq!(err, 0);
if v != expected {
println!("found, expected:");
for m in 0..K::mr() {
for n in 0..K::nr() {
use nu_ansi_term::Color::*;
let f = v[m * K::nr() + n];
let e = expected[m * K::nr() + n];
let color = if f != e { Red } else { Green };
print!("{} ", color.paint(format!("{:4}", f)));
}
print!(" ");
for n in 0..K::nr() {
print!("{:4} ", expected[m * K::nr() + n]);
}
println!();
}
}
assert_eq!(v, expected);
}
pub fn return_c<K, TC, TI>(v: &[TC])
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
fused_ops::<K, TC, TI, _>(v, &[], |_, _, c| c + 1.as_() - 1.as_())
}
pub fn return_c_plus_d<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let d: Vec<TI> = (0..len).map(|f| ((3 * f) % 7).as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::AddUnicast(mmm_stride_storage(&d, K::nr()))],
|row, col, c| c + d[row * K::nr() + col],
);
}
pub fn return_c_min_row<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::mr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerRowMin(bias.as_ptr())],
|row, _, c| {
if c < bias[row] { c } else { bias[row] }
},
)
}
pub fn return_c_max_row<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::mr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerRowMax(bias.as_ptr())],
|row, _, c| {
if c > bias[row] { c } else { bias[row] }
},
)
}
pub fn return_c_add_row<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::mr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerRowAdd(bias.as_ptr())],
|row, _, c| c + bias[row],
)
}
pub fn return_c_mul_row<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::mr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerRowMul(bias.as_ptr())],
|row, _, c| c * bias[row],
)
}
pub fn return_c_sub_row<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::mr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerRowSub(bias.as_ptr())],
|row, _, c| bias[row] - c,
)
}
pub fn return_c_subf_row<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::mr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerRowSubF(bias.as_ptr())],
|row, _, c| c - bias[row],
)
}
pub fn return_c_add_col<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::nr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerColAdd(bias.as_ptr())],
|_, col, c| c + bias[col],
)
}
pub fn return_c_mul_col<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let bias: Vec<TI> = (0..K::nr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::PerColMul(bias.as_ptr())],
|_, col, c| c * bias[col],
)
}
pub fn return_c_add_row_col_product<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let rows: Vec<TI> = (0..K::mr()).map(|f| f.as_()).collect();
let cols: Vec<TI> = (0..K::nr()).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(
&v,
&[FusedKerSpec::AddRowColProducts(
rows.as_ptr(),
cols.as_ptr(),
)],
|row, col, c| c + cols[col] * rows[row],
)
}
pub fn return_c_scalar_min<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(&v, &[FusedKerSpec::ScalarMin(5.as_())], |_, _, c| {
if c > 5.as_() { 5.as_() } else { c }
})
}
pub fn return_c_scalar_max<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(&v, &[FusedKerSpec::ScalarMax(5.as_())], |_, _, c| {
if c < 5.as_() { 5.as_() } else { c }
})
}
pub fn return_c_scalar_add<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(&v, &[FusedKerSpec::ScalarAdd(5.as_())], |_, _, c| {
c + 5.as_()
})
}
pub fn return_c_scalar_mul<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(&v, &[FusedKerSpec::ScalarMul(5.as_())], |_, _, c| {
c * 5.as_()
})
}
pub fn return_c_scalar_sub<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let five: TI = 5.as_();
fused_ops::<K, TC, TI, _>(&v, &[FusedKerSpec::ScalarSub(5.as_())], |_, _, c| five - c)
}
pub fn return_c_scalar_subf<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (0..len).map(|f| f.as_()).collect();
let five: TI = 5.as_();
fused_ops::<K, TC, TI, _>(&v, &[FusedKerSpec::ScalarSubF(5.as_())], |_, _, c| c - five)
}
pub fn return_c_scale_bigpot<K, TC, TI>()
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC> + ScaleShiftAndRound,
isize: AsPrimitive<TC> + AsPrimitive<TI>,
{
let len = K::mr() * K::nr();
let v: Vec<TC> = (-(len as isize) / 2..).take(len).map(|f| f.as_()).collect();
fused_ops::<K, TC, TI, _>(&v, &[FusedKerSpec::ShiftLeft(1)], |_, _, c| c.q_shl(1))
}
#[derive(Debug, new)]
pub struct QScaleProblem<K, TC, TI>
where
K: MatMatMulKer<TI>,
TC: LADatum,
TI: LADatum + AsPrimitive<TC>,
i64: AsPrimitive<TC>,
{
pub c: Vec<TC>,
pub scaler: Scaler,
pub boo: std::marker::PhantomData<(K, TC, TI)>,
}
impl<K, TC, TI> Arbitrary for QScaleProblem<K, TC, TI>
where
K: MatMatMulKer<TI>,
TC: LADatum + Arbitrary,
TI: LADatum + AsPrimitive<TC>,
i64: AsPrimitive<TC>,
{
type Parameters = ();
type Strategy = BoxedStrategy<Self>;
fn arbitrary_with(_p: ()) -> Self::Strategy {
use RoundingPolicy::*;
let len = K::mr() * K::nr();
(
proptest::collection::vec((-20i64..20).prop_map(|i| i.as_()), len..=len),
-5i32..5,
prop_oneof!(Just(1f32), 0f32..1f32),
proptest::prop_oneof![
Just(Zero),
Just(Away),
Just(PlusInf),
Just(MinusInf),
Just(Odd),
Just(Even)
],
)
.prop_map(|(c, scale_pot, scale_mult, policy)| QScaleProblem {
c,
scaler: Scaler::new(scale_mult * 2f32.powi(scale_pot), policy),
boo: std::marker::PhantomData,
})
.boxed()
}
}
impl<K, TC, TI> QScaleProblem<K, TC, TI>
where
K: MatMatMulKer<TI>,
TC: LADatum + AsPrimitive<TI>,
TI: LADatum + AsPrimitive<TC> + ScaleShiftAndRound + AsPrimitive<i64>,
usize: AsPrimitive<TC> + AsPrimitive<TI>,
i64: AsPrimitive<TC>,
{
pub fn run(&self) {
if let FusedSpec::QScale(shift, policy, mult) = self.scaler.as_fused_spec() {
fused_ops::<K, TC, TI, _>(
&self.c,
&[FusedKerSpec::QScale(shift, policy, mult)],
|_, _, c| c.q_scale(self.scaler),
)
} else if let FusedSpec::RoundingShiftRight(shift, policy) = self.scaler.as_fused_spec()
{
fused_ops::<K, TC, TI, _>(
&self.c,
&[FusedKerSpec::RoundingShiftRight(shift, policy)],
|_, _, c| c.q_shr(shift, policy),
)
} else if let FusedSpec::ShiftLeft(shift) = self.scaler.as_fused_spec() {
fused_ops::<K, TC, TI, _>(&self.c, &[FusedKerSpec::ShiftLeft(shift)], |_, _, c| {
c.q_shl(shift)
})
} else {
unreachable!()
}
}
}
pub fn tile<K, TC, TI>() -> BoxedStrategy<Vec<TC>>
where
K: MatMatMulKer<TI>,
TC: LADatum,
TI: LADatum + AsPrimitive<TC>,
i8: AsPrimitive<TC>,
{
let len = K::mr() * K::nr();
proptest::collection::vec(any::<i8>().prop_map(|c| c.as_()), len..=len).boxed()
}
}
@@ -0,0 +1,184 @@
use std::alloc::Layout;
use std::fmt;
use std::ops::Range;
use tract_data::internal::DynHash;
use tract_data::internal::*;
use crate::frame::Packer;
pub trait VirtualInputSpec: DynHash + dyn_clone::DynClone + std::fmt::Debug + Sync + Send {
fn wrap(&self, view: &TensorView) -> Box<dyn VirtualInput>;
}
dyn_clone::clone_trait_object!(VirtualInputSpec);
pub trait VirtualInput: dyn_clone::DynClone + std::fmt::Debug + Sync + Send {
fn input(&self, packer: &Packer, packed_output: *mut u8, k: Range<usize>, mn: Range<usize>);
}
dyn_clone::clone_trait_object!(VirtualInput);
impl std::hash::Hash for Box<dyn VirtualInputSpec> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
use std::any::Any;
std::hash::Hash::hash(&self.type_id(), state);
self.dyn_hash(state)
}
}
#[derive(Clone, Debug, Hash)]
pub enum InputStoreSpec {
Prepacked(PackedStoreSpec),
LatePacking {
packer: Packer,
k_axis: usize,
mn_axis: usize,
},
VirtualPacking {
packer: Packer,
func: Box<dyn VirtualInputSpec>,
k: usize,
},
}
#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)]
pub struct PackedStoreSpec {
pub(crate) panel_bytes: usize,
}
impl InputStoreSpec {
#[inline]
pub unsafe fn wrap(&self, tensor: &TensorView) -> TractResult<InputStore> {
use InputStore::*;
use InputStoreSpec as S;
match self {
S::Prepacked(PackedStoreSpec { panel_bytes }) => Ok(Packed(PackedStore {
ptr: tensor.as_ptr_unchecked::<u8>() as _,
panel_bytes: *panel_bytes as isize,
})),
S::LatePacking {
packer,
k_axis,
mn_axis,
} => Ok(InputStore::LatePacking {
packer: packer.clone(),
ptr: tensor.as_ptr_unchecked::<u8>() as _,
dt: tensor.datum_type(),
k: tensor.shape()[*k_axis],
mn: tensor.shape()[*mn_axis],
k_stride: tensor.strides()[*k_axis],
mn_stride: tensor.strides()[*mn_axis],
}),
S::VirtualPacking { packer, func, k } => Ok(InputStore::VirtualPacking {
packer: packer.clone(),
input: func.wrap(tensor),
k: *k,
dt: tensor.datum_type(),
}),
}
}
}
impl fmt::Display for InputStoreSpec {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
match self {
InputStoreSpec::Prepacked { .. } => write!(fmt, "Packed"),
InputStoreSpec::LatePacking { .. } => write!(fmt, "LatePacking"),
InputStoreSpec::VirtualPacking { .. } => write!(fmt, "VirtualPacking"),
}
}
}
impl PackedStoreSpec {
#[inline]
pub unsafe fn wrap(&self, tensor: &TensorView) -> PackedStore {
PackedStore {
ptr: tensor.as_ptr_unchecked::<u8>() as _,
panel_bytes: self.panel_bytes as isize,
}
}
}
#[derive(Clone, Debug)]
pub enum InputStore {
Packed(PackedStore),
LatePacking {
packer: Packer,
ptr: *const u8,
dt: DatumType,
k: usize,
mn: usize,
k_stride: isize,
mn_stride: isize,
},
VirtualPacking {
packer: Packer,
input: Box<dyn VirtualInput>,
k: usize,
dt: DatumType, // TODO discard me ?
},
}
#[derive(Clone, Copy, Debug)]
pub struct PackedStore {
ptr: *const u8,
panel_bytes: isize,
}
impl InputStore {
pub(super) unsafe fn scratch_panel_buffer_layout(&self) -> Option<Layout> {
match self {
InputStore::Packed(_) => None,
InputStore::LatePacking { packer, dt, k, .. }
| InputStore::VirtualPacking { packer, dt, k, .. } => {
let size = packer.single_panel_len(*k) * dt.size_of();
let align = packer.alignment();
Some(Layout::from_size_align_unchecked(size, align))
}
}
}
#[inline]
pub(super) unsafe fn panel_b(&self, i: usize, buffer: Option<*const u8>) -> *const u8 {
match self {
InputStore::Packed(packed) => packed.panel(i),
InputStore::LatePacking {
packer,
ptr,
dt,
k,
mn,
mn_stride,
k_stride,
} => {
dispatch_copy!(Packer::pack_t(dt)(
packer,
buffer.unwrap() as _,
*ptr as _,
*mn,
*k_stride,
*mn_stride,
0..*k,
packer.r * i..packer.r * (i + 1)
));
buffer.unwrap()
}
InputStore::VirtualPacking {
packer, input, k, ..
} => {
input.input(
packer,
buffer.unwrap() as _,
0..*k,
packer.r * i..packer.r * (i + 1),
);
buffer.unwrap()
}
}
}
}
impl PackedStore {
#[inline]
pub(super) unsafe fn panel(&self, i: usize) -> *const u8 {
self.ptr.offset(self.panel_bytes * i as isize)
}
}
@@ -0,0 +1,399 @@
use std::fmt::Debug;
use crate::LADatum;
use crate::frame::mmm::FusedKerSpec;
use super::{MatMatMul, MatMatMulImpl};
pub trait MatMatMulKer<TI>: Copy + Clone + Debug + Send + Sync + 'static
where
TI: LADatum,
{
fn name() -> &'static str;
fn kernel(op: &[FusedKerSpec<TI>]) -> isize;
fn mr() -> usize;
fn nr() -> usize;
fn alignment_bytes_packed_a() -> usize;
fn end_padding_packed_a() -> usize;
fn alignment_bytes_packed_b() -> usize;
fn end_padding_packed_b() -> usize;
#[allow(unused_variables)]
fn prefetch(ptr: *const u8, len: usize) {}
fn mmm() -> Box<dyn MatMatMul> {
Box::<MatMatMulImpl<Self, TI>>::default()
}
}
#[macro_export]
macro_rules! test_mmm_kernel_f16 {
($k: ident, $cond: expr) => {
paste! {
#[cfg(test)]
#[allow(non_snake_case)]
mod [<test_ $k>] {
mmm_kernel_tests!($cond, $k, f16, f16, f16, f16);
mmm_frame_tests!($cond, $k, f16, f16, f16, f16);
mmm_kernel_fuse_tests!($cond, $k, f16, f16);
}
}
};
}
#[macro_export]
macro_rules! test_mmm_kernel_f32 {
($k: ident, $cond: expr) => {
paste! {
#[cfg(test)]
#[allow(non_snake_case)]
mod [<test_ $k>] {
mmm_kernel_tests!($cond, $k, f32, f32, f32, f32);
mmm_frame_tests!($cond, $k, f32, f32, f32, f32);
mmm_kernel_fuse_tests!($cond, $k, f32, f32);
//qmmm_kernel_fuse_tests!($cond, $k, f32, f32, f32, f32);
}
}
};
}
#[macro_export]
macro_rules! test_mmm_kernel_f64 {
($k: ident, $cond: expr) => {
paste! {
#[cfg(test)]
#[allow(non_snake_case)]
mod [<test_ $k>] {
mmm_kernel_tests!($cond, $k, f64, f64, f64, f64);
mmm_frame_tests!($cond, $k, f64, f64, f64, f64);
mmm_kernel_fuse_tests!($cond, $k, f64, f64);
//qmmm_kernel_fuse_tests!($cond, $k, f64, f64, f64, f64);
}
}
};
}
#[macro_export]
macro_rules! test_mmm_kernel_i32 {
($k: ident, $cond: expr) => {
paste! {
#[cfg(test)]
#[allow(non_snake_case)]
mod [<test_ $k>] {
mmm_kernel_tests!($cond, $k, i8, i8, i8, i32);
mmm_kernel_fuse_tests!($cond, $k, i8, i32);
mmm_frame_tests!($cond, $k, i8, i8, i8, i32);
}
#[cfg(test)]
mod [<test_qi8_ $k>] {
qmmm_kernel_fuse_tests!($cond, $k, i8, i8, i8, i32);
}
#[cfg(test)]
mod [<test_qi32_ $k>] {
qmmm_kernel_fuse_tests!($cond, $k, i8, i8, i32, i32);
}
}
};
}
#[cfg(test)]
#[macro_use]
pub mod test {
use super::*;
use crate::frame::mmm::OutputStoreKer;
use num_traits::{AsPrimitive, One, Zero};
use proptest::collection::vec;
use proptest::prelude::*;
use std::fmt;
use std::marker::PhantomData;
use tract_data::internal::*;
#[macro_export]
macro_rules! mmm_kernel_tests {
($cond:expr, $ker:ident, $ta:ty, $tb:ty, $tc:ty, $ti: ty) => {
mod kernel {
use super::super::$ker;
use num_traits::Zero;
use proptest::prelude::*;
#[allow(unused_imports)]
use tract_data::prelude::f16;
#[allow(unused_imports)]
use $crate::frame::mmm::kernel::test;
use $crate::frame::mmm::kernel::test::PackedPackedProblem;
use $crate::frame::mmm::MatMatMulKer;
proptest::proptest! {
#[test]
fn packed_packed_prop(pb in any::<PackedPackedProblem<$ker, $ta, $tb, $tc, $ti>>()) {
if $cond {
prop_assert_eq!(pb.run(), pb.reference())
}
}
}
#[test]
fn packed_packed_1() {
if $cond {
test::packed_packed::<$ker, $ta, $tb, $tc, $ti>(1)
}
}
#[test]
fn packed_packed_2() {
if $cond {
test::packed_packed::<$ker, $ta, $tb, $tc, $ti>(2)
}
}
#[test]
fn packed_packed_13() {
if $cond {
test::packed_packed::<$ker, $ta, $tb, $tc, $ti>(13)
}
}
#[test]
fn packed_packed_empty() {
if $cond {
let pb = PackedPackedProblem::<$ker, $ta, $tb, $tc, $ti>::new(
0,
vec!(<$ta>::zero(); 0),
vec!(<$tb>::zero(); 0),
false,
false);
assert_eq!(pb.run(), pb.reference())
}
}
#[test]
fn packed_packed_bug_1() {
if $cond {
let pb = PackedPackedProblem::<$ker, $ta, $tb, $tc, $ti>::new(
1,
vec!(<$ta>::zero(); <$ker>::mr()),
vec!(<$tb>::zero(); <$ker>::nr()),
true,
true);
assert_eq!(pb.run(), pb.reference())
}
}
#[test]
fn packed_vec_k1() {
if $cond {
test::packed_vec::<$ker, $ta, $tb, $tc, $ti>(1)
}
}
#[test]
fn packed_vec_k2() {
if $cond {
test::packed_vec::<$ker, $ta, $tb, $tc, $ti>(2)
}
}
#[test]
fn packed_vec_k4() {
if $cond {
test::packed_vec::<$ker, $ta, $tb, $tc, $ti>(4)
}
}
#[test]
fn packed_vec_k13() {
if $cond {
test::packed_vec::<$ker, $ta, $tb, $tc, $ti>(13)
}
}
}
};
}
#[derive(Debug, new)]
pub struct PackedPackedProblem<K, TA, TB, TC, TI>
where
K: MatMatMulKer<TI>,
TA: 'static + Debug + AsPrimitive<TI>,
TB: 'static + Debug + AsPrimitive<TI>,
TC: Copy + PartialEq + 'static + Debug,
TI: LADatum + fmt::Display + AsPrimitive<TC>,
usize: AsPrimitive<TA> + AsPrimitive<TB>,
{
pub k: usize,
pub a: Vec<TA>,
pub b: Vec<TB>,
pub trans_c: bool,
pub add_one: bool,
pub _phantom: PhantomData<(K, TC, TI)>,
}
impl<K, TA, TB, TC, TI> Arbitrary for PackedPackedProblem<K, TA, TB, TC, TI>
where
K: MatMatMulKer<TI>,
TA: 'static + Debug + AsPrimitive<TI>,
TB: 'static + Debug + AsPrimitive<TI>,
TC: Copy + PartialEq + 'static + Debug,
TI: LADatum + fmt::Display + AsPrimitive<TC>,
usize: AsPrimitive<TA> + AsPrimitive<TB>,
{
type Parameters = ();
type Strategy = BoxedStrategy<Self>;
fn arbitrary_with(_: ()) -> Self::Strategy {
(0usize..20, any::<bool>(), any::<bool>())
.prop_flat_map(|(k, trans_c, add_one)| {
let m = k * K::mr();
let n = k * K::nr();
let a = (0usize..10).prop_map(|x| x.as_());
let b = (0usize..10).prop_map(|x| x.as_());
(
Just(k),
Just(trans_c),
Just(add_one),
vec(a, m..=m),
vec(b, n..=n),
)
})
.prop_map(|(k, trans_c, add_one, a, b)| Self {
k,
a,
b,
trans_c,
add_one,
_phantom: PhantomData,
})
.boxed()
}
}
impl<K, TA, TB, TC, TI> PackedPackedProblem<K, TA, TB, TC, TI>
where
K: MatMatMulKer<TI>,
TA: 'static + Debug + AsPrimitive<TI> + Datum,
TB: 'static + Debug + AsPrimitive<TI> + Datum,
TC: Copy + Zero + PartialEq + 'static + Debug,
TI: LADatum + fmt::Display + AsPrimitive<TC>,
usize: AsPrimitive<TA> + AsPrimitive<TB>,
{
pub fn reference(&self) -> Vec<TC> {
let init = if self.add_one { TI::one() } else { TI::zero() };
let mut vi = vec![init; K::mr() * K::nr()];
let mr = K::mr();
let nr = K::nr();
for m in 0..mr {
for n in 0..nr {
for k in 0..self.k {
let a: TI = self.a[m + mr * k].as_();
let b: TI = self.b[n + nr * k].as_();
let offset = if self.trans_c { m + n * mr } else { n + m * nr };
vi[offset] += a * b;
}
}
}
vi.into_iter().map(|ti| ti.as_()).collect()
}
pub fn run(&self) -> Vec<TC> {
unsafe {
let a = self
.a
.iter()
.cloned()
.chain(vec![0.as_(); K::end_padding_packed_a() * K::mr()])
.collect::<Vec<_>>();
let pa = Tensor::from_slice_align(&a, K::alignment_bytes_packed_a()).unwrap();
let b = self
.b
.iter()
.cloned()
.chain(vec![0.as_(); K::end_padding_packed_b() * K::nr()])
.collect::<Vec<_>>();
let pb = Tensor::from_slice_align(&b, K::alignment_bytes_packed_b()).unwrap();
let mut v = vec![TC::zero(); K::mr() * K::nr()];
let c = if self.trans_c {
mmm_stride_storage(&mut v, 1, K::mr())
} else {
mmm_stride_storage(&mut v, K::nr(), 1)
};
let b_store = pb.as_ptr_unchecked::<TB>() as _;
let mut non_linear_ops = tvec!(FusedKerSpec::AddMatMul {
k: self.k,
pa: pa.as_ptr_unchecked::<u8>() as _,
pb: b_store,
cpu_variant: 0,
});
if self.add_one {
non_linear_ops.push(FusedKerSpec::ScalarAdd(TI::one()));
}
non_linear_ops.push(FusedKerSpec::Store(c));
non_linear_ops.push(FusedKerSpec::Done);
non_linear_ops.insert(0, FusedKerSpec::Clear);
let err = K::kernel(&non_linear_ops);
assert_eq!(err, 0);
v
}
}
}
pub fn packed_packed<K, TA, TB, TC, TI>(k: usize)
where
K: MatMatMulKer<TI>,
TA: Copy + One + Datum + AsPrimitive<TI>,
TB: Copy + One + Datum + AsPrimitive<TI>,
TC: Copy + PartialEq + Zero + 'static + Debug,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC> + AsPrimitive<TA> + AsPrimitive<TB>,
{
let a = vec![TA::one(); K::mr() * k];
let b = vec![TB::one(); K::nr() * k];
let pb = PackedPackedProblem::<K, TA, TB, TC, TI>::new(k, a, b, false, false);
assert_eq!(pb.run(), pb.reference())
}
pub fn mmm_stride_storage<T: Copy>(v: &mut [T], rsc: usize, csc: usize) -> OutputStoreKer {
OutputStoreKer {
ptr: v.as_mut_ptr() as _,
row_byte_stride: (std::mem::size_of::<T>() * rsc) as isize,
col_byte_stride: (std::mem::size_of::<T>() * csc) as isize,
item_size: std::mem::size_of::<T>(),
}
}
pub fn packed_vec<K, TA, TB, TC, TI>(k: usize)
where
K: MatMatMulKer<TI>,
TA: Copy + One + AsPrimitive<TI> + Debug + Datum,
TB: Copy + One + AsPrimitive<TI> + Debug + Datum,
TC: Copy + PartialEq + Zero + 'static + Debug,
TI: LADatum + AsPrimitive<TC>,
usize: AsPrimitive<TC>,
{
let pa = unsafe {
Tensor::from_slice_align(
&vec![TA::one(); K::mr() * (k + K::end_padding_packed_a())],
K::alignment_bytes_packed_a(),
)
.unwrap()
};
let b = vec![TB::one(); (k + 1) * K::nr()];
let mut c: Vec<TC> = vec![TC::zero(); K::mr() * K::nr()];
let tile = mmm_stride_storage(&mut c, 1, 0);
let b_store = b.as_ptr() as _;
let non_linear_ops = tvec!(
FusedKerSpec::Clear,
FusedKerSpec::AddMatMul {
pa: unsafe { pa.as_ptr_unchecked::<u8>() as _ },
pb: b_store,
k,
cpu_variant: 0,
},
FusedKerSpec::Store(tile),
FusedKerSpec::Done
);
let err = K::kernel(&non_linear_ops);
assert_eq!(err, 0);
let expected = vec![k.as_(); K::mr()];
assert_eq!(c[..K::mr()], expected);
}
}
@@ -0,0 +1,362 @@
use super::ScratchSpaceFusedNonLinear;
use super::*;
use crate::LADatum;
use crate::frame::Packer;
use anyhow::Context;
use std::fmt;
use std::fmt::Debug;
use std::marker::PhantomData;
use tract_data::anyhow;
use tract_data::internal::*;
pub trait MatMatMul:
Debug + fmt::Display + dyn_clone::DynClone + Send + Sync + std::any::Any
{
fn kernel_name(&self) -> &'static str;
fn mr(&self) -> usize;
fn nr(&self) -> usize;
fn a_pack(&self) -> Packer;
fn b_pack(&self) -> Packer;
fn internal_type(&self) -> DatumType;
unsafe fn a_packed(&self, item_size: usize, k: usize) -> PackedStoreSpec;
unsafe fn b_packed(&self, item_size: usize, k: usize) -> InputStoreSpec;
unsafe fn b_late_packing(&self) -> InputStoreSpec {
self.b_late_packing_with_axes(0, 1)
}
unsafe fn b_late_packing_with_axes(&self, k_axis: usize, n_axis: usize) -> InputStoreSpec;
unsafe fn b_virtual_input(&self, func: Box<dyn VirtualInputSpec>, k: usize) -> InputStoreSpec;
unsafe fn c_view(&self, m_axis: usize, n_axis: usize) -> OutputStoreSpec;
unsafe fn c_from_data_and_strides(
&self,
item_size: usize,
m: usize,
n: usize,
row_stride: isize,
col_stride: isize,
) -> OutputStoreSpec;
unsafe fn run(&self, m: usize, n: usize, non_linear: &[FusedSpec]) -> anyhow::Result<()> {
let mut scratch = self.allocate_scratch_space();
self.run_with_scratch_space(m, n, &mut *scratch, non_linear)
}
unsafe fn allocate_scratch_space(&self) -> Box<dyn ScratchSpace>;
unsafe fn can_use_scratch_space(&self, scratch: &dyn ScratchSpace) -> bool;
unsafe fn run_with_scratch_space(
&self,
m: usize,
n: usize,
scratch: &mut dyn ScratchSpace,
non_linear: &[FusedSpec],
) -> anyhow::Result<()>;
unsafe fn run_with_scratch_space_vec(
&self,
m: usize,
scratch: &mut dyn ScratchSpace,
non_linear: &[FusedSpec],
) -> anyhow::Result<()>;
unsafe fn run_with_scratch_space_col_outer(
&self,
m: usize,
n: usize,
scratch: &mut dyn ScratchSpace,
non_linear: &[FusedSpec],
) -> anyhow::Result<()>;
}
dyn_clone::clone_trait_object!(MatMatMul);
impl PartialEq for Box<dyn MatMatMul> {
fn eq(&self, other: &Box<dyn MatMatMul>) -> bool {
self.as_ref().type_id() == other.as_ref().type_id()
}
}
impl std::hash::Hash for Box<dyn MatMatMul> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.as_ref().type_id().hash(state)
}
}
#[derive(Clone)]
pub struct MatMatMulImpl<K, TI>
where
TI: LADatum,
K: MatMatMulKer<TI> + 'static,
{
phantom: PhantomData<(K, TI)>,
}
unsafe impl<K, TI> Send for MatMatMulImpl<K, TI>
where
TI: LADatum,
K: MatMatMulKer<TI> + 'static,
{
}
unsafe impl<K, TI> Sync for MatMatMulImpl<K, TI>
where
TI: LADatum,
K: MatMatMulKer<TI> + 'static,
{
}
impl<K, TI> Default for MatMatMulImpl<K, TI>
where
TI: LADatum,
K: MatMatMulKer<TI> + 'static,
{
fn default() -> Self {
MatMatMulImpl {
phantom: PhantomData,
}
}
}
impl<K, TI> fmt::Debug for MatMatMulImpl<K, TI>
where
TI: LADatum,
K: MatMatMulKer<TI> + 'static,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "MMM ({} {}x{})", K::name(), K::mr(), K::nr())
}
}
impl<K, TI> MatMatMul for MatMatMulImpl<K, TI>
where
TI: LADatum,
K: MatMatMulKer<TI> + 'static,
{
fn kernel_name(&self) -> &'static str {
K::name()
}
fn mr(&self) -> usize {
K::mr()
}
fn nr(&self) -> usize {
K::nr()
}
fn a_pack(&self) -> Packer {
Packer::new(
K::mr(),
K::alignment_bytes_packed_a(),
K::end_padding_packed_a(),
)
}
fn b_pack(&self) -> Packer {
Packer::new(
K::nr(),
K::alignment_bytes_packed_b(),
K::end_padding_packed_b(),
)
}
fn internal_type(&self) -> DatumType {
TI::datum_type()
}
unsafe fn a_packed(&self, item_size: usize, k: usize) -> PackedStoreSpec {
PackedStoreSpec {
panel_bytes: (k * K::mr() * item_size),
}
}
unsafe fn b_packed(&self, item_size: usize, k: usize) -> InputStoreSpec {
let panel_bytes = k * K::nr() * item_size;
InputStoreSpec::Prepacked(PackedStoreSpec { panel_bytes })
}
unsafe fn b_late_packing_with_axes(&self, k_axis: usize, n_axis: usize) -> InputStoreSpec {
InputStoreSpec::LatePacking {
packer: self.b_pack(),
k_axis,
mn_axis: n_axis,
}
}
unsafe fn b_virtual_input(&self, func: Box<dyn VirtualInputSpec>, k: usize) -> InputStoreSpec {
InputStoreSpec::VirtualPacking {
packer: self.b_pack(),
func,
k,
}
}
unsafe fn c_view(&self, m_axis: usize, n_axis: usize) -> OutputStoreSpec {
OutputStoreSpec::View {
m_axis,
n_axis,
mr: K::mr(),
nr: K::nr(),
}
}
unsafe fn c_from_data_and_strides(
&self,
item_size: usize,
m: usize,
n: usize,
row_stride: isize,
col_stride: isize,
) -> OutputStoreSpec {
OutputStoreSpec::Strides {
row_byte_stride: row_stride * item_size as isize,
col_byte_stride: col_stride * item_size as isize,
mr: K::mr(),
nr: K::nr(),
m,
n,
}
}
unsafe fn allocate_scratch_space(&self) -> Box<dyn ScratchSpace> {
Box::<ScratchSpaceFusedNonLinear<TI>>::default()
}
unsafe fn can_use_scratch_space(&self, scratch: &dyn ScratchSpace) -> bool {
scratch
.downcast_ref::<ScratchSpaceFusedNonLinear<TI>>()
.is_some()
}
unsafe fn run_with_scratch_space_vec(
&self,
m: usize,
scratch: &mut dyn ScratchSpace,
non_linear: &[FusedSpec],
) -> anyhow::Result<()> {
let mr = K::mr();
let scratch = scratch
.downcast_mut::<ScratchSpaceFusedNonLinear<TI>>()
.context("Wrong scratch space type")?;
scratch.prepare::<K>(non_linear)?;
for ia in 0..m / mr {
scratch.for_valid_tile::<K>(non_linear, ia, 0);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
}
if m % mr != 0 {
scratch.for_border_tile::<K>(non_linear, m / mr, 0);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
scratch.postprocess_tile::<K>(non_linear, m / mr, 0, m % mr, 1);
}
Ok(())
}
unsafe fn run_with_scratch_space_col_outer(
&self,
m: usize,
n: usize,
scratch: &mut dyn ScratchSpace,
non_linear: &[FusedSpec],
) -> anyhow::Result<()> {
let mr = K::mr();
let nr = K::nr();
let scratch = scratch
.downcast_mut::<ScratchSpaceFusedNonLinear<TI>>()
.context("Wrong scratch space type")?;
scratch.prepare::<K>(non_linear)?;
for ib in 0..n / nr {
for ia in 0..m / mr {
scratch.for_valid_tile::<K>(non_linear, ia, ib);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
}
if m % mr != 0 {
scratch.for_border_tile::<K>(non_linear, m / mr, ib);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
scratch.postprocess_tile::<K>(non_linear, m / mr, ib, m % mr, nr);
}
}
if n % nr != 0 {
for ia in 0..m / mr {
scratch.for_border_tile::<K>(non_linear, ia, n / nr);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
scratch.postprocess_tile::<K>(non_linear, ia, n / nr, mr, n % nr);
}
if m % mr != 0 {
scratch.for_border_tile::<K>(non_linear, m / mr, n / nr);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
scratch.postprocess_tile::<K>(non_linear, m / mr, n / nr, m % mr, n % nr);
}
}
Ok(())
}
unsafe fn run_with_scratch_space(
&self,
m: usize,
n: usize,
scratch: &mut dyn ScratchSpace,
non_linear: &[FusedSpec],
) -> anyhow::Result<()> {
let mr = K::mr();
let nr = K::nr();
if n == 1 && K::nr() == 1 {
return self.run_with_scratch_space_vec(m, scratch, non_linear);
}
if non_linear.iter().any(|f| f.prefer_col_outer()) {
return self.run_with_scratch_space_col_outer(m, n, scratch, non_linear);
}
let scratch = scratch
.downcast_mut::<ScratchSpaceFusedNonLinear<TI>>()
.context("Wrong scratch space type")?;
scratch.prepare::<K>(non_linear)?;
for ia in 0..m / mr {
for ib in 0..n / nr {
scratch.for_valid_tile::<K>(non_linear, ia, ib);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
}
}
if m % mr != 0 {
for ib in 0..n / nr {
scratch.for_border_tile::<K>(non_linear, m / mr, ib);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
scratch.postprocess_tile::<K>(non_linear, m / mr, ib, m % mr, nr);
}
}
if n % nr != 0 {
for ia in 0..m / mr {
scratch.for_border_tile::<K>(non_linear, ia, n / nr);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
scratch.postprocess_tile::<K>(non_linear, ia, n / nr, mr, n % nr);
}
if m % mr != 0 {
scratch.for_border_tile::<K>(non_linear, m / mr, n / nr);
let err = K::kernel(scratch.uspecs());
debug_assert_eq!(err, 0, "Kernel return error {err}");
scratch.postprocess_tile::<K>(non_linear, m / mr, n / nr, m % mr, n % nr);
}
}
Ok(())
}
}
impl<K, TI> fmt::Display for MatMatMulImpl<K, TI>
where
TI: LADatum,
K: MatMatMulKer<TI>,
{
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
write!(fmt, "({} {}x{})", K::name(), K::mr(), K::nr())
}
}
@@ -0,0 +1,429 @@
use std::alloc::Layout;
use std::fmt::Debug;
use tract_data::internal::*;
use crate::LADatum;
use super::{BinOp, FusedKerSpec, FusedSpec, MatMatMulKer, OutputStoreKer};
use downcast_rs::{Downcast, impl_downcast};
use tract_data::internal::num_integer::Integer;
pub trait ScratchSpace: Downcast + Send {}
impl_downcast!(ScratchSpace);
#[derive(Debug)]
pub struct ScratchSpaceFusedNonLinear<TI: LADatum> {
uspecs: Vec<FusedKerSpec<TI>>,
layout: Layout,
buffer: *const u8,
loc_dependant: TVec<LocDependant>,
}
impl<TI: LADatum> Default for ScratchSpaceFusedNonLinear<TI> {
fn default() -> Self {
ScratchSpaceFusedNonLinear {
uspecs: vec![],
layout: unsafe { Layout::from_size_align_unchecked(0, 1) },
buffer: std::ptr::null(),
loc_dependant: tvec!(),
}
}
}
#[derive(Debug, new)]
struct LocDependant {
spec: usize,
uspec: usize,
loc: *const u8,
buffer: Option<*const u8>,
}
impl<TI: LADatum> ScratchSpace for ScratchSpaceFusedNonLinear<TI> {}
unsafe impl<TI: LADatum> Send for ScratchSpaceFusedNonLinear<TI> {}
impl<TI: LADatum> Drop for ScratchSpaceFusedNonLinear<TI> {
fn drop(&mut self) {
if !self.buffer.is_null() {
unsafe {
std::alloc::dealloc(self.buffer as _, self.layout);
}
}
}
}
struct AddMatMulTemp(*const u8, usize);
impl<TI: LADatum> ScratchSpaceFusedNonLinear<TI> {
pub unsafe fn prepare<K: MatMatMulKer<TI>>(&mut self, specs: &[FusedSpec]) -> TractResult<()> {
use FusedKerSpec as FKS;
use FusedSpec as FS;
self.uspecs.clear();
self.loc_dependant.clear();
self.uspecs.reserve(specs.len() + 2);
self.uspecs.push(FusedKerSpec::Clear);
let mut offset = 0;
let mut align = std::mem::size_of::<*const ()>();
fn ld(spec: usize, uspec: usize, loc: *const u8) -> LocDependant {
LocDependant {
spec,
uspec,
loc,
buffer: None,
}
}
// we're cheating here, storing offset as the buf pointer first
for (ix, spec) in specs.iter().enumerate() {
let uspec = match spec {
FS::BinScalar(t, op) => match op {
BinOp::Min => FKS::ScalarMin(*t.to_scalar()?),
BinOp::Max => FKS::ScalarMax(*t.to_scalar()?),
BinOp::Mul => FKS::ScalarMul(*t.to_scalar()?),
BinOp::Add => FKS::ScalarAdd(*t.to_scalar()?),
BinOp::Sub => FKS::ScalarSub(*t.to_scalar()?),
BinOp::SubF => FKS::ScalarSubF(*t.to_scalar()?),
},
FS::ShiftLeft(s) => FKS::ShiftLeft(*s),
FS::RoundingShiftRight(s, rp) => FKS::RoundingShiftRight(*s, *rp),
FS::QScale(s, rp, m) => FKS::QScale(*s, *rp, *m),
FS::BinPerRow(_, _) => {
self.loc_dependant
.push(ld(ix, self.uspecs.len(), offset as _));
offset += TI::datum_type().size_of() * K::mr();
FusedKerSpec::Done
}
FS::BinPerCol(_, _) => {
self.loc_dependant
.push(ld(ix, self.uspecs.len(), offset as _));
offset += TI::datum_type().size_of() * K::nr();
FusedKerSpec::Done
}
FS::AddRowColProducts(_, _) => {
self.loc_dependant
.push(ld(ix, self.uspecs.len(), offset as _));
offset += TI::datum_type().size_of() * (K::mr() + K::nr());
FusedKerSpec::Done
}
FS::Store(_) | FS::AddUnicast(_) => {
self.loc_dependant
.push(ld(ix, self.uspecs.len(), offset as _));
offset += TI::datum_type().size_of() * K::mr() * K::nr();
FusedKerSpec::Done
}
FS::AddMatMul { b, .. } => {
let mut ld = ld(ix, self.uspecs.len(), offset as _);
offset += std::mem::size_of::<AddMatMulTemp>();
if let Some(tmp) = b.scratch_panel_buffer_layout() {
align = tmp.align().lcm(&align);
offset = Integer::next_multiple_of(&offset, &tmp.align());
ld.buffer = Some(offset as _);
offset += tmp.size();
}
self.loc_dependant.push(ld);
FusedKerSpec::Done
}
};
self.uspecs.push(uspec);
}
self.uspecs.push(FKS::Done);
if offset > self.layout.size() || align > self.layout.align() {
if !self.buffer.is_null() {
std::alloc::dealloc(self.buffer as _, self.layout);
}
self.layout = Layout::from_size_align_unchecked(offset, align);
self.buffer = std::alloc::alloc(self.layout);
assert!(!self.buffer.is_null());
}
for LocDependant {
loc, buffer, spec, ..
} in &mut self.loc_dependant
{
*loc = self.buffer.offset(*loc as _);
if let Some(b) = buffer {
*b = self.buffer.offset(*b as _);
}
let spec = specs.get_unchecked(*spec);
#[allow(clippy::single_match)]
match spec {
FS::AddMatMul { .. } => {
let scratch = *loc as *mut AddMatMulTemp;
(*scratch).1 = usize::MAX;
}
_ => (),
};
}
Ok(())
}
#[inline(always)]
pub unsafe fn for_valid_tile<K: MatMatMulKer<TI>>(
&mut self,
specs: &[FusedSpec],
down: usize,
right: usize,
) {
use FusedKerSpec as FKS;
use FusedSpec as FS;
let ScratchSpaceFusedNonLinear {
uspecs,
loc_dependant,
..
} = self;
debug_assert!(specs.len() + 2 == uspecs.len());
for LocDependant {
spec,
uspec,
loc,
buffer,
} in loc_dependant.iter_mut()
{
let spec = specs.get_unchecked(*spec);
*uspecs.get_unchecked_mut(*uspec) = match spec {
FS::BinPerRow(v, op) => {
let v = v.as_ptr_unchecked::<TI>().add(down * K::mr());
match op {
BinOp::Min => FKS::PerRowMin(v),
BinOp::Max => FKS::PerRowMax(v),
BinOp::Add => FKS::PerRowAdd(v),
BinOp::Mul => FKS::PerRowMul(v),
BinOp::Sub => FKS::PerRowSub(v),
BinOp::SubF => FKS::PerRowSubF(v),
}
}
FS::BinPerCol(v, op) => {
let v = v.as_ptr_unchecked::<TI>().add(right * K::nr());
match op {
BinOp::Min => FKS::PerColMin(v),
BinOp::Max => FKS::PerColMax(v),
BinOp::Add => FKS::PerColAdd(v),
BinOp::Mul => FKS::PerColMul(v),
BinOp::Sub => FKS::PerColSub(v),
BinOp::SubF => FKS::PerColSubF(v),
}
}
FS::AddRowColProducts(rows, cols) => {
let row_ptr = rows.as_ptr_unchecked::<TI>().add(down * K::mr());
let col_ptr = cols.as_ptr_unchecked::<TI>().add(right * K::nr());
FKS::AddRowColProducts(row_ptr, col_ptr)
}
FS::AddUnicast(store) => FKS::AddUnicast(store.tile_c(down, right)),
FS::Store(c_store) => FKS::Store(c_store.tile_c(down, right)),
FS::AddMatMul { k, a, b } => {
let pa = a.panel(down);
K::prefetch(pa as _, 512);
let scratch = *loc as *mut AddMatMulTemp;
if (*scratch).1 != right {
(*scratch).0 = b.panel_b(right, *buffer);
(*scratch).1 = right;
}
FKS::AddMatMul {
k: *k,
pa,
pb: (*scratch).0,
cpu_variant: 0,
}
}
_ => std::hint::unreachable_unchecked(),
};
}
}
#[inline(never)]
pub unsafe fn for_border_tile<K: MatMatMulKer<TI>>(
&mut self,
specs: &[FusedSpec],
down: usize,
right: usize,
) {
use FusedKerSpec as FKS;
use FusedSpec as FS;
let ScratchSpaceFusedNonLinear {
uspecs,
loc_dependant,
..
} = self;
debug_assert!(specs.len() + 2 == uspecs.len());
for LocDependant {
spec,
uspec,
loc,
buffer,
} in loc_dependant.iter_mut()
{
let spec = specs.get_unchecked(*spec);
*uspecs.get_unchecked_mut(*uspec) = match spec {
FS::BinPerRow(v, op) => {
let buf = std::slice::from_raw_parts_mut(*loc as *mut TI, K::mr());
let have = v.len().saturating_sub(down * K::mr()).min(K::mr());
let ptr = if have < K::mr() {
if have > 0 {
buf.get_unchecked_mut(..have).copy_from_slice(
v.as_slice_unchecked()
.get_unchecked(down * K::mr()..)
.get_unchecked(..have),
);
}
if cfg!(debug_assertions) {
buf.get_unchecked_mut(have..)
.iter_mut()
.for_each(|x| *x = TI::zero());
}
buf.as_ptr()
} else {
v.as_ptr_unchecked::<TI>().add(down * K::mr())
};
match op {
BinOp::Min => FKS::PerRowMin(ptr),
BinOp::Max => FKS::PerRowMax(ptr),
BinOp::Add => FKS::PerRowAdd(ptr),
BinOp::Mul => FKS::PerRowMul(ptr),
BinOp::Sub => FKS::PerRowSub(ptr),
BinOp::SubF => FKS::PerRowSubF(ptr),
}
}
FS::BinPerCol(v, op) => {
let buf = std::slice::from_raw_parts_mut(*loc as *mut TI, K::nr());
let have = v.len().saturating_sub(right * K::nr()).min(K::nr());
let ptr = if have < K::nr() {
if have > 0 {
buf.get_unchecked_mut(..have).copy_from_slice(
v.as_slice_unchecked()
.get_unchecked(right * K::nr()..)
.get_unchecked(..have),
);
}
if cfg!(debug_assertions) {
buf.get_unchecked_mut(have..)
.iter_mut()
.for_each(|x| *x = TI::zero());
}
buf.as_ptr()
} else {
v.as_ptr_unchecked::<TI>().add(right * K::nr())
};
match op {
BinOp::Min => FKS::PerColMin(ptr),
BinOp::Max => FKS::PerColMax(ptr),
BinOp::Add => FKS::PerColAdd(ptr),
BinOp::Mul => FKS::PerColMul(ptr),
BinOp::Sub => FKS::PerColSub(ptr),
BinOp::SubF => FKS::PerColSubF(ptr),
}
}
FS::AddRowColProducts(rows, cols) => {
let r = std::slice::from_raw_parts_mut(*loc as *mut TI, K::mr());
let have = rows.len() - down * K::mr();
let row_ptr = if have < K::mr() {
r.get_unchecked_mut(..have).copy_from_slice(
rows.as_slice_unchecked()
.get_unchecked(down * K::mr()..)
.get_unchecked(..have),
);
if cfg!(debug_assertions) {
r.get_unchecked_mut(have..)
.iter_mut()
.for_each(|x| *x = TI::zero());
}
r.as_ptr()
} else {
rows.as_ptr_unchecked::<TI>().add(down * K::mr())
};
let c = std::slice::from_raw_parts_mut((*loc as *mut TI).add(K::mr()), K::nr());
let have = cols.len() - right * K::nr();
let col_ptr = if have < K::nr() {
c.get_unchecked_mut(..have).copy_from_slice(
cols.as_slice_unchecked()
.get_unchecked(right * K::nr()..)
.get_unchecked(..have),
);
if cfg!(debug_assertions) {
r.get_unchecked_mut(have..)
.iter_mut()
.for_each(|x| *x = TI::zero());
}
c.as_ptr()
} else {
cols.as_ptr_unchecked::<TI>().add(right * K::nr())
};
FKS::AddRowColProducts(row_ptr, col_ptr)
}
FS::AddUnicast(store) => {
let row_byte_stride = store.row_byte_stride;
let col_byte_stride = store.col_byte_stride;
let tile_offset = row_byte_stride * down as isize * K::mr() as isize
+ col_byte_stride * right as isize * K::nr() as isize;
let tile_ptr = store.ptr.offset(tile_offset);
let tmp_d_tile =
std::slice::from_raw_parts_mut(*loc as *mut TI, K::mr() * K::nr());
let m = (store.m - down * K::mr()).min(K::mr());
let n = (store.n - right * K::nr()).min(K::nr());
for r in 0..m as isize {
for c in 0..n as isize {
let inner_offset = c * col_byte_stride + r * row_byte_stride;
if inner_offset + tile_offset
< (store.item_size * store.item_count) as isize
{
*tmp_d_tile.get_unchecked_mut(r as usize + c as usize * K::mr()) =
*(tile_ptr.offset(inner_offset) as *const TI);
}
}
}
FKS::AddUnicast(OutputStoreKer {
ptr: tmp_d_tile.as_ptr() as _,
row_byte_stride: std::mem::size_of::<TI>() as isize,
col_byte_stride: (std::mem::size_of::<TI>() * K::mr()) as isize,
item_size: std::mem::size_of::<TI>(),
})
}
FS::Store(c_store) => {
let tmpc = OutputStoreKer {
ptr: *loc as _,
item_size: c_store.item_size,
row_byte_stride: c_store.item_size as isize,
col_byte_stride: (c_store.item_size * K::mr()) as isize,
};
FKS::Store(tmpc)
}
FS::AddMatMul { k, a, b } => {
let pa = a.panel(down);
K::prefetch(pa as _, 512);
let scratch = *loc as *mut AddMatMulTemp;
if (*scratch).1 != right {
(*scratch).0 = b.panel_b(right, *buffer);
(*scratch).1 = right;
}
FKS::AddMatMul {
k: *k,
pa,
pb: (*scratch).0,
cpu_variant: 0,
}
}
_ => std::hint::unreachable_unchecked(),
};
}
}
#[inline]
pub fn uspecs(&self) -> &[FusedKerSpec<TI>] {
&self.uspecs
}
pub unsafe fn postprocess_tile<K: MatMatMulKer<TI>>(
&mut self,
specs: &[FusedSpec],
down: usize,
right: usize,
m_remnant: usize,
n_remnant: usize,
) where
TI: LADatum,
{
for LocDependant { spec, uspec, .. } in self.loc_dependant.iter() {
let spec = specs.get_unchecked(*spec);
let ker_spec = self.uspecs.get_unchecked(*uspec);
if let (FusedSpec::Store(c_store), FusedKerSpec::Store(tmp)) = (spec, ker_spec) {
c_store.set_from_tile(down, right, m_remnant, n_remnant, tmp)
}
}
}
}
@@ -0,0 +1,160 @@
use std::fmt::Debug;
use tract_data::internal::*;
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum OutputStoreSpec {
View {
m_axis: usize,
n_axis: usize,
mr: usize,
nr: usize,
},
Strides {
row_byte_stride: isize,
col_byte_stride: isize,
mr: usize,
nr: usize,
m: usize,
n: usize,
},
}
#[derive(Clone, Copy, Debug)]
pub struct OutputStore {
pub(crate) ptr: *mut u8,
pub(crate) row_byte_stride: isize,
pub(crate) col_byte_stride: isize,
pub(crate) panel_row_byte_stride: isize,
pub(crate) panel_col_byte_stride: isize,
pub(crate) item_size: usize,
pub(crate) item_count: usize,
pub(crate) mr: usize,
pub(crate) m: usize,
pub(crate) n: usize,
}
impl OutputStoreSpec {
#[inline]
pub unsafe fn wrap(&self, tensor: &TensorView) -> OutputStore {
let (mr, nr, row_byte_stride, col_byte_stride) = self.compute_strides(tensor);
let (m, n) = match self {
OutputStoreSpec::View { m_axis, n_axis, .. } => {
(tensor.shape()[*m_axis], tensor.shape()[*n_axis])
}
OutputStoreSpec::Strides { m, n, .. } => (*m, *n),
};
OutputStore {
ptr: tensor.as_ptr_unchecked::<u8>() as _,
row_byte_stride,
col_byte_stride,
panel_row_byte_stride: row_byte_stride * mr as isize,
panel_col_byte_stride: col_byte_stride * nr as isize,
item_size: tensor.datum_type().size_of(),
mr,
item_count: tensor.len(),
m,
n,
}
}
#[inline]
unsafe fn compute_strides(&self, tensor: &TensorView) -> (usize, usize, isize, isize) {
let size_of = tensor.datum_type().size_of() as isize;
match self {
OutputStoreSpec::View {
m_axis,
n_axis,
mr,
nr,
..
} => {
let tensor_strides = tensor.strides();
let row_item_stride = *tensor_strides.get_unchecked(*m_axis);
let col_item_stride = *tensor_strides.get_unchecked(*n_axis);
let row_byte_stride = row_item_stride * size_of;
let col_byte_stride = col_item_stride * size_of;
(*mr, *nr, row_byte_stride, col_byte_stride)
}
OutputStoreSpec::Strides {
row_byte_stride,
col_byte_stride,
mr,
nr,
..
} => (*mr, *nr, *row_byte_stride, *col_byte_stride),
}
}
}
impl OutputStore {
#[inline]
pub(super) unsafe fn tile_c(&self, down: usize, right: usize) -> OutputStoreKer {
let (down, right) = (down as isize, right as isize);
OutputStoreKer {
ptr: self
.ptr
.offset(self.panel_row_byte_stride * down + self.panel_col_byte_stride * right)
as *mut _,
row_byte_stride: self.row_byte_stride,
col_byte_stride: self.col_byte_stride,
item_size: self.item_size,
}
}
#[inline]
pub fn item_size(&self) -> usize {
self.item_size
}
#[inline]
pub(super) unsafe fn set_from_tile(
&self,
down: usize,
right: usize,
height: usize,
width: usize,
tile: &OutputStoreKer,
) {
if self.item_size() == 1 {
self.set_from_tile_t::<i8>(down, right, height, width, tile)
} else if self.item_size() == 2 {
self.set_from_tile_t::<i16>(down, right, height, width, tile)
} else if self.item_size() == 4 {
self.set_from_tile_t::<i32>(down, right, height, width, tile)
} else {
self.set_from_tile_t::<i64>(down, right, height, width, tile)
}
}
#[inline]
unsafe fn set_from_tile_t<T: Datum + Copy>(
&self,
down: usize,
right: usize,
height: usize,
width: usize,
tile: &OutputStoreKer,
) {
let tile = tile.ptr as *mut T;
let dst = self.ptr.add(
self.panel_row_byte_stride as usize * down
+ self.panel_col_byte_stride as usize * right,
);
for y in 0..height as isize {
for x in 0..width as isize {
let value = tile.offset(y + x * self.mr as isize);
let dst = dst.offset(y * self.row_byte_stride + x * self.col_byte_stride);
*(dst as *mut T) = *value;
}
}
}
}
#[repr(C)]
#[derive(PartialEq, Eq, Copy, Clone, Debug)]
pub struct OutputStoreKer {
pub ptr: *mut u8,
pub row_byte_stride: isize,
pub col_byte_stride: isize,
pub item_size: usize,
}
@@ -0,0 +1,601 @@
use super::*;
use crate::LADatum;
use num_traits::AsPrimitive;
use proptest::prelude::*;
use std::ops::Neg;
use tract_data::internal::*;
#[macro_export]
macro_rules! mmm_frame_tests {
($cond:expr, $ker:ident, $ta:ty, $tb:ty, $tc:ty, $ti:ty) => {
mod frame {
#[allow(unused_imports)]
use $crate::frame::mmm::tests::*;
use tract_data::internal::*;
use super::super::$ker;
proptest::proptest! {
#[test]
fn mat_mul_prepacked_prop((m, k, n, ref a, ref b) in strat_mat_mat_mul::<$ta, $tb>()) {
if $cond {
test_mat_mat_mul_prep::<$ker, $ta, $tb, $tc, $ti>(m, k, n, &a, &b)?
}
}
#[test]
fn mat_mul_prepacked_late((m, k, n, ref a, ref b) in strat_mat_mat_mul::<$ta, $tb>()) {
if $cond {
test_mat_mat_mul_late::<$ker, $ta, $tb, $tc, $ti>(m, k, n, &a, &b)?
}
}
#[test]
fn mat_vec_prepacked_prop((m, k, ref a, ref b) in strat_mat_vec_mul::<$ta, $tb>()) {
if $cond {
test_mat_vec_mul_prep::<$ker, $ta, $tb, $tc, $ti>(m, k, &*a, b)?
}
}
}
#[test]
fn mat_mul_1() {
if $cond {
let a = tensor2(&[[-3i32, 3, 5, -5], [6, 0, -6, -5], [0, 0, 9, 7]]).cast_to::<$ta>().unwrap().into_owned();
let b = tensor2(&[[-8i32, 5],[ 5, -3], [5, 7],[ -8, -1]]).cast_to::<$tb>().unwrap().into_owned();
test_mat_mat_mul_prep::<$ker, $ta, $tb, $tc, $ti>(3, 4, 2, &a, &b).unwrap()
}
}
#[test]
fn mat_mul_2() {
if $cond {
let a = tensor2(&[[1i32]]).cast_to::<$ta>().unwrap().into_owned();
let b = tensor2(&[[0i32, 0, 1]]).cast_to::<$tb>().unwrap().into_owned();
test_mat_mat_mul_prep::<$ker, $ta, $tb, $tc, $ti>(1, 1, 3, &a, &b).unwrap()
}
}
#[test]
fn mat_mul_3() {
if $cond {
let a = tensor2(&[[-3i32, 3, 5, -5], [6, 0, -6, -5], [0, 0, 9, 7]]).cast_to::<$ta>().unwrap().into_owned();
let b = tensor2(&[[-8i32, 5],[ 5, -3], [5, 7],[ -8, -1]]).cast_to::<$tb>().unwrap().into_owned();
test_mat_mat_mul_prep::<$ker, $ta, $tb, $tc, $ti>(3, 4, 2, &a, &b).unwrap()
}
}
#[test]
fn mat_mul_4() {
if $cond {
let a = tensor2(&[[122, 82]]).cast_to::<$ta>().unwrap().into_owned();
let b = tensor2(&[[0, 0, 37],[ 0, 0, 57]]).cast_to::<$tb>().unwrap().into_owned();
test_mat_mat_mul_prep::<$ker, $ta, $tb, $tc, $ti>(1, 2, 3, &a, &b).unwrap()
}
}
#[test]
fn mat_mul_1_2_1() {
if $cond {
test_mat_mat_mul_prep::<$ker, $ta, $tb, $tc, $ti>(
1,
2,
1,
&tensor2(&[[0, 1]]).cast_to::<$ta>().unwrap(),
&tensor2(&[[0], [1]]).cast_to::<$tb>().unwrap(),
)
.unwrap()
}
}
#[test]
fn late_packing_1() {
if $cond {
let a = tensor2(&[[1f32, 2f32]]).cast_to::<$ta>().unwrap().into_owned();
let b = tensor2(&[[0f32, 0., 0.], [0., 0., 1.]]).cast_to::<$tb>().unwrap().into_owned();
test_mat_mat_mul_late::<$ker, $ta, $tb, $tc, $ti>(1, 2, 3, &a, &b).unwrap()
}
}
#[test]
fn mat_vec_1() {
if $cond {
let a = tensor2(&[[0], [1]]).cast_to::<$ta>().unwrap().into_owned();
let b = tensor1(&[1]).cast_to::<$tb>().unwrap().into_owned();
test_mat_vec_mul_prep::<$ker, $ta, $tb, $tc, $ti>(2, 1, &a, &b).unwrap()
}
}
#[test]
fn mat_vec_2() {
if $cond {
let a = tensor1(&[0, 0, 0, 0, 0, 0, -4, 1]).into_shape(&[8,1]).unwrap();
let a = a.cast_to::<$ta>().unwrap();
let b = tensor1(&[-64]).cast_to::<$tb>().unwrap().into_owned();
test_mat_vec_mul_prep::<$ker, $ta, $tb, $tc, $ti>(8, 1, &a, &b).unwrap()
}
}
#[test]
fn mat_vec_3() {
if $cond {
let a = tensor1(&[0, 0]).into_shape(&[1, 2]).unwrap();
let a = a.cast_to::<$ta>().unwrap();
let b = tensor1(&[0, 0]).cast_to::<$tb>().unwrap().into_owned();
test_mat_vec_mul_prep::<$ker, $ta, $tb, $tc, $ti>(1, 2, &a, &b).unwrap()
}
}
#[test]
fn row_mul_2_1_3() {
if $cond {
unsafe { row_mul::<$ker, $ta, $tb, $tc, $ti>(2, 3).unwrap() }
}
}
#[test]
fn row_add_2_1_3() {
if $cond {
unsafe { row_add::<$ker, $ta, $tb, $tc, $ti>(2, 3).unwrap() }
}
}
#[test]
fn col_mul_2_1_3() {
if $cond {
unsafe { col_mul::<$ker, $ta, $tb, $tc, $ti>(2, 3).unwrap() }
}
}
#[test]
fn col_add_2_1_3() {
if $cond {
unsafe { col_add::<$ker, $ta, $tb, $tc, $ti>(2, 3).unwrap() }
}
}
#[test]
fn max_2_1_3() {
if $cond {
unsafe { max::<$ker, $ta, $tb, $tc, $ti>(2, 3).unwrap() }
}
}
#[test]
fn min_2_1_3() {
if $cond {
unsafe { min::<$ker, $ta, $tb, $tc, $ti>(2, 3).unwrap() }
}
}
#[test]
fn add_d_2_1_3() {
if $cond {
unsafe { add_d::<$ker, $ta, $tb, $tc, $ti>(2, 3).unwrap() }
}
}
#[test]
fn add_d_big() {
if $cond {
unsafe { add_d::<$ker, $ta, $tb, $tc, $ti>(197, 1).unwrap() }
}
}
}
};
}
fn tensor(dt: DatumType, shape: Vec<usize>) -> BoxedStrategy<Tensor> {
let len = shape.iter().product::<usize>();
// for f16, positive numbers only to avoid worst rounding side effects
// and not too big either to avoid overflow :)
let number = if dt == f16::datum_type() {
(0i16..100).boxed()
} else {
any::<i8>().prop_map(|i| i as i16).boxed()
};
proptest::collection::vec(number, len..=len)
.prop_map(move |vec| {
tract_ndarray::ArrayD::from_shape_vec(shape.clone(), vec)
.unwrap()
.into_tensor()
.cast_to_dt(dt)
.unwrap()
.into_owned()
})
.boxed()
}
pub fn strat_mat_mat_mul<TA: LADatum, TB: LADatum>()
-> BoxedStrategy<(usize, usize, usize, Tensor, Tensor)> {
(1usize..5, 1usize..5, 1usize..5)
.prop_flat_map(move |(m, k, n)| {
(
Just(m),
Just(k),
Just(n),
tensor(TA::datum_type(), vec![m, k]),
tensor(TB::datum_type(), vec![k, n]),
)
})
.boxed()
}
pub fn strat_mat_vec_mul<TA: LADatum, TB: LADatum>() -> BoxedStrategy<(usize, usize, Tensor, Tensor)>
{
(1usize..15, 1usize..15)
.prop_flat_map(move |(m, k)| {
(
Just(m),
Just(k),
tensor(TA::datum_type(), vec![m, k]),
tensor(TB::datum_type(), vec![k, 1]),
)
})
.boxed()
}
pub fn test_mat_mat_mul_prep<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
k: usize,
n: usize,
a: &Tensor,
b: &Tensor,
) -> Result<(), proptest::test_runner::TestCaseError>
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
assert_eq!(a.datum_type(), TA::datum_type());
let op = MatMatMulImpl::<K, TI>::default();
unsafe {
let mut packed_a =
Tensor::uninitialized_aligned::<TA>(&[op.a_pack().len(k, m)], op.a_pack().alignment())
.unwrap();
op.a_pack().pack(packed_a.view_mut(), a.view(), 1, 0);
let mut packed_b =
Tensor::uninitialized_aligned::<TB>(&[op.b_pack().len(k, n)], op.b_pack().alignment())
.unwrap();
op.b_pack().pack(packed_b.view_mut(), b.view(), 0, 1);
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[FusedSpec::AddMatMul {
a: op
.a_packed(TA::datum_type().size_of(), k)
.wrap(&packed_a.view()),
b: op
.b_packed(TB::datum_type().size_of(), k)
.wrap(&packed_b.view())
.unwrap(),
k,
}],
|r, c| {
let mut v: TI = TI::zero();
for i in 0..k {
let a: TI = a.as_slice::<TA>().unwrap()[i + k * r].as_();
let b: TI = b.as_slice::<TB>().unwrap()[c + i * n].as_();
v += a * b;
}
v.as_()
},
)
}
}
pub fn test_mat_mat_mul_late<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
k: usize,
n: usize,
a: &Tensor,
b: &Tensor,
) -> Result<(), proptest::test_runner::TestCaseError>
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
assert_eq!(a.datum_type(), TA::datum_type());
let op = MatMatMulImpl::<K, TI>::default();
unsafe {
let mut packed_a =
Tensor::uninitialized_aligned::<TA>(&[op.a_pack().len(k, m)], op.a_pack().alignment())
.unwrap();
op.a_pack().pack(packed_a.view_mut(), a.view(), 1, 0);
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[FusedSpec::AddMatMul {
a: op
.a_packed(TA::datum_type().size_of(), k)
.wrap(&packed_a.view()),
b: op.b_late_packing().wrap(&b.view()).unwrap(),
k,
}],
|r, c| {
let mut v: TI = TI::zero();
for i in 0..k {
let a: TI = a.as_slice::<TA>().unwrap()[i + k * r].as_();
let b: TI = b.as_slice::<TB>().unwrap()[c + i * n].as_();
v += a * b;
}
v.as_()
},
)
}
}
pub fn test_mat_vec_mul_prep<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
k: usize,
a: &Tensor,
b: &Tensor,
) -> Result<(), proptest::test_runner::TestCaseError>
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
unsafe {
let op = MatMatMulImpl::<K, TI>::default();
let mut packed_a =
Tensor::uninitialized_aligned::<TA>(&[op.a_pack().len(k, m)], op.a_pack().alignment())
.unwrap();
let mut packed_b =
Tensor::uninitialized_aligned::<TB>(&[op.b_pack().len(k, 1)], op.b_pack().alignment())
.unwrap();
op.a_pack().pack(&mut packed_a.view_mut(), &a.view(), 1, 0);
let b = b.clone().into_shape(&[k, 1]).unwrap();
op.b_pack().pack(&mut packed_b.view_mut(), &b.view(), 0, 1);
let pa = op
.a_packed(TA::datum_type().size_of(), k)
.wrap(&packed_a.view());
let pb = op
.b_packed(b.datum_type().size_of(), k)
.wrap(&packed_b.view())
.unwrap();
fused_ops::<K, TA, TB, TC, TI, _>(
m,
1,
&[FusedSpec::AddMatMul { k, a: pa, b: pb }],
|r, _| {
let mut inter = TI::zero();
for i in 0..k {
let a: TI = a.as_slice::<TA>().unwrap()[i + k * r].as_();
let b: TI = b.as_slice::<TB>().unwrap()[i].as_();
inter += a * b;
}
inter.as_()
},
)
}
}
pub unsafe fn fused_ops<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI, F: Fn(usize, usize) -> TC>(
m: usize,
n: usize,
spec: &[FusedSpec],
expect: F,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let op = MatMatMulImpl::<K, TI>::default();
let mut found = Tensor::zero::<TC>(&[m, n]).unwrap();
let c_store = op
.c_from_data_and_strides(TC::datum_type().size_of(), m, n, n as isize, 1)
.wrap(&found.view_mut());
let mut spec: TVec<FusedSpec> = spec.into();
spec.push(FusedSpec::Store(c_store));
op.run(m, n, &spec).unwrap();
let expected =
tract_ndarray::prelude::Array2::from_shape_fn((m, n), |(r, c)| expect(r, c)).into_tensor();
if found.close_enough(&expected, true).is_err() {
println!("found, expected:");
for r in 0..m {
for c in 0..n {
let f = found.as_slice_unchecked::<TC>()[r * n + c];
let e = expected.as_slice_unchecked::<TC>()[r * n + c];
let mut s = format!("{:4} ", f);
if f != e {
s = nu_ansi_term::Color::Red.paint(s).to_string();
}
print!("{:4} ", s);
}
print!(" ");
for c in 0..n {
print!("{:4} ", expected.as_slice_unchecked::<TC>()[r * n + c]);
}
println!();
}
}
found
.close_enough(&expected, true)
.map_err(|e| TestCaseError::Fail(e.to_string().into()))
}
pub unsafe fn row_add<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
n: usize,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let bias = (0..m).map(|i| i.as_()).collect::<Vec<TI>>();
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[FusedSpec::BinPerRow(&tensor1(&bias), BinOp::Add)],
|r, _| bias[r].as_(),
)
}
pub unsafe fn row_mul<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
n: usize,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let bias = (0..m).map(|i| i.as_()).collect::<Vec<TI>>();
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[
FusedSpec::BinScalar(&tensor0(1i32.as_()), BinOp::Add),
FusedSpec::BinPerRow(&tensor1(&bias), BinOp::Mul),
],
|r, _| bias[r].as_(),
)
}
pub unsafe fn col_add<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
n: usize,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let bias = (0..n).map(|i| i.as_()).collect::<Vec<TI>>();
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[FusedSpec::BinPerCol(&tensor1(&bias), BinOp::Add)],
|_, c| bias[c].as_(),
)
}
pub unsafe fn col_mul<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
n: usize,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let bias = (0..n).map(|i| i.as_()).collect::<Vec<TI>>();
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[
FusedSpec::BinScalar(&tensor0(1i32.as_()), BinOp::Add),
FusedSpec::BinPerCol(&tensor1(&bias), BinOp::Mul),
],
|_, c| bias[c].as_(),
)
}
pub unsafe fn add_d<K: MatMatMulKer<TI> + 'static, TA, TB, TC, TI>(
m: usize,
n: usize,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let d = (0..m * n).map(|i| i.as_()).collect::<Vec<TI>>();
let d = tensor1(&d).into_shape(&[m, n]).unwrap();
let store_spec = OutputStoreSpec::View {
m_axis: 0,
n_axis: 1,
mr: K::mr(),
nr: K::nr(),
};
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[FusedSpec::AddUnicast(store_spec.wrap(&d.view()))],
|r, c| {
d.to_array_view_unchecked::<TI>()
.into_dimensionality()
.unwrap()[(r, c)]
.as_()
},
)
}
pub unsafe fn max<K: MatMatMulKer<TI>, TA, TB, TC, TI>(
m: usize,
n: usize,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let five: TI = 5.as_();
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[FusedSpec::BinScalar(&tensor0(five), BinOp::Max)],
|_, _| five.as_(),
)
}
pub unsafe fn min<K: MatMatMulKer<TI>, TA, TB, TC, TI>(
m: usize,
n: usize,
) -> proptest::test_runner::TestCaseResult
where
TA: LADatum + AsPrimitive<TI> + 'static,
TB: LADatum + AsPrimitive<TI> + 'static,
TC: LADatum + AsPrimitive<TI> + 'static,
TI: LADatum + AsPrimitive<TC> + 'static + Neg<Output = TI>,
i32: AsPrimitive<TI>,
usize: AsPrimitive<TI>,
{
let five: TI = 5.as_();
fused_ops::<K, TA, TB, TC, TI, _>(
m,
n,
&[FusedSpec::BinScalar(&tensor0(five), BinOp::Min)],
|_, _| TC::zero(),
)
}
@@ -0,0 +1,629 @@
use std::fmt::Debug;
use std::marker::PhantomData;
use std::ops::Range;
use tract_data::internal::*;
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct Packer {
pub r: usize,
alignment: usize,
end_padding_record: usize,
}
impl Packer {
pub fn new(nr: usize, alignment: usize, end_padding_record: usize) -> Packer {
Packer {
r: nr,
alignment,
end_padding_record,
}
}
pub fn alignment(&self) -> usize {
self.alignment
}
pub fn panel_width(&self) -> usize {
self.r
}
pub fn len<D: DimLike>(&self, k: D, n: D) -> D {
(n.divceil(self.r) * (k + self.end_padding_record)) * self.r
}
pub fn single_panel_len(&self, k: usize) -> usize {
(k + self.end_padding_record) * self.r
}
#[allow(clippy::too_many_arguments)]
pub unsafe fn pack_t<T: Datum + Copy>(
&self,
pb: *mut T,
b: *const T,
mn: usize,
k_stride: isize,
mn_stride: isize,
k_range: Range<usize>,
mn_range: Range<usize>,
) {
if self.r == 1 && k_stride == 1 && mn == 1 {
pb.copy_from_nonoverlapping(b.add(k_range.start), k_range.len())
} else if mn_stride == 1 {
let size_of = T::datum_type().size_of();
let rbytes = self.r * size_of;
let mn_valid_end = mn_range.end.min(mn);
let mn_range_bytes = mn_range.start * size_of..mn_valid_end * size_of;
let k_stride_bytes = k_stride * size_of as isize;
let bb = b as *const u8;
let pbb = pb as *mut u8;
match rbytes {
16 => pack_mn_major::<[u8; 16]>(bb, pbb, k_stride_bytes, mn_range_bytes, k_range),
24 => pack_mn_major::<[u8; 24]>(bb, pbb, k_stride_bytes, mn_range_bytes, k_range),
32 => pack_mn_major::<[u8; 32]>(bb, pbb, k_stride_bytes, mn_range_bytes, k_range),
48 => pack_mn_major::<[u8; 48]>(bb, pbb, k_stride_bytes, mn_range_bytes, k_range),
64 => pack_mn_major::<[u8; 64]>(bb, pbb, k_stride_bytes, mn_range_bytes, k_range),
_ => {
let mut packer = self.write_with_k_outer(pb, k_range.len(), mn_range.len());
for k in k_range {
for x in mn_range.start..mn_valid_end {
packer.write(*b.offset(x as isize + k_stride * k as isize))
}
for _x in mn_valid_end..mn_range.end {
packer.write(T::default())
}
}
}
}
} else if k_stride == 1 {
let mut packer = self.write_with_k_inner(pb, k_range.len(), mn);
let mn_valid_end = mn_range.end.min(mn);
for x in mn_range.start..mn_valid_end {
for k in k_range.clone() {
packer.write(*b.offset(x as isize * mn_stride + k as isize))
}
}
// just ignore invalid mn_range
} else {
let mut packer = self.write_with_k_outer(pb, k_range.len(), mn);
let mn_valid_end = mn_range.end.min(mn);
for k in k_range {
for x in mn_range.start..mn_valid_end {
packer.write(*b.offset(x as isize * mn_stride + k_stride * k as isize))
}
for _x in mn_valid_end..mn_range.end {
packer.write(T::default())
}
}
}
}
pub unsafe fn pack_segment<'a, 'b>(
&self,
mut pb: impl std::borrow::BorrowMut<TensorView<'a>>,
b: impl std::borrow::Borrow<TensorView<'b>>,
k_axis: usize,
mn_axis: usize,
k_range: Range<usize>,
mn_range: Range<usize>,
) {
debug_assert_eq!(pb.borrow().len(), self.len(k_range.len(), mn_range.len()));
let pb = pb.borrow_mut();
let b = b.borrow();
let dt = pb.datum_type();
dispatch_copy!(Self::pack_t(dt)(
self,
pb.as_ptr_mut_unchecked(),
b.as_ptr_unchecked(),
b.shape()[mn_axis],
b.strides()[k_axis],
b.strides()[mn_axis],
k_range,
mn_range
));
}
pub unsafe fn pack<'a, 'b>(
&self,
pb: impl std::borrow::BorrowMut<TensorView<'a>>,
b: impl std::borrow::Borrow<TensorView<'b>>,
k_axis: usize,
mn_axis: usize,
) {
let k = b.borrow().shape()[k_axis];
let mn = b.borrow().shape()[mn_axis];
self.pack_segment(pb, b, k_axis, mn_axis, 0..k, 0..mn);
}
pub fn write_with_k_outer<'p, T: Copy + Debug>(
&self,
pb: *mut T,
k: usize,
mn: usize,
) -> KOutWriter<'p, T> {
KOutWriter::new(pb, self.r, mn, k)
}
pub fn write_single_panel_with_k_outer<'p, T: Copy + Debug>(
&self,
pb: *mut T,
) -> KOutSinglePanelWriter<'p, T> {
KOutSinglePanelWriter::new(pb)
}
pub fn write_with_k_inner<'p, T: Copy + Debug>(
&self,
pb: *mut T,
k: usize,
mn: usize,
) -> KInWriter<'p, T> {
KInWriter::new(pb, self.r, mn, k)
}
}
pub trait PackingWriter<T: Copy> {
fn write(&mut self, t: T);
}
#[derive(Debug)]
pub struct KOutSinglePanelWriter<'p, T>
where
T: Copy + std::fmt::Debug,
{
ptr: *mut T,
_phantom: PhantomData<&'p T>,
}
impl<'p, T> KOutSinglePanelWriter<'p, T>
where
T: Copy + std::fmt::Debug,
{
pub fn new(ptr: *mut T) -> KOutSinglePanelWriter<'p, T> {
KOutSinglePanelWriter {
ptr,
_phantom: PhantomData,
}
}
}
impl<'p, T> PackingWriter<T> for KOutSinglePanelWriter<'p, T>
where
T: Copy + std::fmt::Debug,
{
#[inline(always)]
fn write(&mut self, t: T) {
unsafe {
*self.ptr = t;
self.ptr = self.ptr.offset(1);
}
}
}
#[derive(Debug)]
pub struct KOutWriter<'p, T>
where
T: Copy + std::fmt::Debug,
{
ptr: *mut T,
panels: usize,
panel_width: usize,
last_panel_width: usize,
remain: usize,
current_panel: usize,
next_panel: isize,
next_lane: isize,
_phantom: PhantomData<&'p T>,
}
impl<'p, T> KOutWriter<'p, T>
where
T: Copy + std::fmt::Debug,
{
pub fn new(ptr: *mut T, panel_width: usize, mn: usize, k: usize) -> KOutWriter<'p, T> {
let panels = (mn + panel_width - 1) / panel_width;
let last_panel_width = mn - (panels - 1) * panel_width;
KOutWriter {
ptr,
panels,
panel_width,
last_panel_width,
remain: if panels > 1 {
panel_width
} else {
last_panel_width
},
current_panel: 0,
next_panel: ((k - 1) * panel_width) as isize,
next_lane: panel_width as isize
- ((last_panel_width + (panels - 1) * panel_width * k) as isize),
_phantom: PhantomData,
}
}
}
impl<'p, T> PackingWriter<T> for KOutWriter<'p, T>
where
T: Copy + std::fmt::Debug,
{
#[inline(always)]
fn write(&mut self, t: T) {
unsafe {
*self.ptr = t;
self.remain -= 1;
self.ptr = self.ptr.offset(1);
if self.remain == 0 {
self.current_panel += 1;
if self.current_panel == self.panels {
self.ptr = self.ptr.offset(self.next_lane);
self.current_panel = 0;
} else {
self.ptr = self.ptr.offset(self.next_panel);
}
if self.current_panel == self.panels - 1 {
self.remain = self.last_panel_width;
} else {
self.remain = self.panel_width;
}
}
}
}
}
#[derive(Debug)]
pub struct KInWriter<'p, T>
where
T: Copy + Debug,
{
ptr: *mut T,
k: usize,
panels: usize,
panel_width: usize,
last_panel_width: usize,
remain_on_k: usize,
remain_on_mn: usize,
current_panel: usize,
next_mn_offset: isize,
next_panel_offset: isize,
_phantom: PhantomData<&'p T>,
}
impl<'p, T> KInWriter<'p, T>
where
T: Copy + Debug,
{
pub fn new(ptr: *mut T, panel_width: usize, mn: usize, k: usize) -> KInWriter<'p, T> {
let panels = (mn + panel_width - 1) / panel_width;
let last_panel_width = mn - (panels - 1) * panel_width;
KInWriter {
ptr,
k,
panels,
panel_width,
last_panel_width,
remain_on_k: k,
remain_on_mn: if panels == 1 {
last_panel_width
} else {
panel_width
},
current_panel: 0,
next_mn_offset: 1 - (k * panel_width) as isize,
next_panel_offset: 1 - panel_width as isize,
_phantom: PhantomData,
}
}
}
impl<'p, T> PackingWriter<T> for KInWriter<'p, T>
where
T: Copy + std::fmt::Debug,
{
#[inline(always)]
fn write(&mut self, t: T) {
unsafe {
*self.ptr = t;
self.remain_on_k -= 1;
self.ptr = self.ptr.add(self.panel_width);
if self.remain_on_k == 0 {
self.remain_on_k = self.k;
self.remain_on_mn -= 1;
if self.remain_on_mn > 0 {
self.ptr = self.ptr.offset(self.next_mn_offset);
} else {
self.ptr = self.ptr.offset(self.next_panel_offset);
self.current_panel += 1;
if self.current_panel == self.panels - 1 {
self.remain_on_mn = self.last_panel_width;
} else {
self.remain_on_mn = self.panel_width;
}
}
}
}
}
}
#[inline(never)]
unsafe fn pack_mn_major<Chunk: Copy>(
b: *const u8,
packed: *mut u8,
k_stride_bytes: isize,
mn_range_bytes: Range<usize>,
k_range: Range<usize>,
) {
let mnr = std::mem::size_of::<Chunk>();
let full_panes = mn_range_bytes.len() / mnr;
let partial_pane = mn_range_bytes.len() % mnr;
for k in 0..k_range.len() {
let mut p_row = packed.add(k * mnr);
let mut b_row =
b.offset((k_range.start + k) as isize * k_stride_bytes + mn_range_bytes.start as isize);
for _ in 0..full_panes {
p_row.copy_from_nonoverlapping(b_row, mnr);
p_row = p_row.add(k_range.len() * mnr);
b_row = b_row.add(mnr);
}
if partial_pane > 0 {
p_row.copy_from_nonoverlapping(b_row, partial_pane);
}
}
}
#[cfg(test)]
mod test {
use std::ops::Range;
use proptest::prelude::*;
use tract_data::internal::*;
use tract_ndarray::prelude::*;
#[derive(Debug)]
struct PackProblem {
k: usize,
mn: usize,
is_a: bool,
r: usize,
k_range: Range<usize>,
mn_range: Range<usize>,
}
impl PackProblem {
fn input(&self) -> Array2<u32> {
let shape = if self.is_a {
(self.mn, self.k)
} else {
(self.k, self.mn)
};
let data = (0..(self.k * self.mn) as u32).collect();
Array2::from_shape_vec(shape, data).unwrap()
}
fn packer(&self) -> Array3<u32> {
let panels = self.mn_range.len().divceil(self.r);
let packer = super::Packer::new(self.r, 1, 0);
let input = self.input().into_tensor();
let mut output =
Tensor::zero::<u32>(&[packer.len(self.k_range.len(), self.mn_range.len())])
.unwrap();
unsafe {
packer.pack_segment(
output.view_mut(),
input.view(),
self.is_a as usize,
!self.is_a as usize,
self.k_range.clone(),
self.mn_range.clone(),
)
};
output
.into_array::<u32>()
.unwrap()
.into_shape((panels, self.k_range.len(), self.r))
.unwrap()
}
fn reference(&self) -> Array3<u32> {
let input = self.input();
let panels = self.mn_range.len().divceil(self.r);
Array3::from_shape_fn([panels, self.k_range.len(), self.r], |(panel, k, x)| {
if self.mn_range.start + panel * self.r + x >= self.mn_range.end {
0
} else {
let mn = panel * self.r + x + self.mn_range.start;
let k = k + self.k_range.start;
let coords = if self.is_a { (mn, k) } else { (k, mn) };
*input.get(coords).unwrap_or(&0)
}
})
}
fn check(&self) {
assert_eq!(self.packer(), self.reference())
}
}
impl Arbitrary for PackProblem {
type Parameters = ();
type Strategy = BoxedStrategy<PackProblem>;
fn arbitrary_with(_args: ()) -> Self::Strategy {
(any::<bool>(), 1usize..9, 1usize..20, 1usize..20)
.prop_flat_map(|(is_a, r, k, mn)| {
(
Just((is_a, r, k, mn)),
sub_range_strat(0..k),
sub_range_strat(0..mn),
)
})
.prop_map(|((is_a, r, k, mn), k_range, mn_range)| PackProblem {
k,
mn,
is_a,
r,
k_range,
mn_range,
})
.boxed()
}
}
fn sub_range_strat(range: Range<usize>) -> BoxedStrategy<Range<usize>> {
(0..range.len())
.prop_flat_map(|cropped| (Just(cropped), 0..=cropped))
.prop_map(move |(cropped, left)| range.start + left..range.end - (cropped - left))
.boxed()
}
proptest::proptest! {
#[test]
fn prop(pb in any::<PackProblem>()) {
pb.check();
}
#[test]
fn subrange_prop(_range in sub_range_strat(0..20)) {
}
}
#[test]
fn simple_b_1() {
PackProblem {
k: 2,
mn: 1,
is_a: false,
r: 1,
k_range: 0..2,
mn_range: 0..1,
}
.check();
}
#[test]
fn simple_b_2() {
PackProblem {
k: 2,
mn: 2,
is_a: false,
r: 1,
k_range: 0..2,
mn_range: 0..2,
}
.check()
}
#[test]
fn simple_b_3() {
PackProblem {
k: 2,
mn: 1,
is_a: false,
r: 4,
k_range: 0..2,
mn_range: 0..1,
}
.check();
}
#[test]
fn simple_a_1() {
PackProblem {
k: 2,
mn: 2,
is_a: true,
r: 1,
k_range: 0..2,
mn_range: 0..2,
}
.check();
}
#[test]
fn simple_a_2() {
PackProblem {
k: 2,
mn: 3,
is_a: true,
r: 2,
k_range: 0..2,
mn_range: 0..3,
}
.check();
}
#[test]
fn range_k_0() {
PackProblem {
k: 2,
mn: 1,
is_a: false,
r: 1,
k_range: 1..2,
mn_range: 0..1,
}
.check();
}
#[test]
fn range_k_1() {
PackProblem {
k: 2,
mn: 2,
is_a: false,
r: 1,
k_range: 0..2,
mn_range: 0..1,
}
.check();
}
#[test]
fn range_k_2() {
PackProblem {
k: 2,
mn: 1,
is_a: false,
r: 6,
k_range: 1..2,
mn_range: 0..1,
}
.check();
}
#[test]
fn range_mn_0() {
PackProblem {
k: 1,
mn: 2,
is_a: false,
r: 2,
k_range: 0..1,
mn_range: 0..1,
}
.check();
}
#[test]
fn range_b_4() {
PackProblem {
k: 1,
mn: 2,
is_a: false,
r: 6,
k_range: 0..1,
mn_range: 1..2,
}
.check();
}
#[test]
fn range_b_5() {
PackProblem {
k: 1,
mn: 7,
is_a: false,
r: 6,
k_range: 0..1,
mn_range: 1..7,
}
.check();
}
}
@@ -0,0 +1,96 @@
macro_rules! sigmoid_impl {
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $cond: expr) => {
ew_impl!($ti, $func, $nr, $alignment_items);
#[cfg(test)]
paste! {
mod [<test_ $func>] {
use super::*;
sigmoid_frame_tests!($cond, $ti, $func);
}
}
};
}
#[cfg(test)]
#[macro_use]
pub mod test {
use crate::{LADatum, frame::element_wise::*};
use num_traits::{AsPrimitive, Float};
use proptest::test_runner::TestCaseResult;
#[macro_export]
macro_rules! sigmoid_frame_tests {
($cond:expr, $t: ty, $ker:ty) => {
proptest::proptest! {
#[test]
fn sigmoid(xs in proptest::collection::vec(-25f32..25.0, 0..100)) {
if $cond {
$crate::frame::sigmoid::test::test_sigmoid::<$ker, $t>(&*xs).unwrap()
}
}
}
#[test]
fn sigmoid_4_magic() {
if $cond {
$crate::frame::sigmoid::test::test_sigmoid::<$ker, $t>(&[
0f32, -20.0, 20.0, 0.0,
])
.unwrap()
}
}
#[test]
fn sigmoid_4zeros() {
if $cond {
$crate::frame::sigmoid::test::test_sigmoid::<$ker, $t>(&[0.0; 4]).unwrap();
}
}
#[test]
fn sigmoid_20_ones() {
if $cond {
$crate::frame::sigmoid::test::test_sigmoid::<$ker, $t>(&[1.0; 20]).unwrap();
}
}
#[test]
fn sigmoid_18_zeros() {
if $cond {
$crate::frame::sigmoid::test::test_sigmoid::<$ker, $t>(&[0.0; 18]).unwrap();
}
}
#[test]
fn sigmoid_asymptots() {
use tract_data::internal::*;
use $crate::frame::element_wise::*;
if $cond {
let mut input: Vec<$t> = [-100f32, 100f32]
.iter()
.map(|x| <f32 as num_traits::AsPrimitive<$t>>::as_(*x))
.collect();
let expected: Vec<$t> = [-0f32, 1f32]
.iter()
.map(|x| <f32 as num_traits::AsPrimitive<$t>>::as_(*x))
.collect();
<$ker>::ew().run(&mut input).unwrap();
tensor1(&input)
.close_enough(&tensor1(&expected), Approximation::Close)
.unwrap();
}
}
};
}
pub fn test_sigmoid<K: ElementWiseKer<T>, T: LADatum + Float>(values: &[f32]) -> TestCaseResult
where
f32: AsPrimitive<T>,
T: AsPrimitive<f32>,
{
let values: Vec<T> = values.iter().copied().map(|x| x.as_()).collect();
crate::frame::element_wise::test::test_element_wise::<K, _, _>(&values, |x| {
(1f32).as_() / (1f32.as_() + (-x).exp())
})
}
}
@@ -0,0 +1,101 @@
macro_rules! tanh_impl {
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $cond: expr) => {
ew_impl!($ti, $func, $nr, $alignment_items);
#[cfg(test)]
paste! {
mod [<test_ $func>] {
use super::*;
tanh_frame_tests!($cond, $ti, $func);
}
}
};
}
#[cfg(test)]
#[macro_use]
pub mod test {
use crate::LADatum;
use crate::frame::element_wise::*;
use num_traits::AsPrimitive;
use num_traits::float::Float;
use proptest::test_runner::TestCaseResult;
#[macro_export]
macro_rules! tanh_frame_tests {
($cond:expr, $t:ty, $ker:ty) => {
proptest::proptest! {
#[test]
fn tanh(xs in proptest::collection::vec(-25f32..25.0, 0..100)) {
if $cond {
$crate::frame::tanh::test::test_tanh::<$ker, $t>(&*xs).unwrap()
}
}
}
#[test]
fn tanh_4_magic() {
if $cond {
$crate::frame::tanh::test::test_tanh::<$ker, $t>(&[0f32, -20.0, 20.0, 0.0])
.unwrap()
}
}
#[test]
fn tanh_4zeros() {
if $cond {
$crate::frame::tanh::test::test_tanh::<$ker, $t>(&[0.0; 4]).unwrap();
}
}
#[test]
fn tanh_20_ones() {
if $cond {
$crate::frame::tanh::test::test_tanh::<$ker, $t>(&[1.0; 20]).unwrap();
}
}
#[test]
fn tanh_18_zeros() {
if $cond {
$crate::frame::tanh::test::test_tanh::<$ker, $t>(&[0.0; 18]).unwrap();
}
}
#[test]
fn tanh_foo() {
if $cond {
$crate::frame::tanh::test::test_tanh::<$ker, $t>(&[0.67503357]).unwrap();
}
}
#[test]
fn tanh_asymptots() {
use tract_data::internal::*;
use $crate::frame::element_wise::*;
if $cond {
let mut input: Vec<$t> = [-100f32, 100f32]
.iter()
.map(|x| <f32 as num_traits::AsPrimitive<$t>>::as_(*x))
.collect();
let expected: Vec<$t> = [-1f32, 1f32]
.iter()
.map(|x| <f32 as num_traits::AsPrimitive<$t>>::as_(*x))
.collect();
<$ker>::ew().run(&mut input).unwrap();
tensor1(&input)
.close_enough(&tensor1(&expected), Approximation::Close)
.unwrap();
}
}
};
}
pub fn test_tanh<K: ElementWiseKer<T>, T: LADatum + Float>(values: &[f32]) -> TestCaseResult
where
f32: AsPrimitive<T>,
T: AsPrimitive<f32>,
{
let values: Vec<T> = values.iter().copied().map(|x| x.as_()).collect();
crate::frame::element_wise::test::test_element_wise::<K, _, _>(&values, |x| x.tanh())
}
}
@@ -0,0 +1,14 @@
pub mod erf;
pub mod lut;
pub mod mmm;
pub mod rounding;
pub mod sigmoid;
pub mod tanh;
pub use self::erf::SErf4;
pub use self::lut::GenericLut8;
pub use self::mmm::GenericMmm4x1;
pub use self::mmm::GenericMmm4x4;
pub use self::rounding::{ScaleShiftAndRound, Scaler};
pub use self::sigmoid::{HSigmoid8, SSigmoid4};
pub use self::tanh::{HTanh8, STanh4};
@@ -0,0 +1,51 @@
use crate::element_wise::ElementWiseKer;
#[allow(non_upper_case_globals)]
#[allow(clippy::excessive_precision)]
fn serf(x: &mut f32) {
const a1: f32 = 0.0705230784;
const a2: f32 = 0.0422820123;
const a3: f32 = 0.0092705272;
const a4: f32 = 0.0001520143;
const a5: f32 = 0.0002765672;
const a6: f32 = 0.0000430638;
let signum = x.signum();
let abs = x.abs();
let y = a6 * abs;
let y = (a5 + y) * abs;
let y = (a4 + y) * abs;
let y = (a3 + y) * abs;
let y = (a2 + y) * abs;
let y = (a1 + y) * abs;
let y = 1.0 - (y + 1.0).powi(16).recip();
*x = y.copysign(signum)
}
#[derive(Clone, Debug)]
pub struct SErf4;
impl ElementWiseKer<f32> for SErf4 {
fn name() -> &'static str {
"generic"
}
fn alignment_items() -> usize {
16
}
fn alignment_bytes() -> usize {
16
}
fn nr() -> usize {
4
}
fn run(x: &mut [f32]) {
debug_assert!(x.len() % Self::nr() == 0);
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
x.iter_mut().for_each(serf)
}
}
@@ -0,0 +1,45 @@
use crate::frame::lut::LutKer;
#[derive(Clone, Debug, Hash)]
pub struct GenericLut8;
impl LutKer for GenericLut8 {
fn name() -> &'static str {
"generic"
}
fn input_alignment_bytes() -> usize {
1
}
fn table_alignment_bytes() -> usize {
1
}
fn n() -> usize {
8
}
unsafe fn run(buf: *mut u8, len: usize, table: *const u8) {
debug_assert!(len % Self::n() == 0);
debug_assert!(buf as usize % Self::input_alignment_bytes() == 0);
debug_assert!(table as usize % Self::table_alignment_bytes() == 0);
for i in 0..((len / 8) as isize) {
let ptr = buf.offset(8 * i);
*ptr.offset(0) = *table.offset(*ptr.offset(0) as isize);
*ptr.offset(1) = *table.offset(*ptr.offset(1) as isize);
*ptr.offset(2) = *table.offset(*ptr.offset(2) as isize);
*ptr.offset(3) = *table.offset(*ptr.offset(3) as isize);
*ptr.offset(4) = *table.offset(*ptr.offset(4) as isize);
*ptr.offset(5) = *table.offset(*ptr.offset(5) as isize);
*ptr.offset(6) = *table.offset(*ptr.offset(6) as isize);
*ptr.offset(7) = *table.offset(*ptr.offset(7) as isize);
}
}
}
#[cfg(test)]
#[macro_use]
pub mod test {
lut_frame_tests!(true, crate::generic::GenericLut8);
}
@@ -0,0 +1,594 @@
#![allow(clippy::needless_range_loop)]
use num_traits::AsPrimitive;
use std::marker::PhantomData;
use std::{fmt, ops};
use tract_data::prelude::*;
use super::*;
use crate::LADatum;
use crate::frame::mmm::*;
#[derive(Copy, Clone, Debug)]
pub struct GenericMmm4x4<TA, TB, TI>(PhantomData<(TA, TB, TI)>)
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound;
unsafe impl<TA, TB, TI> Send for GenericMmm4x4<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
{
}
unsafe impl<TA, TB, TI> Sync for GenericMmm4x4<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
{
}
macro_rules! scalar {
($ab: expr, $m: expr, $f: expr) => {
for i in 0..$ab.len() {
for j in 0..$ab[0].len() {
$ab[i][j] = $f($m, $ab[i][j])
}
}
};
}
macro_rules! per_row {
($ab: expr, $m: expr, $f: expr) => {
for i in 0..$ab.len() {
for j in 0..$ab[0].len() {
$ab[i][j] = $f(*$m.add(i), $ab[i][j])
}
}
};
}
macro_rules! per_col {
($ab: expr, $m: expr, $f: expr) => {
for i in 0..$ab.len() {
for j in 0..$ab[0].len() {
$ab[i][j] = $f(*$m.add(j), $ab[i][j])
}
}
};
}
impl<TA, TB, TI> MatMatMulKer<TI> for GenericMmm4x4<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
usize: AsPrimitive<TI>,
{
#[inline(always)]
fn name() -> &'static str {
match TI::datum_type() {
DatumType::F16 => "generic_f16_4x4",
DatumType::F32 => "generic_f32_4x4",
DatumType::I32 => "generic_i32_4x4",
DatumType::F64 => "generic_f64_4x4",
_ => panic!(),
}
}
#[inline(always)]
fn mr() -> usize {
4
}
#[inline(always)]
fn nr() -> usize {
4
}
fn end_padding_packed_a() -> usize {
0
}
fn end_padding_packed_b() -> usize {
0
}
#[inline(always)]
fn alignment_bytes_packed_a() -> usize {
std::mem::size_of::<TA>()
}
#[inline(always)]
fn alignment_bytes_packed_b() -> usize {
std::mem::size_of::<TB>()
}
#[inline(never)]
fn kernel(spec: &[FusedKerSpec<TI>]) -> isize {
unsafe {
let mut ab = [[TI::zero(); 4]; 4];
let mut pnl = spec.as_ptr();
loop {
if pnl.is_null() {
break;
}
match *pnl {
FusedKerSpec::Done => break,
FusedKerSpec::Clear => ab = std::mem::zeroed(),
FusedKerSpec::ScalarAdd(a) => scalar!(ab, a, |a, b| a + b),
FusedKerSpec::ScalarMul(a) => scalar!(ab, a, |a, b| a * b),
FusedKerSpec::ScalarMin(m) => scalar!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::ScalarMax(m) => scalar!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::ScalarSub(m) => scalar!(ab, m, |a, b| a - b),
FusedKerSpec::ScalarSubF(m) => scalar!(ab, m, |a, b| b - a),
FusedKerSpec::PerRowMin(m) => per_row!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::PerRowMax(m) => per_row!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::PerRowAdd(m) => per_row!(ab, m, |a, b| a + b),
FusedKerSpec::PerRowMul(m) => per_row!(ab, m, |a, b| a * b),
FusedKerSpec::PerRowSub(m) => per_row!(ab, m, |a, b| a - b),
FusedKerSpec::PerRowSubF(m) => per_row!(ab, m, |a, b| b - a),
FusedKerSpec::PerColMin(m) => per_col!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::PerColMax(m) => per_col!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::PerColAdd(m) => per_col!(ab, m, |a, b| a + b),
FusedKerSpec::PerColMul(m) => per_col!(ab, m, |a, b| a * b),
FusedKerSpec::PerColSub(m) => per_col!(ab, m, |a, b| a - b),
FusedKerSpec::PerColSubF(m) => per_col!(ab, m, |a, b| b - a),
FusedKerSpec::AddRowColProducts(rows, cols) => {
for i in 0..4 {
for j in 0..4 {
ab[i][j] += *rows.add(i) * *cols.add(j);
}
}
}
FusedKerSpec::AddUnicast(tile) => add_unicast::<TI, _>(&tile, &mut ab),
FusedKerSpec::ShiftLeft(shift) => {
for i in 0..4 {
for j in 0..4 {
ab[i][j] = ab[i][j].q_shl(shift);
}
}
}
FusedKerSpec::RoundingShiftRight(shift, rp) => {
for i in 0..4 {
for j in 0..4 {
ab[i][j] = ab[i][j].q_shr(shift, rp);
}
}
}
FusedKerSpec::QScale(shift, rp, mult) => {
for i in 0..4 {
for j in 0..4 {
ab[i][j] =
ab[i][j].q_scale(Scaler::from_fuse_params(shift, rp, mult));
}
}
}
FusedKerSpec::AddMatMul { k, pa, pb, .. } => {
let a = pa as *const TA;
let b = pb as *const TB;
for i in 0..k {
let a = std::slice::from_raw_parts(a.offset(4 * i as isize), 4);
let b = std::slice::from_raw_parts(b.offset(4 * i as isize), 4);
ab[0][0] += a[0].as_() * b[0].as_();
ab[0][1] += a[0].as_() * b[1].as_();
ab[0][2] += a[0].as_() * b[2].as_();
ab[0][3] += a[0].as_() * b[3].as_();
ab[1][0] += a[1].as_() * b[0].as_();
ab[1][1] += a[1].as_() * b[1].as_();
ab[1][2] += a[1].as_() * b[2].as_();
ab[1][3] += a[1].as_() * b[3].as_();
ab[2][0] += a[2].as_() * b[0].as_();
ab[2][1] += a[2].as_() * b[1].as_();
ab[2][2] += a[2].as_() * b[2].as_();
ab[2][3] += a[2].as_() * b[3].as_();
ab[3][0] += a[3].as_() * b[0].as_();
ab[3][1] += a[3].as_() * b[1].as_();
ab[3][2] += a[3].as_() * b[2].as_();
ab[3][3] += a[3].as_() * b[3].as_();
}
}
FusedKerSpec::Store(tile) => store(&tile, &ab),
};
pnl = pnl.add(1);
}
}
0
}
}
#[derive(Copy, Clone, Debug)]
pub struct GenericMmm4x1<TA, TB, TI>(PhantomData<(TA, TB, TI)>)
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound;
unsafe impl<TA, TB, TI> Send for GenericMmm4x1<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
{
}
unsafe impl<TA, TB, TI> Sync for GenericMmm4x1<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
{
}
impl<TA, TB, TI> MatMatMulKer<TI> for GenericMmm4x1<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
usize: AsPrimitive<TI>,
{
#[inline(always)]
fn name() -> &'static str {
match TI::datum_type() {
DatumType::F16 => "generic_f16_4x1",
DatumType::F32 => "generic_f32_4x1",
DatumType::I32 => "generic_i32_4x1",
DatumType::F64 => "generic_f64_4x1",
_ => panic!(),
}
}
#[inline(always)]
fn mr() -> usize {
4
}
#[inline(always)]
fn nr() -> usize {
1
}
fn end_padding_packed_a() -> usize {
0
}
fn end_padding_packed_b() -> usize {
0
}
#[inline(always)]
fn alignment_bytes_packed_a() -> usize {
std::mem::size_of::<TA>()
}
#[inline(always)]
fn alignment_bytes_packed_b() -> usize {
std::mem::size_of::<TB>()
}
#[inline(never)]
fn kernel(spec: &[FusedKerSpec<TI>]) -> isize {
unsafe {
let mut ab = [[TI::zero(); 1]; 4];
let mut pnl = spec.as_ptr();
loop {
if pnl.is_null() {
break;
}
match *pnl {
FusedKerSpec::Done => break,
FusedKerSpec::Clear => ab = std::mem::zeroed(),
FusedKerSpec::ScalarAdd(a) => scalar!(ab, a, |a, b| a + b),
FusedKerSpec::ScalarMul(a) => scalar!(ab, a, |a, b| a * b),
FusedKerSpec::ScalarMin(m) => scalar!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::ScalarMax(m) => scalar!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::ScalarSub(m) => scalar!(ab, m, |a, b| a - b),
FusedKerSpec::ScalarSubF(m) => scalar!(ab, m, |a, b| b - a),
FusedKerSpec::PerRowMul(m) => per_row!(ab, m, |a, b| a * b),
FusedKerSpec::PerRowMin(m) => per_row!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::PerRowMax(m) => per_row!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::PerRowAdd(m) => per_row!(ab, m, |a, b| a + b),
FusedKerSpec::PerRowSub(m) => per_row!(ab, m, |a, b| a - b),
FusedKerSpec::PerRowSubF(m) => per_row!(ab, m, |a, b| b - a),
FusedKerSpec::PerColMin(m) => per_col!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::PerColMax(m) => per_col!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::PerColAdd(m) => per_col!(ab, m, |a, b| a + b),
FusedKerSpec::PerColMul(m) => per_col!(ab, m, |a, b| a * b),
FusedKerSpec::PerColSub(m) => per_col!(ab, m, |a, b| a - b),
FusedKerSpec::PerColSubF(m) => per_col!(ab, m, |a, b| b - a),
FusedKerSpec::AddRowColProducts(rows, cols) => {
let col = *cols;
for i in 0..4 {
ab[i][0] += *rows.add(i) * col;
}
}
FusedKerSpec::AddUnicast(tile) => add_unicast::<TI, _>(
&tile,
&mut [
std::slice::from_raw_parts_mut(ab.as_ptr().offset(0) as _, 1),
std::slice::from_raw_parts_mut(ab.as_ptr().offset(1) as _, 1),
std::slice::from_raw_parts_mut(ab.as_ptr().offset(2) as _, 1),
std::slice::from_raw_parts_mut(ab.as_ptr().offset(3) as _, 1),
],
),
FusedKerSpec::ShiftLeft(shift) => {
for i in 0..4 {
ab[i][0] = ab[i][0].q_shl(shift);
}
}
FusedKerSpec::RoundingShiftRight(shift, rp) => {
for i in 0..4 {
ab[i][0] = ab[i][0].q_shr(shift, rp);
}
}
FusedKerSpec::QScale(shift, rp, mult) => {
for i in 0..4 {
ab[i][0] = ab[i][0].q_scale(Scaler::from_fuse_params(shift, rp, mult));
}
}
FusedKerSpec::AddMatMul { k, pa, pb, .. } => {
let a = pa as *const TA;
let b = pb as *const TB;
for i in 0..k {
let a = std::slice::from_raw_parts(a.offset(4 * i as isize), 4);
let b = *b.add(i);
ab[0][0] += a[0].as_() * b.as_();
ab[1][0] += a[1].as_() * b.as_();
ab[2][0] += a[2].as_() * b.as_();
ab[3][0] += a[3].as_() * b.as_();
}
}
FusedKerSpec::Store(tile) => store(
&tile,
&[
std::slice::from_raw_parts(ab.as_ptr().offset(0) as _, 1),
std::slice::from_raw_parts(ab.as_ptr().offset(1) as _, 1),
std::slice::from_raw_parts(ab.as_ptr().offset(2) as _, 1),
std::slice::from_raw_parts(ab.as_ptr().offset(3) as _, 1),
],
),
}
pnl = pnl.add(1);
}
}
0
}
}
#[cfg(test)]
#[derive(Copy, Clone, Debug)]
pub struct GenericMmmTest3x2<TA, TB, TI>(PhantomData<(TA, TB, TI)>)
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound;
#[cfg(test)]
unsafe impl<TA, TB, TI> Send for GenericMmmTest3x2<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
{
}
#[cfg(test)]
unsafe impl<TA, TB, TI> Sync for GenericMmmTest3x2<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
{
}
#[cfg(test)]
impl<TA, TB, TI> MatMatMulKer<TI> for GenericMmmTest3x2<TA, TB, TI>
where
TA: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TB: Datum + Copy + fmt::Debug + AsPrimitive<TI>,
TI: LADatum + ScaleShiftAndRound,
usize: AsPrimitive<TI>,
{
#[inline(always)]
fn name() -> &'static str {
match TI::datum_type() {
DatumType::F16 => "generic_f16_3x2",
DatumType::F32 => "generic_f32_3x2",
DatumType::I32 => "generic_i32_3x2",
DatumType::F64 => "generic_f64_3x2",
_ => panic!(),
}
}
#[inline(always)]
fn mr() -> usize {
3
}
#[inline(always)]
fn nr() -> usize {
2
}
fn end_padding_packed_a() -> usize {
0
}
fn end_padding_packed_b() -> usize {
0
}
#[inline(always)]
fn alignment_bytes_packed_a() -> usize {
std::mem::size_of::<TA>()
}
#[inline(always)]
fn alignment_bytes_packed_b() -> usize {
std::mem::size_of::<TB>()
}
#[inline(never)]
fn kernel(spec: &[FusedKerSpec<TI>]) -> isize {
unsafe {
let mut ab = [[TI::zero(); 2]; 3];
let mut pnl = spec.as_ptr();
loop {
if pnl.is_null() {
break;
}
match *pnl {
FusedKerSpec::Done => break,
FusedKerSpec::Clear => ab = std::mem::zeroed(),
FusedKerSpec::ScalarAdd(a) => scalar!(ab, a, |a, b| a + b),
FusedKerSpec::ScalarMul(a) => scalar!(ab, a, |a, b| a * b),
FusedKerSpec::ScalarMin(m) => scalar!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::ScalarMax(m) => scalar!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::ScalarSub(m) => scalar!(ab, m, |a, b| a - b),
FusedKerSpec::ScalarSubF(m) => scalar!(ab, m, |a, b| b - a),
FusedKerSpec::PerRowMin(m) => per_row!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::PerRowMax(m) => per_row!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::PerRowAdd(m) => per_row!(ab, m, |a, b| a + b),
FusedKerSpec::PerRowMul(m) => per_row!(ab, m, |a, b| a * b),
FusedKerSpec::PerRowSub(m) => per_row!(ab, m, |a, b| a - b),
FusedKerSpec::PerRowSubF(m) => per_row!(ab, m, |a, b| b - a),
FusedKerSpec::PerColMin(m) => per_col!(ab, m, |a, b| if a < b { a } else { b }),
FusedKerSpec::PerColMax(m) => per_col!(ab, m, |a, b| if a > b { a } else { b }),
FusedKerSpec::PerColAdd(m) => per_col!(ab, m, |a, b| a + b),
FusedKerSpec::PerColMul(m) => per_col!(ab, m, |a, b| a * b),
FusedKerSpec::PerColSub(m) => per_col!(ab, m, |a, b| a - b),
FusedKerSpec::PerColSubF(m) => per_col!(ab, m, |a, b| b - a),
FusedKerSpec::AddRowColProducts(rows, cols) => {
for i in 0..3 {
for j in 0..2 {
ab[i][j] += *rows.add(i) * *cols.add(j);
}
}
}
FusedKerSpec::AddUnicast(tile) => add_unicast::<TI, _>(&tile, &mut ab),
FusedKerSpec::ShiftLeft(shift) => {
for i in 0..3 {
for j in 0..2 {
ab[i][j] = ab[i][j].q_shl(shift);
}
}
}
FusedKerSpec::RoundingShiftRight(shift, rp) => {
for i in 0..3 {
for j in 0..2 {
ab[i][j] = ab[i][j].q_shr(shift, rp)
}
}
}
FusedKerSpec::QScale(shift, rp, mult) => {
for i in 0..3 {
for j in 0..2 {
ab[i][j] =
ab[i][j].q_scale(Scaler::from_fuse_params(shift, rp, mult));
}
}
}
FusedKerSpec::AddMatMul { k, pa, pb, .. } => {
let a = pa as *const TA;
let b = pb as *const TB;
for i in 0..k {
let a = std::slice::from_raw_parts(a.offset(3 * i as isize), 3);
let b = std::slice::from_raw_parts(b.offset(2 * i as isize), 2);
ab[0][0] += a[0].as_() * b[0].as_();
ab[0][1] += a[0].as_() * b[1].as_();
ab[1][0] += a[1].as_() * b[0].as_();
ab[1][1] += a[1].as_() * b[1].as_();
ab[2][0] += a[2].as_() * b[0].as_();
ab[2][1] += a[2].as_() * b[1].as_();
}
}
FusedKerSpec::Store(tile) => store(&tile, &ab),
}
pnl = pnl.add(1);
}
}
0
}
}
unsafe fn store_t<TC, TI, AB>(tile: &OutputStoreKer, ab: &[AB])
where
TC: Copy,
AB: AsRef<[TI]> + fmt::Debug,
{
for i in 0usize..ab.len() {
for j in 0usize..ab[0].as_ref().len() {
let loc: *mut TC = tile
.ptr
.offset(tile.row_byte_stride * i as isize + tile.col_byte_stride * j as isize)
as _;
let val: *const TC = (&ab[i].as_ref()[j]) as *const TI as _;
*loc = *val
}
}
}
unsafe fn store<TI, AB>(tile: &OutputStoreKer, ab: &[AB])
where
AB: AsRef<[TI]> + fmt::Debug,
{
match tile.item_size {
1 => store_t::<u8, _, _>(tile, ab),
2 => store_t::<u16, _, _>(tile, ab),
4 => store_t::<u32, _, _>(tile, ab),
8 => store_t::<f64, _, _>(tile, ab),
_ => unimplemented!(),
}
}
unsafe fn add_unicast<TI, AB>(tile: &OutputStoreKer, ab: &mut [AB])
where
TI: LADatum + ops::AddAssign<TI>,
AB: AsMut<[TI]> + fmt::Debug,
{
if tile.item_size == TI::datum_type().size_of() {
for i in 0usize..ab.len() {
for j in 0usize..ab[0].as_mut().len() {
let value: *const TI = tile
.ptr
.offset(tile.row_byte_stride * i as isize + tile.col_byte_stride * j as isize)
as _;
ab[i].as_mut()[j] += *value;
}
}
} else if TI::datum_type() == i32::datum_type() && tile.item_size == 1 {
for i in 0usize..ab.len() {
for j in 0usize..ab[0].as_mut().len() {
let value: i8 = *(tile
.ptr
.offset(tile.row_byte_stride * i as isize + tile.col_byte_stride * j as isize)
as *const i8);
let acc: *mut i32 = ab[i].as_mut().as_mut_ptr().add(j) as *mut i32;
*acc += value as i32;
}
}
} else {
unimplemented!("Missing AddUnicast type");
}
}
#[allow(non_camel_case_types)]
pub type generic_f16_4x4 = GenericMmm4x4<f16, f16, f16>;
test_mmm_kernel_f16!(generic_f16_4x4, true);
#[allow(non_camel_case_types)]
pub type generic_f32_4x4 = GenericMmm4x4<f32, f32, f32>;
test_mmm_kernel_f32!(generic_f32_4x4, true);
#[allow(non_camel_case_types)]
pub type generic_f64_4x4 = GenericMmm4x4<f64, f64, f64>;
test_mmm_kernel_f64!(generic_f64_4x4, true);
#[allow(non_camel_case_types)]
pub type generic_i32_4x4 = GenericMmm4x4<i8, i8, i32>;
test_mmm_kernel_i32!(generic_i32_4x4, true);
#[allow(non_camel_case_types)]
pub type generic_f32_4x1 = GenericMmm4x1<f32, f32, f32>;
test_mmm_kernel_f32!(generic_f32_4x1, true);
#[allow(non_camel_case_types)]
pub type generic_f64_4x1 = GenericMmm4x1<f64, f64, f64>;
test_mmm_kernel_f64!(generic_f64_4x1, true);
#[allow(non_camel_case_types)]
pub type generic_i32_4x1 = GenericMmm4x1<i8, i8, i32>;
test_mmm_kernel_i32!(generic_i32_4x1, true);
#[cfg(test)]
#[allow(non_camel_case_types)]
type generic_f32_3x2 = GenericMmmTest3x2<f32, f32, f32>;
test_mmm_kernel_f32!(generic_f32_3x2, true);
#[cfg(test)]
#[allow(non_camel_case_types)]
type generic_i32_3x2 = GenericMmmTest3x2<i8, i8, i32>;
test_mmm_kernel_i32!(generic_i32_3x2, true);
@@ -0,0 +1,534 @@
use crate::frame::mmm::*;
use std::hash::{Hash, Hasher};
use std::ops::Mul;
use tract_data::prelude::f16;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Scaler {
pub scale: f32,
pub mult: Option<i32>,
pub shift: isize,
pub policy: RoundingPolicy,
}
impl Eq for Scaler {}
#[allow(clippy::derived_hash_with_manual_eq)]
impl Hash for Scaler {
fn hash<H>(&self, state: &mut H)
where
H: Hasher,
{
Hash::hash(&self.scale.to_bits(), state)
}
}
impl Scaler {
pub fn new(scale: f32, policy: RoundingPolicy) -> Self {
let (mult, shift) = Self::convert_scale_to_mult_shift(scale);
Self {
scale,
mult,
shift,
policy,
}
}
pub fn as_fused_spec(&self) -> FusedSpec<'_> {
if let Some(multiplier) = self.mult {
FusedSpec::QScale(self.shift, self.policy, multiplier)
} else if self.shift > 0 {
FusedSpec::RoundingShiftRight(self.shift as usize, self.policy)
} else {
FusedSpec::ShiftLeft((-self.shift) as usize)
}
}
// FIXME: Only to avoid fused op breaking
pub fn from_fuse_params(shift: isize, policy: RoundingPolicy, mult: i32) -> Self {
let scale = mult as f32 * 2f32.powi(-(31 + shift as i32));
Self {
scale,
mult: Some(mult),
shift,
policy,
}
}
#[inline]
// This function convert a scale (actually a fraction of two integers Q/D)
// into an integer multiplier and a shift (the multiplier being 1/2D in Q0_31).
fn convert_scale_to_mult_shift(scale: f32) -> (Option<i32>, isize) {
// Zero is a special case to handle
if scale == 0.0 {
return (None, 0);
}
// Convert f32 to bits representation with the following pattern
// Bit | 31 | 30-23 | 22-0 |
// | Sign | Exponent | Fraction |
let scale_bits = scale.to_bits();
// Get actual value of the exponent
let current_exponent = (scale_bits >> 23) & 0xff;
// Extract fractional part of the float with:
// - 0x007fffff that represents the mask of the 23 lower bits (fractional part)
// (partial because it doesn't include the hidden bit (24) of the float representation)
let partial_frac = scale_bits & 0x007fffff;
if partial_frac == 0 {
let shift = 127 - current_exponent as isize;
(None, shift)
} else {
// We add 0x800000 that represents the hidden bit set to one.
// Here the frac is encoded as a Q8_23.
let frac = partial_frac | 0x800000;
// We rescale the result to be in Q0_31
// We should have shifted the result by 8 but the frac value is in [1.0, 2.0)
// so we cannot do that (we would need one bit for the integer).
// Instead we devide the frac by two to be in [0.5, 1.0) in Q0_31
// which lead to a shift of (8-1 = 7).
let half_frac = (frac << 7) as i32;
// Compute the actual value of the shift
// Here, we remove one as half_frac needs to be multiplied by 2.
let shift = 127 - current_exponent as isize - 1;
(Some(half_frac), shift)
}
}
}
impl Mul<f16> for Scaler {
type Output = f16;
#[inline]
fn mul(self, rhs: f16) -> Self::Output {
f16::from_f32(self.scale) * rhs
}
}
impl Mul<f32> for Scaler {
type Output = f32;
#[inline]
fn mul(self, rhs: f32) -> Self::Output {
self.scale * rhs
}
}
impl Mul<f64> for Scaler {
type Output = f64;
#[inline]
fn mul(self, rhs: f64) -> Self::Output {
self.scale as f64 * rhs
}
}
impl Mul<Scaler> for f16 {
type Output = f16;
#[inline]
fn mul(self, rhs: Scaler) -> Self::Output {
rhs * self
}
}
impl Mul<Scaler> for f32 {
type Output = f32;
#[inline]
fn mul(self, rhs: Scaler) -> Self::Output {
rhs * self
}
}
impl Mul<Scaler> for f64 {
type Output = f64;
#[inline]
fn mul(self, rhs: Scaler) -> Self::Output {
rhs * self
}
}
impl Mul<i32> for Scaler {
type Output = i32;
#[inline]
fn mul(self, rhs: i32) -> Self::Output {
let (val, shift) = if let Some(multiplier) = self.mult {
(multiplier as i64 * rhs as i64, self.shift + 31)
} else {
(rhs as i64, self.shift)
};
// Round according to rounding policy
use RoundingPolicy::*;
if shift > 0 {
let half: i64 = 1 << (shift - 1);
let nudge: i64 = match self.policy {
Zero => -1,
MinusInf => -((val >= 0) as i64),
PlusInf => -((val <= 0) as i64),
Away => 0,
Even => ((val.abs() >> shift) & 0x1) - 1,
Odd => -((val.abs() >> shift) & 0x1),
_ => panic!(),
};
(val.signum() * ((val.abs() + half + nudge) >> shift)) as i32
} else {
(val << -shift) as i32
}
}
}
impl Mul<Scaler> for i32 {
type Output = i32;
#[inline]
fn mul(self, rhs: Scaler) -> Self::Output {
rhs * self
}
}
pub trait ScaleShiftAndRound {
fn q_scale(self, scaler: Scaler) -> Self;
fn q_shl(self, shift: usize) -> Self;
fn q_shr(self, shift: usize, rp: RoundingPolicy) -> Self;
}
impl ScaleShiftAndRound for f64 {
fn q_scale(self, scaler: Scaler) -> Self {
self * scaler
}
fn q_shl(self, shift: usize) -> Self {
self * 2f64.powi(shift as i32)
}
fn q_shr(self, shift: usize, _rp: RoundingPolicy) -> Self {
self * 2f64.powi(-(shift as i32))
}
}
impl ScaleShiftAndRound for f32 {
fn q_scale(self, scaler: Scaler) -> Self {
self * scaler
}
fn q_shl(self, shift: usize) -> Self {
self * 2f32.powi(shift as i32)
}
fn q_shr(self, shift: usize, _rp: RoundingPolicy) -> Self {
self * 2f32.powi(-(shift as i32))
}
}
impl ScaleShiftAndRound for f16 {
fn q_scale(self, scaler: Scaler) -> Self {
self * scaler
}
fn q_shl(self, shift: usize) -> Self {
self * f16::from_f32(2f32.powi(shift as i32))
}
fn q_shr(self, shift: usize, _rp: RoundingPolicy) -> Self {
self * f16::from_f32(2f32.powi(-(shift as i32)))
}
}
impl ScaleShiftAndRound for i32 {
fn q_scale(self, scaler: Scaler) -> Self {
self * scaler
}
fn q_shr(self, shift: usize, rp: RoundingPolicy) -> Self {
use RoundingPolicy::*;
let half: i32 = 1 << (shift - 1);
let nudge: i32 = match rp {
Zero => -1,
MinusInf => -((self >= 0) as i32),
PlusInf => -((self <= 0) as i32),
Away => 0,
Even => ((self.abs() >> shift) & 0x1) - 1,
Odd => -((self.abs() >> shift) & 0x1),
_ => panic!(),
};
self.signum() * ((self.abs() + half + nudge) >> shift)
}
fn q_shl(self, shift: usize) -> Self {
self << shift
}
}
// 6 / 4 -> 1.5 -> arrondi: 2. rien a faire
// 2 / 4 -> 0.5 -> arrondi: 1. veut 0 -> nudge = -1
#[cfg(test)]
mod test {
use super::RoundingPolicy::*;
use super::*;
#[test]
fn test_scale_rounding_f32() {
assert_eq!(0f32.q_scale(Scaler::new(0.5, Zero)), 0.0);
assert_eq!(1f32.q_scale(Scaler::new(0.5, Zero)), 0.5);
assert_eq!(2f32.q_scale(Scaler::new(0.5, Zero)), 1.0);
assert_eq!(3f32.q_scale(Scaler::new(0.5, Zero)), 1.5);
assert_eq!((-1f32).q_scale(Scaler::new(0.5, Zero)), -0.5);
assert_eq!((-2f32).q_scale(Scaler::new(0.5, Zero)), -1.0);
assert_eq!((-3f32).q_scale(Scaler::new(0.5, Zero)), -1.5);
}
#[test]
fn test_shift_rounding_zero() {
assert_eq!(0i32.q_shr(1, Zero), 0);
assert_eq!(1i32.q_shr(1, Zero), 0);
assert_eq!(2i32.q_shr(1, Zero), 1);
assert_eq!(3i32.q_shr(1, Zero), 1);
assert_eq!(0i32.q_shr(2, Zero), 0);
assert_eq!(1i32.q_shr(2, Zero), 0);
assert_eq!(2i32.q_shr(2, Zero), 0);
assert_eq!(3i32.q_shr(2, Zero), 1);
assert_eq!(4i32.q_shr(2, Zero), 1);
assert_eq!(5i32.q_shr(2, Zero), 1);
assert_eq!(6i32.q_shr(2, Zero), 1);
assert_eq!((-1i32).q_shr(2, Zero), 0);
assert_eq!((-2i32).q_shr(2, Zero), 0);
assert_eq!((-3i32).q_shr(2, Zero), -1);
assert_eq!((-4i32).q_shr(2, Zero), -1);
assert_eq!((-5i32).q_shr(2, Zero), -1);
assert_eq!((-6i32).q_shr(2, Zero), -1);
}
#[test]
fn test_scale_rounding_zero() {
assert_eq!(0i32.q_scale(Scaler::new(0.5, Zero)), 0);
assert_eq!(1i32.q_scale(Scaler::new(0.5, Zero)), 0);
assert_eq!(2i32.q_scale(Scaler::new(0.5, Zero)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.5, Zero)), 1);
assert_eq!((-1i32).q_scale(Scaler::new(0.5, Zero)), 0);
assert_eq!((-2i32).q_scale(Scaler::new(0.5, Zero)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.5, Zero)), -1);
assert_eq!(2i32.q_scale(Scaler::new(0.25, Zero)), 0);
assert_eq!(3i32.q_scale(Scaler::new(0.25, Zero)), 1);
assert_eq!(4i32.q_scale(Scaler::new(0.25, Zero)), 1);
assert_eq!(5i32.q_scale(Scaler::new(0.25, Zero)), 1);
assert_eq!(6i32.q_scale(Scaler::new(0.25, Zero)), 1);
assert_eq!((-2i32).q_scale(Scaler::new(0.25, Zero)), 0);
assert_eq!((-3i32).q_scale(Scaler::new(0.25, Zero)), -1);
assert_eq!((-4i32).q_scale(Scaler::new(0.25, Zero)), -1);
assert_eq!((-5i32).q_scale(Scaler::new(0.25, Zero)), -1);
assert_eq!((-6i32).q_scale(Scaler::new(0.25, Zero)), -1);
}
#[test]
fn test_shift_rounding_away() {
assert_eq!(0i32.q_shr(1, Away), 0);
assert_eq!(1i32.q_shr(1, Away), 1);
assert_eq!(2i32.q_shr(1, Away), 1);
assert_eq!(3i32.q_shr(1, Away), 2);
assert_eq!(0i32.q_shr(2, Away), 0);
assert_eq!(1i32.q_shr(2, Away), 0);
assert_eq!(2i32.q_shr(2, Away), 1);
assert_eq!(3i32.q_shr(2, Away), 1);
assert_eq!(4i32.q_shr(2, Away), 1);
assert_eq!(5i32.q_shr(2, Away), 1);
assert_eq!(6i32.q_shr(2, Away), 2);
assert_eq!((-1i32).q_shr(2, Away), 0);
assert_eq!((-2i32).q_shr(2, Away), -1);
assert_eq!((-3i32).q_shr(2, Away), -1);
assert_eq!((-4i32).q_shr(2, Away), -1);
assert_eq!((-5i32).q_shr(2, Away), -1);
assert_eq!((-6i32).q_shr(2, Away), -2);
}
#[test]
fn test_scale_rounding_away() {
assert_eq!(0i32.q_scale(Scaler::new(0.5, Away)), 0);
assert_eq!(1i32.q_scale(Scaler::new(0.5, Away)), 1);
assert_eq!(2i32.q_scale(Scaler::new(0.5, Away)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.5, Away)), 2);
assert_eq!((-1i32).q_scale(Scaler::new(0.5, Away)), -1);
assert_eq!((-2i32).q_scale(Scaler::new(0.5, Away)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.5, Away)), -2);
assert_eq!(2i32.q_scale(Scaler::new(0.25, Away)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.25, Away)), 1);
assert_eq!(4i32.q_scale(Scaler::new(0.25, Away)), 1);
assert_eq!(5i32.q_scale(Scaler::new(0.25, Away)), 1);
assert_eq!(6i32.q_scale(Scaler::new(0.25, Away)), 2);
assert_eq!((-2i32).q_scale(Scaler::new(0.25, Away)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.25, Away)), -1);
assert_eq!((-4i32).q_scale(Scaler::new(0.25, Away)), -1);
assert_eq!((-5i32).q_scale(Scaler::new(0.25, Away)), -1);
assert_eq!((-6i32).q_scale(Scaler::new(0.25, Away)), -2);
}
#[test]
fn test_shift_rounding_plus_inf() {
assert_eq!(0i32.q_shr(1, PlusInf), 0);
assert_eq!(1i32.q_shr(1, PlusInf), 1);
assert_eq!(2i32.q_shr(1, PlusInf), 1);
assert_eq!(3i32.q_shr(1, PlusInf), 2);
assert_eq!(0i32.q_shr(2, PlusInf), 0);
assert_eq!(1i32.q_shr(2, PlusInf), 0);
assert_eq!(2i32.q_shr(2, PlusInf), 1);
assert_eq!(3i32.q_shr(2, PlusInf), 1);
assert_eq!(4i32.q_shr(2, PlusInf), 1);
assert_eq!(5i32.q_shr(2, PlusInf), 1);
assert_eq!(6i32.q_shr(2, PlusInf), 2);
assert_eq!((-1i32).q_shr(2, PlusInf), 0);
assert_eq!((-2i32).q_shr(2, PlusInf), 0);
assert_eq!((-3i32).q_shr(2, PlusInf), -1);
assert_eq!((-4i32).q_shr(2, PlusInf), -1);
assert_eq!((-5i32).q_shr(2, PlusInf), -1);
assert_eq!((-6i32).q_shr(2, PlusInf), -1);
}
#[test]
fn test_scale_rounding_plus_inf() {
assert_eq!(0i32.q_scale(Scaler::new(0.5, PlusInf)), 0);
assert_eq!(1i32.q_scale(Scaler::new(0.5, PlusInf)), 1);
assert_eq!(2i32.q_scale(Scaler::new(0.5, PlusInf)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.5, PlusInf)), 2);
assert_eq!((-1i32).q_scale(Scaler::new(0.5, PlusInf)), 0);
assert_eq!((-2i32).q_scale(Scaler::new(0.5, PlusInf)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.5, PlusInf)), -1);
assert_eq!(2i32.q_scale(Scaler::new(0.25, PlusInf)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.25, PlusInf)), 1);
assert_eq!(4i32.q_scale(Scaler::new(0.25, PlusInf)), 1);
assert_eq!(5i32.q_scale(Scaler::new(0.25, PlusInf)), 1);
assert_eq!(6i32.q_scale(Scaler::new(0.25, PlusInf)), 2);
assert_eq!((-2i32).q_scale(Scaler::new(0.25, PlusInf)), 0);
assert_eq!((-3i32).q_scale(Scaler::new(0.25, PlusInf)), -1);
assert_eq!((-4i32).q_scale(Scaler::new(0.25, PlusInf)), -1);
assert_eq!((-5i32).q_scale(Scaler::new(0.25, PlusInf)), -1);
assert_eq!((-6i32).q_scale(Scaler::new(0.25, PlusInf)), -1);
}
#[test]
fn test_shift_rounding_minus_inf() {
assert_eq!(0i32.q_shr(1, MinusInf), 0);
assert_eq!(1i32.q_shr(1, MinusInf), 0);
assert_eq!(2i32.q_shr(1, MinusInf), 1);
assert_eq!(3i32.q_shr(1, MinusInf), 1);
assert_eq!(0i32.q_shr(2, MinusInf), 0);
assert_eq!(1i32.q_shr(2, MinusInf), 0);
assert_eq!(2i32.q_shr(2, MinusInf), 0);
assert_eq!(3i32.q_shr(2, MinusInf), 1);
assert_eq!(4i32.q_shr(2, MinusInf), 1);
assert_eq!(5i32.q_shr(2, MinusInf), 1);
assert_eq!(6i32.q_shr(2, MinusInf), 1);
assert_eq!((-1i32).q_shr(2, MinusInf), 0);
assert_eq!((-2i32).q_shr(2, MinusInf), -1);
assert_eq!((-3i32).q_shr(2, MinusInf), -1);
assert_eq!((-4i32).q_shr(2, MinusInf), -1);
assert_eq!((-5i32).q_shr(2, MinusInf), -1);
assert_eq!((-6i32).q_shr(2, MinusInf), -2);
}
#[test]
fn test_scale_rounding_minus_inf() {
assert_eq!(0i32.q_scale(Scaler::new(0.5, MinusInf)), 0);
assert_eq!(1i32.q_scale(Scaler::new(0.5, MinusInf)), 0);
assert_eq!(2i32.q_scale(Scaler::new(0.5, MinusInf)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.5, MinusInf)), 1);
assert_eq!((-1i32).q_scale(Scaler::new(0.5, MinusInf)), -1);
assert_eq!((-2i32).q_scale(Scaler::new(0.5, MinusInf)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.5, MinusInf)), -2);
assert_eq!(2i32.q_scale(Scaler::new(0.25, MinusInf)), 0);
assert_eq!(3i32.q_scale(Scaler::new(0.25, MinusInf)), 1);
assert_eq!(4i32.q_scale(Scaler::new(0.25, MinusInf)), 1);
assert_eq!(5i32.q_scale(Scaler::new(0.25, MinusInf)), 1);
assert_eq!(6i32.q_scale(Scaler::new(0.25, MinusInf)), 1);
assert_eq!((-2i32).q_scale(Scaler::new(0.25, MinusInf)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.25, MinusInf)), -1);
assert_eq!((-4i32).q_scale(Scaler::new(0.25, MinusInf)), -1);
assert_eq!((-5i32).q_scale(Scaler::new(0.25, MinusInf)), -1);
assert_eq!((-6i32).q_scale(Scaler::new(0.25, MinusInf)), -2);
//assert_eq!((-9i32).q_scale(ONE_OVER_TWO_IN_Q0_30, 5, MinusInf), 0);
}
#[test]
fn test_shift_rounding_even() {
assert_eq!(0i32.q_shr(1, Even), 0);
assert_eq!(1i32.q_shr(1, Even), 0);
assert_eq!(2i32.q_shr(1, Even), 1);
assert_eq!(3i32.q_shr(1, Even), 2);
assert_eq!(0i32.q_shr(2, Even), 0);
assert_eq!(1i32.q_shr(2, Even), 0);
assert_eq!(2i32.q_shr(2, Even), 0);
assert_eq!(3i32.q_shr(2, Even), 1);
assert_eq!(4i32.q_shr(2, Even), 1);
assert_eq!(5i32.q_shr(2, Even), 1);
assert_eq!(6i32.q_shr(2, Even), 2);
assert_eq!((-1i32).q_shr(2, Even), 0);
assert_eq!((-2i32).q_shr(2, Even), 0);
assert_eq!((-3i32).q_shr(2, Even), -1);
assert_eq!((-4i32).q_shr(2, Even), -1);
assert_eq!((-5i32).q_shr(2, Even), -1);
assert_eq!((-6i32).q_shr(2, Even), -2);
}
#[test]
fn test_scale_rounding_even() {
assert_eq!(0i32.q_scale(Scaler::new(0.5, Even)), 0);
assert_eq!(1i32.q_scale(Scaler::new(0.5, Even)), 0);
assert_eq!(2i32.q_scale(Scaler::new(0.5, Even)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.5, Even)), 2);
assert_eq!((-1i32).q_scale(Scaler::new(0.5, Even)), 0);
assert_eq!((-2i32).q_scale(Scaler::new(0.5, Even)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.5, Even)), -2);
assert_eq!(2i32.q_scale(Scaler::new(0.25, Even)), 0);
assert_eq!(3i32.q_scale(Scaler::new(0.25, Even)), 1);
assert_eq!(4i32.q_scale(Scaler::new(0.25, Even)), 1);
assert_eq!(5i32.q_scale(Scaler::new(0.25, Even)), 1);
assert_eq!(6i32.q_scale(Scaler::new(0.25, Even)), 2);
assert_eq!((-2i32).q_scale(Scaler::new(0.25, Even)), 0);
assert_eq!((-3i32).q_scale(Scaler::new(0.25, Even)), -1);
assert_eq!((-4i32).q_scale(Scaler::new(0.25, Even)), -1);
assert_eq!((-5i32).q_scale(Scaler::new(0.25, Even)), -1);
assert_eq!((-6i32).q_scale(Scaler::new(0.25, Even)), -2);
}
#[test]
fn test_shift_rounding_odd() {
assert_eq!(0i32.q_shr(1, Odd), 0);
assert_eq!(1i32.q_shr(1, Odd), 1);
assert_eq!(2i32.q_shr(1, Odd), 1);
assert_eq!(3i32.q_shr(1, Odd), 1);
assert_eq!(0i32.q_shr(2, Odd), 0);
assert_eq!(1i32.q_shr(2, Odd), 0);
assert_eq!(2i32.q_shr(2, Odd), 1);
assert_eq!(3i32.q_shr(2, Odd), 1);
assert_eq!(4i32.q_shr(2, Odd), 1);
assert_eq!(5i32.q_shr(2, Odd), 1);
assert_eq!(6i32.q_shr(2, Odd), 1);
assert_eq!((-1i32).q_shr(2, Odd), 0);
assert_eq!((-2i32).q_shr(2, Odd), -1);
assert_eq!((-3i32).q_shr(2, Odd), -1);
assert_eq!((-4i32).q_shr(2, Odd), -1);
assert_eq!((-5i32).q_shr(2, Odd), -1);
assert_eq!((-6i32).q_shr(2, Odd), -1);
}
#[test]
fn test_scale_rounding_odd() {
assert_eq!(0i32.q_scale(Scaler::new(0.5, Odd)), 0);
assert_eq!(1i32.q_scale(Scaler::new(0.5, Odd)), 1);
assert_eq!(2i32.q_scale(Scaler::new(0.5, Odd)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.5, Odd)), 1);
assert_eq!((-1i32).q_scale(Scaler::new(0.5, Odd)), -1);
assert_eq!((-2i32).q_scale(Scaler::new(0.5, Odd)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.5, Odd)), -1);
assert_eq!(2i32.q_scale(Scaler::new(0.25, Odd)), 1);
assert_eq!(3i32.q_scale(Scaler::new(0.25, Odd)), 1);
assert_eq!(4i32.q_scale(Scaler::new(0.25, Odd)), 1);
assert_eq!(5i32.q_scale(Scaler::new(0.25, Odd)), 1);
assert_eq!(6i32.q_scale(Scaler::new(0.25, Odd)), 1);
assert_eq!((-2i32).q_scale(Scaler::new(0.25, Odd)), -1);
assert_eq!((-3i32).q_scale(Scaler::new(0.25, Odd)), -1);
assert_eq!((-4i32).q_scale(Scaler::new(0.25, Odd)), -1);
assert_eq!((-5i32).q_scale(Scaler::new(0.25, Odd)), -1);
assert_eq!((-6i32).q_scale(Scaler::new(0.25, Odd)), -1);
}
}
@@ -0,0 +1,142 @@
#![allow(clippy::excessive_precision)]
use crate::frame::element_wise::ElementWiseKer;
use tract_data::internal::*;
pub fn ssigmoid(x: f32) -> f32 {
const LOW: f32 = -18.6;
const HIGH: f32 = -LOW;
const ALPHA_13: f32 = -4.433153405e-18;
const ALPHA_11: f32 = 1.169974371e-14;
const ALPHA_9: f32 = -1.875289645e-11;
const ALPHA_7: f32 = 4.257889523e-8;
const ALPHA_5: f32 = 0.00004811817576;
const ALPHA_3: f32 = 0.008163842030;
const ALPHA_1: f32 = 0.2499999971;
const BETA_6: f32 = 3.922935744e-6;
const BETA_4: f32 = 0.001524872358;
const BETA_2: f32 = 0.1159886749;
const BETA_0: f32 = 1.0;
let x = x.clamp(LOW, HIGH);
let x2 = x * x;
let p = ALPHA_13;
let p = x2 * p + ALPHA_11;
let p = x2 * p + ALPHA_9;
let p = x2 * p + ALPHA_7;
let p = x2 * p + ALPHA_5;
let p = x2 * p + ALPHA_3;
let p = x2 * p + ALPHA_1;
let p = p * x;
let q = BETA_6;
let q = x2 * q + BETA_4;
let q = x2 * q + BETA_2;
let q = x2 * q + BETA_0;
p / q + 0.5
}
pub fn hsigmoid(x: f16) -> f16 {
/*
* (x (0.249895 + x^2 (0.00400222 - 0.0000124702 x^2)))
* /
* (1. + 0.098734 x^2)
*/
const LOW: f16 = f16::from_f32_const(-6.92);
const HIGH: f16 = f16::from_f32_const(6.92);
const ALPHA_5: f16 = f16::from_f32_const(-0.0000124702);
const ALPHA_3: f16 = f16::from_f32_const(0.00400222);
const ALPHA_1: f16 = f16::from_f32_const(0.249895);
const BETA_2: f16 = f16::from_f32_const(0.098734);
const BETA_0: f16 = f16::from_f32_const(1.0);
let x = x.clamp(LOW, HIGH);
let x2 = x * x;
let p = ALPHA_5;
let p = x2 * p + ALPHA_3;
let p = x2 * p + ALPHA_1;
let p = p * x;
let q = BETA_2;
let q = x2 * q + BETA_0;
p / q + f16::from_f32_const(0.5)
}
#[derive(Clone, Debug)]
pub struct SSigmoid4;
impl ElementWiseKer<f32> for SSigmoid4 {
fn name() -> &'static str {
"generic"
}
fn alignment_bytes() -> usize {
16
}
fn alignment_items() -> usize {
4
}
fn nr() -> usize {
4
}
fn run(x: &mut [f32]) {
debug_assert!(x.len() % Self::nr() == 0);
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
x.iter_mut().for_each(|px| *px = ssigmoid(*px))
}
}
#[derive(Clone, Debug)]
pub struct HSigmoid8;
impl ElementWiseKer<f16> for HSigmoid8 {
fn name() -> &'static str {
"generic"
}
fn alignment_bytes() -> usize {
16
}
fn alignment_items() -> usize {
4
}
fn nr() -> usize {
8
}
fn run(x: &mut [f16]) {
debug_assert!(x.len() % Self::nr() == 0);
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
x.iter_mut().for_each(|px| *px = hsigmoid(*px))
}
}
#[cfg(test)]
#[macro_use]
pub mod s {
sigmoid_frame_tests!(true, f32, crate::generic::sigmoid::SSigmoid4);
}
#[cfg(test)]
#[macro_use]
pub mod h {
sigmoid_frame_tests!(
true,
tract_data::internal::f16,
crate::generic::sigmoid::HSigmoid8
);
}
@@ -0,0 +1,137 @@
#![allow(clippy::excessive_precision)]
use crate::frame::element_wise::ElementWiseKer;
use tract_data::internal::*;
pub fn stanh(x: f32) -> f32 {
const LOW: f32 = -8.9;
const HIGH: f32 = 8.9;
const ALPHA_13: f32 = -8.488492677e-14;
const ALPHA_11: f32 = 5.277853000e-11;
const ALPHA_9: f32 = -2.022500419e-8;
const ALPHA_7: f32 = 0.00001115424833;
const ALPHA_5: f32 = 0.003103950131;
const ALPHA_3: f32 = 0.1308400453;
const ALPHA_1: f32 = 0.9999999934;
const BETA_6: f32 = 0.0002546136580;
const BETA_4: f32 = 0.02449515379;
const BETA_2: f32 = 0.4641733162;
const BETA_0: f32 = 1.0;
let x = x.clamp(LOW, HIGH);
let x2 = x * x;
let p = ALPHA_13;
let p = x2 * p + ALPHA_11;
let p = x2 * p + ALPHA_9;
let p = x2 * p + ALPHA_7;
let p = x2 * p + ALPHA_5;
let p = x2 * p + ALPHA_3;
let p = x2 * p + ALPHA_1;
let p = p * x;
let q = BETA_6;
let q = x2 * q + BETA_4;
let q = x2 * q + BETA_2;
let q = x2 * q + BETA_0;
p / q
}
pub fn htanh(x: f16) -> f16 {
const LOW: f16 = f16::from_f32_const(-3.84);
const HIGH: f16 = f16::from_f32_const(3.84);
const ALPHA_3: f16 = f16::from_f32_const(0.082654955);
const ALPHA_1: f16 = f16::from_f32_const(0.99963124);
const BETA_4: f16 = f16::from_f32_const(0.0065383179);
const BETA_2: f16 = f16::from_f32_const(0.41401828);
const BETA_0: f16 = f16::from_f32_const(1.0);
let x = x.clamp(LOW, HIGH);
let x2 = x * x;
let p = ALPHA_3;
let p = x2 * p + ALPHA_1;
let p = p * x;
let q = BETA_4;
let q = x2 * q + BETA_2;
let q = x2 * q + BETA_0;
p / q
}
#[derive(Clone, Debug)]
pub struct STanh4;
impl ElementWiseKer<f32> for STanh4 {
fn name() -> &'static str {
"generic"
}
fn alignment_items() -> usize {
16
}
fn alignment_bytes() -> usize {
16
}
fn nr() -> usize {
4
}
fn run(x: &mut [f32]) {
debug_assert!(x.len() % Self::nr() == 0);
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
x.iter_mut().for_each(|px| *px = stanh(*px))
}
}
#[cfg(test)]
#[macro_use]
pub mod s {
tanh_frame_tests!(true, f32, crate::generic::tanh::STanh4);
}
#[derive(Clone, Debug)]
pub struct HTanh8;
impl ElementWiseKer<f16> for HTanh8 {
fn name() -> &'static str {
"generic"
}
fn alignment_items() -> usize {
16
}
fn alignment_bytes() -> usize {
16
}
fn nr() -> usize {
8
}
fn run(x: &mut [f16]) {
debug_assert!(x.len() % Self::nr() == 0);
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
x.iter_mut().for_each(|px| *px = htanh(*px))
}
}
#[cfg(test)]
#[macro_use]
pub mod h {
tanh_frame_tests!(
true,
tract_data::internal::f16,
crate::generic::tanh::HTanh8
);
}
@@ -0,0 +1,219 @@
#![allow(clippy::missing_safety_doc)]
#[macro_use]
extern crate derive_new;
extern crate lazy_static;
extern crate log;
extern crate num_traits;
#[macro_use]
extern crate paste;
#[cfg(test)]
extern crate proptest;
include!(concat!(env!("OUT_DIR"), "/extern_kernel_macro.rs"));
#[macro_use]
pub mod frame;
pub mod generic;
use frame::MatMatMul;
use frame::element_wise::ElementWiseKer;
pub use generic::{ScaleShiftAndRound, Scaler};
#[cfg(target_arch = "x86_64")]
pub mod x86_64_fma;
#[cfg(target_arch = "aarch64")]
pub mod arm64;
#[cfg(target_arch = "arm")]
pub mod arm32;
pub use self::frame::{element_wise, lut, mmm};
use crate::frame::mmm::kernel::MatMatMulKer;
use tract_data::prelude::*;
type MMMImpl = Box<
dyn Fn(Option<usize>, Option<usize>, Option<usize>) -> Box<dyn mmm::MatMatMul> + Send + Sync,
>;
type MMVImpl = Box<dyn Fn(Option<usize>, Option<usize>) -> Box<dyn mmm::MatMatMul> + Send + Sync>;
#[allow(clippy::type_complexity)]
pub struct Ops {
mmm_f64: MMMImpl,
mmv_f64: MMVImpl,
mmm_f32_impls: Vec<Box<dyn MatMatMul>>,
mmm_f32: MMMImpl,
mmv_f32: MMVImpl,
mmm_f16: MMMImpl,
mmv_f16: MMVImpl,
qmmm_i32: MMMImpl,
qmmv_i32: MMVImpl,
pub sigmoid_f16: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f16>> + Send + Sync>,
pub sigmoid_f32: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f32>> + Send + Sync>,
pub tanh_f16: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f16>> + Send + Sync>,
pub tanh_f32: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f32>> + Send + Sync>,
pub erf_f32: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f32>> + Send + Sync>,
pub lut_u8: Box<dyn Fn(&[u8]) -> Box<dyn lut::Lut> + Send + Sync>,
}
impl Ops {
pub fn mmm_f32_impls(&self) -> &[Box<dyn MatMatMul>] {
&self.mmm_f32_impls
}
pub fn mmm(
&self,
a: DatumType,
b: DatumType,
c: DatumType,
m: Option<usize>,
k: Option<usize>,
n: Option<usize>,
) -> Option<Box<dyn mmm::MatMatMul>> {
use DatumType::*;
match (a.unquantized(), b.unquantized(), c.unquantized()) {
(F64, F64, F64) => Some(if n == Some(1) {
(self.mmv_f64)(m, k)
} else {
(self.mmm_f64)(m, k, n)
}),
(F32, F32, F32) => Some(if n == Some(1) {
(self.mmv_f32)(m, k)
} else {
(self.mmm_f32)(m, k, n)
}),
(F16, F16, F16) => Some(if n == Some(1) {
(self.mmv_f16)(m, k)
} else {
(self.mmm_f16)(m, k, n)
}),
(I8, I8, I32) => Some(if n == Some(1) {
(self.qmmv_i32)(m, k)
} else {
(self.qmmm_i32)(m, k, n)
}),
(I8, I8, I8) => Some(if n == Some(1) {
(self.qmmv_i32)(m, k)
} else {
(self.qmmm_i32)(m, k, n)
}),
_ => None,
}
}
}
pub fn generic() -> Ops {
Ops {
mmm_f64: Box::new(|_, _, _| generic::GenericMmm4x4::<f64, f64, f64>::mmm()),
mmv_f64: Box::new(|_, _| generic::GenericMmm4x1::<f64, f64, f64>::mmm()),
mmm_f32_impls: vec![generic::GenericMmm4x4::<f32, f32, f32>::mmm()],
mmm_f32: Box::new(|_, _, _| generic::GenericMmm4x4::<f32, f32, f32>::mmm()),
mmv_f32: Box::new(|_, _| generic::GenericMmm4x1::<f32, f32, f32>::mmm()),
mmm_f16: Box::new(|_, _, _| generic::GenericMmm4x4::<f16, f16, f16>::mmm()),
mmv_f16: Box::new(|_, _| generic::GenericMmm4x1::<f16, f16, f16>::mmm()),
qmmm_i32: Box::new(|_, _, _| generic::GenericMmm4x4::<i8, i8, i32>::mmm()),
qmmv_i32: Box::new(|_, _| generic::GenericMmm4x1::<i8, i8, i32>::mmm()),
sigmoid_f16: Box::new(|| generic::HSigmoid8::ew()),
sigmoid_f32: Box::new(|| generic::SSigmoid4::ew()),
tanh_f16: Box::new(|| generic::HTanh8::ew()),
tanh_f32: Box::new(|| generic::STanh4::ew()),
erf_f32: Box::new(|| generic::SErf4::ew()),
lut_u8: Box::new(|table: &[u8]| Box::new(lut::LutImpl::<generic::GenericLut8>::new(table))),
}
}
#[allow(unreachable_code, unused_mut)]
pub fn best() -> Ops {
let mut ops = generic();
#[cfg(target_arch = "x86_64")]
x86_64_fma::plug(&mut ops);
#[cfg(target_arch = "arm")]
arm32::plug(&mut ops);
#[cfg(target_arch = "aarch64")]
arm64::plug(&mut ops);
ops
}
lazy_static::lazy_static! {
static ref OPS: Ops = {
best()
};
}
pub fn ops() -> &'static Ops {
&OPS
}
use num_traits::*;
use std::fmt::Debug;
use std::ops::*;
pub trait LADatum:
Sized
+ std::fmt::Display
+ Debug
+ Copy
+ Clone
+ Zero
+ One
+ 'static
+ Add<Output = Self>
+ Sub<Output = Self>
+ Mul
+ AddAssign
+ PartialOrd
+ Bounded
+ tract_data::prelude::Datum
{
#[cfg(test)]
fn strat() -> proptest::prelude::BoxedStrategy<Self>;
}
#[cfg(test)]
use proptest::prelude::*;
impl LADatum for f16 {
#[cfg(test)]
fn strat() -> BoxedStrategy<Self> {
f32::strat().prop_map(|f| f.as_()).boxed()
}
}
impl LADatum for f32 {
#[cfg(test)]
fn strat() -> BoxedStrategy<Self> {
(-1000isize..1000).prop_map(|i| i as f32 / 1000.0).boxed()
}
}
impl LADatum for f64 {
#[cfg(test)]
fn strat() -> BoxedStrategy<Self> {
(-1000isize..1000).prop_map(|i| i as f64 / 1000.0).boxed()
}
}
impl LADatum for u8 {
#[cfg(test)]
fn strat() -> BoxedStrategy<Self> {
any::<u8>().boxed()
}
}
impl LADatum for i8 {
#[cfg(test)]
fn strat() -> BoxedStrategy<Self> {
any::<i8>().boxed()
}
}
impl LADatum for i32 {
#[cfg(test)]
fn strat() -> BoxedStrategy<Self> {
any::<i32>().boxed()
}
}
@@ -0,0 +1,93 @@
use crate::Ops;
use crate::frame::element_wise::ElementWiseKer;
use crate::frame::mmm::kernel::MatMatMulKer;
pub mod mmm;
mod intel;
tanh_impl!(f32, fma_tanh_f32, 8, 8, is_x86_feature_detected!("fma"));
sigmoid_impl!(f32, fma_sigmoid_f32, 8, 8, is_x86_feature_detected!("fma"));
pub fn plug(ops: &mut Ops) {
if is_x86_feature_detected!("fma") {
ops.mmv_f32 = Box::new(|_, _| mmm::fma_mmm_f32_64x1::mmm());
ops.mmm_f32 = Box::new(|_, _, n| {
if n.is_none() {
return mmm::fma_mmm_f32_16x6::mmm();
}
let n = n.unwrap();
match n {
1 => unreachable!("should've been mmv"),
2 => return mmm::fma_mmm_f32_40x2::mmm(),
3 => return mmm::fma_mmm_f32_32x3::mmm(),
4 => return mmm::fma_mmm_f32_24x4::mmm(),
5 => return mmm::fma_mmm_f32_16x5::mmm(),
6 => return mmm::fma_mmm_f32_16x6::mmm(),
8 => return mmm::fma_mmm_f32_8x8::mmm(),
_ => {}
};
let scaling_baseline = 60.0;
let kernel_normalized_perf = [
44.0 / scaling_baseline, // 8x8
54.0 / scaling_baseline, // 2x6
54.0 / scaling_baseline, // 2x5
54.0 / scaling_baseline, // 3x4
54.0 / scaling_baseline, // 4x3
54.0 / scaling_baseline, // 5x2
];
fn compute_efficiency(n: usize, kernel_width: usize, scale: f32) -> f32 {
let kernel_width = kernel_width as f32;
let n = n as f32;
let batch_count = (n / kernel_width).ceil();
let actual_count = batch_count * kernel_width;
let multi_batch_penalty = 1.0 - batch_count / 100.0;
n / actual_count * scale * multi_batch_penalty
}
let efficiencies = [
compute_efficiency(n, 8, kernel_normalized_perf[0]),
compute_efficiency(n, 6, kernel_normalized_perf[1]),
compute_efficiency(n, 5, kernel_normalized_perf[2]),
compute_efficiency(n, 4, kernel_normalized_perf[3]),
compute_efficiency(n, 3, kernel_normalized_perf[4]),
compute_efficiency(n, 2, kernel_normalized_perf[5]),
];
let best_idx = efficiencies
.iter()
.copied()
.enumerate()
.fold((0, 0.0), |max, val| if val.1 > max.1 { val } else { max });
match best_idx.0 {
0 => mmm::fma_mmm_f32_8x8::mmm(),
1 => mmm::fma_mmm_f32_16x6::mmm(),
2 => mmm::fma_mmm_f32_16x5::mmm(),
3 => mmm::fma_mmm_f32_24x4::mmm(),
4 => mmm::fma_mmm_f32_32x3::mmm(),
5 => mmm::fma_mmm_f32_40x2::mmm(),
_ => unreachable!("not a valid index"),
}
});
ops.mmm_f32_impls.push(mmm::fma_mmm_f32_16x6::mmm());
ops.mmm_f32_impls.push(mmm::fma_mmm_f32_16x5::mmm());
ops.mmm_f32_impls.push(mmm::fma_mmm_f32_24x4::mmm());
ops.mmm_f32_impls.push(mmm::fma_mmm_f32_32x3::mmm());
ops.mmm_f32_impls.push(mmm::fma_mmm_f32_40x2::mmm());
ops.mmm_f32_impls.push(mmm::fma_mmm_f32_8x8::mmm());
ops.sigmoid_f32 = Box::new(|| fma_sigmoid_f32::ew());
ops.tanh_f32 = Box::new(|| fma_tanh_f32::ew());
log::info!("mmm_f32, sigmoid_f32, tanh_f32: x86_64/fma activated");
}
if is_x86_feature_detected!("avx2") {
ops.qmmm_i32 = Box::new(|_, _, _| mmm::avx2_mmm_i32_8x8::mmm());
log::info!("mmm_i8_i8 and mmm_i8_i32: x86_64/avx2 activated");
}
}
@@ -0,0 +1,5 @@
use crate::frame::mmm::cost_model::CostModel;
#[allow(dead_code)]
pub fn models() -> Vec<(&'static str, CostModel<'static>)> {
vec![]
}
@@ -0,0 +1,11 @@
use crate::frame::mmm::*;
MMMKernel!(f32, fma_mmm_f32_8x8; 8, 8; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("fma"));
MMMKernel!(f32, fma_mmm_f32_16x6; 16, 6; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("fma"));
MMMKernel!(f32, fma_mmm_f32_16x5; 16, 5; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("fma"));
MMMKernel!(f32, fma_mmm_f32_24x4; 24, 4; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("fma"));
MMMKernel!(f32, fma_mmm_f32_32x3; 32, 3; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("fma"));
MMMKernel!(f32, fma_mmm_f32_40x2; 40, 2; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("fma"));
MMMKernel!(f32, fma_mmm_f32_64x1; 64, 1; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("fma"));
MMMKernel!(i32, avx2_mmm_i32_8x8; 8, 8; 32, 4; 0, 0; no_prefetch, is_x86_feature_detected!("avx2"));