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:
@@ -0,0 +1,58 @@
|
||||
#[cfg(any(target_os = "macos", all(target_os = "ios", feature = "apple-amx-ios")))]
|
||||
|
||||
#[cfg(target_os = "ios")]
|
||||
lazy_static::lazy_static! {
|
||||
static ref IPHONE_MODEL_MAJOR:Option<usize> = {
|
||||
use std::ffi::{c_char, c_void, CStr, CString};
|
||||
use std::ptr::null_mut;
|
||||
|
||||
extern "C" {
|
||||
fn sysctlbyname(
|
||||
name: *const c_char,
|
||||
oldp: *mut c_void,
|
||||
oldlenp: *mut isize,
|
||||
newp: *mut c_void,
|
||||
newlen: isize,
|
||||
);
|
||||
}
|
||||
|
||||
unsafe {
|
||||
let mut len: isize = 0;
|
||||
let name = CString::new("hw.machine").unwrap();
|
||||
sysctlbyname(name.as_ptr(), null_mut(), &mut len, null_mut(), 0);
|
||||
let mut buf = vec![0u8; len as _];
|
||||
sysctlbyname(name.as_ptr(), buf.as_mut_ptr() as _, &mut len, null_mut(), 0);
|
||||
let version = CStr::from_bytes_with_nul(&buf).unwrap().to_string_lossy().into_owned();
|
||||
let Some((major, _)) = version.trim_start_matches("iPhone").split_once(",") else { return None };
|
||||
major.parse::<usize>().ok()
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
pub fn has_amx() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
#[cfg(all(target_os = "ios", feature = "apple-amx-ios"))]
|
||||
fn has_amx() -> bool {
|
||||
// iPhone12,1 is the one branded "iPhone 11", with Apple A13 bionic, first CPU featuring amx
|
||||
IPHONE_MODEL_MAJOR.map(|it| it >= 12).unwrap_or(false)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
#[cfg(target_os = "ios")]
|
||||
pub fn has_fp16() -> bool {
|
||||
// iPhone10,1 is the one branded "iPhone 8", with Apple A11 bionic, first CPU featuring fp16
|
||||
IPHONE_MODEL_MAJOR.map(|it| it >= 10).unwrap_or(false)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
#[cfg(not(target_os = "ios"))]
|
||||
pub fn has_fp16() -> bool {
|
||||
cfg!(target_os = "macos")
|
||||
|| cfg!(feature_cpu = "fp16")
|
||||
|| *KIND == Kind::CortexA55
|
||||
|| *KIND == Kind::CortexA75
|
||||
|| *HAS_FP16
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
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::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) {
|
||||
if has_neon() {
|
||||
log::info!("armv7neon activated (smmm, ssigmoid), stanh)");
|
||||
armv7neon::plug(ops);
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
let cost_managed_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.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(&cost_managed_impls, m, k, n))
|
||||
}
|
||||
0xc09 => {
|
||||
let model = cortex_a9::model();
|
||||
Box::new(move |m, k, n| model.pick(&cost_managed_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 {
|
||||
armvfpv2::plug(ops);
|
||||
}
|
||||
}
|
||||
|
||||
#[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());
|
||||
}
|
||||
}
|
||||
}
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
use crate::Ops;
|
||||
use crate::frame::mmm::ImplementationQuality::ManuallyOptimized;
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
const NEON: fn() -> bool = || crate::arm32::has_neon();
|
||||
|
||||
MMMExternKernel!(armv7neon_mmm_f32_8x4_cortexa7 <f32>( 8, 4)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_8x4_cortexa9 <f32>( 8, 4)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_8x4_generic <f32>( 8, 4)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_8x6_cortexa7 <f32>( 8, 6)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_8x6_cortexa9 <f32>( 8, 6)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_8x6_generic <f32>( 8, 6)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_8x1_generic <f32>( 8, 1)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_32x1_cortexa7<f32>(32, 1)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_32x1_cortexa9<f32>(32, 1)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(armv7neon_mmm_f32_32x1_generic <f32>(32, 1)@(16, 4) where(NEON) quality(ManuallyOptimized));
|
||||
|
||||
MMMExternKernel!(armv7neon_mmm_i32_8x4<i32>(8, 4)@(32, 4) where(NEON)
|
||||
packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 8, 32), PackedFormat::new(DatumType::I8, 4, 32));
|
||||
quality(ManuallyOptimized)
|
||||
store(i8)
|
||||
);
|
||||
|
||||
MMMExternKernel!(armv7neon_mmm_i32_32x1<i32>(32, 1)@(32, 4) where(NEON)
|
||||
packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 32, 32), PackedFormat::new(DatumType::I8, 1, 4));
|
||||
quality(ManuallyOptimized)
|
||||
store(i8)
|
||||
);
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
ops.mmm_impls.extend_from_slice(&[
|
||||
armv7neon_mmm_f32_8x4_cortexa7.mmm(),
|
||||
armv7neon_mmm_f32_8x4_cortexa9.mmm(),
|
||||
armv7neon_mmm_f32_8x4_generic.mmm(),
|
||||
armv7neon_mmm_f32_8x6_cortexa7.mmm(),
|
||||
armv7neon_mmm_f32_8x6_cortexa9.mmm(),
|
||||
armv7neon_mmm_f32_8x6_generic.mmm(),
|
||||
armv7neon_mmm_f32_8x1_generic.mmm(),
|
||||
armv7neon_mmm_f32_32x1_cortexa7.mmm(),
|
||||
armv7neon_mmm_f32_32x1_cortexa9.mmm(),
|
||||
armv7neon_mmm_f32_32x1_generic.mmm(),
|
||||
]);
|
||||
}
|
||||
|
||||
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());
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
use crate::Ops;
|
||||
use crate::frame::mmm::ImplementationQuality::ManuallyOptimized;
|
||||
use crate::frame::mmm::*;
|
||||
|
||||
MMMExternKernel!(armvfpv2_mmm_f32_4x4<f32>(4, 4)@(4, 4) quality(ManuallyOptimized));
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
log::info!("armvfpv2 activated for smmm");
|
||||
ops.mmm_f32 = Box::new(|_, _, _| armvfpv2_mmm_f32_4x4.mmm());
|
||||
ops.mmm_impls.push(armvfpv2_mmm_f32_4x4.mmm());
|
||||
}
|
||||
+861
@@ -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,
|
||||
],
|
||||
}
|
||||
}
|
||||
+1701
File diff suppressed because it is too large
Load Diff
+861
@@ -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,
|
||||
],
|
||||
}
|
||||
}
|
||||
+1701
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,542 @@
|
||||
#![allow(clippy::excessive_precision)]
|
||||
#[cfg(any(target_os = "macos", all(target_os = "ios", feature = "apple-amx-ios")))]
|
||||
mod apple_amx;
|
||||
mod arm64simd;
|
||||
pub mod cortex_a53;
|
||||
mod cortex_a55;
|
||||
// `tract_sme` is set by build.rs only when the assembler can assemble SME
|
||||
// (gates out e.g. the old Debian stretch aarch64 toolchain).
|
||||
#[cfg(all(any(target_os = "macos", target_os = "linux"), tract_sme))]
|
||||
mod sme;
|
||||
mod sve;
|
||||
//mod cortex_a72;
|
||||
//mod cortex_a73;
|
||||
pub use arm64simd::*;
|
||||
|
||||
#[cfg(not(feature = "no_fp16"))]
|
||||
pub mod arm64fp16;
|
||||
#[cfg(not(feature = "no_fp16"))]
|
||||
pub use arm64fp16::*;
|
||||
|
||||
use crate::f16;
|
||||
use crate::{BinOp, DatumType, LinalgRegistry, Ops};
|
||||
|
||||
use crate::frame::by_scalar::ByScalarKer;
|
||||
use crate::frame::element_wise::ElementWiseKer;
|
||||
use crate::frame::reduce::{MapReduceKer, ReduceKer};
|
||||
use crate::frame::unicast::UnicastKer;
|
||||
|
||||
// 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";
|
||||
#[allow(dead_code)]
|
||||
const PART_NEOVERSE_N1: &str = "0xd0c";
|
||||
#[allow(dead_code)]
|
||||
const PART_NEOVERSE_N2: &str = "0xd49";
|
||||
#[allow(dead_code)]
|
||||
const PART_NEOVERSE_N3: &str = "0xd8e";
|
||||
#[allow(dead_code)]
|
||||
const PART_NEOVERSE_V1: &str = "0xd40";
|
||||
#[allow(dead_code)]
|
||||
const PART_NEOVERSE_V2: &str = "0xd4f";
|
||||
#[allow(dead_code)]
|
||||
const PART_NEOVERSE_V3: &str = "0xd83";
|
||||
|
||||
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())
|
||||
}
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref KIND: Kind = Kind::choose();
|
||||
|
||||
static ref CPU_FEATURES: Vec<String> = {
|
||||
#[cfg(test)] crate::setup_test_logger();
|
||||
let Ok(cpu_info) = std::fs::read_to_string("/proc/cpuinfo") else {
|
||||
log::warn!("Could not read /proc/cpuinfo. CPU Features detection may be impaired.");
|
||||
return vec!();
|
||||
};
|
||||
if let Some(line) = cpu_info
|
||||
.lines()
|
||||
.find(|line| line.starts_with("Features")) {
|
||||
line.split_once(':').unwrap().1.split_whitespace().map(|s| s.to_string()).collect()
|
||||
} else {
|
||||
log::warn!("Could not find \"Features :\" lines in /proc/cpuinfo. CPU Features detection may be impaired.");
|
||||
vec!()
|
||||
}
|
||||
};
|
||||
|
||||
static ref HAS_FP16: bool = {
|
||||
CPU_FEATURES.iter().any(|s| &**s == "asimdhp")
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "macos", target_os = "ios"))]
|
||||
fn apple_string_from_c_bytes(buf: &[u8]) -> String {
|
||||
use std::ffi::CStr;
|
||||
|
||||
CStr::from_bytes_until_nul(buf)
|
||||
.ok()
|
||||
.map(|s| s.to_string_lossy().into_owned())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "macos", target_os = "ios"))]
|
||||
fn apple_get_syscall(key: &str) -> String {
|
||||
use std::ffi::{CString, c_char, c_int, c_void};
|
||||
use std::ptr::null_mut;
|
||||
|
||||
unsafe extern "C" {
|
||||
fn sysctlbyname(
|
||||
name: *const c_char,
|
||||
oldp: *mut c_void,
|
||||
oldlenp: *mut usize,
|
||||
newp: *mut c_void,
|
||||
newlen: usize,
|
||||
) -> c_int;
|
||||
}
|
||||
|
||||
let Ok(name) = CString::new(key) else {
|
||||
return String::new();
|
||||
};
|
||||
|
||||
unsafe {
|
||||
let mut len_needed: usize = 0;
|
||||
if sysctlbyname(name.as_ptr(), null_mut(), &mut len_needed, null_mut(), 0) != 0 {
|
||||
return String::new();
|
||||
}
|
||||
|
||||
let mut buf = vec![0u8; len_needed.saturating_add(1)];
|
||||
let mut len: usize = buf.len();
|
||||
if sysctlbyname(
|
||||
name.as_ptr(),
|
||||
buf.as_mut_ptr() as _,
|
||||
&mut len,
|
||||
null_mut(),
|
||||
0,
|
||||
) != 0
|
||||
{
|
||||
return String::new();
|
||||
}
|
||||
|
||||
buf.truncate(len.min(buf.len()));
|
||||
if buf.last().copied() != Some(0) {
|
||||
buf.push(0);
|
||||
}
|
||||
|
||||
apple_string_from_c_bytes(&buf)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, any(target_os = "macos", target_os = "ios")))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn apple_string_from_c_bytes_returns_empty_without_nul() {
|
||||
assert_eq!(apple_string_from_c_bytes(b"hello"), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apple_string_from_c_bytes_stops_at_first_nul() {
|
||||
assert_eq!(apple_string_from_c_bytes(b"hello\0world\0"), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apple_get_syscall_does_not_panic() {
|
||||
let _ = apple_get_syscall("machdep.cpu.brand_string");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
pub fn has_amx() -> bool {
|
||||
!apple_get_syscall("machdep.cpu.brand_string").contains("(Virtual)")
|
||||
}
|
||||
|
||||
#[cfg(target_os = "ios")]
|
||||
lazy_static::lazy_static! {
|
||||
static ref IPHONE_MODEL_MAJOR:Option<usize> = {
|
||||
let version = apple_get_syscall("hw.machine");
|
||||
let Some((major, _)) = version.trim_start_matches("iPhone").split_once(",") else { return None };
|
||||
major.parse::<usize>().ok()
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(all(target_os = "ios", feature = "apple-amx-ios"))]
|
||||
fn has_amx() -> bool {
|
||||
// iPhone12,1 is the one branded "iPhone 11", with Apple A13 bionic, first CPU featuring amx
|
||||
IPHONE_MODEL_MAJOR.map(|it| it >= 12).unwrap_or(false)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
#[cfg(target_os = "ios")]
|
||||
pub fn has_fp16() -> bool {
|
||||
// iPhone10,1 is the one branded "iPhone 8", with Apple A11 bionic, first CPU featuring fp16
|
||||
IPHONE_MODEL_MAJOR.map(|it| it >= 10).unwrap_or(false)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
#[cfg(not(target_os = "ios"))]
|
||||
pub fn has_fp16() -> bool {
|
||||
cfg!(target_os = "macos")
|
||||
|| cfg!(feature_cpu = "fp16")
|
||||
|| *KIND == Kind::CortexA55
|
||||
|| *KIND == Kind::CortexA75
|
||||
|| *HAS_FP16
|
||||
}
|
||||
|
||||
// FEAT_DotProd (SDOT/UDOT), ARMv8.2. TRACT_DOTPROD_DISABLE=1 forces it off so
|
||||
// callers can A/B the SDOT kernel against the SMLAL 8x8 fallback on one binary.
|
||||
#[cfg(target_os = "macos")]
|
||||
pub fn has_dotprod() -> bool {
|
||||
// Every Apple arm64 CPU (M1+/A11+) implements FEAT_DotProd.
|
||||
std::env::var_os("TRACT_DOTPROD_DISABLE").is_none()
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
pub fn has_dotprod() -> bool {
|
||||
if std::env::var_os("TRACT_DOTPROD_DISABLE").is_some() {
|
||||
return false;
|
||||
}
|
||||
// HWCAP_ASIMDDP = 1 << 20 on aarch64.
|
||||
const HWCAP_ASIMDDP: u64 = 1 << 20;
|
||||
const AT_HWCAP: u64 = 16;
|
||||
unsafe extern "C" {
|
||||
fn getauxval(t: u64) -> u64;
|
||||
}
|
||||
unsafe { (getauxval(AT_HWCAP) & HWCAP_ASIMDDP) != 0 }
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "ios")))]
|
||||
pub fn has_dotprod() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
#[cfg(target_os = "ios")]
|
||||
pub fn has_dotprod() -> bool {
|
||||
// A11+ (iPhone10,1+) implement FEAT_DotProd.
|
||||
std::env::var_os("TRACT_DOTPROD_DISABLE").is_none()
|
||||
&& IPHONE_MODEL_MAJOR.map(|it| it >= 10).unwrap_or(false)
|
||||
}
|
||||
|
||||
#[target_feature(enable = "fp16")]
|
||||
#[inline]
|
||||
pub unsafe fn add_f16(a: f16, b: f16) -> f16 {
|
||||
unsafe {
|
||||
let result: u16;
|
||||
std::arch::asm!(
|
||||
"fadd {0:h}, {1:h}, {2:h}",
|
||||
lateout(vreg) result,
|
||||
in(vreg) a.to_bits(),
|
||||
in(vreg) b.to_bits(),
|
||||
options(pure, nomem, nostack, preserves_flags));
|
||||
f16::from_bits(result)
|
||||
}
|
||||
}
|
||||
|
||||
#[target_feature(enable = "fp16")]
|
||||
#[inline]
|
||||
pub unsafe fn mul_f16(a: f16, b: f16) -> f16 {
|
||||
unsafe {
|
||||
let result: u16;
|
||||
std::arch::asm!(
|
||||
"fmul {0:h}, {1:h}, {2:h}",
|
||||
lateout(vreg) result,
|
||||
in(vreg) a.to_bits(),
|
||||
in(vreg) b.to_bits(),
|
||||
options(pure, nomem, nostack, preserves_flags));
|
||||
f16::from_bits(result)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
|
||||
pub enum Kind {
|
||||
Generic,
|
||||
AppleM,
|
||||
Neoverse,
|
||||
CortexA53,
|
||||
CortexA55,
|
||||
CortexA72,
|
||||
CortexA73,
|
||||
CortexA75,
|
||||
}
|
||||
|
||||
impl Kind {
|
||||
pub fn choose() -> Kind {
|
||||
#[cfg(test)]
|
||||
crate::setup_test_logger();
|
||||
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("neoverse") {
|
||||
Kind::Neoverse
|
||||
} 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,
|
||||
PART_NEOVERSE_N1 | PART_NEOVERSE_N2 | PART_NEOVERSE_N3 | PART_NEOVERSE_V1
|
||||
| PART_NEOVERSE_V2 | PART_NEOVERSE_V3 => Kind::Neoverse,
|
||||
_ => Kind::Generic,
|
||||
}
|
||||
};
|
||||
log::info!("CPU optimisation: {:?}", kind);
|
||||
kind
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn register_all_unicast(registry: &mut LinalgRegistry) {
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F32),
|
||||
Box::new(|| arm64simd_unicast_mul_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F16),
|
||||
Box::new(|| arm64fp16_unicast_mul_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F32),
|
||||
Box::new(|| arm64simd_unicast_add_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F16),
|
||||
Box::new(|| arm64fp16_unicast_add_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F32),
|
||||
Box::new(|| arm64simd_unicast_sub_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F16),
|
||||
Box::new(|| arm64fp16_unicast_sub_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F32),
|
||||
Box::new(|| arm64simd_unicast_subf_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F16),
|
||||
Box::new(|| arm64fp16_unicast_subf_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F32),
|
||||
Box::new(|| arm64simd_unicast_min_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F16),
|
||||
Box::new(|| arm64fp16_unicast_min_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F32),
|
||||
Box::new(|| arm64simd_unicast_max_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F16),
|
||||
Box::new(|| arm64fp16_unicast_max_f16_32n::bin()),
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn register_all_by_scalar(registry: &mut LinalgRegistry) {
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F32),
|
||||
Box::new(|| arm64simd_mul_by_scalar_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F16),
|
||||
Box::new(|| arm64fp16_mul_by_scalar_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F32),
|
||||
Box::new(|| arm64simd_add_by_scalar_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F16),
|
||||
Box::new(|| arm64fp16_add_by_scalar_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F32),
|
||||
Box::new(|| arm64simd_sub_by_scalar_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F16),
|
||||
Box::new(|| arm64fp16_sub_by_scalar_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F32),
|
||||
Box::new(|| arm64simd_subf_by_scalar_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F16),
|
||||
Box::new(|| arm64fp16_subf_by_scalar_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F32),
|
||||
Box::new(|| arm64simd_min_by_scalar_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F16),
|
||||
Box::new(|| arm64fp16_min_by_scalar_f16_32n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F32),
|
||||
Box::new(|| arm64simd_max_by_scalar_f32_16n::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F16),
|
||||
Box::new(|| arm64fp16_max_by_scalar_f16_32n::bin()),
|
||||
);
|
||||
}
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
arm64simd::plug(ops);
|
||||
|
||||
#[cfg(not(feature = "no_fp16"))]
|
||||
if has_fp16() {
|
||||
arm64fp16::plug(ops);
|
||||
}
|
||||
|
||||
// SDOT (~4x the SMLAL 8x8) when FEAT_DotProd is present, else the SMLAL 8x8 fallback.
|
||||
// The SDOT kernel only exists when the assembler could encode `sdot`
|
||||
// (`tract_arm64_dotprod`, set by build.rs); otherwise always use the SMLAL 8x8.
|
||||
#[cfg(tract_arm64_dotprod)]
|
||||
if has_dotprod() {
|
||||
ops.qmmm_i32 = Box::new(|_, _, _| arm64simd_mmm_i32_8x8_dot.mmm());
|
||||
} else {
|
||||
ops.qmmm_i32 = Box::new(|_, _, _| arm64simd_mmm_i32_8x8.mmm());
|
||||
}
|
||||
#[cfg(not(tract_arm64_dotprod))]
|
||||
{
|
||||
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,
|
||||
};
|
||||
let impls = ops.mmm_impls.clone();
|
||||
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 {
|
||||
log::info!("Cortex-A55 mmm_f16 and mmv_f16 activated");
|
||||
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 {
|
||||
log::info!("ARMv8.2 mmm_f16 and mmv_f16 activated");
|
||||
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.leaky_relu_f32 = Box::new(|| arm64simd_leaky_relu_f32_8n::ew());
|
||||
ops.hardswish_f32 = Box::new(|| arm64simd_hardswish_f32_8n::ew());
|
||||
ops.silu_f32 = Box::new(|| arm64simd_silu_f32_4n_fused::ew());
|
||||
ops.gelu_f32 = Box::new(|| arm64simd_gelu_f32_4n_fused::ew());
|
||||
ops.sigmoid_f32 = Box::new(|| arm64simd_sigmoid_f32_4n::ew());
|
||||
ops.tanh_f32 = Box::new(|| arm64simd_tanh_f32_4n::ew());
|
||||
ops.max_f32 = Box::new(|| arm64simd_max_f32_16n::red());
|
||||
ops.sum_f32 = Box::new(|| arm64simd_sum_f32_16n::red());
|
||||
ops.mul_by_scalar_f32 = Box::new(|| arm64simd_mul_by_scalar_f32_16n::ew());
|
||||
ops.softmax2_fastcompact_f32 = Box::new(|| arm64simd_softmax2_fastcompact_f32_16n::red());
|
||||
ops.rms_norm_f32 = Box::new(arm64simd_rms_norm_f32);
|
||||
#[cfg(not(feature = "no_fp16"))]
|
||||
if has_fp16() {
|
||||
log::info!("ARMv8.2 tanh_f16 and sigmoid_f16 activated");
|
||||
ops.leaky_relu_f16 = Box::new(|| arm64fp16_leaky_relu_f16_16n::ew());
|
||||
ops.tanh_f16 = Box::new(|| arm64fp16_tanh_f16_8n::ew());
|
||||
ops.sigmoid_f16 = Box::new(|| arm64fp16_sigmoid_f16_8n::ew());
|
||||
ops.max_f16 = Box::new(|| arm64fp16_max_f16_32n::red());
|
||||
ops.sum_f16 = Box::new(|| arm64fp16_sum_f16_32n::red());
|
||||
ops.mul_by_scalar_f16 = Box::new(|| arm64fp16_mul_by_scalar_f16_32n::ew());
|
||||
} else {
|
||||
log::info!("No native fp16 support");
|
||||
}
|
||||
#[cfg(any(target_os = "macos", all(target_os = "ios", feature = "apple-amx-ios")))]
|
||||
{
|
||||
apple_amx::plug(ops);
|
||||
}
|
||||
#[cfg(all(any(target_os = "macos", target_os = "linux"), tract_sme))]
|
||||
{
|
||||
sme::plug(ops);
|
||||
}
|
||||
sve::plug(ops);
|
||||
}
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
use crate::Ops;
|
||||
use crate::frame::mmm::ImplementationQuality::ManuallyOptimized;
|
||||
use crate::mmm::*;
|
||||
use tract_data::prelude::*;
|
||||
|
||||
use super::has_amx;
|
||||
use super::{arm64fp16_mmm_f16_16x8_gen, arm64simd_mmm_f32_8x8_gen, arm64simd_mmm_f32_64x1_gen};
|
||||
|
||||
const AMX: fn() -> bool = crate::arm64::has_amx;
|
||||
const CAN_FUSE: fn(&FusedSpec) -> bool = |f| !matches!(f, &FusedSpec::LeakyRelu(_));
|
||||
|
||||
MMMExternKernel!(apple_amx_mmm_f32_32x32<f32>(32, 32)@(128, 128) where(AMX) can_fuse(CAN_FUSE) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(apple_amx_mmm_f32_32x1<f32>(32, 1)@(128, 128) where(AMX) can_fuse(CAN_FUSE) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(apple_amx_mmm_f16_64x32<f16>(64, 32)@(128, 128) where(AMX) can_fuse(CAN_FUSE) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(apple_amx_mmm_f16_64x1<f16>(64, 1)@(128, 128) where(AMX) can_fuse(CAN_FUSE) quality(ManuallyOptimized));
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
if has_amx() {
|
||||
log::info!(
|
||||
"AMX optimisation activated (A7v2: AMX only for f32 mmm with M>=32 AND N>=32; \
|
||||
smaller shapes + all f32 mmv route to NEON kernels)"
|
||||
);
|
||||
// ----- A7v2 dispatch logic (data-driven) -----
|
||||
//
|
||||
// Empirical finding from /tmp/amx_vs_neon.md microbench (Apple M1 Pro):
|
||||
// the AMX 32x32 kernel beats NEON 8x8 only when BOTH M and N are at
|
||||
// least 32 — the AMX tile dimensions. At smaller shapes the per-tile
|
||||
// padding waste + AMX dispatch overhead make NEON faster.
|
||||
//
|
||||
// Predicate validation: 88.3% accuracy on 512-shape sweep.
|
||||
//
|
||||
// Canary impact (measured 2026-05-13, see notes/tract-amx-low-m-investigation.md):
|
||||
// turning AMX off entirely yielded:
|
||||
// df_dec 1.55× faster mobilenetv2 1.59× faster
|
||||
// erb_dec 1.49× squeezenet 1.22×
|
||||
// enc 1.17× yolov8n 1.15× SLOWER
|
||||
// inception_v3 1.43× SLOWER sam2_tiny 1.54× SLOWER
|
||||
// The shape-aware predicate keeps the AMX wins for the heavy models
|
||||
// (Inception, YOLO, SAM2) while routing small shapes to NEON.
|
||||
ops.mmm_f32 = Box::new(|m, _, n| {
|
||||
let big_enough = m.is_some_and(|m| m >= 32) && n.is_some_and(|n| n >= 32);
|
||||
if big_enough {
|
||||
apple_amx_mmm_f32_32x32.mmm()
|
||||
} else {
|
||||
arm64simd_mmm_f32_8x8_gen.mmm()
|
||||
}
|
||||
});
|
||||
// mmv (n=1) f32: AMX 32x1 is dominated by NEON 64x1 across the entire
|
||||
// shape sweep — confirmed by canary deltas on DFN3 (which is mmv-heavy).
|
||||
// Always use NEON.
|
||||
ops.mmv_f32 = Box::new(|_, _| arm64simd_mmm_f32_64x1_gen.mmm());
|
||||
|
||||
// ----- f16 paths kept conservative for now -----
|
||||
//
|
||||
// We didn't run the f16 microbench yet, so retain the original logic
|
||||
// and the previous low-M-routes-to-NEON heuristic.
|
||||
ops.mmm_f16 = Box::new(|m, _, _| {
|
||||
if m.is_some_and(|m| m <= 16) {
|
||||
arm64fp16_mmm_f16_16x8_gen.mmm()
|
||||
} else {
|
||||
apple_amx_mmm_f16_64x32.mmm()
|
||||
}
|
||||
});
|
||||
ops.mmv_f16 = Box::new(|_, _| apple_amx_mmm_f16_64x1.mmm());
|
||||
ops.mmm_impls.extend_from_slice(&[
|
||||
apple_amx_mmm_f32_32x32.mmm(),
|
||||
apple_amx_mmm_f32_32x1.mmm(),
|
||||
apple_amx_mmm_f16_64x32.mmm(),
|
||||
apple_amx_mmm_f16_64x1.mmm(),
|
||||
]);
|
||||
} else {
|
||||
log::info!("No AMX optimisation");
|
||||
}
|
||||
}
|
||||
+70
@@ -0,0 +1,70 @@
|
||||
use tract_data::half::f16;
|
||||
|
||||
mod by_scalar;
|
||||
mod leaky_relu;
|
||||
mod max;
|
||||
pub mod panel_extract;
|
||||
mod sum;
|
||||
mod unicast;
|
||||
pub use by_scalar::*;
|
||||
pub use leaky_relu::*;
|
||||
pub use max::*;
|
||||
pub use sum::*;
|
||||
pub use unicast::*;
|
||||
|
||||
use crate::Ops;
|
||||
use crate::block_quant::PackedBlockQuantFormat;
|
||||
use crate::block_quant::Q4_0;
|
||||
use crate::frame::mmm::ImplementationQuality::ManuallyOptimized;
|
||||
|
||||
const FP16: fn() -> bool = crate::arm64::has_fp16;
|
||||
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_16x8_gen<f16>(16, 8)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_16x8_a55<f16>(16, 8)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_32x4_gen<f16>(32, 4)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_32x4_a55<f16>(32, 4)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_128x1_gen<f16>(128,1)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_128x1_a55<f16>(128,1)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_64x3_gen<f16>(64, 3)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64fp16_mmm_f16_32x6_gen<f16>(32, 6)@(16, 16) where(FP16) quality(ManuallyOptimized));
|
||||
|
||||
MMMExternKernel! { arm64fp16_mmm_f16_64x1_gen<f16>(64, 1)@(16, 16) where(FP16)
|
||||
packing[1] = q40f16z16se => |k| k.with_packing_a(PackedBlockQuantFormat::new(&Q4_0, 64, 16, true));
|
||||
packing[2] = q40f16z16 => |k| k.with_packing_a(PackedBlockQuantFormat::new(&Q4_0, 64, 16, false));
|
||||
quality(ManuallyOptimized)
|
||||
}
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
panel_extract::plug(ops);
|
||||
ops.mmm_impls.extend_from_slice(&[
|
||||
arm64fp16_mmm_f16_16x8_a55.mmm(),
|
||||
arm64fp16_mmm_f16_16x8_gen.mmm(),
|
||||
arm64fp16_mmm_f16_32x4_a55.mmm(),
|
||||
arm64fp16_mmm_f16_32x4_gen.mmm(),
|
||||
arm64fp16_mmm_f16_128x1_a55.mmm(),
|
||||
arm64fp16_mmm_f16_128x1_gen.mmm(),
|
||||
arm64fp16_mmm_f16_64x3_gen.mmm(),
|
||||
arm64fp16_mmm_f16_32x6_gen.mmm(),
|
||||
arm64fp16_mmm_f16_64x1_gen.mmm(),
|
||||
]);
|
||||
}
|
||||
|
||||
tanh_impl!(f16, arm64fp16_tanh_f16_8n, 8, 8, crate::arm64::has_fp16());
|
||||
sigmoid_impl!(
|
||||
f16,
|
||||
arm64fp16_sigmoid_f16_8n,
|
||||
8,
|
||||
8,
|
||||
crate::arm64::has_fp16()
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
|
||||
#[test]
|
||||
fn kits() {
|
||||
let mut ops = crate::generic();
|
||||
super::plug(&mut ops);
|
||||
}
|
||||
}
|
||||
Vendored
+258
@@ -0,0 +1,258 @@
|
||||
use crate::f16;
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_mul_by_scalar_f16_32n,
|
||||
32,
|
||||
4,
|
||||
f16,
|
||||
fn run(buf: &mut [f16], s: f16) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &mut [f16], s: f16) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.8h, v0.h[0]
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}]
|
||||
fmul v4.8h, v4.8h, v0.8h
|
||||
fmul v5.8h, v5.8h, v0.8h
|
||||
fmul v6.8h, v6.8h, v0.8h
|
||||
fmul v7.8h, v7.8h, v0.8h
|
||||
st1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s.to_bits(),
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(buf, s) }
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_add_by_scalar_f16_32n,
|
||||
32,
|
||||
4,
|
||||
f16,
|
||||
fn run(buf: &mut [f16], s: f16) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &mut [f16], s: f16) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.8h, v0.h[0]
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}]
|
||||
fadd v4.8h, v4.8h, v0.8h
|
||||
fadd v5.8h, v5.8h, v0.8h
|
||||
fadd v6.8h, v6.8h, v0.8h
|
||||
fadd v7.8h, v7.8h, v0.8h
|
||||
st1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s.to_bits(),
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(buf, s) }
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_sub_by_scalar_f16_32n,
|
||||
32,
|
||||
4,
|
||||
f16,
|
||||
fn run(buf: &mut [f16], s: f16) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &mut [f16], s: f16) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.8h, v0.h[0]
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}]
|
||||
fsub v4.8h, v4.8h, v0.8h
|
||||
fsub v5.8h, v5.8h, v0.8h
|
||||
fsub v6.8h, v6.8h, v0.8h
|
||||
fsub v7.8h, v7.8h, v0.8h
|
||||
st1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s.to_bits(),
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(buf, s) }
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_subf_by_scalar_f16_32n,
|
||||
32,
|
||||
4,
|
||||
f16,
|
||||
fn run(buf: &mut [f16], s: f16) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &mut [f16], s: f16) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.8h, v0.h[0]
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}]
|
||||
fsub v4.8h, v0.8h, v4.8h
|
||||
fsub v5.8h, v0.8h, v5.8h
|
||||
fsub v6.8h, v0.8h, v6.8h
|
||||
fsub v7.8h, v0.8h, v7.8h
|
||||
st1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s.to_bits(),
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(buf, s) }
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_min_by_scalar_f16_32n,
|
||||
32,
|
||||
4,
|
||||
f16,
|
||||
fn run(buf: &mut [f16], s: f16) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &mut [f16], s: f16) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.8h, v0.h[0]
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}]
|
||||
fmin v4.8h, v4.8h, v0.8h
|
||||
fmin v5.8h, v5.8h, v0.8h
|
||||
fmin v6.8h, v6.8h, v0.8h
|
||||
fmin v7.8h, v7.8h, v0.8h
|
||||
st1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s.to_bits(),
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(buf, s) }
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_max_by_scalar_f16_32n,
|
||||
32,
|
||||
4,
|
||||
f16,
|
||||
fn run(buf: &mut [f16], s: f16) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &mut [f16], s: f16) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.8h, v0.h[0]
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}]
|
||||
fmax v4.8h, v4.8h, v0.8h
|
||||
fmax v5.8h, v5.8h, v0.8h
|
||||
fmax v6.8h, v6.8h, v0.8h
|
||||
fmax v7.8h, v7.8h, v0.8h
|
||||
st1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s.to_bits(),
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(buf, s) }
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64fp16_mul_by_scalar_f16_32n {
|
||||
use super::*;
|
||||
by_scalar_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_mul_by_scalar_f16_32n,
|
||||
|a, b| a * b
|
||||
);
|
||||
by_scalar_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_add_by_scalar_f16_32n,
|
||||
|a, b| a + b
|
||||
);
|
||||
by_scalar_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_sub_by_scalar_f16_32n,
|
||||
|a, b| a - b
|
||||
);
|
||||
by_scalar_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_subf_by_scalar_f16_32n,
|
||||
|a, b| b - a
|
||||
);
|
||||
by_scalar_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_min_by_scalar_f16_32n,
|
||||
|a, b| a.min(b)
|
||||
);
|
||||
by_scalar_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_max_by_scalar_f16_32n,
|
||||
|a, b| a.max(b)
|
||||
);
|
||||
}
|
||||
Vendored
+56
@@ -0,0 +1,56 @@
|
||||
use tract_data::internal::f16;
|
||||
|
||||
ew_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_leaky_relu_f16_16n,
|
||||
16,
|
||||
8,
|
||||
f16,
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f16], alpha: f16) {
|
||||
assert!(buf.len() % 8 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &mut [f16], alpha: f16) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.8h, {alpha:v}.h[0]
|
||||
dup v1.8h, {one:v}.h[0]
|
||||
2:
|
||||
ldp q3, q4, [{ptr}]
|
||||
|
||||
fcmgt v5.8h, v3.8h, #0.0
|
||||
fcmgt v6.8h, v4.8h, #0.0
|
||||
bsl v5.16b, v1.16b, v0.16b
|
||||
bsl v6.16b, v1.16b, v0.16b
|
||||
fmul v3.8h, v3.8h, v5.8h
|
||||
fmul v4.8h, v4.8h, v6.8h
|
||||
|
||||
stp q3, q4, [{ptr}], #32
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
one = in(vreg) f16::from_f32(1.0f32).to_bits(),
|
||||
alpha = in(vreg) alpha.to_bits(),
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
out("v0") _,
|
||||
out("v1") _,
|
||||
out("q3") _,
|
||||
out("q4") _,
|
||||
out("q5") _,
|
||||
out("q6") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
unsafe { run(buf, alpha) }
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_arm64simd_leaky_relu_f16_16n {
|
||||
use super::*;
|
||||
leaky_relu_frame_tests!(crate::arm64::has_fp16(), f16, arm64fp16_leaky_relu_f16_16n);
|
||||
}
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
use tract_data::half::f16;
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_max_f16_32n,
|
||||
32,
|
||||
8,
|
||||
(),
|
||||
f16::MIN,
|
||||
#[inline(never)]
|
||||
fn run(buf: &[f16], _: ()) -> f16 {
|
||||
assert!(buf.len() % 32 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &[f16]) -> f16 {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
let mut out: u16;
|
||||
std::arch::asm!("
|
||||
ins v0.h[0], {min:w}
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
fmax v0.8h, v0.8h, v4.8h
|
||||
fmax v1.8h, v1.8h, v5.8h
|
||||
fmax v2.8h, v2.8h, v6.8h
|
||||
fmax v3.8h, v3.8h, v7.8h
|
||||
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
|
||||
fmax v0.8h, v0.8h, v1.8h
|
||||
fmax v2.8h, v2.8h, v3.8h
|
||||
fmax v0.8h, v0.8h, v2.8h
|
||||
fmaxv h0, v0.8h
|
||||
",
|
||||
// using v0 as inout triggers https://github.com/rust-lang/rust/issues/120374
|
||||
min = in(reg) f16::MIN.to_bits(),
|
||||
ptr = inout(reg) ptr => _,
|
||||
len = inout(reg) len => _,
|
||||
out("v0") out, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
f16::from_bits(out)
|
||||
}
|
||||
}
|
||||
unsafe { run(buf) }
|
||||
},
|
||||
#[inline(never)]
|
||||
fn reduce_two(a: f16, b: f16) -> f16 {
|
||||
a.max(b)
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64fp16_max_f16_32n {
|
||||
use super::*;
|
||||
crate::max_frame_tests!(crate::arm64::has_fp16(), f16, arm64fp16_max_f16_32n);
|
||||
}
|
||||
Vendored
+94
@@ -0,0 +1,94 @@
|
||||
use super::FP16;
|
||||
use crate::Ops;
|
||||
use crate::block_quant::{PackedBlockQuantFormat, Q4_0};
|
||||
use crate::pack::Packing;
|
||||
use tract_data::internal::*;
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
ops.panel_extractors.push(packed_64_q40_to_f16.clone());
|
||||
}
|
||||
|
||||
panel_extractor!(kernel_packed_64_q40_to_f16 as packed_64_q40_to_f16(
|
||||
Box::new(PackedBlockQuantFormat::new(&Q4_0, 64, 16, true)),
|
||||
f16::packing(64).align(16)
|
||||
) where(FP16));
|
||||
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn kernel_packed_64_q40_to_f16(input: *const u8, output: *mut u8, k: usize) {
|
||||
unsafe {
|
||||
if k == 0 {
|
||||
return;
|
||||
}
|
||||
let lookup_table: [u8; 16] = [
|
||||
0xc8, 0xc7, 0xc6, 0xc5, 0xc4, 0xc2, 0xc0, 0xbc, 0x00, 0x3c, 0x40, 0x42, 0x44, 0x45,
|
||||
0x46, 0x47,
|
||||
];
|
||||
std::arch::asm!("
|
||||
ld1 {{v13.16b}}, [{lookup_table}]
|
||||
movi v15.16b, 15
|
||||
eor v12.16b, v12.16b, v12.16b
|
||||
|
||||
2:
|
||||
add {scales}, {i}, 1024 // scales at end: 32 (cols) * 64 (rows) / 2 (half byte)
|
||||
ld1 {{v16.16b-v19.16b}}, [{scales}], #64
|
||||
ld1 {{v20.16b-v23.16b}}, [{scales}]
|
||||
|
||||
mov {k2}, 32
|
||||
3:
|
||||
ld1 {{ v9.16b-v10.16b }}, [{i}], #32
|
||||
|
||||
and v0.16b, v9.16b, v15.16b
|
||||
ushr v2.16b, v9.16b, 4
|
||||
|
||||
and v4.16b, v10.16b, v15.16b
|
||||
ushr v6.16b, v10.16b, 4
|
||||
|
||||
tbl v0.16b, {{ v13.16b }}, v0.16b
|
||||
tbl v2.16b, {{ v13.16b }}, v2.16b
|
||||
tbl v4.16b, {{ v13.16b }}, v4.16b
|
||||
tbl v6.16b, {{ v13.16b }}, v6.16b
|
||||
|
||||
zip2 v1.16b, v12.16b, v0.16b
|
||||
zip2 v3.16b, v12.16b, v2.16b
|
||||
zip2 v5.16b, v12.16b, v4.16b
|
||||
zip2 v7.16b, v12.16b, v6.16b
|
||||
|
||||
zip1 v0.16b, v12.16b, v0.16b
|
||||
zip1 v2.16b, v12.16b, v2.16b
|
||||
zip1 v4.16b, v12.16b, v4.16b
|
||||
zip1 v6.16b, v12.16b, v6.16b
|
||||
|
||||
fmul v0.8h, v0.8h, v16.8h
|
||||
fmul v1.8h, v1.8h, v17.8h
|
||||
fmul v2.8h, v2.8h, v18.8h
|
||||
fmul v3.8h, v3.8h, v19.8h
|
||||
fmul v4.8h, v4.8h, v20.8h
|
||||
fmul v5.8h, v5.8h, v21.8h
|
||||
fmul v6.8h, v6.8h, v22.8h
|
||||
fmul v7.8h, v7.8h, v23.8h
|
||||
|
||||
st1 {{v0.16b-v3.16b}}, [{o}], #64
|
||||
st1 {{v4.16b-v7.16b}}, [{o}], #64
|
||||
|
||||
subs {k2}, {k2}, #1
|
||||
bne 3b
|
||||
|
||||
add {i}, {i}, 128 // skip scales
|
||||
subs {k}, {k}, 32
|
||||
bne 2b
|
||||
",
|
||||
lookup_table = in(reg) &lookup_table,
|
||||
k = inout(reg) k => _,
|
||||
k2 = out(reg) _,
|
||||
scales = out(reg) _,
|
||||
i = inout(reg) input => _,
|
||||
o = inout(reg) output => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,
|
||||
out("v8") _, out("v9") _, out("v10") _, out("v11") _,
|
||||
out("v12") _, out("v13") _, out("v14") _, out("v15") _,
|
||||
out("v16") _, out("v17") _, out("v18") _, out("v19") _,
|
||||
out("v20") _, out("v21") _, out("v22") _, out("v23") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
+62
@@ -0,0 +1,62 @@
|
||||
use crate::num_traits::Zero;
|
||||
use tract_data::half::f16;
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_sum_f16_32n,
|
||||
32,
|
||||
8,
|
||||
(),
|
||||
f16::zero(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &[f16], _: ()) -> f16 {
|
||||
assert!(buf.len() % 32 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(buf: &[f16]) -> f16 {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
let mut out: u16;
|
||||
std::arch::asm!("
|
||||
movi v0.8h, #0
|
||||
movi v1.8h, #0
|
||||
movi v2.8h, #0
|
||||
movi v3.8h, #0
|
||||
2:
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{ptr}], 64
|
||||
fadd v0.8h, v0.8h, v4.8h
|
||||
fadd v1.8h, v1.8h, v5.8h
|
||||
fadd v2.8h, v2.8h, v6.8h
|
||||
fadd v3.8h, v3.8h, v7.8h
|
||||
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
|
||||
fadd v0.8h, v0.8h, v1.8h
|
||||
fadd v2.8h, v2.8h, v3.8h
|
||||
fadd v0.8h, v0.8h, v2.8h
|
||||
faddp v0.8h, v0.8h, v0.8h
|
||||
faddp v0.8h, v0.8h, v0.8h
|
||||
faddp v0.8h, v0.8h, v0.8h
|
||||
",
|
||||
ptr = inout(reg) ptr => _,
|
||||
len = inout(reg) len => _,
|
||||
out("s0") out, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
f16::from_bits(out)
|
||||
}
|
||||
}
|
||||
unsafe { run(buf) }
|
||||
},
|
||||
#[inline(never)]
|
||||
fn reduce_two(a: f16, b: f16) -> f16 {
|
||||
a + b
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64fp16_sum_f16_32n {
|
||||
use super::*;
|
||||
crate::sum_frame_tests!(crate::arm64::has_fp16(), f16, arm64fp16_sum_f16_32n);
|
||||
}
|
||||
Vendored
+271
@@ -0,0 +1,271 @@
|
||||
use tract_data::half::f16;
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_unicast_mul_f16_32n,
|
||||
32,
|
||||
8,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 32 == 0);
|
||||
assert!(a.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(a: &mut [f16], b: &[f16]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}]
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{b_ptr}], 64
|
||||
fmul v0.8h, v0.8h, v4.8h
|
||||
fmul v1.8h, v1.8h, v5.8h
|
||||
fmul v2.8h, v2.8h, v6.8h
|
||||
fmul v3.8h, v3.8h, v7.8h
|
||||
st1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_unicast_add_f16_32n,
|
||||
32,
|
||||
8,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 32 == 0);
|
||||
assert!(a.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(a: &mut [f16], b: &[f16]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}]
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{b_ptr}], 64
|
||||
fadd v0.8h, v0.8h, v4.8h
|
||||
fadd v1.8h, v1.8h, v5.8h
|
||||
fadd v2.8h, v2.8h, v6.8h
|
||||
fadd v3.8h, v3.8h, v7.8h
|
||||
st1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_unicast_sub_f16_32n,
|
||||
32,
|
||||
8,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 32 == 0);
|
||||
assert!(a.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(a: &mut [f16], b: &[f16]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}]
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{b_ptr}], 64
|
||||
fsub v0.8h, v0.8h, v4.8h
|
||||
fsub v1.8h, v1.8h, v5.8h
|
||||
fsub v2.8h, v2.8h, v6.8h
|
||||
fsub v3.8h, v3.8h, v7.8h
|
||||
st1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_unicast_subf_f16_32n,
|
||||
32,
|
||||
8,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 32 == 0);
|
||||
assert!(a.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(a: &mut [f16], b: &[f16]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}]
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{b_ptr}], 64
|
||||
fsub v0.8h, v4.8h, v0.8h
|
||||
fsub v1.8h, v5.8h, v1.8h
|
||||
fsub v2.8h, v6.8h, v2.8h
|
||||
fsub v3.8h, v7.8h, v3.8h
|
||||
st1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_unicast_min_f16_32n,
|
||||
32,
|
||||
8,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 32 == 0);
|
||||
assert!(a.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(a: &mut [f16], b: &[f16]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}]
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{b_ptr}], 64
|
||||
fmin v0.8h, v0.8h, v4.8h
|
||||
fmin v1.8h, v1.8h, v5.8h
|
||||
fmin v2.8h, v2.8h, v6.8h
|
||||
fmin v3.8h, v3.8h, v7.8h
|
||||
st1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
arm64fp16_unicast_max_f16_32n,
|
||||
32,
|
||||
8,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 32 == 0);
|
||||
assert!(a.len() > 0);
|
||||
#[target_feature(enable = "fp16")]
|
||||
unsafe fn run(a: &mut [f16], b: &[f16]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}]
|
||||
ld1 {{v4.8h, v5.8h, v6.8h, v7.8h}}, [{b_ptr}], 64
|
||||
fmax v0.8h, v0.8h, v4.8h
|
||||
fmax v1.8h, v1.8h, v5.8h
|
||||
fmax v2.8h, v2.8h, v6.8h
|
||||
fmax v3.8h, v3.8h, v7.8h
|
||||
st1 {{v0.8h, v1.8h, v2.8h, v3.8h}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 32
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64fp16_unicast_mul_f16_32n {
|
||||
use super::*;
|
||||
use proptest::strategy::Strategy;
|
||||
crate::unicast_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_unicast_mul_f16_32n,
|
||||
|a, b| a * b
|
||||
);
|
||||
crate::unicast_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_unicast_add_f16_32n,
|
||||
|a, b| a + b
|
||||
);
|
||||
crate::unicast_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_unicast_sub_f16_32n,
|
||||
|a, b| a - b
|
||||
);
|
||||
crate::unicast_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_unicast_subf_f16_32n,
|
||||
|a, b| b - a
|
||||
);
|
||||
crate::unicast_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_unicast_min_f16_32n,
|
||||
|a, b| a.min(b)
|
||||
);
|
||||
crate::unicast_frame_tests!(
|
||||
crate::arm64::has_fp16(),
|
||||
f16,
|
||||
arm64fp16_unicast_max_f16_32n,
|
||||
|a, b| a.max(b)
|
||||
);
|
||||
}
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
mod by_scalar;
|
||||
mod gelu;
|
||||
mod gelu_fused;
|
||||
mod hardswish;
|
||||
mod leaky_relu;
|
||||
mod max;
|
||||
mod panel_extract;
|
||||
mod rms_norm;
|
||||
mod silu;
|
||||
mod silu_fused;
|
||||
mod softmax;
|
||||
mod sum;
|
||||
mod unicast;
|
||||
|
||||
pub use by_scalar::*;
|
||||
pub use gelu::arm64simd_gelu_f32_4n;
|
||||
pub use gelu_fused::arm64simd_gelu_f32_4n_fused;
|
||||
pub use hardswish::arm64simd_hardswish_f32_8n;
|
||||
pub use leaky_relu::arm64simd_leaky_relu_f32_8n;
|
||||
pub use max::arm64simd_max_f32_16n;
|
||||
pub use rms_norm::rms_norm_f32 as arm64simd_rms_norm_f32;
|
||||
pub use silu::arm64simd_silu_f32_4n;
|
||||
pub use silu_fused::arm64simd_silu_f32_4n_fused;
|
||||
pub use softmax::arm64simd_softmax2_fastcompact_f32_16n;
|
||||
pub use sum::arm64simd_sum_f32_16n;
|
||||
pub use unicast::*;
|
||||
|
||||
use crate::Ops;
|
||||
use crate::block_quant::{PackedBlockQuantFormat, Q4_0};
|
||||
use crate::frame::mmm::ImplementationQuality::ManuallyOptimized;
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
use super::Kind;
|
||||
|
||||
fn a55() -> isize {
|
||||
if *super::KIND == Kind::CortexA55 {
|
||||
1
|
||||
} else {
|
||||
-1
|
||||
}
|
||||
}
|
||||
|
||||
fn a53() -> isize {
|
||||
if *super::KIND == Kind::CortexA53 {
|
||||
1
|
||||
} else {
|
||||
-1
|
||||
}
|
||||
}
|
||||
|
||||
MMMExternKernel!(arm64simd_mmm_f32_8x8_a55 <f32>(8, 8)@(16, 16) quality(ManuallyOptimized) boost(a55));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_12x8_a55<f32>(12, 8)@(16, 16) quality(ManuallyOptimized) boost(a55));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_16x4_a55<f32>(16, 4)@(16, 16) quality(ManuallyOptimized) boost(a55));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_24x4_a55<f32>(24, 4)@(16, 16) quality(ManuallyOptimized) boost(a55));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_64x1_a55<f32>(64, 1)@(16, 16) quality(ManuallyOptimized) boost(a55));
|
||||
|
||||
MMMExternKernel!(arm64simd_mmm_f32_16x4_a53<f32>(16, 4)@(16, 16) quality(ManuallyOptimized) boost(a53));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_24x4_a53<f32>(24, 4)@(16, 16) quality(ManuallyOptimized) boost(a53));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_8x8_a53 <f32>(8, 8)@(16, 16) quality(ManuallyOptimized) boost(a53));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_12x8_a53<f32>(12, 8)@(16, 16) quality(ManuallyOptimized) boost(a53));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_64x1_a53<f32>(64, 1)@(16, 16) quality(ManuallyOptimized) boost(a53));
|
||||
|
||||
MMMExternKernel!(arm64simd_mmm_f32_16x4_gen<f32>(16, 4)@(16, 16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_24x4_gen<f32>(24, 4)@(16, 16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_8x8_gen <f32>(8, 8)@(16, 16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_12x8_gen<f32>(12, 8)@(16, 16) quality(ManuallyOptimized));
|
||||
MMMExternKernel!(arm64simd_mmm_f32_64x1_gen<f32>(64, 1)@(16, 16) quality(ManuallyOptimized));
|
||||
|
||||
fn q40p32z16se() -> PackedBlockQuantFormat {
|
||||
PackedBlockQuantFormat::new(&Q4_0, 32, 16, true)
|
||||
}
|
||||
|
||||
MMMExternKernel!(arm64simd_mmm_f32_32x1_gen<f32>(32, 1)@(16, 16)
|
||||
packing[1] = q40f16 => |k| k.with_packing(q40p32z16se(), f16::packing(1));
|
||||
packing[2] = q40f32 => |k| k.with_packing(q40p32z16se(), f32::packing(1));
|
||||
packing[3] = f16f16 => |k| k.with_packing(f16::packing(32), f16::packing(1));
|
||||
packing[4] = f32f16 => |k| k.with_packing(f32::packing(32), f16::packing(1));
|
||||
packing[5] = f16f32 => |k| k.with_packing(f16::packing(32), f32::packing(1));
|
||||
quality(ManuallyOptimized)
|
||||
store(f16)
|
||||
);
|
||||
|
||||
MMMExternKernel!(arm64simd_mmm_f32_32x3_gen<f32>(32, 3)@(16, 16)
|
||||
packing[1] = f32f16 => |k| k.with_packing(f32::packing(32), f16::packing(3));
|
||||
packing[2] = f16f32 => |k| k.with_packing(f16::packing(32), f32::packing(3));
|
||||
packing[3] = f16f16 => |k| k.with_packing(f16::packing(32), f16::packing(3));
|
||||
quality(ManuallyOptimized)
|
||||
store(f16)
|
||||
);
|
||||
|
||||
MMMExternKernel!(arm64simd_mmm_i32_8x8<i32>(8, 8)@(16, 16)
|
||||
packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 8, 16), PackedFormat::new(DatumType::I8, 8, 16));
|
||||
quality(ManuallyOptimized)
|
||||
store(i8)
|
||||
);
|
||||
|
||||
// SDOT (FEAT_DotProd) variant: 4-K reduction per instruction (~4x the SMLAL
|
||||
// 8x8 above). Uses the K=4-inner PackedI8K4 packing; identical v16..v31 tile
|
||||
// layout, so it reuses all the i32 fuse/store/q_scale machinery.
|
||||
//
|
||||
// Gated on `tract_arm64_dotprod` (set by build.rs when the assembler can encode
|
||||
// `sdot`; binutils < 2.30 cannot). On old toolchains the kernel is omitted and
|
||||
// dispatch falls back to the SMLAL 8x8 i32 kernel.
|
||||
#[cfg(tract_arm64_dotprod)]
|
||||
MMMExternKernel!(arm64simd_mmm_i32_8x8_dot<i32>(8, 8)@(16, 16)
|
||||
where(super::has_dotprod)
|
||||
packing[1] = i8i8 => |k| k.with_packing(crate::pack::PackedI8K4::new(8), crate::pack::PackedI8K4::new(8));
|
||||
quality(ManuallyOptimized)
|
||||
store(i8)
|
||||
);
|
||||
|
||||
MMMExternKernel!(arm64simd_mmm_i32_64x1<i32>(64, 1)@(16, 1)
|
||||
packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 64,16), PackedFormat::new(DatumType::I8, 1, 1));
|
||||
quality(ManuallyOptimized)
|
||||
store(i8)
|
||||
);
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
ops.mmm_impls.extend([
|
||||
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(),
|
||||
arm64simd_mmm_f32_32x1_gen.mmm(),
|
||||
arm64simd_mmm_f32_32x3_gen.mmm(),
|
||||
arm64simd_mmm_f32_64x1_gen.mmm(),
|
||||
arm64simd_mmm_f32_64x1_a53.mmm(),
|
||||
arm64simd_mmm_f32_64x1_a55.mmm(),
|
||||
arm64simd_mmm_i32_8x8.mmm(),
|
||||
arm64simd_mmm_i32_64x1.mmm(),
|
||||
]);
|
||||
#[cfg(tract_arm64_dotprod)]
|
||||
ops.mmm_impls.push(arm64simd_mmm_i32_8x8_dot.mmm());
|
||||
panel_extract::plug(ops);
|
||||
}
|
||||
|
||||
tanh_impl!(f32, arm64simd_tanh_f32_4n, 4, 4, true);
|
||||
sigmoid_impl!(f32, arm64simd_sigmoid_f32_4n, 4, 4, true);
|
||||
Vendored
+202
@@ -0,0 +1,202 @@
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_mul_by_scalar_f32_16n,
|
||||
16,
|
||||
4,
|
||||
f32,
|
||||
fn run(buf: &mut [f32], s: f32) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, v0.s[0]
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}]
|
||||
fmul v4.4s, v4.4s, v0.4s
|
||||
fmul v5.4s, v5.4s, v0.4s
|
||||
fmul v6.4s, v6.4s, v0.4s
|
||||
fmul v7.4s, v7.4s, v0.4s
|
||||
st1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_add_by_scalar_f32_16n,
|
||||
16,
|
||||
4,
|
||||
f32,
|
||||
fn run(buf: &mut [f32], s: f32) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, v0.s[0]
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}]
|
||||
fadd v4.4s, v4.4s, v0.4s
|
||||
fadd v5.4s, v5.4s, v0.4s
|
||||
fadd v6.4s, v6.4s, v0.4s
|
||||
fadd v7.4s, v7.4s, v0.4s
|
||||
st1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_sub_by_scalar_f32_16n,
|
||||
16,
|
||||
4,
|
||||
f32,
|
||||
fn run(buf: &mut [f32], s: f32) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, v0.s[0]
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}]
|
||||
fsub v4.4s, v4.4s, v0.4s
|
||||
fsub v5.4s, v5.4s, v0.4s
|
||||
fsub v6.4s, v6.4s, v0.4s
|
||||
fsub v7.4s, v7.4s, v0.4s
|
||||
st1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_subf_by_scalar_f32_16n,
|
||||
16,
|
||||
4,
|
||||
f32,
|
||||
fn run(buf: &mut [f32], s: f32) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, v0.s[0]
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}]
|
||||
fsub v4.4s, v0.4s, v4.4s
|
||||
fsub v5.4s, v0.4s, v5.4s
|
||||
fsub v6.4s, v0.4s, v6.4s
|
||||
fsub v7.4s, v0.4s, v7.4s
|
||||
st1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_min_by_scalar_f32_16n,
|
||||
16,
|
||||
4,
|
||||
f32,
|
||||
fn run(buf: &mut [f32], s: f32) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, v0.s[0]
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}]
|
||||
fmin v4.4s, v4.4s, v0.4s
|
||||
fmin v5.4s, v5.4s, v0.4s
|
||||
fmin v6.4s, v6.4s, v0.4s
|
||||
fmin v7.4s, v7.4s, v0.4s
|
||||
st1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_max_by_scalar_f32_16n,
|
||||
16,
|
||||
4,
|
||||
f32,
|
||||
fn run(buf: &mut [f32], s: f32) {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, v0.s[0]
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}]
|
||||
fmax v4.4s, v4.4s, v0.4s
|
||||
fmax v5.4s, v5.4s, v0.4s
|
||||
fmax v6.4s, v6.4s, v0.4s
|
||||
fmax v7.4s, v7.4s, v0.4s
|
||||
st1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
in("v0") s,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64simd_mul_by_scalar_f32_16n {
|
||||
use super::*;
|
||||
by_scalar_frame_tests!(true, f32, arm64simd_mul_by_scalar_f32_16n, |a, b| a * b);
|
||||
by_scalar_frame_tests!(true, f32, arm64simd_add_by_scalar_f32_16n, |a, b| a + b);
|
||||
by_scalar_frame_tests!(true, f32, arm64simd_sub_by_scalar_f32_16n, |a, b| a - b);
|
||||
by_scalar_frame_tests!(true, f32, arm64simd_subf_by_scalar_f32_16n, |a, b| b - a);
|
||||
by_scalar_frame_tests!(true, f32, arm64simd_min_by_scalar_f32_16n, |a, b| a.min(b));
|
||||
by_scalar_frame_tests!(true, f32, arm64simd_max_by_scalar_f32_16n, |a, b| a.max(b));
|
||||
}
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
// Tanh-form GELU (pow=3) matching tract's GeluApproximate:
|
||||
// gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
|
||||
//
|
||||
// Composed at the kernel level: save the original x, compute the tanh
|
||||
// argument, call tract's NEON tanh kernel in place, then finish with the
|
||||
// 0.5 * x * (1 + tanh) multiply. Chunked to keep the scratch buffer L1-resident.
|
||||
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_gelu_f32_4n,
|
||||
4,
|
||||
4,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
const SQRT_2_OVER_PI: f32 = 0.7978845608028654;
|
||||
const COEF: f32 = 0.044715;
|
||||
const CHUNK: usize = 256;
|
||||
let mut scratch = [0f32; CHUNK];
|
||||
let mut start = 0;
|
||||
while start < buf.len() {
|
||||
let end = (start + CHUNK).min(buf.len());
|
||||
let chunk = &mut buf[start..end];
|
||||
let n = chunk.len();
|
||||
// Save original x and pre-compute the tanh argument in place.
|
||||
for i in 0..n {
|
||||
let x = chunk[i];
|
||||
scratch[i] = x;
|
||||
chunk[i] = SQRT_2_OVER_PI * (x + COEF * x * x * x);
|
||||
}
|
||||
super::arm64simd_tanh_f32_4n::run(chunk, ());
|
||||
// chunk now holds tanh(arg). Combine with saved x.
|
||||
for i in 0..n {
|
||||
chunk[i] = 0.5 * scratch[i] * (1.0 + chunk[i]);
|
||||
}
|
||||
start = end;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_arm64simd_gelu_f32_4n {
|
||||
use super::*;
|
||||
gelu_frame_tests!(true, f32, arm64simd_gelu_f32_4n);
|
||||
}
|
||||
Vendored
+292
@@ -0,0 +1,292 @@
|
||||
// Fused GELU (tanh-form, pow=3):
|
||||
// gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
|
||||
//
|
||||
// loop4 (16 lanes per iter) + loop1 (4-lane tail). Clones the tanh Padé
|
||||
// polynomial from arm64simd_tanh_f32_4n.S.j2, with pre-tanh argument
|
||||
// computation up front and the final 0.5*x*(1+tanh) combined via fmla.
|
||||
// Single memory pass (load + store), no scratch buffer.
|
||||
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_gelu_f32_4n_fused,
|
||||
4,
|
||||
4,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
// Tanh Padé coefficients (matches arm64simd_tanh_f32_4n.S.j2) +
|
||||
// 3 GELU constants packed into the last vector lanes:
|
||||
// index 13: 0.5
|
||||
// index 14: sqrt(2/pi) ≈ 0.7978846
|
||||
// index 15: 0.044715 * sqrt(2/pi) ≈ 0.0356774
|
||||
static COEFFS: [f32; 16] = [
|
||||
-8.9,
|
||||
8.9,
|
||||
-8.488492677e-14,
|
||||
5.277853000e-11,
|
||||
-2.022500419e-8,
|
||||
0.00001115424833,
|
||||
0.003103950131,
|
||||
0.1308400453,
|
||||
0.9999999934,
|
||||
0.0002546136580,
|
||||
0.02449515379,
|
||||
0.4641733162,
|
||||
1.0,
|
||||
0.5,
|
||||
0.7978845608028654,
|
||||
0.03567739613,
|
||||
];
|
||||
|
||||
assert!(buf.len() % 4 == 0);
|
||||
if buf.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_mut_ptr();
|
||||
let coef_ptr = COEFFS.as_ptr();
|
||||
|
||||
// Register layout (loop4):
|
||||
// v0-v3: coefficients
|
||||
// v4: sqrt(2/pi) (broadcast of v3.s[2])
|
||||
// v5: tanh clamp low (-8.9, dup v0.s[0])
|
||||
// v6: tanh clamp high (8.9, dup v0.s[1])
|
||||
// v7: 0.5 (broadcast of v3.s[1])
|
||||
// v8-v11: 0.5 * original x (saved after load)
|
||||
// v16-v19: working (load -> pre_tanh -> clamped -> numerator)
|
||||
// v20-v23: x² for tanh polynomial
|
||||
// v24-v27: polynomial intermediates (denominator at end)
|
||||
// v28-v31: polynomial intermediates (also x³ temp before tanh)
|
||||
std::arch::asm!("
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [{coef}]
|
||||
dup v5.4s, v0.s[0]
|
||||
dup v6.4s, v0.s[1]
|
||||
dup v7.4s, v3.s[1]
|
||||
dup v4.4s, v3.s[2]
|
||||
|
||||
cmp {len}, #16
|
||||
blt 9f
|
||||
|
||||
1:
|
||||
ld1 {{ v16.4s, v17.4s, v18.4s, v19.4s }}, [{ptr}]
|
||||
|
||||
// Save 0.5 * original x into v8-v11
|
||||
fmul v8.4s, v16.4s, v7.4s
|
||||
fmul v9.4s, v17.4s, v7.4s
|
||||
fmul v10.4s, v18.4s, v7.4s
|
||||
fmul v11.4s, v19.4s, v7.4s
|
||||
|
||||
// Compute x^3 into v28-v31
|
||||
fmul v28.4s, v16.4s, v16.4s
|
||||
fmul v29.4s, v17.4s, v17.4s
|
||||
fmul v30.4s, v18.4s, v18.4s
|
||||
fmul v31.4s, v19.4s, v19.4s
|
||||
|
||||
fmul v28.4s, v28.4s, v16.4s
|
||||
fmul v29.4s, v29.4s, v17.4s
|
||||
fmul v30.4s, v30.4s, v18.4s
|
||||
fmul v31.4s, v31.4s, v19.4s
|
||||
|
||||
// pre_tanh = sqrt(2/pi)*x + 0.0356774 * x^3
|
||||
fmul v16.4s, v16.4s, v4.4s
|
||||
fmul v17.4s, v17.4s, v4.4s
|
||||
fmul v18.4s, v18.4s, v4.4s
|
||||
fmul v19.4s, v19.4s, v4.4s
|
||||
|
||||
fmla v16.4s, v28.4s, v3.s[3]
|
||||
fmla v17.4s, v29.4s, v3.s[3]
|
||||
fmla v18.4s, v30.4s, v3.s[3]
|
||||
fmla v19.4s, v31.4s, v3.s[3]
|
||||
|
||||
// Clamp pre_tanh argument
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmax v17.4s, v17.4s, v5.4s
|
||||
fmax v18.4s, v18.4s, v5.4s
|
||||
fmax v19.4s, v19.4s, v5.4s
|
||||
|
||||
fmin v16.4s, v16.4s, v6.4s
|
||||
fmin v17.4s, v17.4s, v6.4s
|
||||
fmin v18.4s, v18.4s, v6.4s
|
||||
fmin v19.4s, v19.4s, v6.4s
|
||||
|
||||
// Tanh Padé polynomial (cloned from arm64simd_tanh_f32_4n.S.j2)
|
||||
fmul v20.4s, v16.4s, v16.4s
|
||||
fmul v21.4s, v17.4s, v17.4s
|
||||
fmul v22.4s, v18.4s, v18.4s
|
||||
fmul v23.4s, v19.4s, v19.4s
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v25.4s, v0.s[3]
|
||||
fmla v25.4s, v21.4s, v0.s[2]
|
||||
dup v26.4s, v0.s[3]
|
||||
fmla v26.4s, v22.4s, v0.s[2]
|
||||
dup v27.4s, v0.s[3]
|
||||
fmla v27.4s, v23.4s, v0.s[2]
|
||||
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[1]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[1]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[1]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[2]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[2]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[2]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[3]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[3]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[3]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
fmul v16.4s, v16.4s, v28.4s
|
||||
fmul v17.4s, v17.4s, v29.4s
|
||||
fmul v18.4s, v18.4s, v30.4s
|
||||
fmul v19.4s, v19.4s, v31.4s
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v25.4s, v2.s[2]
|
||||
fmla v25.4s, v21.4s, v2.s[1]
|
||||
dup v26.4s, v2.s[2]
|
||||
fmla v26.4s, v22.4s, v2.s[1]
|
||||
dup v27.4s, v2.s[2]
|
||||
fmla v27.4s, v23.4s, v2.s[1]
|
||||
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[3]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[3]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[3]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v3.s[0]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v3.s[0]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v3.s[0]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
// tanh(pre_arg) = num/denom
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
fdiv v17.4s, v17.4s, v25.4s
|
||||
fdiv v18.4s, v18.4s, v26.4s
|
||||
fdiv v19.4s, v19.4s, v27.4s
|
||||
|
||||
// result = 0.5*x * (1 + tanh) = (0.5*x) + (0.5*x) * tanh
|
||||
fmla v8.4s, v8.4s, v16.4s
|
||||
fmla v9.4s, v9.4s, v17.4s
|
||||
fmla v10.4s, v10.4s, v18.4s
|
||||
fmla v11.4s, v11.4s, v19.4s
|
||||
|
||||
st1 {{ v8.4s, v9.4s, v10.4s, v11.4s }}, [{ptr}], #64
|
||||
sub {len}, {len}, #16
|
||||
cmp {len}, #16
|
||||
bge 1b
|
||||
|
||||
9:
|
||||
cbz {len}, 3f
|
||||
|
||||
2:
|
||||
ld1 {{ v16.4s }}, [{ptr}]
|
||||
fmul v8.4s, v16.4s, v7.4s
|
||||
|
||||
fmul v28.4s, v16.4s, v16.4s
|
||||
fmul v28.4s, v28.4s, v16.4s
|
||||
|
||||
fmul v16.4s, v16.4s, v4.4s
|
||||
fmla v16.4s, v28.4s, v3.s[3]
|
||||
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmin v16.4s, v16.4s, v6.4s
|
||||
|
||||
fmul v20.4s, v16.4s, v16.4s
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
fmul v16.4s, v16.4s, v28.4s
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
|
||||
fmla v8.4s, v8.4s, v16.4s
|
||||
|
||||
st1 {{ v8.4s }}, [{ptr}], #16
|
||||
subs {len}, {len}, 4
|
||||
bne 2b
|
||||
|
||||
3:
|
||||
",
|
||||
coef = in(reg) coef_ptr,
|
||||
ptr = inout(reg) ptr => _,
|
||||
len = inout(reg) len => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,
|
||||
out("v8") _, out("v9") _, out("v10") _, out("v11") _,
|
||||
out("v16") _, out("v17") _, out("v18") _, out("v19") _,
|
||||
out("v20") _, out("v21") _, out("v22") _, out("v23") _,
|
||||
out("v24") _, out("v25") _, out("v26") _, out("v27") _,
|
||||
out("v28") _, out("v29") _, out("v30") _, out("v31") _,
|
||||
options(nostack),
|
||||
);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_arm64simd_gelu_f32_4n_fused {
|
||||
use super::*;
|
||||
gelu_frame_tests!(true, f32, arm64simd_gelu_f32_4n_fused);
|
||||
}
|
||||
Vendored
+63
@@ -0,0 +1,63 @@
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_hardswish_f32_8n,
|
||||
8,
|
||||
4,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
assert!(buf.len() % 8 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, {three:v}.s[0]
|
||||
dup v1.4s, {six:v}.s[0]
|
||||
dup v2.4s, {inv6:v}.s[0]
|
||||
movi v3.4s, #0
|
||||
2:
|
||||
ldp q4, q5, [{ptr}]
|
||||
|
||||
fadd v6.4s, v4.4s, v0.4s
|
||||
fadd v7.4s, v5.4s, v0.4s
|
||||
|
||||
fmin v6.4s, v6.4s, v1.4s
|
||||
fmin v7.4s, v7.4s, v1.4s
|
||||
|
||||
fmax v6.4s, v6.4s, v3.4s
|
||||
fmax v7.4s, v7.4s, v3.4s
|
||||
|
||||
fmul v6.4s, v6.4s, v4.4s
|
||||
fmul v7.4s, v7.4s, v5.4s
|
||||
|
||||
fmul v6.4s, v6.4s, v2.4s
|
||||
fmul v7.4s, v7.4s, v2.4s
|
||||
|
||||
stp q6, q7, [{ptr}], #32
|
||||
subs {len}, {len}, 8
|
||||
bne 2b
|
||||
",
|
||||
three = in(vreg) 3.0f32,
|
||||
six = in(vreg) 6.0f32,
|
||||
inv6 = in(vreg) 1.0f32 / 6.0f32,
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
out("v0") _,
|
||||
out("v1") _,
|
||||
out("v2") _,
|
||||
out("v3") _,
|
||||
out("q4") _,
|
||||
out("q5") _,
|
||||
out("q6") _,
|
||||
out("q7") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_arm64simd_hardswish_f32_8n {
|
||||
use super::*;
|
||||
hardswish_frame_tests!(true, f32, arm64simd_hardswish_f32_8n);
|
||||
}
|
||||
Vendored
+50
@@ -0,0 +1,50 @@
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_leaky_relu_f32_8n,
|
||||
8,
|
||||
4,
|
||||
f32,
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], alpha: f32) {
|
||||
assert!(buf.len() % 8 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
dup v0.4s, {alpha:v}.s[0]
|
||||
dup v1.4s, {one:v}.s[0]
|
||||
2:
|
||||
ldp q3, q4, [{ptr}]
|
||||
|
||||
fcmgt v5.4s, v3.4s, #0.0
|
||||
fcmgt v6.4s, v4.4s, #0.0
|
||||
bsl v5.16b, v1.16b, v0.16b
|
||||
bsl v6.16b, v1.16b, v0.16b
|
||||
fmul v3.4s, v3.4s, v5.4s
|
||||
fmul v4.4s, v4.4s, v6.4s
|
||||
|
||||
stp q3, q4, [{ptr}], #32
|
||||
subs {len}, {len}, 8
|
||||
bne 2b
|
||||
",
|
||||
one = in(vreg) 1.0f32,
|
||||
alpha = in(vreg) alpha,
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
out("v0") _,
|
||||
out("v1") _,
|
||||
out("q3") _,
|
||||
out("q4") _,
|
||||
out("q5") _,
|
||||
out("q6") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_arm64simd_leaky_relu_f32_8n {
|
||||
use super::*;
|
||||
leaky_relu_frame_tests!(true, f32, arm64simd_leaky_relu_f32_8n);
|
||||
}
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
use std::arch::aarch64::{float32x4_t, vdupq_n_f32, vgetq_lane_f32};
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_max_f32_16n,
|
||||
16,
|
||||
4,
|
||||
(),
|
||||
f32::MIN,
|
||||
#[inline(never)]
|
||||
fn run(buf: &[f32], _: ()) -> f32 {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
let mut out: float32x4_t = vdupq_n_f32(f32::MIN);
|
||||
std::arch::asm!("
|
||||
and v1.16b, v0.16b, v0.16b
|
||||
and v2.16b, v0.16b, v0.16b
|
||||
and v3.16b, v0.16b, v0.16b
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
fmax v0.4s, v0.4s, v4.4s
|
||||
fmax v1.4s, v1.4s, v5.4s
|
||||
fmax v2.4s, v2.4s, v6.4s
|
||||
fmax v3.4s, v3.4s, v7.4s
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
fmax v0.4s, v0.4s, v1.4s
|
||||
fmax v2.4s, v2.4s, v3.4s
|
||||
fmax v0.4s, v0.4s, v2.4s
|
||||
fmaxv s0, v0.4s
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
inout("v0") out, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
vgetq_lane_f32(out, 0)
|
||||
}
|
||||
},
|
||||
#[inline(never)]
|
||||
fn reduce_two(a: f32, b: f32) -> f32 {
|
||||
a.max(b)
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64simd_max_f32_16n {
|
||||
use super::*;
|
||||
crate::max_frame_tests!(true, f32, arm64simd_max_f32_16n);
|
||||
}
|
||||
Vendored
+98
@@ -0,0 +1,98 @@
|
||||
use crate::Ops;
|
||||
use crate::pack::Packing;
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
ops.panel_extractors.push(packed_32_q40_to_f32.clone());
|
||||
}
|
||||
|
||||
panel_extractor!(kernel_packed_32_q40_to_f32 as packed_32_q40_to_f32(
|
||||
Box::new(super::q40p32z16se()),
|
||||
f32::packing(32).align(16)
|
||||
));
|
||||
|
||||
unsafe fn kernel_packed_32_q40_to_f32(input: *const u8, output: *mut u8, k: usize) {
|
||||
unsafe {
|
||||
if k == 0 {
|
||||
return;
|
||||
}
|
||||
let lookup_table: [u8; 16] = [
|
||||
0xc8, 0xc7, 0xc6, 0xc5, 0xc4, 0xc2, 0xc0, 0xbc, 0x00, 0x3c, 0x40, 0x42, 0x44, 0x45,
|
||||
0x46, 0x47,
|
||||
];
|
||||
std::arch::asm!("
|
||||
ld1 {{v13.16b}}, [{lookup_table}]
|
||||
movi v15.16b, 15
|
||||
eor v12.16b, v12.16b, v12.16b
|
||||
|
||||
2:
|
||||
add {scales}, {i}, 512 // scales at end: 32 (cols) * 32 (rows) / 2 (half byte)
|
||||
ld1 {{v0.8h-v3.8h}}, [{scales}]
|
||||
|
||||
fcvtl v16.4s, v0.4h
|
||||
fcvtl2 v17.4s, v0.8h
|
||||
fcvtl v18.4s, v1.4h
|
||||
fcvtl2 v19.4s, v1.8h
|
||||
fcvtl v20.4s, v2.4h
|
||||
fcvtl2 v21.4s, v2.8h
|
||||
fcvtl v22.4s, v3.4h
|
||||
fcvtl2 v23.4s, v3.8h
|
||||
|
||||
mov {k2}, 32
|
||||
3:
|
||||
ld1 {{ v9.16b }}, [{i}], #16
|
||||
|
||||
and v0.16b, v9.16b, v15.16b
|
||||
ushr v4.16b, v9.16b, 4
|
||||
|
||||
tbl v0.16b, {{ v13.16b }}, v0.16b
|
||||
tbl v4.16b, {{ v13.16b }}, v4.16b
|
||||
|
||||
zip2 v2.16b, v12.16b, v0.16b
|
||||
zip2 v6.16b, v12.16b, v4.16b
|
||||
|
||||
zip1 v0.16b, v12.16b, v0.16b
|
||||
zip1 v4.16b, v12.16b, v4.16b
|
||||
|
||||
fcvtl2 v1.4s, v0.8h
|
||||
fcvtl v0.4s, v0.4h
|
||||
fcvtl2 v3.4s, v2.8h
|
||||
fcvtl v2.4s, v2.4h
|
||||
fcvtl2 v5.4s, v4.8h
|
||||
fcvtl v4.4s, v4.4h
|
||||
fcvtl2 v7.4s, v6.8h
|
||||
fcvtl v6.4s, v6.4h
|
||||
|
||||
fmul v0.4s, v0.4s, v16.4s
|
||||
fmul v1.4s, v1.4s, v17.4s
|
||||
fmul v2.4s, v2.4s, v18.4s
|
||||
fmul v3.4s, v3.4s, v19.4s
|
||||
fmul v4.4s, v4.4s, v20.4s
|
||||
fmul v5.4s, v5.4s, v21.4s
|
||||
fmul v6.4s, v6.4s, v22.4s
|
||||
fmul v7.4s, v7.4s, v23.4s
|
||||
|
||||
st1 {{v0.16b-v3.16b}}, [{o}], #64
|
||||
st1 {{v4.16b-v7.16b}}, [{o}], #64
|
||||
|
||||
subs {k2}, {k2}, #1
|
||||
bne 3b
|
||||
|
||||
add {i}, {i}, 64 // skip scales
|
||||
subs {k}, {k}, 32
|
||||
bne 2b
|
||||
",
|
||||
lookup_table = in(reg) &lookup_table,
|
||||
k = inout(reg) k => _,
|
||||
k2 = out(reg) _,
|
||||
scales = out(reg) _,
|
||||
i = inout(reg) input => _,
|
||||
o = inout(reg) output => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,
|
||||
out("v8") _, out("v9") _, out("v10") _, out("v11") _,
|
||||
out("v12") _, out("v13") _, out("v14") _, out("v15") _,
|
||||
out("v16") _, out("v17") _, out("v18") _, out("v19") _,
|
||||
out("v20") _, out("v21") _, out("v22") _, out("v23") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
Vendored
+164
@@ -0,0 +1,164 @@
|
||||
// NEON (aarch64, 128-bit, 4 f32 lanes) fused row-wise RmsNorm.
|
||||
//
|
||||
// Mirrors the AVX-512 kernel structure from `x86_64_fma/rms_norm.rs`:
|
||||
//
|
||||
// Pass 1 (sum of squares): acc += x² over 4 v-registers (4 lanes each →
|
||||
// 16 f32 / iter), horizontal reduce to a scalar,
|
||||
// then rsqrt(mean + eps) in scalar.
|
||||
// Pass 2 (multiply-back): broadcast inv_std into v0, multiply each
|
||||
// 4-v-register chunk in place.
|
||||
// Scalar tail handles the (len % 16 != 0) remainder.
|
||||
//
|
||||
// Drops into `Ops::rms_norm_f32` (added by the parent PR) — the core-side
|
||||
// dispatcher in `core::ops::nn::RmsNorm::eval` is already arch-neutral and
|
||||
// will pick this up automatically.
|
||||
|
||||
#[target_feature(enable = "neon")]
|
||||
unsafe fn rms_norm_f32_inner(buf: &mut [f32], eps: f32) {
|
||||
use std::arch::aarch64::*;
|
||||
let n = buf.len();
|
||||
let chunks = n / 16;
|
||||
let tail_start = chunks * 16;
|
||||
let ptr = buf.as_mut_ptr();
|
||||
|
||||
// --- Pass 1: sum of squares ---
|
||||
let mut sum_sq: f32 = 0.0;
|
||||
if chunks > 0 {
|
||||
let p = ptr;
|
||||
let c = chunks;
|
||||
let mut sum_v: float32x4_t = vdupq_n_f32(0.0);
|
||||
unsafe {
|
||||
std::arch::asm!("
|
||||
movi v1.4s, 0
|
||||
movi v2.4s, 0
|
||||
movi v3.4s, 0
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{p}], 64
|
||||
fmla v0.4s, v4.4s, v4.4s
|
||||
fmla v1.4s, v5.4s, v5.4s
|
||||
fmla v2.4s, v6.4s, v6.4s
|
||||
fmla v3.4s, v7.4s, v7.4s
|
||||
subs {c}, {c}, 1
|
||||
bne 2b
|
||||
fadd v0.4s, v0.4s, v1.4s
|
||||
fadd v2.4s, v2.4s, v3.4s
|
||||
fadd v0.4s, v0.4s, v2.4s
|
||||
",
|
||||
p = inout(reg) p => _,
|
||||
c = inout(reg) c => _,
|
||||
inout("v0") sum_v,
|
||||
out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,
|
||||
);
|
||||
}
|
||||
// horizontal sum across the 4 surviving lanes
|
||||
sum_sq = vaddvq_f32(sum_v);
|
||||
}
|
||||
// scalar tail
|
||||
for i in tail_start..n {
|
||||
let x = unsafe { *buf.get_unchecked(i) };
|
||||
sum_sq += x * x;
|
||||
}
|
||||
|
||||
// --- Compute inv_std (scalar) ---
|
||||
let mean_sq = sum_sq / (n as f32);
|
||||
let inv_std = (mean_sq + eps).sqrt().recip();
|
||||
|
||||
// --- Pass 2: multiply by inv_std ---
|
||||
if chunks > 0 {
|
||||
let p = ptr;
|
||||
let c = chunks;
|
||||
let inv_v: float32x4_t = vdupq_n_f32(inv_std);
|
||||
unsafe {
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{p}]
|
||||
fmul v4.4s, v4.4s, v0.4s
|
||||
fmul v5.4s, v5.4s, v0.4s
|
||||
fmul v6.4s, v6.4s, v0.4s
|
||||
fmul v7.4s, v7.4s, v0.4s
|
||||
st1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{p}], 64
|
||||
subs {c}, {c}, 1
|
||||
bne 2b
|
||||
",
|
||||
p = inout(reg) p => _,
|
||||
c = inout(reg) c => _,
|
||||
in("v0") inv_v,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
for i in tail_start..n {
|
||||
unsafe {
|
||||
*buf.get_unchecked_mut(i) *= inv_std;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn rms_norm_f32(buf: &mut [f32], eps: f32) {
|
||||
if buf.is_empty() {
|
||||
return;
|
||||
}
|
||||
unsafe { rms_norm_f32_inner(buf, eps) }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn ref_rms_norm(buf: &mut [f32], eps: f32) {
|
||||
let n = buf.len() as f32;
|
||||
let sum_sq: f32 = buf.iter().map(|x| x * x).sum();
|
||||
let mean_sq = sum_sq / n;
|
||||
let inv_std = (mean_sq + eps).sqrt().recip();
|
||||
for x in buf.iter_mut() {
|
||||
*x *= inv_std;
|
||||
}
|
||||
}
|
||||
|
||||
fn close_enough(got: &[f32], want: &[f32], tol: f32) {
|
||||
assert_eq!(got.len(), want.len());
|
||||
for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() {
|
||||
let diff = (g - w).abs();
|
||||
assert!(diff <= tol, "lane {i}: got {g}, want {w}, diff {diff}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_reference_16() {
|
||||
// 16 = exactly one inner iteration, no tail.
|
||||
let mut x: Vec<f32> = (0..16).map(|i| (i as f32 * 0.13).sin() * 5.0).collect();
|
||||
let mut y = x.clone();
|
||||
rms_norm_f32(&mut x, 1e-5);
|
||||
ref_rms_norm(&mut y, 1e-5);
|
||||
close_enough(&x, &y, 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_reference_1024_with_tail() {
|
||||
// 1024 + 7 = exercises the scalar tail loop (len % 16 = 7).
|
||||
let n = 1024 + 7;
|
||||
let mut x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.07).cos() * 3.0).collect();
|
||||
let mut y = x.clone();
|
||||
rms_norm_f32(&mut x, 1e-5);
|
||||
ref_rms_norm(&mut y, 1e-5);
|
||||
close_enough(&x, &y, 1e-4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_reference_short_below_chunk() {
|
||||
// 8 elements — shorter than one NEON iteration; all scalar tail.
|
||||
let mut x: Vec<f32> = vec![0.5, -1.5, 2.5, -3.5, 0.0, 4.0, -4.0, 1.0];
|
||||
let mut y = x.clone();
|
||||
rms_norm_f32(&mut x, 1e-5);
|
||||
ref_rms_norm(&mut y, 1e-5);
|
||||
close_enough(&x, &y, 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_is_noop() {
|
||||
let mut x: Vec<f32> = vec![];
|
||||
rms_norm_f32(&mut x, 1e-5);
|
||||
assert!(x.is_empty());
|
||||
}
|
||||
}
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_silu_f32_4n,
|
||||
4,
|
||||
4,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
// SiLU(x) = x * sigmoid(x). Compose by saving the input chunk to a
|
||||
// stack scratch buffer, running tract's NEON sigmoid kernel in place,
|
||||
// then multiplying back by the saved original. Multiply loop
|
||||
// auto-vectorises on aarch64.
|
||||
const CHUNK: usize = 256;
|
||||
let mut scratch = [0f32; CHUNK];
|
||||
let mut start = 0;
|
||||
while start < buf.len() {
|
||||
let end = (start + CHUNK).min(buf.len());
|
||||
let chunk = &mut buf[start..end];
|
||||
let n = chunk.len();
|
||||
scratch[..n].copy_from_slice(chunk);
|
||||
super::arm64simd_sigmoid_f32_4n::run(chunk, ());
|
||||
for i in 0..n {
|
||||
chunk[i] *= scratch[i];
|
||||
}
|
||||
start = end;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_arm64simd_silu_f32_4n {
|
||||
use super::*;
|
||||
silu_frame_tests!(true, f32, arm64simd_silu_f32_4n);
|
||||
}
|
||||
Vendored
+246
@@ -0,0 +1,246 @@
|
||||
// Fused SiLU: x * sigmoid(x).
|
||||
// loop4 (16 lanes per iter) + loop1 (4-lane tail).
|
||||
// Clones the sigmoid Padé polynomial from arm64simd_sigmoid_f32_4n.S.j2,
|
||||
// with the input saved before clamp (in v8-v11) and multiplied back at the
|
||||
// end. Single memory pass (load + store), no scratch buffer.
|
||||
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_silu_f32_4n_fused,
|
||||
4,
|
||||
4,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
// Sigmoid Padé coefficients (matches arm64simd_sigmoid_f32_4n.S.j2).
|
||||
static COEFFS: [f32; 16] = [
|
||||
-18.6,
|
||||
18.6,
|
||||
-4.433153405e-18,
|
||||
1.169974371e-14,
|
||||
-1.875289645e-11,
|
||||
4.257889523e-8,
|
||||
0.00004811817576,
|
||||
0.008163842030,
|
||||
0.2499999971,
|
||||
3.922935744e-6,
|
||||
0.001524872358,
|
||||
0.1159886749,
|
||||
1.0,
|
||||
0.5,
|
||||
0.0,
|
||||
0.0,
|
||||
];
|
||||
|
||||
assert!(buf.len() % 4 == 0);
|
||||
if buf.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_mut_ptr();
|
||||
let coef_ptr = COEFFS.as_ptr();
|
||||
|
||||
std::arch::asm!("
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [{coef}]
|
||||
dup v5.4s, v0.s[0]
|
||||
dup v6.4s, v0.s[1]
|
||||
dup v7.4s, v3.s[1]
|
||||
|
||||
cmp {len}, #16
|
||||
blt 9f
|
||||
|
||||
1:
|
||||
ld1 {{ v16.4s, v17.4s, v18.4s, v19.4s }}, [{ptr}]
|
||||
|
||||
mov v8.16b, v16.16b
|
||||
mov v9.16b, v17.16b
|
||||
mov v10.16b, v18.16b
|
||||
mov v11.16b, v19.16b
|
||||
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmax v17.4s, v17.4s, v5.4s
|
||||
fmax v18.4s, v18.4s, v5.4s
|
||||
fmax v19.4s, v19.4s, v5.4s
|
||||
|
||||
fmin v16.4s, v16.4s, v6.4s
|
||||
fmin v17.4s, v17.4s, v6.4s
|
||||
fmin v18.4s, v18.4s, v6.4s
|
||||
fmin v19.4s, v19.4s, v6.4s
|
||||
|
||||
fmul v20.4s, v16.4s, v16.4s
|
||||
fmul v21.4s, v17.4s, v17.4s
|
||||
fmul v22.4s, v18.4s, v18.4s
|
||||
fmul v23.4s, v19.4s, v19.4s
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v25.4s, v0.s[3]
|
||||
fmla v25.4s, v21.4s, v0.s[2]
|
||||
dup v26.4s, v0.s[3]
|
||||
fmla v26.4s, v22.4s, v0.s[2]
|
||||
dup v27.4s, v0.s[3]
|
||||
fmla v27.4s, v23.4s, v0.s[2]
|
||||
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[1]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[1]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[1]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[2]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[2]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[2]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[3]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[3]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[3]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
fmul v16.4s, v16.4s, v28.4s
|
||||
fmul v17.4s, v17.4s, v29.4s
|
||||
fmul v18.4s, v18.4s, v30.4s
|
||||
fmul v19.4s, v19.4s, v31.4s
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v25.4s, v2.s[2]
|
||||
fmla v25.4s, v21.4s, v2.s[1]
|
||||
dup v26.4s, v2.s[2]
|
||||
fmla v26.4s, v22.4s, v2.s[1]
|
||||
dup v27.4s, v2.s[2]
|
||||
fmla v27.4s, v23.4s, v2.s[1]
|
||||
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[3]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[3]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[3]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v3.s[0]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v3.s[0]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v3.s[0]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
fdiv v17.4s, v17.4s, v25.4s
|
||||
fdiv v18.4s, v18.4s, v26.4s
|
||||
fdiv v19.4s, v19.4s, v27.4s
|
||||
|
||||
fadd v16.4s, v16.4s, v7.4s
|
||||
fadd v17.4s, v17.4s, v7.4s
|
||||
fadd v18.4s, v18.4s, v7.4s
|
||||
fadd v19.4s, v19.4s, v7.4s
|
||||
|
||||
fmul v16.4s, v16.4s, v8.4s
|
||||
fmul v17.4s, v17.4s, v9.4s
|
||||
fmul v18.4s, v18.4s, v10.4s
|
||||
fmul v19.4s, v19.4s, v11.4s
|
||||
|
||||
st1 {{ v16.4s, v17.4s, v18.4s, v19.4s }}, [{ptr}], #64
|
||||
sub {len}, {len}, #16
|
||||
cmp {len}, #16
|
||||
bge 1b
|
||||
|
||||
9:
|
||||
cbz {len}, 3f
|
||||
|
||||
2:
|
||||
ld1 {{ v16.4s }}, [{ptr}]
|
||||
mov v8.16b, v16.16b
|
||||
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmin v16.4s, v16.4s, v6.4s
|
||||
fmul v20.4s, v16.4s, v16.4s
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
fmul v16.4s, v16.4s, v28.4s
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
fadd v16.4s, v16.4s, v7.4s
|
||||
|
||||
fmul v16.4s, v16.4s, v8.4s
|
||||
|
||||
st1 {{ v16.4s }}, [{ptr}], #16
|
||||
subs {len}, {len}, 4
|
||||
bne 2b
|
||||
|
||||
3:
|
||||
",
|
||||
coef = in(reg) coef_ptr,
|
||||
ptr = inout(reg) ptr => _,
|
||||
len = inout(reg) len => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v5") _, out("v6") _, out("v7") _,
|
||||
out("v8") _, out("v9") _, out("v10") _, out("v11") _,
|
||||
out("v16") _, out("v17") _, out("v18") _, out("v19") _,
|
||||
out("v20") _, out("v21") _, out("v22") _, out("v23") _,
|
||||
out("v24") _, out("v25") _, out("v26") _, out("v27") _,
|
||||
out("v28") _, out("v29") _, out("v30") _, out("v31") _,
|
||||
options(nostack),
|
||||
);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_arm64simd_silu_f32_4n_fused {
|
||||
use super::*;
|
||||
silu_frame_tests!(true, f32, arm64simd_silu_f32_4n_fused);
|
||||
}
|
||||
Vendored
+110
@@ -0,0 +1,110 @@
|
||||
map_reduce_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_softmax2_fastcompact_f32_16n,
|
||||
16,
|
||||
4,
|
||||
f32,
|
||||
f32::MIN,
|
||||
0f32,
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], max: f32) -> f32 {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
let mut acc;
|
||||
const MLN2: f32 = 0.6931471805f32;
|
||||
const A: f32 = 8388608.0f32;
|
||||
const B: f32 = 1065353216.0f32;
|
||||
const C: f32 = 60801.0f32;
|
||||
const SLOPE: f32 = A / MLN2;
|
||||
const OFFSET: f32 = B - C;
|
||||
unsafe {
|
||||
std::arch::asm!("
|
||||
// v0-v3 sum acc
|
||||
eor v0.16b, v0.16b, v0.16b
|
||||
eor v1.16b, v1.16b, v1.16b
|
||||
eor v2.16b, v2.16b, v2.16b
|
||||
eor v3.16b, v3.16b, v3.16b
|
||||
|
||||
dup v4.4s, v4.s[0] // max
|
||||
dup v5.4s, v5.s[0] // slope
|
||||
dup v6.4s, v6.s[0] // offset
|
||||
eor v7.16b, v7.16b, v7.16b // zero for max
|
||||
2:
|
||||
ld1 {{v8.4s, v9.4s, v10.4s, v11.4s}}, [{ptr}]
|
||||
|
||||
fsub v8.4s, v8.4s, v4.4s
|
||||
fsub v9.4s, v9.4s, v4.4s
|
||||
fsub v10.4s, v10.4s, v4.4s
|
||||
fsub v11.4s, v11.4s, v4.4s
|
||||
|
||||
fmul v8.4s, v8.4s, v5.4s
|
||||
fmul v9.4s, v9.4s, v5.4s
|
||||
fmul v10.4s, v10.4s, v5.4s
|
||||
fmul v11.4s, v11.4s, v5.4s
|
||||
|
||||
fadd v8.4s, v8.4s, v6.4s
|
||||
fadd v9.4s, v9.4s, v6.4s
|
||||
fadd v10.4s, v10.4s, v6.4s
|
||||
fadd v11.4s, v11.4s, v6.4s
|
||||
|
||||
fmax v8.4s, v8.4s, v7.4s
|
||||
fmax v9.4s, v9.4s, v7.4s
|
||||
fmax v10.4s, v10.4s, v7.4s
|
||||
fmax v11.4s, v11.4s, v7.4s
|
||||
|
||||
fcvtnu v8.4s, v8.4s
|
||||
fcvtnu v9.4s, v9.4s
|
||||
fcvtnu v10.4s, v10.4s
|
||||
fcvtnu v11.4s, v11.4s
|
||||
|
||||
fadd v0.4s, v0.4s, v8.4s
|
||||
fadd v1.4s, v1.4s, v9.4s
|
||||
fadd v2.4s, v2.4s, v10.4s
|
||||
fadd v3.4s, v3.4s, v11.4s
|
||||
|
||||
st1 {{v8.4s, v9.4s, v10.4s, v11.4s}}, [{ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
|
||||
fadd v0.4s, v0.4s, v1.4s
|
||||
fadd v2.4s, v2.4s, v3.4s
|
||||
fadd v0.4s, v0.4s, v2.4s
|
||||
|
||||
ext v1.16b, v0.16b, v0.16b, 4
|
||||
ext v2.16b, v0.16b, v0.16b, 8
|
||||
ext v3.16b, v0.16b, v0.16b, 12
|
||||
fadd v0.4s, v0.4s, v1.4s
|
||||
fadd v2.4s, v2.4s, v3.4s
|
||||
fadd v0.4s, v0.4s, v2.4s
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
out("v0") acc,
|
||||
out("v1") _,
|
||||
out("v2") _,
|
||||
out("v3") _,
|
||||
inout("v4") max => _,
|
||||
inout("v5") SLOPE => _,
|
||||
inout("v6") OFFSET => _,
|
||||
out("v7") _,
|
||||
out("v8") _,
|
||||
out("v9") _,
|
||||
out("v10") _,
|
||||
out("v11") _,
|
||||
);
|
||||
}
|
||||
acc
|
||||
},
|
||||
#[inline(never)]
|
||||
fn reduce_two(a: f32, b: f32) -> f32 {
|
||||
a + b
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64simd_softmax2_fastcompact_f32_16n {
|
||||
use super::*;
|
||||
crate::softmax_l2_frame_tests!(true, f32, arm64simd_softmax2_fastcompact_f32_16n);
|
||||
}
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
use crate::num_traits::Zero;
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_sum_f32_16n,
|
||||
16,
|
||||
4,
|
||||
(),
|
||||
f32::zero(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &[f32], _: ()) -> f32 {
|
||||
assert!(buf.len() % 16 == 0);
|
||||
assert!(buf.len() > 0);
|
||||
unsafe fn run(buf: &[f32]) -> f32 {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
let mut out: u32;
|
||||
std::arch::asm!("
|
||||
movi v0.4s, #0
|
||||
movi v1.4s, #0
|
||||
movi v2.4s, #0
|
||||
movi v3.4s, #0
|
||||
2:
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{ptr}], 64
|
||||
fadd v0.4s, v0.4s, v4.4s
|
||||
fadd v1.4s, v1.4s, v5.4s
|
||||
fadd v2.4s, v2.4s, v6.4s
|
||||
fadd v3.4s, v3.4s, v7.4s
|
||||
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
|
||||
fadd v0.4s, v0.4s, v1.4s
|
||||
fadd v2.4s, v2.4s, v3.4s
|
||||
fadd v0.4s, v0.4s, v2.4s
|
||||
faddp v0.4s, v0.4s, v0.4s
|
||||
faddp v0.4s, v0.4s, v0.4s
|
||||
",
|
||||
ptr = inout(reg) ptr => _,
|
||||
len = inout(reg) len => _,
|
||||
out("s0") out, out("v1") _, out("v2") _, out("v3") _,
|
||||
out("v4") _, out("v5") _, out("v6") _, out("v7") _,);
|
||||
f32::from_bits(out)
|
||||
}
|
||||
}
|
||||
unsafe { run(buf) }
|
||||
},
|
||||
#[inline(never)]
|
||||
fn reduce_two(a: f32, b: f32) -> f32 {
|
||||
a + b
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64simd_sum_f32_16n {
|
||||
use super::*;
|
||||
crate::sum_frame_tests!(true, f32, arm64simd_sum_f32_16n);
|
||||
}
|
||||
Vendored
+233
@@ -0,0 +1,233 @@
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_unicast_mul_f32_16n,
|
||||
16,
|
||||
4,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 16 == 0);
|
||||
assert!(a.len() > 0);
|
||||
unsafe fn run(a: &mut [f32], b: &[f32]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}]
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{b_ptr}], 64
|
||||
fmul v0.4s, v0.4s, v4.4s
|
||||
fmul v1.4s, v1.4s, v5.4s
|
||||
fmul v2.4s, v2.4s, v6.4s
|
||||
fmul v3.4s, v3.4s, v7.4s
|
||||
st1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_unicast_add_f32_16n,
|
||||
16,
|
||||
4,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 16 == 0);
|
||||
assert!(a.len() > 0);
|
||||
unsafe fn run(a: &mut [f32], b: &[f32]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}]
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{b_ptr}], 64
|
||||
fadd v0.4s, v0.4s, v4.4s
|
||||
fadd v1.4s, v1.4s, v5.4s
|
||||
fadd v2.4s, v2.4s, v6.4s
|
||||
fadd v3.4s, v3.4s, v7.4s
|
||||
st1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_unicast_sub_f32_16n,
|
||||
16,
|
||||
4,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 16 == 0);
|
||||
assert!(a.len() > 0);
|
||||
unsafe fn run(a: &mut [f32], b: &[f32]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}]
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{b_ptr}], 64
|
||||
fsub v0.4s, v0.4s, v4.4s
|
||||
fsub v1.4s, v1.4s, v5.4s
|
||||
fsub v2.4s, v2.4s, v6.4s
|
||||
fsub v3.4s, v3.4s, v7.4s
|
||||
st1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_unicast_subf_f32_16n,
|
||||
16,
|
||||
4,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 16 == 0);
|
||||
assert!(a.len() > 0);
|
||||
unsafe fn run(a: &mut [f32], b: &[f32]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}]
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{b_ptr}], 64
|
||||
fsub v0.4s, v4.4s, v0.4s
|
||||
fsub v1.4s, v5.4s, v1.4s
|
||||
fsub v2.4s, v6.4s, v2.4s
|
||||
fsub v3.4s, v7.4s, v3.4s
|
||||
st1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_unicast_max_f32_16n,
|
||||
16,
|
||||
4,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 16 == 0);
|
||||
assert!(a.len() > 0);
|
||||
unsafe fn run(a: &mut [f32], b: &[f32]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}]
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{b_ptr}], 64
|
||||
fmax v0.4s, v0.4s, v4.4s
|
||||
fmax v1.4s, v1.4s, v5.4s
|
||||
fmax v2.4s, v2.4s, v6.4s
|
||||
fmax v3.4s, v3.4s, v7.4s
|
||||
st1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
arm64simd_unicast_min_f32_16n,
|
||||
16,
|
||||
4,
|
||||
#[inline(never)]
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
assert!(a.len() == b.len());
|
||||
assert!(a.len() % 16 == 0);
|
||||
assert!(a.len() > 0);
|
||||
unsafe fn run(a: &mut [f32], b: &[f32]) {
|
||||
unsafe {
|
||||
let len = a.len();
|
||||
let a_ptr = a.as_ptr();
|
||||
let b_ptr = b.as_ptr();
|
||||
std::arch::asm!("
|
||||
2:
|
||||
ld1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}]
|
||||
ld1 {{v4.4s, v5.4s, v6.4s, v7.4s}}, [{b_ptr}], 64
|
||||
fmin v0.4s, v0.4s, v4.4s
|
||||
fmin v1.4s, v1.4s, v5.4s
|
||||
fmin v2.4s, v2.4s, v6.4s
|
||||
fmin v3.4s, v3.4s, v7.4s
|
||||
st1 {{v0.4s, v1.4s, v2.4s, v3.4s}}, [{a_ptr}], 64
|
||||
subs {len}, {len}, 16
|
||||
bne 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
a_ptr = inout(reg) a_ptr => _,
|
||||
b_ptr = inout(reg) b_ptr => _,
|
||||
out("v0") _, out("v1") _, out("v2") _, out("v3") _,);
|
||||
}
|
||||
}
|
||||
unsafe { run(a, b) }
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_arm64simd_unicast_mul_f32_16n {
|
||||
use super::*;
|
||||
use proptest::strategy::Strategy;
|
||||
crate::unicast_frame_tests!(true, f32, arm64simd_unicast_mul_f32_16n, |a, b| a * b);
|
||||
crate::unicast_frame_tests!(true, f32, arm64simd_unicast_add_f32_16n, |a, b| a + b);
|
||||
crate::unicast_frame_tests!(true, f32, arm64simd_unicast_sub_f32_16n, |a, b| a - b);
|
||||
crate::unicast_frame_tests!(true, f32, arm64simd_unicast_subf_f32_16n, |a, b| b - a);
|
||||
crate::unicast_frame_tests!(true, f32, arm64simd_unicast_min_f32_16n, |a, b| a.min(b));
|
||||
crate::unicast_frame_tests!(true, f32, arm64simd_unicast_max_f32_16n, |a, b| a.max(b));
|
||||
}
|
||||
+1197
File diff suppressed because it is too large
Load Diff
+1365
File diff suppressed because it is too large
Load Diff
+4
@@ -0,0 +1,4 @@
|
||||
use crate::frame::mmm::cost_model::CostModel;
|
||||
pub fn models() -> Vec<(&'static str, CostModel<'static>)> {
|
||||
vec![]
|
||||
}
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
use crate::frame::mmm::cost_model::CostModel;
|
||||
pub fn models() -> Vec<(&'static str, CostModel<'static>)> {
|
||||
vec![]
|
||||
}
|
||||
+309
@@ -0,0 +1,309 @@
|
||||
use crate::Ops;
|
||||
use crate::frame::mmm::ImplementationQuality::ManuallyOptimized;
|
||||
use crate::mmm::*;
|
||||
|
||||
// CAN_FUSE: everything except LeakyRelu / QScale / RoundingShiftRight /
|
||||
// ShiftLeft. LoadTile, AddUnicast, AddRowColProducts, per-row/col/scalar
|
||||
// arithmetic, Clear, Store, AddMatMul are all in. (Matches AMX
|
||||
// `apple_amx.rs` CAN_FUSE, minus the i32-only quantization ops.)
|
||||
const CAN_FUSE: fn(&FusedSpec) -> bool = |f| {
|
||||
!matches!(
|
||||
f,
|
||||
FusedSpec::LeakyRelu(_)
|
||||
| FusedSpec::QScale(_, _, _)
|
||||
| FusedSpec::RoundingShiftRight(_, _)
|
||||
| FusedSpec::ShiftLeft(_)
|
||||
)
|
||||
};
|
||||
|
||||
const SME: fn() -> bool = has_sme;
|
||||
const SME2: fn() -> bool = has_sme2;
|
||||
// The SMOPA i32 kernel implements the quant fuse ops (QScale / RoundingShiftRight
|
||||
// / ShiftLeft) bit-exactly; only LeakyRelu is unsupported (kernel returns 1).
|
||||
const CAN_FUSE_I32: fn(&FusedSpec) -> bool = |f| !matches!(f, FusedSpec::LeakyRelu(_));
|
||||
|
||||
MMMExternKernel!(sme_qmmm_i32_32x32<i32>(32,32)@(128,128) where(SME2) can_fuse(CAN_FUSE_I32)
|
||||
packing[1] = i8i8 => |k| k.with_packing(crate::pack::PackedI8K4::new(32), crate::pack::PackedI8K4::new(32));
|
||||
quality(ManuallyOptimized) store(i8));
|
||||
|
||||
// Streaming vector length in bytes, read via `RDSVL x0, #1` (encoding
|
||||
// 0x04bf5820). RDSVL is legal in non-streaming mode, but is UNDEFINED
|
||||
// unless FEAT_SME is implemented — callers MUST confirm FEAT_SME first
|
||||
// (sysctl on macOS, HWCAP2 on Linux) or this SIGILLs.
|
||||
#[cfg(any(target_os = "macos", target_os = "linux"))]
|
||||
unsafe fn streaming_vector_bytes() -> u64 {
|
||||
let svl: u64;
|
||||
unsafe {
|
||||
std::arch::asm!(
|
||||
".inst 0x04bf5820", // rdsvl x0, #1
|
||||
out("x0") svl,
|
||||
options(nomem, nostack, preserves_flags),
|
||||
);
|
||||
}
|
||||
svl
|
||||
}
|
||||
|
||||
// Our SME kernels hardcode a 512-bit streaming vector length (16 f32 lanes
|
||||
// per ZA.S slice — the 32x32 and 64x1 tile geometries depend on it). A host
|
||||
// that advertises FEAT_SME with a different SVL would run the kernels with
|
||||
// mismatched geometry and produce silently-wrong results. The prime offender
|
||||
// is qemu-aarch64 user-mode emulation, which sets HWCAP2_SME / HWCAP2_SME2
|
||||
// but uses a non-512 SVL — that is exactly what makes the cross-compiled
|
||||
// aarch64 CI jobs (run under QEMU) fail. Reject any non-512 SVL here so we
|
||||
// fall back to the portable path. MUST only be called once FEAT_SME is known
|
||||
// present.
|
||||
#[cfg(any(target_os = "macos", target_os = "linux"))]
|
||||
fn sme_geometry_supported() -> bool {
|
||||
// SVL = 512 bits = 64 bytes.
|
||||
unsafe { streaming_vector_bytes() == 64 }
|
||||
}
|
||||
|
||||
MMMExternKernel!(
|
||||
sme_mmm_f32_32x32<f32>(32, 32)@(128, 128)
|
||||
where(SME)
|
||||
can_fuse(CAN_FUSE)
|
||||
quality(ManuallyOptimized)
|
||||
);
|
||||
|
||||
MMMExternKernel!(
|
||||
sme_mmv_f32_64x1<f32>(64, 1)@(128, 128)
|
||||
where(SME2)
|
||||
can_fuse(CAN_FUSE)
|
||||
quality(ManuallyOptimized)
|
||||
);
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
pub fn has_sme() -> bool {
|
||||
// TRACT_SME_DISABLE=1 forces the SME path off so callers can A/B
|
||||
// against the AMX path on the same binary.
|
||||
if std::env::var_os("TRACT_SME_DISABLE").is_some() {
|
||||
return false;
|
||||
}
|
||||
// hw.optional.arm.FEAT_SME is an INTEGER sysctl, not a string. The
|
||||
// generic apple_get_syscall reads bytes-as-C-string which fails here
|
||||
// (`\x01\x00\x00\x00` would compare against the ASCII "1"), so we
|
||||
// read it as a u64 directly.
|
||||
use std::ffi::{CString, c_char, c_int, c_void};
|
||||
use std::ptr::null_mut;
|
||||
unsafe extern "C" {
|
||||
fn sysctlbyname(
|
||||
name: *const c_char,
|
||||
oldp: *mut c_void,
|
||||
oldlenp: *mut usize,
|
||||
newp: *mut c_void,
|
||||
newlen: usize,
|
||||
) -> c_int;
|
||||
}
|
||||
let Ok(name) = CString::new("hw.optional.arm.FEAT_SME") else {
|
||||
return false;
|
||||
};
|
||||
let mut value: u64 = 0;
|
||||
let mut len: usize = std::mem::size_of::<u64>();
|
||||
unsafe {
|
||||
if sysctlbyname(
|
||||
name.as_ptr(),
|
||||
&mut value as *mut _ as *mut c_void,
|
||||
&mut len,
|
||||
null_mut(),
|
||||
0,
|
||||
) != 0
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// FEAT_SME present AND the streaming vector length matches our kernels'
|
||||
// hardcoded 512-bit geometry.
|
||||
value != 0 && sme_geometry_supported()
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
pub fn has_sme() -> bool {
|
||||
// HWCAP2_SME = 1 << 23 on aarch64 (kernel ABI).
|
||||
const HWCAP2_SME: u64 = 1 << 23;
|
||||
unsafe extern "C" {
|
||||
fn getauxval(t: u64) -> u64;
|
||||
}
|
||||
const AT_HWCAP2: u64 = 26;
|
||||
let feat = unsafe { (getauxval(AT_HWCAP2) & HWCAP2_SME) != 0 };
|
||||
// FEAT_SME present AND the streaming vector length matches our kernels'
|
||||
// hardcoded 512-bit geometry (rejects qemu-user, which advertises SME
|
||||
// with a non-512 SVL — the cause of the cross-compiled CI failures).
|
||||
feat && sme_geometry_supported()
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
|
||||
pub fn has_sme() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
pub fn has_sme2() -> bool {
|
||||
// TRACT_SME_DISABLE=1 disables both SME and SME2 dispatch on the same
|
||||
// binary so end users can A/B the entire SME backend.
|
||||
if std::env::var_os("TRACT_SME_DISABLE").is_some() {
|
||||
return false;
|
||||
}
|
||||
use std::ffi::{CString, c_char, c_int, c_void};
|
||||
use std::ptr::null_mut;
|
||||
unsafe extern "C" {
|
||||
fn sysctlbyname(
|
||||
name: *const c_char,
|
||||
oldp: *mut c_void,
|
||||
oldlenp: *mut usize,
|
||||
newp: *mut c_void,
|
||||
newlen: usize,
|
||||
) -> c_int;
|
||||
}
|
||||
let Ok(name) = CString::new("hw.optional.arm.FEAT_SME2") else {
|
||||
return false;
|
||||
};
|
||||
let mut value: u64 = 0;
|
||||
let mut len: usize = std::mem::size_of::<u64>();
|
||||
unsafe {
|
||||
if sysctlbyname(
|
||||
name.as_ptr(),
|
||||
&mut value as *mut _ as *mut c_void,
|
||||
&mut len,
|
||||
null_mut(),
|
||||
0,
|
||||
) != 0
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// FEAT_SME2 present AND the streaming vector length matches our kernels'
|
||||
// hardcoded 512-bit geometry.
|
||||
value != 0 && sme_geometry_supported()
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
pub fn has_sme2() -> bool {
|
||||
// HWCAP2_SME2 = 1 << 37 on aarch64 (kernel ABI).
|
||||
const HWCAP2_SME2: u64 = 1 << 37;
|
||||
unsafe extern "C" {
|
||||
fn getauxval(t: u64) -> u64;
|
||||
}
|
||||
const AT_HWCAP2: u64 = 26;
|
||||
let feat = unsafe { (getauxval(AT_HWCAP2) & HWCAP2_SME2) != 0 };
|
||||
// FEAT_SME2 present AND the streaming vector length matches our kernels'
|
||||
// hardcoded 512-bit geometry (rejects qemu-user, which advertises SME2
|
||||
// with a non-512 SVL — the cause of the cross-compiled CI failures).
|
||||
feat && sme_geometry_supported()
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
|
||||
pub fn has_sme2() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
if has_sme() {
|
||||
log::info!("SME optimisation activated");
|
||||
ops.mmm_f32 = Box::new(|_, _, _| sme_mmm_f32_32x32.mmm());
|
||||
ops.mmm_impls.extend_from_slice(&[sme_mmm_f32_32x32.mmm()]);
|
||||
}
|
||||
if has_sme2() {
|
||||
log::info!("SME2 GEMV optimisation activated");
|
||||
ops.mmv_f32 = Box::new(|_, _| sme_mmv_f32_64x1.mmm());
|
||||
ops.qmmm_i32 = Box::new(|_, _, _| sme_qmmm_i32_32x32.mmm());
|
||||
ops.mmm_impls
|
||||
.extend_from_slice(&[sme_mmv_f32_64x1.mmm(), sme_qmmm_i32_32x32.mmm()]);
|
||||
}
|
||||
if !has_sme() && !has_sme2() {
|
||||
log::info!("No SME optimisation");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::frame::mmm::tests::packed_packed::PackedPackedProblem;
|
||||
use tract_data::internal::Approximation;
|
||||
|
||||
// Phase 1A correctness: AddMatMul + Clear + Store + Done on a few
|
||||
// shapes. Bypasses auto-tests (SME_OFF) by calling run/reference
|
||||
// directly. Skipped if hardware lacks SME.
|
||||
fn check_shape(m_tile: usize, k: usize, n_tile: usize) {
|
||||
const MR: usize = 32;
|
||||
const NR: usize = 32;
|
||||
let m = m_tile * MR;
|
||||
let n = n_tile * NR;
|
||||
let a: Vec<f32> = (0..m * k).map(|i| (i as f32 * 0.013) - 1.5).collect();
|
||||
let b: Vec<f32> = (0..k * n).map(|i| (i as f32 * 0.017) + 0.25).collect();
|
||||
let pb = PackedPackedProblem::kernel(&*sme_mmm_f32_32x32, 0, a, b);
|
||||
let expected = pb.reference().expect("scalar reference");
|
||||
let found = pb.run().expect("SME kernel run");
|
||||
found
|
||||
.close_enough(&expected, Approximation::Approximate)
|
||||
.unwrap_or_else(|e| panic!("SME mmm mismatch at k={k}: {e}"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sme_mmm_f32_32x32_k1() {
|
||||
if !has_sme() {
|
||||
eprintln!("SME not present, skipping");
|
||||
return;
|
||||
}
|
||||
check_shape(1, 1, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sme_mmm_f32_32x32_k8() {
|
||||
if !has_sme() {
|
||||
return;
|
||||
}
|
||||
check_shape(1, 8, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sme_mmm_f32_32x32_k128() {
|
||||
if !has_sme() {
|
||||
return;
|
||||
}
|
||||
check_shape(1, 128, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sme_mmm_f32_32x32_multi_tile() {
|
||||
if !has_sme() {
|
||||
return;
|
||||
}
|
||||
// 64x64 output (2x2 tiles), K=64 — exercises the framework
|
||||
// iterating across multiple kernel calls.
|
||||
check_shape(2, 64, 2);
|
||||
}
|
||||
|
||||
// Strided store path: hand-built Clear + Store chain with non-contig C.
|
||||
#[test]
|
||||
fn sme_store_non_contiguous() {
|
||||
if !has_sme() {
|
||||
return;
|
||||
}
|
||||
use crate::frame::mmm::{FusedKerSpec, OutputStoreKer};
|
||||
const MR: usize = 32;
|
||||
const NR: usize = 32;
|
||||
let mut v: Vec<f32> = vec![f32::MAX; MR * 5 * NR * 3];
|
||||
let c = OutputStoreKer {
|
||||
ptr: v.as_mut_ptr() as _,
|
||||
row_byte_stride: (4 * 3 * NR * 5) as isize,
|
||||
col_byte_stride: 4 * 3,
|
||||
item_size: 4,
|
||||
};
|
||||
let non_linear = [
|
||||
FusedKerSpec::<f32>::Clear,
|
||||
FusedKerSpec::Store(c),
|
||||
FusedKerSpec::Done,
|
||||
];
|
||||
let err = unsafe { (sme_mmm_f32_32x32.kernel)(&non_linear) };
|
||||
assert_eq!(err, 0, "kernel returned non-zero error code");
|
||||
let mut expected = vec![f32::MAX; v.len()];
|
||||
for col in 0..NR {
|
||||
for row in 0..MR {
|
||||
expected[col * 3 + row * 3 * 5 * NR] = 0.0;
|
||||
}
|
||||
}
|
||||
for (i, (got, exp)) in v.iter().zip(expected.iter()).enumerate() {
|
||||
assert_eq!(got, exp, "mismatch at idx {i}: got {got} expected {exp}");
|
||||
}
|
||||
}
|
||||
}
|
||||
+230
@@ -0,0 +1,230 @@
|
||||
use crate::Ops;
|
||||
|
||||
// `tract_sve` is set by build.rs only on aarch64-linux when the C compiler
|
||||
// supports SVE intrinsics. The kernel registration + extern live behind it so
|
||||
// non-SVE builds never reference the (absent) C symbol.
|
||||
#[cfg(tract_sve)]
|
||||
use crate::frame::mmm::ImplementationQuality::ManuallyOptimized;
|
||||
#[cfg(tract_sve)]
|
||||
use crate::mmm::*;
|
||||
#[cfg(tract_sve)]
|
||||
use crate::pack::PackedFormat;
|
||||
// Explicit import so `f16` is tract's half::f16 (LADatum), not rustc's builtin
|
||||
// primitive f16 — a glob import would not shadow the primitive.
|
||||
#[cfg(tract_sve)]
|
||||
use tract_data::prelude::f16;
|
||||
|
||||
// f32 SVE kernel can't do LeakyRelu or the i32 quantization ops (matches the
|
||||
// arm64simd / SME f32 CAN_FUSE).
|
||||
#[cfg(tract_sve)]
|
||||
const CAN_FUSE: fn(&FusedSpec) -> bool = |f| {
|
||||
!matches!(
|
||||
f,
|
||||
FusedSpec::LeakyRelu(_)
|
||||
| FusedSpec::QScale(_, _, _)
|
||||
| FusedSpec::RoundingShiftRight(_, _)
|
||||
| FusedSpec::ShiftLeft(_)
|
||||
)
|
||||
};
|
||||
|
||||
// The i32 quantized kernel keeps the quantization fuse ops (QScale /
|
||||
// RoundingShiftRight / ShiftLeft) — they are the whole point of a quantized
|
||||
// kernel — and excludes only LeakyRelu (matches arm64simd's i32 surface; i32
|
||||
// LeakyRelu has no practical use and the C kernel does not implement it).
|
||||
#[cfg(tract_sve)]
|
||||
const CAN_FUSE_I32: fn(&FusedSpec) -> bool = |f| !matches!(f, FusedSpec::LeakyRelu(_));
|
||||
|
||||
#[cfg(tract_sve)]
|
||||
const SVE2: fn() -> bool = has_sve2;
|
||||
|
||||
// The f16 kernels need FEAT_SVE2 AND FEAT_FP16 (native f16 arithmetic).
|
||||
#[cfg(tract_sve)]
|
||||
const SVE2_FP16: fn() -> bool = || has_sve2() && crate::arm64::has_fp16();
|
||||
|
||||
// The VLA SVE f32 GEMM kernel, implemented in C (arm64/sve/sve_mmm_f32.c) since
|
||||
// Rust has no stable SVE intrinsics. Broadcast-A rank-1 update, N-tile walked in
|
||||
// svcntw() chunks → correct and full-width at any VL.
|
||||
#[cfg(tract_sve)]
|
||||
mod sve_sys {
|
||||
use crate::frame::mmm::FusedKerSpec;
|
||||
use tract_data::prelude::f16;
|
||||
unsafe extern "C" {
|
||||
pub fn sve_mmm_f32_kernel(ops: *const FusedKerSpec<f32>) -> isize;
|
||||
pub fn sve_mmv_f32_64x1_kernel(ops: *const FusedKerSpec<f32>) -> isize;
|
||||
pub fn sve_mmm_i32_kernel(ops: *const FusedKerSpec<i32>) -> isize;
|
||||
pub fn sve_mmm_i32_64x1_kernel(ops: *const FusedKerSpec<i32>) -> isize;
|
||||
pub fn sve_mmm_f16_kernel(ops: *const FusedKerSpec<f16>) -> isize;
|
||||
pub fn sve_mmv_f16_64x1_kernel(ops: *const FusedKerSpec<f16>) -> isize;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(tract_sve)]
|
||||
MMMRustKernel!(sve_sys::sve_mmm_f32_kernel => sve_mmm_f32_8x8<f32>(8, 8)
|
||||
where(SVE2)
|
||||
can_fuse(CAN_FUSE)
|
||||
quality(ManuallyOptimized)
|
||||
);
|
||||
|
||||
// The VLA SVE f32 GEMV kernel (arm64/sve/sve_mmv_f32_64x1.c), MR=64 NR=1,
|
||||
// dispatched when N == 1 (matrix x f32 column vector). Wired to ops.mmv_f32.
|
||||
#[cfg(tract_sve)]
|
||||
MMMRustKernel!(sve_sys::sve_mmv_f32_64x1_kernel => sve_mmv_f32_64x1<f32>(64, 1)
|
||||
where(SVE2)
|
||||
can_fuse(CAN_FUSE)
|
||||
quality(ManuallyOptimized)
|
||||
);
|
||||
|
||||
// The VLA SVE int8 -> int32 GEMM kernel (arm64/sve/sve_mmm_i32.c). Consumes
|
||||
// tract's native i8i8 K-major packing via the widening rank-1 update (svld1sb +
|
||||
// svmla), and supports the i32 quantization fuse ops. Wired to ops.qmmm_i32.
|
||||
#[cfg(tract_sve)]
|
||||
MMMRustKernel!(sve_sys::sve_mmm_i32_kernel => sve_mmm_i32_8x8<i32>(8, 8)
|
||||
where(SVE2)
|
||||
can_fuse(CAN_FUSE_I32)
|
||||
packing[1] = i8i8 => |k| k.with_packing(
|
||||
PackedFormat::new(DatumType::I8, 8, 16),
|
||||
PackedFormat::new(DatumType::I8, 8, 16),
|
||||
);
|
||||
quality(ManuallyOptimized)
|
||||
store(i8)
|
||||
);
|
||||
|
||||
// The VLA SVE int8 -> int32 GEMV kernel (arm64/sve/sve_mmm_i32_64x1.c), MR=64
|
||||
// NR=1, dispatched when N == 1. Same widening update vectorized over M. Wired to
|
||||
// ops.qmmv_i32.
|
||||
#[cfg(tract_sve)]
|
||||
MMMRustKernel!(sve_sys::sve_mmm_i32_64x1_kernel => sve_mmm_i32_64x1<i32>(64, 1)
|
||||
where(SVE2)
|
||||
can_fuse(CAN_FUSE_I32)
|
||||
packing[1] = i8i8 => |k| k.with_packing(
|
||||
PackedFormat::new(DatumType::I8, 64, 16),
|
||||
PackedFormat::new(DatumType::I8, 1, 1),
|
||||
);
|
||||
quality(ManuallyOptimized)
|
||||
store(i8)
|
||||
);
|
||||
|
||||
// The VLA SVE f16 GEMM kernel (arm64/sve/sve_mmm_f16.c), native f16 FMA, gated on
|
||||
// SVE2 + FP16. Wired to ops.mmm_f16 when has_fp16().
|
||||
#[cfg(tract_sve)]
|
||||
MMMRustKernel!(sve_sys::sve_mmm_f16_kernel => sve_mmm_f16_8x8<f16>(8, 8)
|
||||
where(SVE2_FP16)
|
||||
can_fuse(CAN_FUSE)
|
||||
quality(ManuallyOptimized)
|
||||
);
|
||||
|
||||
// The VLA SVE f16 GEMV kernel (arm64/sve/sve_mmv_f16_64x1.c), MR=64 NR=1,
|
||||
// dispatched when N == 1. Wired to ops.mmv_f16 when has_fp16().
|
||||
#[cfg(tract_sve)]
|
||||
MMMRustKernel!(sve_sys::sve_mmv_f16_64x1_kernel => sve_mmv_f16_64x1<f16>(64, 1)
|
||||
where(SVE2_FP16)
|
||||
can_fuse(CAN_FUSE)
|
||||
quality(ManuallyOptimized)
|
||||
);
|
||||
|
||||
// SVE / SVE2 backend.
|
||||
//
|
||||
// Unlike SME (Apple M4) and AMX (Apple), SVE/SVE2 is NOT present on any Apple
|
||||
// silicon — it lives on ARMv9 server/mobile cores (Neoverse V1+/N2+, Cortex-X2+
|
||||
// / A510+, Graviton 3/4). So detection is Linux-only in practice; macOS always
|
||||
// returns false.
|
||||
//
|
||||
// The kernels are vector-length-agnostic (VLA): they read the vector width at
|
||||
// runtime via `whilelt` predication and `svcntw()`, so a single kernel is
|
||||
// correct at every VL (128..2048-bit). That means — unlike the SME kernels,
|
||||
// which hardcoded SVL=512 and needed an RDSVL gate — the SVE kernels need NO
|
||||
// vector-length gate for correctness. `rdvl_bytes()` is provided only for
|
||||
// optional VL-matched dispatch (selecting a wider-tiled variant when the
|
||||
// hardware VL is large), not for correctness.
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
pub fn has_sve() -> bool {
|
||||
if std::env::var_os("TRACT_SVE_DISABLE").is_some() {
|
||||
return false;
|
||||
}
|
||||
// HWCAP_SVE = 1 << 22 on aarch64 (kernel ABI).
|
||||
const HWCAP_SVE: u64 = 1 << 22;
|
||||
unsafe extern "C" {
|
||||
fn getauxval(t: u64) -> u64;
|
||||
}
|
||||
const AT_HWCAP: u64 = 16;
|
||||
unsafe { (getauxval(AT_HWCAP) & HWCAP_SVE) != 0 }
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
pub fn has_sve() -> bool {
|
||||
// No Apple silicon implements SVE; no SVE on non-Linux targets we support.
|
||||
false
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
pub fn has_sve2() -> bool {
|
||||
if std::env::var_os("TRACT_SVE_DISABLE").is_some() {
|
||||
return false;
|
||||
}
|
||||
// HWCAP2_SVE2 = 1 << 1 on aarch64 (kernel ABI).
|
||||
const HWCAP2_SVE2: u64 = 1 << 1;
|
||||
unsafe extern "C" {
|
||||
fn getauxval(t: u64) -> u64;
|
||||
}
|
||||
const AT_HWCAP2: u64 = 26;
|
||||
unsafe { (getauxval(AT_HWCAP2) & HWCAP2_SVE2) != 0 }
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
pub fn has_sve2() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// SVE vector length in bytes, via `RDVL x0, #1` (encoding 0x04bf5020).
|
||||
/// Legal whenever FEAT_SVE is implemented; callers MUST confirm `has_sve()`
|
||||
/// first (RDVL is UNDEFINED without SVE and would SIGILL). Used only for
|
||||
/// optional VL-matched kernel selection — VLA kernels do not need it.
|
||||
#[cfg(target_os = "linux")]
|
||||
#[allow(dead_code)]
|
||||
pub fn rdvl_bytes() -> u64 {
|
||||
let vl: u64;
|
||||
unsafe {
|
||||
std::arch::asm!(
|
||||
".inst 0x04bf5020", // rdvl x0, #1
|
||||
out("x0") vl,
|
||||
options(nomem, nostack, preserves_flags),
|
||||
);
|
||||
}
|
||||
vl
|
||||
}
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
let _ = ops;
|
||||
if has_sve2() {
|
||||
#[cfg(target_os = "linux")]
|
||||
log::info!("SVE2 optimisation available (VL = {} bytes)", rdvl_bytes());
|
||||
#[cfg(tract_sve)]
|
||||
{
|
||||
// Force the SVE kernels for f32 mmm and i32 quantized mmm (mirrors the
|
||||
// SME backend) and also register them as candidates. TRACT_SVE_DISABLE=1
|
||||
// already turns the whole thing off via has_sve2().
|
||||
ops.mmm_f32 = Box::new(|_, _, _| sve_mmm_f32_8x8.mmm());
|
||||
ops.mmv_f32 = Box::new(|_, _| sve_mmv_f32_64x1.mmm());
|
||||
ops.qmmm_i32 = Box::new(|_, _, _| sve_mmm_i32_8x8.mmm());
|
||||
ops.qmmv_i32 = Box::new(|_, _| sve_mmm_i32_64x1.mmm());
|
||||
ops.mmm_impls.extend_from_slice(&[
|
||||
sve_mmm_f32_8x8.mmm(),
|
||||
sve_mmv_f32_64x1.mmm(),
|
||||
sve_mmm_i32_8x8.mmm(),
|
||||
sve_mmm_i32_64x1.mmm(),
|
||||
]);
|
||||
// f16 kernels additionally require FEAT_FP16.
|
||||
if crate::arm64::has_fp16() {
|
||||
ops.mmm_f16 = Box::new(|_, _, _| sve_mmm_f16_8x8.mmm());
|
||||
ops.mmv_f16 = Box::new(|_, _| sve_mmv_f16_64x1.mmm());
|
||||
ops.mmm_impls
|
||||
.extend_from_slice(&[sve_mmm_f16_8x8.mmm(), sve_mmv_f16_64x1.mmm()]);
|
||||
}
|
||||
}
|
||||
} else if has_sve() {
|
||||
log::info!("SVE (v1) present; SVE2 kernels not enabled");
|
||||
} else {
|
||||
log::info!("No SVE optimisation");
|
||||
}
|
||||
}
|
||||
Vendored
+81
@@ -0,0 +1,81 @@
|
||||
use byteorder::{LE, ReadBytesExt, WriteBytesExt};
|
||||
use std::io::{Cursor, Read, Write};
|
||||
use tract_data::internal::*;
|
||||
|
||||
pub struct NibbleReader<R> {
|
||||
second_half: Option<i8>,
|
||||
reader: R,
|
||||
}
|
||||
|
||||
impl<'s> NibbleReader<Cursor<&'s [u8]>> {
|
||||
pub fn for_slice(slice: &'s [u8]) -> Self {
|
||||
NibbleReader::new(Cursor::new(slice))
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: Read> NibbleReader<R> {
|
||||
pub fn new(reader: R) -> NibbleReader<R> {
|
||||
NibbleReader {
|
||||
reader,
|
||||
second_half: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_f16(&mut self) -> f16 {
|
||||
assert!(self.second_half.is_none());
|
||||
f16::from_bits(self.reader.read_u16::<LE>().unwrap())
|
||||
}
|
||||
|
||||
pub fn read_i4(&mut self) -> i8 {
|
||||
if let Some(second) = self.second_half.take() {
|
||||
second
|
||||
} else {
|
||||
let byte = self.reader.read_u8().unwrap();
|
||||
self.second_half = Some((byte >> 4) as i8);
|
||||
(byte & 0x0F) as i8
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_i8(&mut self) -> i8 {
|
||||
self.reader.read_i8().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct NibbleWriter<W> {
|
||||
first_half: Option<i8>,
|
||||
writer: W,
|
||||
}
|
||||
|
||||
impl<'s> NibbleWriter<Cursor<&'s mut [u8]>> {
|
||||
pub fn for_slice(slice: &'s mut [u8]) -> Self {
|
||||
NibbleWriter::new(Cursor::new(slice))
|
||||
}
|
||||
}
|
||||
|
||||
impl<W: Write> NibbleWriter<W> {
|
||||
pub fn new(writer: W) -> NibbleWriter<W> {
|
||||
NibbleWriter {
|
||||
writer,
|
||||
first_half: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn write_f16(&mut self, f: f16) {
|
||||
assert!(self.first_half.is_none());
|
||||
self.writer.write_u16::<LE>(f.to_bits()).unwrap()
|
||||
}
|
||||
|
||||
pub fn write_i4(&mut self, q: i8) {
|
||||
if let Some(first) = self.first_half.take() {
|
||||
self.writer
|
||||
.write_u8(first as u8 | ((q as u8) << 4))
|
||||
.unwrap()
|
||||
} else {
|
||||
self.first_half = Some(q);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn write_i8(&mut self, q: i8) {
|
||||
self.writer.write_i8(q).unwrap()
|
||||
}
|
||||
}
|
||||
Vendored
+357
@@ -0,0 +1,357 @@
|
||||
use downcast_rs::{Downcast, impl_downcast};
|
||||
use dyn_clone::{DynClone, clone_box};
|
||||
use dyn_eq::DynEq;
|
||||
use dyn_hash::DynHash;
|
||||
use num_traits::Zero;
|
||||
use tract_data::internal::*;
|
||||
use tract_data::itertools::Itertools;
|
||||
|
||||
use std::alloc::Layout;
|
||||
use std::borrow::Cow;
|
||||
use std::fmt::{Debug, Display};
|
||||
use std::hash::Hash;
|
||||
use std::sync::Arc;
|
||||
|
||||
mod helpers;
|
||||
mod q4_0;
|
||||
mod q8_1;
|
||||
mod storage;
|
||||
mod value;
|
||||
|
||||
pub use helpers::{NibbleReader, NibbleWriter};
|
||||
pub use q4_0::Q4_0;
|
||||
pub use q8_1::Q8_1;
|
||||
pub use storage::{BlockQuantStorage, block_quant_slice};
|
||||
pub use value::{BlockQuantFact, PackedBlockQuantFact};
|
||||
|
||||
use crate::mmm::{EagerPackedInput, MMMInputFormat};
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
use crate::WeightType;
|
||||
|
||||
use super::mmm::MMMInputValue;
|
||||
|
||||
pub trait BlockQuant:
|
||||
Debug + Display + Send + Sync + DynClone + DynHash + dyn_eq::DynEq + Downcast
|
||||
{
|
||||
fn block_len(&self) -> usize;
|
||||
|
||||
fn block_bytes(&self) -> usize;
|
||||
|
||||
fn dequant_block_f32(&self, quant: &[u8], block: &mut [f32]);
|
||||
fn dequant_block_f16(&self, quant: &[u8], block: &mut [f16]);
|
||||
fn quant_block_f16(&self, block: &[f16], quant: &mut [u8]);
|
||||
fn quant_block_f32(&self, block: &[f32], quant: &mut [u8]);
|
||||
|
||||
fn quant_f16(&self, input: &[f16]) -> TractResult<Blob> {
|
||||
unsafe {
|
||||
let blocks = input.len() / self.block_len();
|
||||
let mut quant = Blob::for_layout(
|
||||
Layout::from_size_align(blocks * self.block_bytes(), 128).unwrap(),
|
||||
);
|
||||
for b in 0..blocks {
|
||||
let block = &input[b * self.block_len()..][..self.block_len()];
|
||||
let qblock = &mut quant[b * self.block_bytes()..][..self.block_bytes()];
|
||||
self.quant_block_f16(block, qblock);
|
||||
}
|
||||
Ok(quant)
|
||||
}
|
||||
}
|
||||
|
||||
fn quant_f32(&self, input: &[f32]) -> TractResult<Blob> {
|
||||
unsafe {
|
||||
let blocks = input.len() / self.block_len();
|
||||
let mut quant = Blob::for_layout(
|
||||
Layout::from_size_align(blocks * self.block_bytes(), 128).unwrap(),
|
||||
);
|
||||
for b in 0..blocks {
|
||||
let block = &input[b * self.block_len()..][..self.block_len()];
|
||||
let qblock = &mut quant[b * self.block_bytes()..][..self.block_bytes()];
|
||||
self.quant_block_f32(block, qblock);
|
||||
}
|
||||
Ok(quant)
|
||||
}
|
||||
}
|
||||
|
||||
fn dequant_f32(&self, input: &[u8]) -> TractResult<Tensor> {
|
||||
unsafe {
|
||||
let blocks = input.len() / self.block_bytes();
|
||||
let mut tensor = Tensor::uninitialized::<f32>(&[blocks * self.block_len()])?;
|
||||
let mut tensor_plain = tensor.try_as_plain_mut()?;
|
||||
let slice = tensor_plain.as_slice_mut::<f32>()?;
|
||||
for b in 0..blocks {
|
||||
let block = &mut slice[b * self.block_len()..][..self.block_len()];
|
||||
let qblock = &input[b * self.block_bytes()..][..self.block_bytes()];
|
||||
self.dequant_block_f32(qblock, block);
|
||||
}
|
||||
Ok(tensor)
|
||||
}
|
||||
}
|
||||
|
||||
fn dequant_f16(&self, input: &[u8]) -> TractResult<Tensor> {
|
||||
unsafe {
|
||||
let blocks = input.len() / self.block_bytes();
|
||||
let mut tensor = Tensor::uninitialized::<f16>(&[blocks * self.block_len()])?;
|
||||
let mut tensor_plain = tensor.try_as_plain_mut()?;
|
||||
let slice = tensor_plain.as_slice_mut::<f16>()?;
|
||||
for b in 0..blocks {
|
||||
let block = &mut slice[b * self.block_len()..][..self.block_len()];
|
||||
let qblock = &input[b * self.block_bytes()..][..self.block_bytes()];
|
||||
self.dequant_block_f16(qblock, block);
|
||||
}
|
||||
Ok(tensor)
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_at_offset_f16(&self, input: &[u8], offset: usize) -> f16 {
|
||||
let len = self.block_len();
|
||||
let block_id = offset / len;
|
||||
let mut block = vec![f16::zero(); self.block_len()];
|
||||
self.dequant_block_f16(
|
||||
&input[block_id * self.block_bytes()..][..self.block_bytes()],
|
||||
&mut block,
|
||||
);
|
||||
block[offset % len]
|
||||
}
|
||||
|
||||
fn extract_at_offset_f32(&self, input: &[u8], offset: usize) -> f32 {
|
||||
let len = self.block_len();
|
||||
let block_id = offset / len;
|
||||
let mut block = vec![f32::zero(); self.block_len()];
|
||||
self.dequant_block_f32(
|
||||
&input[block_id * self.block_bytes()..][..self.block_bytes()],
|
||||
&mut block,
|
||||
);
|
||||
block[offset % len]
|
||||
}
|
||||
|
||||
fn simulate_precision_loss(
|
||||
&self,
|
||||
mut tensor: Tensor,
|
||||
block_axis: usize,
|
||||
) -> TractResult<Tensor> {
|
||||
ensure!(block_axis == tensor.rank() - 1);
|
||||
ensure!(tensor.shape()[block_axis] % self.block_len() == 0);
|
||||
let mut scratch = vec![0u8; self.block_bytes()];
|
||||
if tensor.datum_type() == f32::datum_type() {
|
||||
let mut tensor_plain = tensor.try_as_plain_mut()?;
|
||||
for block in tensor_plain
|
||||
.as_slice_mut::<f32>()?
|
||||
.chunks_mut(self.block_len())
|
||||
{
|
||||
self.quant_block_f32(block, &mut scratch);
|
||||
self.dequant_block_f32(&scratch, block);
|
||||
}
|
||||
drop(tensor_plain);
|
||||
Ok(tensor)
|
||||
} else if tensor.datum_type() == f16::datum_type() {
|
||||
let mut tensor_plain = tensor.try_as_plain_mut()?;
|
||||
for block in tensor_plain
|
||||
.as_slice_mut::<f16>()?
|
||||
.chunks_mut(self.block_len())
|
||||
{
|
||||
self.quant_block_f16(block, &mut scratch);
|
||||
self.dequant_block_f16(&scratch, block);
|
||||
}
|
||||
drop(tensor_plain);
|
||||
Ok(tensor)
|
||||
} else {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
||||
fn pack(
|
||||
&self,
|
||||
input: &[u8],
|
||||
k: usize,
|
||||
r: usize,
|
||||
zip: usize,
|
||||
scales_at_end: bool,
|
||||
) -> TractResult<EagerPackedInput>;
|
||||
|
||||
unsafe fn extract_packed_panel(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
target: &PackedFormat,
|
||||
panel: usize,
|
||||
scratch: *mut u8,
|
||||
) -> TractResult<()>;
|
||||
|
||||
fn extract_at_mn_f16(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [f16],
|
||||
) -> TractResult<()>;
|
||||
|
||||
fn extract_at_mn_f32(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [f32],
|
||||
) -> TractResult<()>;
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(BlockQuant);
|
||||
dyn_hash::hash_trait_object!(BlockQuant);
|
||||
dyn_eq::eq_trait_object!(BlockQuant);
|
||||
impl_downcast!(BlockQuant);
|
||||
|
||||
#[allow(clippy::derived_hash_with_manual_eq)]
|
||||
#[derive(Clone, Hash)]
|
||||
pub struct PackedBlockQuantFormat {
|
||||
pub bq: Box<dyn BlockQuant>,
|
||||
pub r: usize,
|
||||
pub zip: usize,
|
||||
pub scales_at_end: bool,
|
||||
}
|
||||
|
||||
impl PartialEq for PackedBlockQuantFormat {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
*self.bq == *other.bq
|
||||
&& self.r == other.r
|
||||
&& self.zip == other.zip
|
||||
&& self.scales_at_end == other.scales_at_end
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for PackedBlockQuantFormat {}
|
||||
|
||||
impl Display for PackedBlockQuantFormat {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "Packed{}[{}]", &*self.bq, self.r)?;
|
||||
if self.zip != 0 {
|
||||
write!(f, "Z{}", self.zip)?;
|
||||
}
|
||||
if self.scales_at_end {
|
||||
write!(f, "Se")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Debug for PackedBlockQuantFormat {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
<Self as Display>::fmt(self, f)
|
||||
}
|
||||
}
|
||||
|
||||
impl PackedBlockQuantFormat {
|
||||
pub fn new(bq: &dyn BlockQuant, r: usize, zip: usize, scales_at_end: bool) -> Self {
|
||||
PackedBlockQuantFormat {
|
||||
bq: clone_box(bq),
|
||||
r,
|
||||
zip,
|
||||
scales_at_end,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn simulate_precision_loss(
|
||||
&self,
|
||||
tensor: Tensor,
|
||||
block_axis: usize,
|
||||
) -> TractResult<Tensor> {
|
||||
self.bq.simulate_precision_loss(tensor, block_axis)
|
||||
}
|
||||
|
||||
pub fn pack(&self, input: &[u8], k: usize) -> TractResult<EagerPackedInput> {
|
||||
self.bq.pack(input, k, self.r, self.zip, self.scales_at_end)
|
||||
}
|
||||
}
|
||||
|
||||
impl MMMInputFormat for PackedBlockQuantFormat {
|
||||
fn prepare_tensor(&self, t: &Tensor, _k_axis: usize, _mn_axis: usize) -> TractResult<Tensor> {
|
||||
let bqs = t.try_storage_as::<BlockQuantStorage>()?;
|
||||
let num_groups: usize = if t.rank() > 2 {
|
||||
t.shape()[..t.rank() - 2].iter().product()
|
||||
} else {
|
||||
1
|
||||
};
|
||||
let m_per_group = t.shape()[t.rank().saturating_sub(2)];
|
||||
let k = *t.shape().last().unwrap();
|
||||
let values = (0..num_groups)
|
||||
.map(|g| {
|
||||
let slice = block_quant_slice(bqs.value(), &*self.bq, m_per_group, k, g);
|
||||
let packed = self.pack(slice, k)?;
|
||||
Ok(Box::new(packed) as Box<dyn MMMInputValue>)
|
||||
})
|
||||
.collect::<TractResult<Vec<_>>>()?;
|
||||
let leading_shape = &t.shape()[..t.rank().saturating_sub(2)];
|
||||
Ok(
|
||||
crate::mmm::PackedMatrixStorage::new_batched(leading_shape, values)
|
||||
.into_tensor(t.datum_type()),
|
||||
)
|
||||
}
|
||||
|
||||
fn prepare_one(
|
||||
&self,
|
||||
t: &Tensor,
|
||||
k_axis: usize,
|
||||
mn_axis: usize,
|
||||
) -> TractResult<Box<dyn MMMInputValue>> {
|
||||
// this code path is essentially there for test scenarios
|
||||
let t = if t.is_plain() && t.datum_type().is_number() {
|
||||
let k = t.shape()[k_axis];
|
||||
let m = t.shape()[mn_axis];
|
||||
assert!(k % self.bq.block_len() == 0);
|
||||
let t: Cow<Tensor> = if k_axis == 1 && mn_axis == 0 {
|
||||
Cow::Borrowed(t)
|
||||
} else {
|
||||
Cow::Owned(t.clone().move_axis(1, 0)?)
|
||||
};
|
||||
let quant = if t.datum_type() == f32::datum_type() {
|
||||
self.bq.quant_f32(t.try_as_plain()?.as_slice()?)?
|
||||
} else if t.datum_type() == f16::datum_type() {
|
||||
self.bq.quant_f16(t.try_as_plain()?.as_slice()?)?
|
||||
} else {
|
||||
todo!()
|
||||
};
|
||||
Cow::Owned(
|
||||
BlockQuantStorage::new(self.bq.clone(), m, k, Arc::new(quant))?
|
||||
.into_tensor_with_shape(t.datum_type(), &[1, m, k]),
|
||||
)
|
||||
} else {
|
||||
Cow::Borrowed(t)
|
||||
};
|
||||
ensure!(mn_axis == 0);
|
||||
ensure!(k_axis == 1);
|
||||
let bqs = t.try_storage_as::<BlockQuantStorage>()?;
|
||||
let k = *t.shape().last().unwrap();
|
||||
let packed = self.pack(bqs.value(), k)?;
|
||||
Ok(Box::new(packed))
|
||||
}
|
||||
|
||||
fn precursor(&self) -> WeightType {
|
||||
WeightType::BlockQuant(self.bq.clone())
|
||||
}
|
||||
|
||||
fn k_alignment(&self) -> usize {
|
||||
self.bq.block_len()
|
||||
}
|
||||
|
||||
fn r(&self) -> usize {
|
||||
self.r
|
||||
}
|
||||
|
||||
fn mem_size(&self, k: TDim, mn: TDim) -> TDim {
|
||||
k * mn * self.bq.block_bytes() / self.bq.block_len()
|
||||
}
|
||||
|
||||
fn extract_at_mn_f16(
|
||||
&self,
|
||||
data: &EagerPackedInput,
|
||||
mn: usize,
|
||||
slice: &mut [f16],
|
||||
) -> TractResult<()> {
|
||||
self.bq.extract_at_mn_f16(data, mn, slice)
|
||||
}
|
||||
|
||||
fn extract_at_mn_f32(
|
||||
&self,
|
||||
data: &EagerPackedInput,
|
||||
mn: usize,
|
||||
slice: &mut [f32],
|
||||
) -> TractResult<()> {
|
||||
self.bq.extract_at_mn_f32(data, mn, slice)
|
||||
}
|
||||
}
|
||||
Vendored
+538
@@ -0,0 +1,538 @@
|
||||
use crate::mmm::PackedExoticFact;
|
||||
|
||||
use super::*;
|
||||
use num_traits::{AsPrimitive, Float, Zero};
|
||||
use std::alloc::Layout;
|
||||
|
||||
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
|
||||
pub struct BaseQ4_0<const QK: usize = 32>;
|
||||
|
||||
pub const Q4_0: BaseQ4_0 = BaseQ4_0::<32>;
|
||||
|
||||
impl<const QK: usize> Debug for BaseQ4_0<QK> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
if QK == 32 {
|
||||
write!(f, "Q4_0")
|
||||
} else {
|
||||
write!(f, "BaseQ4_0<{QK}>")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<const QK: usize> BaseQ4_0<QK> {
|
||||
fn quant_block<T>(&self, block: &[T], quant: &mut [u8])
|
||||
where
|
||||
f32: AsPrimitive<i8> + From<T>,
|
||||
T: Debug + Float,
|
||||
{
|
||||
assert!(quant.len() == self.block_bytes());
|
||||
assert!(block.len() == self.block_len());
|
||||
let mut writer = NibbleWriter::for_slice(quant);
|
||||
let mut amax = T::zero();
|
||||
let mut max = T::zero();
|
||||
for v in block {
|
||||
if amax < v.abs() {
|
||||
amax = v.abs();
|
||||
max = *v;
|
||||
}
|
||||
}
|
||||
let scale = f32::from(max) / -8f32;
|
||||
let r_scale = if scale.is_zero() { 0f32 } else { scale.recip() };
|
||||
writer.write_f16(f16::from_f32(scale));
|
||||
|
||||
for idx in 0..block.len() {
|
||||
// Quant block in GGML nibble order
|
||||
let ggml_idx = (block.len() / 2) * (idx % 2) + (idx / 2);
|
||||
let i: i8 = (f32::from(block[ggml_idx]) * r_scale + 8.5f32).as_();
|
||||
writer.write_i4(i.min(15));
|
||||
}
|
||||
}
|
||||
|
||||
fn dequant_block<T: Float + 'static>(&self, quant: &[u8], block: &mut [T])
|
||||
where
|
||||
f16: AsPrimitive<T>,
|
||||
i8: AsPrimitive<T>,
|
||||
{
|
||||
assert!(quant.len() == self.block_bytes());
|
||||
assert!(block.len() == self.block_len());
|
||||
let mut nibbles = NibbleReader::for_slice(quant);
|
||||
let d: T = nibbles.read_f16().as_();
|
||||
for idx in 0..block.len() {
|
||||
let ggml_idx = (block.len() / 2) * (idx % 2) + (idx / 2);
|
||||
block[ggml_idx] = (nibbles.read_i4() - 8).as_() * d;
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn extract_panel_t<T: Float + Debug + 'static>(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
target: &PackedFormat,
|
||||
panel: usize,
|
||||
scratch: *mut u8,
|
||||
) -> TractResult<()>
|
||||
where
|
||||
f16: AsPrimitive<T>,
|
||||
i8: AsPrimitive<T>,
|
||||
{
|
||||
let pbqf: &PackedBlockQuantFormat =
|
||||
value.fact.format.downcast_ref().with_context(|| {
|
||||
format!(
|
||||
"Expecing PackedBlockQuantFormat, found {:?}",
|
||||
value.fact.format
|
||||
)
|
||||
})?;
|
||||
ensure!(pbqf.r == target.r);
|
||||
ensure!(value.fact.k % self.block_len() == 0);
|
||||
ensure!(*pbqf.bq == *(self as &dyn BlockQuant));
|
||||
let scratch =
|
||||
unsafe { std::slice::from_raw_parts_mut(scratch as *mut T, value.fact.k * target.r) };
|
||||
let blocks_for_k = value.fact.k / self.block_len();
|
||||
let row_bytes = blocks_for_k * self.block_bytes();
|
||||
let input = &value.packed[panel * target.r * row_bytes..];
|
||||
let mut scales = vec![T::zero(); target.r];
|
||||
let mut scratch = scratch.iter_mut();
|
||||
let zipped_order = zipped_order(pbqf.r, pbqf.zip);
|
||||
let mut weights = vec![0i8; pbqf.r];
|
||||
let panel_block_bytes = target.r * self.block_bytes();
|
||||
let (scale_offset, weights_offset) = if pbqf.scales_at_end {
|
||||
(
|
||||
panel_block_bytes - target.r * f16::datum_type().size_of(),
|
||||
0,
|
||||
)
|
||||
} else {
|
||||
(0, target.r * f16::datum_type().size_of())
|
||||
};
|
||||
for block in 0..blocks_for_k {
|
||||
let block = &input[block * panel_block_bytes..][..panel_block_bytes];
|
||||
let mut s_reader = NibbleReader::for_slice(&block[scale_offset..]);
|
||||
let mut w_reader = NibbleReader::for_slice(&block[weights_offset..]);
|
||||
for s in &mut scales {
|
||||
*s = s_reader.read_f16().as_();
|
||||
}
|
||||
for _ in 0..self.block_len() {
|
||||
for &o in &zipped_order {
|
||||
weights[o] = w_reader.read_i4();
|
||||
}
|
||||
for (w, s) in weights.iter().zip(scales.iter()) {
|
||||
*scratch.next().unwrap() = *s * (*w - 8).as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn extract_at_mn_t<T: Float + Debug + 'static>(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [T],
|
||||
) -> TractResult<()>
|
||||
where
|
||||
f16: AsPrimitive<T>,
|
||||
i8: AsPrimitive<T>,
|
||||
{
|
||||
let pbqf: &PackedBlockQuantFormat =
|
||||
value.fact.format.downcast_ref().with_context(|| {
|
||||
format!(
|
||||
"Expecing PackedBlockQuantFormat, found {:?}",
|
||||
value.fact.format
|
||||
)
|
||||
})?;
|
||||
ensure!(value.fact.k % self.block_len() == 0);
|
||||
ensure!(*pbqf.bq == *(self as &dyn BlockQuant));
|
||||
ensure!(
|
||||
value
|
||||
.fact
|
||||
.mn
|
||||
.to_usize()
|
||||
.ok()
|
||||
.map(|it| mn < it)
|
||||
.unwrap_or(true)
|
||||
);
|
||||
ensure!(value.fact.k == target.len());
|
||||
let blocks_for_k = value.fact.k / self.block_len();
|
||||
let row_bytes = blocks_for_k * self.block_bytes();
|
||||
let panel = mn / pbqf.r;
|
||||
let value = &value.packed[panel * pbqf.r * row_bytes..];
|
||||
let mut target = target.iter_mut();
|
||||
let zipped_order = zipped_order(pbqf.r, pbqf.zip)
|
||||
.iter()
|
||||
.position(|x| *x == mn % pbqf.r)
|
||||
.unwrap();
|
||||
|
||||
let panel_block_bytes = pbqf.r * self.block_bytes();
|
||||
let (scale_offset, weights_offset) = if pbqf.scales_at_end {
|
||||
(panel_block_bytes - pbqf.r * f16::datum_type().size_of(), 0)
|
||||
} else {
|
||||
(0, pbqf.r * f16::datum_type().size_of())
|
||||
};
|
||||
unsafe {
|
||||
for block in 0..blocks_for_k {
|
||||
let block = value.as_ptr().add(block * panel_block_bytes);
|
||||
let scale = *((block.add(scale_offset) as *const f16).add(mn % pbqf.r));
|
||||
let scale: T = scale.as_();
|
||||
for i in 0..self.block_len() {
|
||||
let byte = *block.add(weights_offset + i * pbqf.r / 2 + zipped_order / 2);
|
||||
let nib = if zipped_order % 2 == 0 {
|
||||
byte & 0x0F
|
||||
} else {
|
||||
byte >> 4
|
||||
};
|
||||
*target.next().unwrap() = scale * ((nib as i8) - 8).as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn zipped_order(r: usize, zip: usize) -> Vec<usize> {
|
||||
if zip == 0 {
|
||||
(0..r).collect_vec()
|
||||
} else {
|
||||
(0..r)
|
||||
.map(|i| {
|
||||
let vec_pair_ix = i / (2 * zip);
|
||||
let lane = (i % (2 * zip)) / 2;
|
||||
let side = i % 2;
|
||||
vec_pair_ix * 2 * zip + side * zip + lane
|
||||
})
|
||||
.collect_vec()
|
||||
}
|
||||
}
|
||||
|
||||
impl<const QK: usize> BlockQuant for BaseQ4_0<QK> {
|
||||
fn block_len(&self) -> usize {
|
||||
QK
|
||||
}
|
||||
|
||||
fn block_bytes(&self) -> usize {
|
||||
2 + self.block_len() / 2
|
||||
}
|
||||
|
||||
fn quant_block_f32(&self, block: &[f32], quant: &mut [u8]) {
|
||||
self.quant_block(block, quant)
|
||||
}
|
||||
|
||||
fn quant_block_f16(&self, block: &[f16], quant: &mut [u8]) {
|
||||
self.quant_block(block, quant)
|
||||
}
|
||||
|
||||
fn dequant_block_f32(&self, quant: &[u8], block: &mut [f32]) {
|
||||
self.dequant_block(quant, block)
|
||||
}
|
||||
|
||||
fn dequant_block_f16(&self, quant: &[u8], block: &mut [f16]) {
|
||||
self.dequant_block(quant, block)
|
||||
}
|
||||
|
||||
// s0_0 n0_0 n0_1 n0_2 n0_3 ... n0_30n0_31 s0_32 n0_32n0_33 ...
|
||||
// s1_0 n1_0 n1_1 n1_2 n1_3 ... n1_30n1_31 s1_32 n1_32n1_33 ...
|
||||
//
|
||||
// becomes (with r=4)
|
||||
//
|
||||
// s0_0 s1_0 s2_0 s3_0 n0_0 n1_0 n2_0 n3_0 n0_1 n1_1 n2_1 n3_1 ... n0_33 n1_33 n2_33 n3_33
|
||||
// s0_32 s1_32 s2_32 s3_32 n0_0 n1_0 n2_0 n3_0 n0_1 n1_1 n2_1 n3_1 ... n0_33 n1_33 n2_33 n3_33
|
||||
// ...
|
||||
fn pack(
|
||||
&self,
|
||||
input: &[u8],
|
||||
k: usize,
|
||||
r: usize,
|
||||
zip: usize,
|
||||
scales_at_end: bool,
|
||||
) -> TractResult<EagerPackedInput> {
|
||||
ensure!(input.len() % self.block_bytes() == 0);
|
||||
ensure!(k % self.block_len() == 0);
|
||||
// ensure!(input.len() == k * r / self.block_len() * self.block_bytes());
|
||||
ensure!(zip < r);
|
||||
let m = if input.len() == 0 {
|
||||
0
|
||||
} else {
|
||||
input.len() / self.block_bytes() * self.block_len() / k
|
||||
};
|
||||
let panels = m.divceil(r);
|
||||
let blocks_for_k = k / self.block_len();
|
||||
let row_bytes = blocks_for_k * self.block_bytes();
|
||||
let panel_bytes = row_bytes * r;
|
||||
let mut blob =
|
||||
unsafe { Blob::for_layout(Layout::from_size_align(panel_bytes * panels, 128)?) };
|
||||
let mut writer = NibbleWriter::for_slice(&mut blob);
|
||||
let order = zipped_order(r, zip);
|
||||
let mut scales = vec![f16::zero(); r];
|
||||
for p in 0..panels {
|
||||
let input = &input[(r * p) * row_bytes..];
|
||||
let mut readers = (0..r)
|
||||
.map(|r| {
|
||||
// manage partial panel
|
||||
let offset = if r * row_bytes < input.len() {
|
||||
r * row_bytes
|
||||
} else {
|
||||
0
|
||||
};
|
||||
NibbleReader::for_slice(&input[offset..])
|
||||
})
|
||||
.collect_vec();
|
||||
let mut temp_nibbles = vec![vec![0i8; self.block_len()]; r];
|
||||
for _ in 0..blocks_for_k {
|
||||
for (row, reader) in readers.iter_mut().enumerate() {
|
||||
scales[row] = reader.read_f16();
|
||||
temp_nibbles[row] = (0..self.block_len())
|
||||
.map(|_| reader.read_i4())
|
||||
.collect_vec();
|
||||
}
|
||||
if !scales_at_end {
|
||||
scales.iter().for_each(|s| writer.write_f16(*s))
|
||||
}
|
||||
for pos in 0..self.block_len() {
|
||||
for &row in &order {
|
||||
let ggml_idx = pos / (self.block_len() / 2) + (2 * pos) % self.block_len();
|
||||
let nib = temp_nibbles[row][ggml_idx];
|
||||
writer.write_i4(nib);
|
||||
}
|
||||
}
|
||||
if scales_at_end {
|
||||
scales.iter().for_each(|s| writer.write_f16(*s))
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(EagerPackedInput {
|
||||
fact: PackedExoticFact {
|
||||
format: Box::new(PackedBlockQuantFormat {
|
||||
bq: Box::new(*self),
|
||||
r,
|
||||
zip,
|
||||
scales_at_end,
|
||||
}),
|
||||
mn: m.to_dim(),
|
||||
k,
|
||||
},
|
||||
packed: blob.into(),
|
||||
panel_bytes,
|
||||
mn: m,
|
||||
})
|
||||
}
|
||||
|
||||
unsafe fn extract_packed_panel(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
target: &PackedFormat,
|
||||
panel: usize,
|
||||
scratch: *mut u8,
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
dispatch_floatlike!(Self::extract_panel_t(target.dt)(
|
||||
self, value, target, panel, scratch
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_at_mn_f16(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [f16],
|
||||
) -> TractResult<()> {
|
||||
self.extract_at_mn_t(value, mn, target)
|
||||
}
|
||||
|
||||
fn extract_at_mn_f32(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [f32],
|
||||
) -> TractResult<()> {
|
||||
self.extract_at_mn_t(value, mn, target)
|
||||
}
|
||||
}
|
||||
|
||||
impl<const QK: usize> Display for BaseQ4_0<QK> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "Q4_0")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use num_traits::Zero;
|
||||
use tract_data::internal::tract_ndarray::Array2;
|
||||
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn test_loop_f32(b: impl BlockQuant, data: &[f32]) {
|
||||
let mut input = data.to_vec();
|
||||
while input.len() % b.block_len() != 0 {
|
||||
input.push(0f32);
|
||||
}
|
||||
let quant = b.quant_f32(&input).unwrap();
|
||||
let result = b.dequant_f32(&quant).unwrap();
|
||||
let view = &result.try_as_plain().unwrap().as_slice::<f32>().unwrap()[..data.len()];
|
||||
assert_eq!(data, view);
|
||||
}
|
||||
|
||||
fn test_loop_f16(b: impl BlockQuant, data: &[f32]) {
|
||||
let mut input = data.iter().map(|f| f16::from_f32(*f)).collect_vec();
|
||||
while input.len() % b.block_len() != 0 {
|
||||
input.push(f16::zero());
|
||||
}
|
||||
let quant = b.quant_f16(&input).unwrap();
|
||||
let result = b.dequant_f16(&quant).unwrap();
|
||||
let view = &result.try_as_plain().unwrap().as_slice::<f16>().unwrap();
|
||||
assert_eq!(&input, view);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q4f32_pos() {
|
||||
test_loop_f32(Q4_0, &[1.0, 2.0, 3.0, 4.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q4f16_pos() {
|
||||
test_loop_f16(Q4_0, &[1.0, 2.0, 3.0, 4.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q4f32_neg() {
|
||||
test_loop_f32(Q4_0, &[-1.0, -2.0, -3.0, -4.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q4f16_neg() {
|
||||
test_loop_f16(Q4_0, &[-1.0, -2.0, -3.0, -4.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q4_big_pos() {
|
||||
test_loop_f32(Q4_0, &[1234.0]);
|
||||
test_loop_f16(Q4_0, &[1234.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q4_big_neg() {
|
||||
test_loop_f32(Q4_0, &[-1234.0]);
|
||||
test_loop_f16(Q4_0, &[-1234.0]);
|
||||
}
|
||||
|
||||
fn test_extract_f32(b: impl BlockQuant, data: &[f32]) {
|
||||
let mut input = data.to_vec();
|
||||
while input.len() % b.block_len() != 0 {
|
||||
input.push(0f32);
|
||||
}
|
||||
let quant = b.quant_f32(&input).unwrap();
|
||||
for (ix, v) in data.iter().enumerate() {
|
||||
assert_eq!(b.extract_at_offset_f32(&quant, ix).round(), *v);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_q40f32_pos() {
|
||||
let data = (1..)
|
||||
.map(|i| ((i % 14) - 6) as f32)
|
||||
.take(5 * Q4_0.block_len())
|
||||
.collect_vec();
|
||||
test_extract_f32(Q4_0, &data);
|
||||
}
|
||||
|
||||
fn test_pack_then_extract_panel(
|
||||
q: impl BlockQuant,
|
||||
k: usize,
|
||||
m: usize,
|
||||
r: usize,
|
||||
zip: usize,
|
||||
scales_at_end: bool,
|
||||
) -> TractResult<()> {
|
||||
let weights_orig =
|
||||
Array2::from_shape_fn((m, k), |(m, k)| ((m * 31 + k * 17) % 20) as f32 - 10.)
|
||||
.into_tensor();
|
||||
let weights_f32 = q
|
||||
.dequant_f32(&q.quant_f32(weights_orig.try_as_plain()?.as_slice::<f32>()?)?)?
|
||||
.into_shape(&[m, k])?;
|
||||
let packer = PackedFormat::new(f32::datum_type(), r, 128);
|
||||
let packed_f32 = packer.pack_tensor(&weights_f32, 1, 0)?;
|
||||
|
||||
let q4 = q.quant_f32(weights_f32.try_as_plain()?.as_slice::<f32>()?)?;
|
||||
let packed_q4 = q.pack(&q4, k, r, zip, scales_at_end)?;
|
||||
|
||||
for panel in 0..packed_f32.panels_count() {
|
||||
unsafe {
|
||||
let panel_f32 = packed_f32.panel_bytes(panel, None)?;
|
||||
let panel_f32 = std::slice::from_raw_parts(panel_f32 as *const f32, k * r);
|
||||
let mut panel_q4 = Tensor::zero::<f32>(&[k * r])?;
|
||||
q.extract_packed_panel(
|
||||
&packed_q4,
|
||||
&packer,
|
||||
panel,
|
||||
panel_q4.as_bytes_mut().as_mut_ptr(),
|
||||
)?;
|
||||
assert_eq!(panel_q4.try_as_plain()?.as_slice::<f32>()?, panel_f32);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_panel() -> TractResult<()> {
|
||||
test_pack_then_extract_panel(BaseQ4_0::<2>, 4, 4, 2, 0, false)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_panel_with_zip() -> TractResult<()> {
|
||||
test_pack_then_extract_panel(BaseQ4_0::<2>, 2, 8, 8, 4, false)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_panel_with_scales_at_end() -> TractResult<()> {
|
||||
test_pack_then_extract_panel(BaseQ4_0::<2>, 2, 4, 4, 0, true)
|
||||
}
|
||||
|
||||
fn test_pack_then_extract_row(
|
||||
q: impl BlockQuant,
|
||||
k: usize,
|
||||
m: usize,
|
||||
r: usize,
|
||||
zip: usize,
|
||||
scales_at_end: bool,
|
||||
) -> TractResult<()> {
|
||||
let weights_orig =
|
||||
Array2::from_shape_fn((m, k), |(m, k)| ((m * 31 + k * 17) % 20) as f32 - 10.)
|
||||
.into_tensor();
|
||||
let weights_f32 = q
|
||||
.dequant_f32(&q.quant_f32(weights_orig.try_as_plain()?.as_slice::<f32>()?)?)?
|
||||
.into_shape(&[m, k])?;
|
||||
let packer = PackedFormat::new(f32::datum_type(), r, 128);
|
||||
let packed_f32 = packer.pack_tensor(&weights_f32, 1, 0)?;
|
||||
|
||||
let q4 = q.quant_f32(weights_f32.try_as_plain()?.as_slice::<f32>()?)?;
|
||||
let packed_q4 = q.pack(&q4, k, r, zip, scales_at_end)?;
|
||||
|
||||
for row in 0..packed_f32.mn() {
|
||||
unsafe {
|
||||
let panel_f32 = packed_f32.panel_bytes(row / r, None)?;
|
||||
let panel_f32 = std::slice::from_raw_parts(panel_f32 as *const f32, k * r);
|
||||
let row_f32 = (0..k).map(|ix| panel_f32[row % r + r * ix]).collect_vec();
|
||||
|
||||
let mut q4 = vec![0f32; k];
|
||||
q.extract_at_mn_f32(&packed_q4, row, &mut q4)?;
|
||||
assert_eq!(q4, row_f32);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_row() -> TractResult<()> {
|
||||
test_pack_then_extract_row(BaseQ4_0::<2>, 4, 4, 2, 0, false)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_row_with_zip() -> TractResult<()> {
|
||||
test_pack_then_extract_row(BaseQ4_0::<2>, 2, 8, 8, 4, false)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_row_with_scales_at_end() -> TractResult<()> {
|
||||
test_pack_then_extract_row(BaseQ4_0::<2>, 2, 4, 4, 0, true)
|
||||
}
|
||||
}
|
||||
Vendored
+523
@@ -0,0 +1,523 @@
|
||||
use crate::mmm::PackedExoticFact;
|
||||
|
||||
use super::*;
|
||||
use num_traits::{AsPrimitive, Float, Zero};
|
||||
use std::alloc::Layout;
|
||||
use std::ops::AddAssign;
|
||||
|
||||
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
|
||||
pub struct BaseQ8_1<const QK: usize = 32>;
|
||||
|
||||
pub const Q8_1: BaseQ8_1 = BaseQ8_1::<32>;
|
||||
|
||||
impl<const QK: usize> Debug for BaseQ8_1<QK> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
if QK == 32 {
|
||||
write!(f, "Q8_1")
|
||||
} else {
|
||||
write!(f, "BaseQ8_1<{QK}>")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<const QK: usize> BaseQ8_1<QK> {
|
||||
fn quant_block<T>(&self, block: &[T], quant: &mut [u8])
|
||||
where
|
||||
f32: AsPrimitive<i8> + From<T>,
|
||||
T: Debug + Float + AsPrimitive<f16> + AddAssign + 'static,
|
||||
{
|
||||
assert!(quant.len() == self.block_bytes());
|
||||
assert!(block.len() == self.block_len());
|
||||
let mut writer = NibbleWriter::for_slice(quant);
|
||||
let mut amax = T::zero();
|
||||
let mut max = T::zero();
|
||||
let mut sum = T::zero();
|
||||
for v in block {
|
||||
if amax < v.abs() {
|
||||
amax = v.abs();
|
||||
max = *v;
|
||||
}
|
||||
sum += *v;
|
||||
}
|
||||
|
||||
let scale = f32::from(max) / 127f32;
|
||||
let r_scale = if scale.is_zero() { 0f32 } else { scale.recip() };
|
||||
writer.write_f16(f16::from_f32(scale));
|
||||
writer.write_f16(sum.as_());
|
||||
|
||||
for val_f in block {
|
||||
let i: i8 = (f32::from(*val_f) * r_scale).round().as_();
|
||||
writer.write_i8(i);
|
||||
}
|
||||
}
|
||||
|
||||
fn dequant_block<T: Float + 'static>(&self, quant: &[u8], block: &mut [T])
|
||||
where
|
||||
f16: AsPrimitive<T>,
|
||||
i8: AsPrimitive<T>,
|
||||
{
|
||||
assert!(quant.len() == self.block_bytes());
|
||||
assert!(block.len() == self.block_len());
|
||||
let mut quants = NibbleReader::for_slice(quant);
|
||||
let d: T = quants.read_f16().as_();
|
||||
let _sum: T = quants.read_f16().as_();
|
||||
for val_f in block {
|
||||
*val_f = (quants.read_i8()).as_() * d;
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn extract_panel_t<T: Float + Debug + 'static>(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
target: &PackedFormat,
|
||||
panel: usize,
|
||||
scratch: *mut u8,
|
||||
) -> TractResult<()>
|
||||
where
|
||||
f16: AsPrimitive<T>,
|
||||
i8: AsPrimitive<T>,
|
||||
{
|
||||
let pbqf: &PackedBlockQuantFormat =
|
||||
value.fact.format.downcast_ref().with_context(|| {
|
||||
format!(
|
||||
"Expecing PackedBlockQuantFormat, found {:?}",
|
||||
value.fact.format
|
||||
)
|
||||
})?;
|
||||
ensure!(pbqf.r == target.r);
|
||||
ensure!(value.fact.k % self.block_len() == 0);
|
||||
ensure!(*pbqf.bq == *(self as &dyn BlockQuant));
|
||||
let scratch =
|
||||
unsafe { std::slice::from_raw_parts_mut(scratch as *mut T, value.fact.k * target.r) };
|
||||
let blocks_for_k = value.fact.k / self.block_len();
|
||||
let row_bytes = blocks_for_k * self.block_bytes();
|
||||
let input = &value.packed[panel * target.r * row_bytes..];
|
||||
let mut scales = vec![T::zero(); target.r];
|
||||
let mut scratch = scratch.iter_mut();
|
||||
let mut weights = vec![0i8; pbqf.r];
|
||||
let panel_block_bytes = target.r * self.block_bytes();
|
||||
let (params_offset, weights_offset) = if pbqf.scales_at_end {
|
||||
(
|
||||
panel_block_bytes - target.r * 2 * f16::datum_type().size_of(),
|
||||
0,
|
||||
)
|
||||
} else {
|
||||
(0, target.r * 2 * f16::datum_type().size_of())
|
||||
};
|
||||
for block in 0..blocks_for_k {
|
||||
let block = &input[block * panel_block_bytes..][..panel_block_bytes];
|
||||
let mut s_reader = NibbleReader::for_slice(&block[params_offset..]);
|
||||
let mut w_reader = NibbleReader::for_slice(&block[weights_offset..]);
|
||||
// Layout: [scale_0, sum_0, scale_1, sum_1, .., weights]
|
||||
for s in &mut scales {
|
||||
*s = s_reader.read_f16().as_();
|
||||
// Unused sums
|
||||
s_reader.read_f16();
|
||||
}
|
||||
|
||||
for _ in 0..self.block_len() {
|
||||
for w in &mut weights {
|
||||
*w = w_reader.read_i8();
|
||||
}
|
||||
for (w, s) in weights.iter().zip(scales.iter()) {
|
||||
*scratch.next().unwrap() = *s * (*w).as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn extract_at_mn_t<T: Float + Debug + 'static>(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [T],
|
||||
) -> TractResult<()>
|
||||
where
|
||||
f16: AsPrimitive<T>,
|
||||
i8: AsPrimitive<T>,
|
||||
{
|
||||
let pbqf: &PackedBlockQuantFormat =
|
||||
value.fact.format.downcast_ref().with_context(|| {
|
||||
format!(
|
||||
"Expecing PackedBlockQuantFormat, found {:?}",
|
||||
value.fact.format
|
||||
)
|
||||
})?;
|
||||
ensure!(value.fact.k % self.block_len() == 0);
|
||||
ensure!(*pbqf.bq == *(self as &dyn BlockQuant));
|
||||
ensure!(
|
||||
value
|
||||
.fact
|
||||
.mn
|
||||
.to_usize()
|
||||
.ok()
|
||||
.map(|it| mn < it)
|
||||
.unwrap_or(true)
|
||||
);
|
||||
ensure!(value.fact.k == target.len());
|
||||
let blocks_for_k = value.fact.k / self.block_len();
|
||||
let row_bytes = blocks_for_k * self.block_bytes();
|
||||
let panel = mn / pbqf.r;
|
||||
let value = &value.packed[panel * pbqf.r * row_bytes..];
|
||||
let mut target = target.iter_mut();
|
||||
let panel_block_bytes = pbqf.r * self.block_bytes();
|
||||
let (scale_offset, weights_offset) = if pbqf.scales_at_end {
|
||||
(
|
||||
panel_block_bytes - pbqf.r * 2 * f16::datum_type().size_of(),
|
||||
0,
|
||||
)
|
||||
} else {
|
||||
(0, pbqf.r * 2 * f16::datum_type().size_of())
|
||||
};
|
||||
unsafe {
|
||||
for block in 0..blocks_for_k {
|
||||
let block = value.as_ptr().add(block * panel_block_bytes);
|
||||
let scale = *((block.add(scale_offset) as *const f16).add(2 * (mn % pbqf.r)));
|
||||
let scale: T = scale.as_();
|
||||
for i in 0..self.block_len() {
|
||||
let byte = *block.add(weights_offset + i * pbqf.r + mn % pbqf.r);
|
||||
*target.next().unwrap() = scale * (byte as i8).as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<const QK: usize> BlockQuant for BaseQ8_1<QK> {
|
||||
fn block_len(&self) -> usize {
|
||||
QK
|
||||
}
|
||||
|
||||
fn block_bytes(&self) -> usize {
|
||||
4 + self.block_len()
|
||||
}
|
||||
|
||||
fn quant_block_f32(&self, block: &[f32], quant: &mut [u8]) {
|
||||
self.quant_block(block, quant)
|
||||
}
|
||||
|
||||
fn quant_block_f16(&self, block: &[f16], quant: &mut [u8]) {
|
||||
self.quant_block(block, quant)
|
||||
}
|
||||
|
||||
fn dequant_block_f32(&self, quant: &[u8], block: &mut [f32]) {
|
||||
self.dequant_block(quant, block)
|
||||
}
|
||||
|
||||
fn dequant_block_f16(&self, quant: &[u8], block: &mut [f16]) {
|
||||
self.dequant_block(quant, block)
|
||||
}
|
||||
|
||||
// s0_0 sum_0_0 n0_0 n0_1 n0_2 n0_3 ... n0_30n0_31 s0_32 sum0_32 n0_32n0_33 ...
|
||||
// s1_0 sum_1_0 n1_0 n1_1 n1_2 n1_3 ... n1_30n1_31 s1_32 sum_1_32 n1_32n1_33 ...
|
||||
//
|
||||
// becomes (with r=4)
|
||||
//
|
||||
// s0_0 sum0_0 s1_0 sum 1_0 s2_0 sum2_0 s3_0 sum3_0 n0_0 n1_0 n2_0 n3_0 n0_1 n1_1 n2_1 n3_1 ... n0_33 n1_33 n2_33 n3_33
|
||||
// s0_32 sum0_32 s1_32 sum 1_32 s2_32 sum2_32 s3_32 sum3_32 n0_0 n1_0 n2_0 n3_0 n0_1 n1_1 n2_1 n3_1 ... n0_33 n1_33 n2_33 n3_33
|
||||
// ...
|
||||
fn pack(
|
||||
&self,
|
||||
input: &[u8],
|
||||
k: usize,
|
||||
r: usize,
|
||||
zip: usize,
|
||||
scales_at_end: bool,
|
||||
) -> TractResult<EagerPackedInput> {
|
||||
ensure!(input.len() % self.block_bytes() == 0);
|
||||
ensure!(k % self.block_len() == 0);
|
||||
ensure!(zip == 0, "No zipping required for Q8_1");
|
||||
let m = if input.len() == 0 {
|
||||
0
|
||||
} else {
|
||||
input.len() / self.block_bytes() * self.block_len() / k
|
||||
};
|
||||
let panels = m.divceil(r);
|
||||
let blocks_for_k = k / self.block_len();
|
||||
let row_bytes = blocks_for_k * self.block_bytes();
|
||||
let panel_bytes = row_bytes * r;
|
||||
let mut blob =
|
||||
unsafe { Blob::for_layout(Layout::from_size_align(panel_bytes * panels, 128)?) };
|
||||
let mut writer = NibbleWriter::for_slice(&mut blob);
|
||||
let mut scales = vec![f16::zero(); r];
|
||||
let mut sums = vec![f16::zero(); r];
|
||||
for p in 0..panels {
|
||||
let input = &input[(r * p) * row_bytes..];
|
||||
let mut readers = (0..r)
|
||||
.map(|r| {
|
||||
// manage partial panel
|
||||
let offset = if r * row_bytes < input.len() {
|
||||
r * row_bytes
|
||||
} else {
|
||||
0
|
||||
};
|
||||
NibbleReader::for_slice(&input[offset..])
|
||||
})
|
||||
.collect_vec();
|
||||
let mut temp_quants = vec![vec![0i8; self.block_len()]; r];
|
||||
for _ in 0..blocks_for_k {
|
||||
for (row, reader) in readers.iter_mut().enumerate() {
|
||||
scales[row] = reader.read_f16();
|
||||
sums[row] = reader.read_f16();
|
||||
temp_quants[row] = (0..self.block_len())
|
||||
.map(|_| reader.read_i8())
|
||||
.collect_vec();
|
||||
}
|
||||
if !scales_at_end {
|
||||
scales.iter().zip(&sums).for_each(|(scale, sum)| {
|
||||
writer.write_f16(*scale);
|
||||
writer.write_f16(*sum);
|
||||
});
|
||||
}
|
||||
for pos in 0..self.block_len() {
|
||||
for row in &temp_quants {
|
||||
let q = row[pos];
|
||||
writer.write_i8(q);
|
||||
}
|
||||
}
|
||||
if scales_at_end {
|
||||
scales.iter().zip(&sums).for_each(|(scale, sum)| {
|
||||
writer.write_f16(*scale);
|
||||
writer.write_f16(*sum);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(EagerPackedInput {
|
||||
fact: PackedExoticFact {
|
||||
format: Box::new(PackedBlockQuantFormat {
|
||||
bq: Box::new(*self),
|
||||
r,
|
||||
zip,
|
||||
scales_at_end,
|
||||
}),
|
||||
mn: m.to_dim(),
|
||||
k,
|
||||
},
|
||||
packed: blob.into(),
|
||||
panel_bytes,
|
||||
mn: m,
|
||||
})
|
||||
}
|
||||
|
||||
unsafe fn extract_packed_panel(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
target: &PackedFormat,
|
||||
panel: usize,
|
||||
scratch: *mut u8,
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
dispatch_floatlike!(Self::extract_panel_t(target.dt)(
|
||||
self, value, target, panel, scratch
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_at_mn_f16(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [f16],
|
||||
) -> TractResult<()> {
|
||||
self.extract_at_mn_t(value, mn, target)
|
||||
}
|
||||
|
||||
fn extract_at_mn_f32(
|
||||
&self,
|
||||
value: &EagerPackedInput,
|
||||
mn: usize,
|
||||
target: &mut [f32],
|
||||
) -> TractResult<()> {
|
||||
self.extract_at_mn_t(value, mn, target)
|
||||
}
|
||||
}
|
||||
|
||||
impl<const QK: usize> Display for BaseQ8_1<QK> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "Q8_1")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use num_traits::Zero;
|
||||
use tract_data::internal::tract_ndarray::Array2;
|
||||
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn test_loop_f32(b: impl BlockQuant, data: &[f32]) -> TractResult<()> {
|
||||
let mut input = data.to_vec();
|
||||
while input.len() % b.block_len() != 0 {
|
||||
input.push(0f32);
|
||||
}
|
||||
let ref_tensor = unsafe { Tensor::from_slice_align(&input, vector_size())? };
|
||||
|
||||
let quant = b.quant_f32(&input).unwrap();
|
||||
let result = b.dequant_f32(&quant).unwrap();
|
||||
result.close_enough(&ref_tensor, Approximation::VeryApproximate)
|
||||
}
|
||||
|
||||
fn test_loop_f16(b: impl BlockQuant, data: &[f32]) -> TractResult<()> {
|
||||
let mut input = data.iter().map(|f| f16::from_f32(*f)).collect_vec();
|
||||
while input.len() % b.block_len() != 0 {
|
||||
input.push(f16::zero());
|
||||
}
|
||||
let ref_tensor = unsafe { Tensor::from_slice_align(&input, vector_size())? };
|
||||
|
||||
let quant = b.quant_f16(&input).unwrap();
|
||||
let result = b.dequant_f16(&quant).unwrap();
|
||||
result.close_enough(&ref_tensor, Approximation::VeryApproximate)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q81f32_pos() -> TractResult<()> {
|
||||
test_loop_f32(Q8_1, &[1.0, 2.0, 3.0, 4.0])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q81f16_pos() -> TractResult<()> {
|
||||
test_loop_f16(Q8_1, &[1.0, 2.0, 3.0, 4.0])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q81f32_neg() -> TractResult<()> {
|
||||
test_loop_f32(Q8_1, &[-1.0, -2.0, -3.0, -4.0])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q81f16_neg() -> TractResult<()> {
|
||||
test_loop_f16(Q8_1, &[-1.0, -2.0, -3.0, -4.0])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q81_big_pos() -> TractResult<()> {
|
||||
test_loop_f32(Q8_1, &[1234.0])?;
|
||||
test_loop_f16(Q8_1, &[1234.0])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loop_q81_big_neg() -> TractResult<()> {
|
||||
test_loop_f32(Q8_1, &[-1234.0])?;
|
||||
test_loop_f16(Q8_1, &[-1234.0])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn test_extract_f32(b: impl BlockQuant, data: &[f32]) {
|
||||
let mut input = data.to_vec();
|
||||
while input.len() % b.block_len() != 0 {
|
||||
input.push(0f32);
|
||||
}
|
||||
let quant = b.quant_f32(&input).unwrap();
|
||||
for (ix, v) in data.iter().enumerate() {
|
||||
assert_eq!(b.extract_at_offset_f32(&quant, ix).round(), *v);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_q81f32_pos() {
|
||||
let data = (1..)
|
||||
.map(|i| ((i % 14) - 6) as f32)
|
||||
.take(5 * Q8_1.block_len())
|
||||
.collect_vec();
|
||||
test_extract_f32(Q8_1, &data);
|
||||
}
|
||||
|
||||
fn test_pack_then_extract_panel(
|
||||
q: impl BlockQuant,
|
||||
k: usize,
|
||||
m: usize,
|
||||
r: usize,
|
||||
scales_at_end: bool,
|
||||
) -> TractResult<()> {
|
||||
let weights_orig =
|
||||
Array2::from_shape_fn((m, k), |(m, k)| ((m * 31 + k * 17) % 20) as f32 - 10.)
|
||||
.into_tensor();
|
||||
let weights_f32 = q
|
||||
.dequant_f32(&q.quant_f32(weights_orig.try_as_plain()?.as_slice::<f32>()?)?)?
|
||||
.into_shape(&[m, k])?;
|
||||
let packer = PackedFormat::new(f32::datum_type(), r, 128);
|
||||
let packed_f32 = packer.pack_tensor(&weights_f32, 1, 0)?;
|
||||
|
||||
let q81 = q.quant_f32(weights_f32.try_as_plain()?.as_slice::<f32>()?)?;
|
||||
let packed_q81 = q.pack(&q81, k, r, 0, scales_at_end)?;
|
||||
|
||||
for panel in 0..packed_f32.panels_count() {
|
||||
unsafe {
|
||||
let panel_f32 = packed_f32.panel_bytes(panel, None)?;
|
||||
let panel_f32 = std::slice::from_raw_parts(panel_f32 as *const f32, k * r);
|
||||
let mut panel_q81 = Tensor::zero::<f32>(&[k * r])?;
|
||||
q.extract_packed_panel(
|
||||
&packed_q81,
|
||||
&packer,
|
||||
panel,
|
||||
panel_q81.as_bytes_mut().as_mut_ptr(),
|
||||
)?;
|
||||
assert_eq!(panel_q81.try_as_plain()?.as_slice::<f32>()?, panel_f32);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_panel() -> TractResult<()> {
|
||||
test_pack_then_extract_panel(BaseQ8_1::<2>, 4, 4, 2, false)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_panel_with_scales_at_end() -> TractResult<()> {
|
||||
test_pack_then_extract_panel(BaseQ8_1::<2>, 2, 4, 4, true)
|
||||
}
|
||||
|
||||
fn test_pack_then_extract_row(
|
||||
q: impl BlockQuant,
|
||||
k: usize,
|
||||
m: usize,
|
||||
r: usize,
|
||||
scales_at_end: bool,
|
||||
) -> TractResult<()> {
|
||||
let weights_orig =
|
||||
Array2::from_shape_fn((m, k), |(m, k)| ((m * 31 + k * 17) % 20) as f32 - 10.)
|
||||
.into_tensor();
|
||||
let weights_f32 = q
|
||||
.dequant_f32(&q.quant_f32(weights_orig.try_as_plain()?.as_slice::<f32>()?)?)?
|
||||
.into_shape(&[m, k])?;
|
||||
let packer = PackedFormat::new(f32::datum_type(), r, 128);
|
||||
let packed_f32 = packer.pack_tensor(&weights_f32, 1, 0)?;
|
||||
|
||||
let q81 = q.quant_f32(weights_f32.try_as_plain()?.as_slice::<f32>()?)?;
|
||||
let packed_q81 = q.pack(&q81, k, r, 0, scales_at_end)?;
|
||||
|
||||
for row in 0..packed_f32.mn() {
|
||||
unsafe {
|
||||
let panel_f32 = packed_f32.panel_bytes(row / r, None)?;
|
||||
let panel_f32 = std::slice::from_raw_parts(panel_f32 as *const f32, k * r);
|
||||
let row_f32 = (0..k).map(|ix| panel_f32[row % r + r * ix]).collect_vec();
|
||||
|
||||
let mut q81 = vec![0f32; k];
|
||||
q.extract_at_mn_f32(&packed_q81, row, &mut q81)?;
|
||||
assert_eq!(q81, row_f32);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_row() -> TractResult<()> {
|
||||
test_pack_then_extract_row(BaseQ8_1::<2>, 4, 4, 2, false)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_then_extract_row_with_scales_at_end() -> TractResult<()> {
|
||||
test_pack_then_extract_row(BaseQ8_1::<2>, 2, 4, 4, true)
|
||||
}
|
||||
}
|
||||
Vendored
+134
@@ -0,0 +1,134 @@
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use tract_data::internal::*;
|
||||
|
||||
use super::BlockQuant;
|
||||
use super::BlockQuantFact;
|
||||
|
||||
/// Concrete tensor storage for block-quantized weights.
|
||||
///
|
||||
/// Stores a single contiguous `Arc<Blob>` of quantized data along with the
|
||||
/// block-quant format. Shape lives on the tensor, not here.
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct BlockQuantStorage {
|
||||
format: Box<dyn BlockQuant>,
|
||||
data: Arc<Blob>,
|
||||
}
|
||||
|
||||
impl BlockQuantStorage {
|
||||
fn expected_bytes(format: &dyn BlockQuant, m: usize, k: usize) -> usize {
|
||||
m * k / format.block_len() * format.block_bytes()
|
||||
}
|
||||
|
||||
pub fn new(
|
||||
format: Box<dyn BlockQuant>,
|
||||
m: usize,
|
||||
k: usize,
|
||||
data: Arc<Blob>,
|
||||
) -> TractResult<Self> {
|
||||
let expected = Self::expected_bytes(&*format, m, k);
|
||||
ensure!(
|
||||
data.len() == expected,
|
||||
"BlockQuantStorage::new: blob length {} does not match expected {} (m={}, k={}, format={})",
|
||||
data.len(),
|
||||
expected,
|
||||
m,
|
||||
k,
|
||||
format,
|
||||
);
|
||||
Ok(Self { format, data })
|
||||
}
|
||||
|
||||
pub fn format(&self) -> &dyn BlockQuant {
|
||||
&*self.format
|
||||
}
|
||||
|
||||
/// Returns the single contiguous blob.
|
||||
pub fn value(&self) -> &Arc<Blob> {
|
||||
&self.data
|
||||
}
|
||||
|
||||
/// Converts this storage into a `Tensor` with the given shape.
|
||||
///
|
||||
/// `dt` is the logical element type (e.g. f32, f16) — the type these
|
||||
/// weights represent when dequantized.
|
||||
pub fn into_tensor_with_shape(self, dt: DatumType, shape: &[usize]) -> Tensor {
|
||||
Tensor::from_storage(dt, shape, self)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a byte slice for a single group within contiguous block-quant data.
|
||||
pub fn block_quant_slice<'a>(
|
||||
data: &'a [u8],
|
||||
format: &dyn BlockQuant,
|
||||
m_per_group: usize,
|
||||
k: usize,
|
||||
g: usize,
|
||||
) -> &'a [u8] {
|
||||
let row_bytes = k / format.block_len() * format.block_bytes();
|
||||
let group_bytes = m_per_group * row_bytes;
|
||||
let start = g * group_bytes;
|
||||
&data[start..start + group_bytes]
|
||||
}
|
||||
|
||||
impl fmt::Debug for BlockQuantStorage {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"BlockQuantStorage({}, bytes={})",
|
||||
self.format,
|
||||
self.data.len()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for BlockQuantStorage {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"BlockQuantStorage({}, bytes={})",
|
||||
self.format,
|
||||
self.data.len()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl TensorStorage for BlockQuantStorage {
|
||||
fn byte_len(&self) -> usize {
|
||||
self.data.len()
|
||||
}
|
||||
|
||||
fn is_empty(&self) -> bool {
|
||||
self.data.is_empty()
|
||||
}
|
||||
|
||||
fn deep_clone(&self) -> Box<dyn TensorStorage> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
|
||||
fn as_plain(&self) -> Option<&PlainStorage> {
|
||||
None
|
||||
}
|
||||
|
||||
fn as_plain_mut(&mut self) -> Option<&mut PlainStorage> {
|
||||
None
|
||||
}
|
||||
|
||||
fn into_plain(self: Box<Self>) -> Option<PlainStorage> {
|
||||
None
|
||||
}
|
||||
|
||||
fn dyn_hash(&self, state: &mut dyn std::hash::Hasher) {
|
||||
state.write_u8(1);
|
||||
self.format.dyn_hash(state);
|
||||
state.write(self.data.as_bytes());
|
||||
}
|
||||
|
||||
fn exotic_fact(&self, shape: &[usize]) -> TractResult<Option<Box<dyn ExoticFact>>> {
|
||||
Ok(Some(Box::new(BlockQuantFact::new(
|
||||
dyn_clone::clone_box(&*self.format),
|
||||
shape.into(),
|
||||
))))
|
||||
}
|
||||
}
|
||||
Vendored
+83
@@ -0,0 +1,83 @@
|
||||
use super::{BlockQuant, PackedBlockQuantFormat};
|
||||
use tract_data::TVec;
|
||||
use tract_data::internal::*;
|
||||
|
||||
#[allow(clippy::derived_hash_with_manual_eq)]
|
||||
#[derive(Clone, Hash)]
|
||||
pub struct BlockQuantFact {
|
||||
pub format: Box<dyn BlockQuant>,
|
||||
shape: TVec<usize>,
|
||||
}
|
||||
impl BlockQuantFact {
|
||||
pub fn new(format: Box<dyn BlockQuant>, shape: TVec<usize>) -> Self {
|
||||
Self { format, shape }
|
||||
}
|
||||
|
||||
/// Product of all leading dims except the last two (M, K).
|
||||
/// For rank <= 2, returns 1.
|
||||
pub fn num_groups(&self) -> usize {
|
||||
if self.shape.len() <= 2 {
|
||||
1
|
||||
} else {
|
||||
self.shape[..self.shape.len() - 2].iter().product()
|
||||
}
|
||||
}
|
||||
|
||||
/// Product of all dims except the last (K). This is the flat M
|
||||
/// dimension (groups * m_per_group).
|
||||
pub fn m(&self) -> usize {
|
||||
self.shape[..self.shape.len() - 1].iter().product()
|
||||
}
|
||||
|
||||
/// Last dimension.
|
||||
pub fn k(&self) -> usize {
|
||||
*self.shape.last().unwrap()
|
||||
}
|
||||
|
||||
pub fn shape(&self) -> &[usize] {
|
||||
&self.shape
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for BlockQuantFact {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}({:?})", self.format, self.shape)
|
||||
}
|
||||
}
|
||||
|
||||
impl ExoticFact for BlockQuantFact {
|
||||
fn buffer_sizes(&self) -> TVec<TDim> {
|
||||
let total = self.m() * self.k() / self.format.block_len() * self.format.block_bytes();
|
||||
tvec!(total.to_dim())
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for BlockQuantFact {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
*self.format == *other.format && self.shape == other.shape
|
||||
}
|
||||
}
|
||||
impl Eq for BlockQuantFact {}
|
||||
|
||||
#[derive(Clone, Hash, PartialEq)]
|
||||
pub struct PackedBlockQuantFact {
|
||||
pub format: PackedBlockQuantFormat,
|
||||
pub shape: TVec<usize>,
|
||||
}
|
||||
impl Eq for PackedBlockQuantFact {}
|
||||
|
||||
impl std::fmt::Debug for PackedBlockQuantFact {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}({:?})", self.format, self.shape)
|
||||
}
|
||||
}
|
||||
|
||||
impl ExoticFact for PackedBlockQuantFact {
|
||||
fn buffer_sizes(&self) -> TVec<TDim> {
|
||||
tvec!(
|
||||
(self.shape.iter().product::<usize>() / self.format.bq.block_len()
|
||||
* self.format.bq.block_bytes())
|
||||
.to_dim()
|
||||
)
|
||||
}
|
||||
}
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
use std::fmt::Debug;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use crate::element_wise::{ElementWise, ElementWiseKer};
|
||||
use crate::element_wise_helper::map_slice_with_alignment;
|
||||
use crate::{LADatum, LinalgFn};
|
||||
use tract_data::internal::*;
|
||||
|
||||
/// Generic implementation struct that unify all by scalar kernels.
|
||||
/// A by scalar operation is an ElementWise operation with a scalar paramerer.
|
||||
#[derive(Debug, Clone, new)]
|
||||
pub struct ByScalarImpl<K, T>
|
||||
where
|
||||
T: LADatum,
|
||||
K: ByScalarKer<T> + Clone,
|
||||
{
|
||||
phantom: PhantomData<(K, T)>,
|
||||
}
|
||||
|
||||
impl<K, T> ElementWise<T, T> for ByScalarImpl<K, T>
|
||||
where
|
||||
T: LADatum,
|
||||
K: ByScalarKer<T> + Clone,
|
||||
{
|
||||
fn name(&self) -> &'static str {
|
||||
K::name()
|
||||
}
|
||||
fn run_with_params(&self, vec: &mut [T], params: T) -> TractResult<()> {
|
||||
map_slice_with_alignment(
|
||||
vec,
|
||||
|data| K::run(data, params),
|
||||
K::nr(),
|
||||
K::alignment_bytes(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ByScalarKer<T>: ElementWiseKer<T, T>
|
||||
where
|
||||
T: LADatum,
|
||||
{
|
||||
fn bin() -> Box<LinalgFn> {
|
||||
Box::new(|a: &mut TensorView, b: &TensorView| {
|
||||
let a_slice = a.as_slice_mut()?;
|
||||
let b = b.as_slice()?[0];
|
||||
(Self::ew()).run_with_params(a_slice, b)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! by_scalar_impl_wrap {
|
||||
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $run: item) => {
|
||||
paste! {
|
||||
ew_impl_wrap!($ti, $func, $nr, $alignment_items, $ti, $run);
|
||||
|
||||
impl crate::frame::by_scalar::ByScalarKer<$ti> for $func {}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::element_wise::ElementWiseKer;
|
||||
use num_traits::{AsPrimitive, Float};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! by_scalar_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty, $func:expr) => {
|
||||
pastey::paste! {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn [<prop_ $ker:snake>](xs in proptest::collection::vec(-25f32..25.0, 0..100), scalar in -25f32..25f32) {
|
||||
if $cond {
|
||||
$crate::frame::by_scalar::test::test_by_scalar::<$ker, $t>(&*xs, scalar, $func).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_by_scalar<K: ElementWiseKer<T, T>, T: LADatum + Float>(
|
||||
values: &[f32],
|
||||
scalar: f32,
|
||||
func: impl Fn(T, T) -> T,
|
||||
) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
let values: Vec<T> = values.iter().copied().map(|x| x.as_()).collect();
|
||||
crate::frame::element_wise::test::test_element_wise_params::<K, T, _, T>(
|
||||
&values,
|
||||
|a| (func)(a, scalar.as_()),
|
||||
scalar.as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
use std::fmt::Debug;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use tract_data::TractResult;
|
||||
|
||||
use crate::LADatum;
|
||||
|
||||
use super::element_wise_helper::map_slice_with_alignment;
|
||||
|
||||
macro_rules! ew_impl_wrap {
|
||||
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $run: item) => {
|
||||
paste! {
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
#[allow(non_camel_case_types)]
|
||||
pub struct $func;
|
||||
|
||||
impl crate::frame::element_wise::ElementWiseKer<$ti, $params> for $func {
|
||||
#[inline(always)]
|
||||
fn name() -> &'static str {
|
||||
stringify!($func)
|
||||
}
|
||||
#[inline(always)]
|
||||
fn nr() -> usize {
|
||||
$nr
|
||||
}
|
||||
#[inline(always)]
|
||||
fn alignment_items() -> usize {
|
||||
$alignment_items
|
||||
}
|
||||
$run
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
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) -> ());
|
||||
}
|
||||
ew_impl_wrap!($ti, $func, $nr, $alignment_items, (),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [$ti], _params: ()) {
|
||||
unsafe { [<sys_ $func>]::$func(buf.as_mut_ptr(), buf.len()) }
|
||||
}
|
||||
);
|
||||
}
|
||||
};
|
||||
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty) => {
|
||||
paste! {
|
||||
mod [<sys_ $func>] {
|
||||
#[allow(unused_imports)]
|
||||
use tract_data::prelude::f16;
|
||||
extern_kernel!(fn $func(ptr: *mut $ti, count: usize, params: $params) -> ());
|
||||
}
|
||||
ew_impl_wrap!($ti, $func, $nr, $alignment_items, $params,
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [$ti], params: $params) {
|
||||
unsafe { [<sys_ $func>]::$func(buf.as_mut_ptr(), buf.len(), params) }
|
||||
}
|
||||
);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub trait ElementWise<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
T: Copy + Debug + PartialEq + Send + Sync,
|
||||
{
|
||||
fn name(&self) -> &'static str;
|
||||
fn run(&self, vec: &mut [T]) -> TractResult<()> {
|
||||
self.run_with_params(vec, Params::default())
|
||||
}
|
||||
fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<()>;
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(<T, Params> ElementWise<T, Params> where T: Copy, Params: Copy);
|
||||
|
||||
#[derive(Debug, Clone, new)]
|
||||
pub struct ElementWiseImpl<K, T, Params = ()>
|
||||
where
|
||||
T: LADatum,
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
K: ElementWiseKer<T, Params> + Clone,
|
||||
{
|
||||
phantom: PhantomData<(K, T, Params)>,
|
||||
}
|
||||
|
||||
impl<K, T, Params> ElementWise<T, Params> for ElementWiseImpl<K, T, Params>
|
||||
where
|
||||
T: LADatum,
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
K: ElementWiseKer<T, Params> + Clone,
|
||||
{
|
||||
fn name(&self) -> &'static str {
|
||||
K::name()
|
||||
}
|
||||
fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<()> {
|
||||
map_slice_with_alignment(
|
||||
vec,
|
||||
|data| K::run(data, params),
|
||||
K::nr(),
|
||||
K::alignment_bytes(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ElementWiseKer<T, Params = ()>:
|
||||
Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
T: LADatum,
|
||||
{
|
||||
fn name() -> &'static str;
|
||||
fn alignment_bytes() -> usize {
|
||||
Self::alignment_items() * T::datum_type().size_of()
|
||||
}
|
||||
fn alignment_items() -> usize;
|
||||
fn nr() -> usize;
|
||||
fn run(vec: &mut [T], params: Params);
|
||||
fn ew() -> Box<dyn ElementWise<T, Params>> {
|
||||
Box::new(ElementWiseImpl::<Self, T, Params>::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 {
|
||||
test_element_wise_params::<K, T, F, ()>(values, reference, ())
|
||||
}
|
||||
|
||||
pub fn test_element_wise_params<
|
||||
K: ElementWiseKer<T, Params>,
|
||||
T: LADatum,
|
||||
F: Fn(T) -> T,
|
||||
Params,
|
||||
>(
|
||||
values: &[T],
|
||||
reference: F,
|
||||
params: Params,
|
||||
) -> TestCaseResult
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
let op = ElementWiseImpl::<K, T, Params>::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_with_params(&mut found, params).unwrap();
|
||||
tensor1(&found)
|
||||
.close_enough(&tensor1(&expected), true)
|
||||
.map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Vendored
+172
@@ -0,0 +1,172 @@
|
||||
use crate::LADatum;
|
||||
use std::alloc::*;
|
||||
use tract_data::TractResult;
|
||||
|
||||
pub(crate) fn map_slice_with_alignment<T>(
|
||||
vec: &mut [T],
|
||||
f: impl Fn(&mut [T]),
|
||||
nr: usize,
|
||||
alignment_bytes: usize,
|
||||
) -> TractResult<()>
|
||||
where
|
||||
T: LADatum,
|
||||
{
|
||||
if vec.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
unsafe {
|
||||
TMP.with(|buffer| {
|
||||
let mut buffer = buffer.borrow_mut();
|
||||
buffer.ensure(nr * T::datum_type().size_of(), alignment_bytes);
|
||||
let tmp = std::slice::from_raw_parts_mut(buffer.buffer as *mut T, nr);
|
||||
let mut compute_via_temp_buffer = |slice: &mut [T]| {
|
||||
tmp[..slice.len()].copy_from_slice(slice);
|
||||
f(tmp);
|
||||
slice.copy_from_slice(&tmp[..slice.len()])
|
||||
};
|
||||
let prefix_len = vec.as_ptr().align_offset(alignment_bytes).min(vec.len());
|
||||
if prefix_len > 0 {
|
||||
compute_via_temp_buffer(&mut vec[..prefix_len]);
|
||||
}
|
||||
let aligned_len = (vec.len() - prefix_len) / nr * nr;
|
||||
if aligned_len > 0 {
|
||||
f(&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(crate) fn reduce_slice_with_alignment<T>(
|
||||
vec: &[T],
|
||||
f: impl Fn(&[T]) -> T,
|
||||
nr: usize,
|
||||
alignment_bytes: usize,
|
||||
neutral: T,
|
||||
reduce: impl Fn(T, T) -> T,
|
||||
) -> TractResult<T>
|
||||
where
|
||||
T: LADatum,
|
||||
{
|
||||
if vec.is_empty() {
|
||||
return Ok(neutral);
|
||||
}
|
||||
let mut red = neutral;
|
||||
unsafe {
|
||||
TMP.with(|buffer| {
|
||||
let mut buffer = buffer.borrow_mut();
|
||||
buffer.ensure(nr * T::datum_type().size_of(), alignment_bytes);
|
||||
let tmp = std::slice::from_raw_parts_mut(buffer.buffer as *mut T, nr);
|
||||
let mut compute_via_temp_buffer = |slice: &[T], red: &mut T| {
|
||||
tmp[..slice.len()].copy_from_slice(slice);
|
||||
tmp[slice.len()..].fill(neutral);
|
||||
*red = reduce(*red, f(tmp));
|
||||
};
|
||||
let prefix_len = vec.as_ptr().align_offset(alignment_bytes).min(vec.len());
|
||||
if prefix_len > 0 {
|
||||
compute_via_temp_buffer(&vec[..prefix_len], &mut red);
|
||||
}
|
||||
let aligned_len = (vec.len() - prefix_len) / nr * nr;
|
||||
if aligned_len > 0 {
|
||||
let t = f(&vec[prefix_len..][..aligned_len]);
|
||||
red = reduce(red, t);
|
||||
}
|
||||
if prefix_len + aligned_len < vec.len() {
|
||||
compute_via_temp_buffer(&vec[prefix_len + aligned_len..], &mut red);
|
||||
}
|
||||
})
|
||||
}
|
||||
Ok(red)
|
||||
}
|
||||
|
||||
pub(crate) fn map_reduce_slice_with_alignment<T>(
|
||||
vec: &mut [T],
|
||||
f: impl Fn(&mut [T]) -> T,
|
||||
nr: usize,
|
||||
alignment_bytes: usize,
|
||||
map_neutral: T,
|
||||
neutral: T,
|
||||
reduce: impl Fn(T, T) -> T,
|
||||
) -> TractResult<T>
|
||||
where
|
||||
T: LADatum,
|
||||
{
|
||||
if vec.is_empty() {
|
||||
return Ok(neutral);
|
||||
}
|
||||
let mut red = neutral;
|
||||
unsafe {
|
||||
TMP.with(|buffer| {
|
||||
let mut buffer = buffer.borrow_mut();
|
||||
buffer.ensure(nr * T::datum_type().size_of(), alignment_bytes);
|
||||
let tmp = std::slice::from_raw_parts_mut(buffer.buffer as *mut T, nr);
|
||||
let mut compute_via_temp_buffer = |slice: &mut [T], red: &mut T| {
|
||||
tmp[..slice.len()].copy_from_slice(slice);
|
||||
tmp[slice.len()..].fill(map_neutral);
|
||||
*red = reduce(*red, f(tmp));
|
||||
slice.copy_from_slice(&tmp[..slice.len()]);
|
||||
};
|
||||
let prefix_len = vec.as_ptr().align_offset(alignment_bytes).min(vec.len());
|
||||
if prefix_len > 0 {
|
||||
compute_via_temp_buffer(&mut vec[..prefix_len], &mut red);
|
||||
}
|
||||
let aligned_len = (vec.len() - prefix_len) / nr * nr;
|
||||
if aligned_len > 0 {
|
||||
let t = f(&mut vec[prefix_len..][..aligned_len]);
|
||||
red = reduce(red, t);
|
||||
}
|
||||
if prefix_len + aligned_len < vec.len() {
|
||||
compute_via_temp_buffer(&mut vec[prefix_len + aligned_len..], &mut red);
|
||||
}
|
||||
})
|
||||
}
|
||||
Ok(red)
|
||||
}
|
||||
|
||||
std::thread_local! {
|
||||
static TMP: std::cell::RefCell<TempBuffer> = std::cell::RefCell::new(TempBuffer::default());
|
||||
}
|
||||
|
||||
pub struct TempBuffer {
|
||||
pub layout: Layout,
|
||||
pub buffer: *mut u8,
|
||||
}
|
||||
|
||||
impl Default for TempBuffer {
|
||||
fn default() -> Self {
|
||||
TempBuffer {
|
||||
layout: Layout::new::<()>(),
|
||||
buffer: std::ptr::null_mut(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TempBuffer {
|
||||
pub 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
#[allow(unused_macros)]
|
||||
macro_rules! erf_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::*;
|
||||
erf_frame_tests!($cond, $ti, $func);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::element_wise::*;
|
||||
use num_traits::{AsPrimitive, Float};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! erf_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-5f32..5.0, 0..100)) {
|
||||
if $cond {
|
||||
$crate::frame::erf::test::test_erf::<$ker, $t>(&*xs).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn trivial() {
|
||||
if $cond {
|
||||
$crate::frame::erf::test::test_erf::<$ker, $t>(&[
|
||||
-5f32, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 5.0,
|
||||
])
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn zeros() {
|
||||
if $cond {
|
||||
$crate::frame::erf::test::test_erf::<$ker, $t>(&[0.0; 16]).unwrap();
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_erf<K: ElementWiseKer<T>, T: LADatum + Float>(values: &[f32]) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
T: AsPrimitive<f32>,
|
||||
{
|
||||
let data = tract_data::prelude::tensor1(values);
|
||||
let data = data.cast_to::<T>().unwrap();
|
||||
let data = data.try_as_plain().unwrap().as_slice::<T>().unwrap();
|
||||
crate::frame::element_wise::test::test_element_wise::<K, T, _>(data, |x: T| {
|
||||
// Abramowitz & Stegun 7.1.26 six-coefficient approximation, mirroring
|
||||
// generic/erf.rs::serf so the test reference matches the production scalar path.
|
||||
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 x: f32 = x.as_();
|
||||
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();
|
||||
y.copysign(signum).as_()
|
||||
})
|
||||
}
|
||||
}
|
||||
+61
@@ -0,0 +1,61 @@
|
||||
#[allow(unused_macros)]
|
||||
macro_rules! gelu_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::*;
|
||||
gelu_frame_tests!($cond, $ti, $func);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::element_wise::*;
|
||||
use num_traits::{AsPrimitive, Float};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! gelu_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-10f32..10.0, 0..100)) {
|
||||
if $cond {
|
||||
$crate::frame::gelu::test::test_gelu::<$ker, $t>(&*xs).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn trivial() {
|
||||
if $cond {
|
||||
$crate::frame::gelu::test::test_gelu::<$ker, $t>(&[-5f32, -1.0, 0.0, 1.0, 5.0])
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_gelu<K: ElementWiseKer<T>, T: LADatum + Float>(values: &[f32]) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
let data = tract_data::prelude::tensor1(values);
|
||||
let data = data.cast_to::<T>().unwrap();
|
||||
let data = data.try_as_plain().unwrap().as_slice::<T>().unwrap();
|
||||
// Tanh-form GELU (pow=3): 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
|
||||
crate::frame::element_wise::test::test_element_wise::<K, T, _>(data, |x: T| {
|
||||
let half: T = 0.5f32.as_();
|
||||
let one: T = 1f32.as_();
|
||||
let coef: T = 0.044715f32.as_();
|
||||
let sqrt_2_over_pi: T = 0.7978845608028654f32.as_();
|
||||
let inner = sqrt_2_over_pi * (x + coef * x * x * x);
|
||||
half * x * (one + inner.tanh())
|
||||
})
|
||||
}
|
||||
}
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
#[allow(unused_macros)]
|
||||
macro_rules! hardswish_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::*;
|
||||
hardswish_frame_tests!($cond, $ti, $func);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::element_wise::*;
|
||||
use num_traits::{AsPrimitive, Float, Zero};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! hardswish_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-25f32..25.0, 0..100)) {
|
||||
if $cond {
|
||||
$crate::frame::hardswish::test::test_hardswish::<$ker, $t>(&*xs).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn trivial() {
|
||||
if $cond {
|
||||
$crate::frame::hardswish::test::test_hardswish::<$ker, $t>(&[
|
||||
-10f32, -3.0, -1.0, 0.0, 1.0, 3.0, 6.0, 10.0,
|
||||
])
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_hardswish<K: ElementWiseKer<T>, T: LADatum + Float>(
|
||||
values: &[f32],
|
||||
) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
let data = tract_data::prelude::tensor1(values);
|
||||
let data = data.cast_to::<T>().unwrap();
|
||||
let data = data.try_as_plain().unwrap().as_slice::<T>().unwrap();
|
||||
crate::frame::element_wise::test::test_element_wise::<K, T, _>(data, |x: T| {
|
||||
let three: T = 3f32.as_();
|
||||
let six: T = 6f32.as_();
|
||||
let zero: T = T::zero();
|
||||
let inv6: T = (1f32 / 6f32).as_();
|
||||
let relu6 = ((x + three).min(six)).max(zero);
|
||||
x * relu6 * inv6
|
||||
})
|
||||
}
|
||||
}
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
#[allow(unused_macros)]
|
||||
macro_rules! leaky_relu_impl {
|
||||
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $cond: expr) => {
|
||||
ew_impl!($ti, $func, $nr, $alignment_items, $ti);
|
||||
#[cfg(test)]
|
||||
paste! {
|
||||
mod [<test_ $func>] {
|
||||
use super::*;
|
||||
leaky_relu_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! leaky_relu_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-25f32..25.0, 0..100), alpha in 0f32..1f32) {
|
||||
if $cond {
|
||||
$crate::frame::leaky_relu::test::test_leaky_relu::<$ker, $t>(&*xs, alpha).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn trivial() {
|
||||
if $cond {
|
||||
$crate::frame::leaky_relu::test::test_leaky_relu::<$ker, $t>(&[-10f32], 0.0496).unwrap();
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_leaky_relu<K: ElementWiseKer<T, T>, T: LADatum + Float>(
|
||||
values: &[f32],
|
||||
alpha: f32,
|
||||
) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
let data = tract_data::prelude::tensor1(values);
|
||||
let data = data.cast_to::<T>().unwrap();
|
||||
let data = data.try_as_plain().unwrap().as_slice::<T>().unwrap();
|
||||
let alpha: T = tract_data::prelude::tensor0(alpha)
|
||||
.cast_to_scalar::<T>()
|
||||
.unwrap();
|
||||
crate::frame::element_wise::test::test_element_wise_params::<K, T, _, T>(
|
||||
data,
|
||||
|x: T| {
|
||||
if x > T::zero() { x } else { alpha * x }
|
||||
},
|
||||
alpha,
|
||||
)
|
||||
}
|
||||
}
|
||||
+151
@@ -0,0 +1,151 @@
|
||||
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 {
|
||||
fn table(&self) -> &[u8];
|
||||
fn run(&self, buf: &mut [u8]);
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(Lut);
|
||||
|
||||
impl PartialEq for dyn Lut {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.table() == other.table()
|
||||
}
|
||||
}
|
||||
impl Eq for dyn Lut {}
|
||||
|
||||
#[derive(Debug, Clone, Hash)]
|
||||
pub struct LutImpl<K: LutKer> {
|
||||
table: Tensor,
|
||||
_boo: PhantomData<K>,
|
||||
}
|
||||
|
||||
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.try_as_plain().unwrap().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).next_multiple_of(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())
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
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 CostModel<'_> {
|
||||
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.name() == choice).unwrap().clone()
|
||||
} else {
|
||||
impls
|
||||
.iter()
|
||||
.find(|k| k.name() == self.big_product_kernel_choice)
|
||||
.unwrap()
|
||||
.clone()
|
||||
}
|
||||
}
|
||||
}
|
||||
+135
@@ -0,0 +1,135 @@
|
||||
use std::fmt::Debug;
|
||||
use std::ops::Deref;
|
||||
|
||||
use crate::BinOp;
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
use super::{MMMInputValue, OutputStore, OutputStoreKer};
|
||||
use tract_data::internal::*;
|
||||
|
||||
#[repr(usize)]
|
||||
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum RoundingPolicy {
|
||||
Native,
|
||||
Zero,
|
||||
Away,
|
||||
MinusInf,
|
||||
PlusInf,
|
||||
Even,
|
||||
Odd,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum AsInputValue<'t> {
|
||||
Owned(Box<dyn MMMInputValue>),
|
||||
Borrowed(&'t dyn MMMInputValue),
|
||||
}
|
||||
|
||||
impl Deref for AsInputValue<'_> {
|
||||
type Target = dyn MMMInputValue;
|
||||
fn deref(&self) -> &Self::Target {
|
||||
match self {
|
||||
AsInputValue::Owned(b) => &**b,
|
||||
AsInputValue::Borrowed(r) => *r,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum FusedSpec<'t> {
|
||||
BinScalar(&'t Tensor, BinOp),
|
||||
BinPerRow(TensorView<'t>, BinOp),
|
||||
BinPerCol(TensorView<'t>, BinOp),
|
||||
AddRowColProducts(&'t Tensor, &'t Tensor),
|
||||
AddUnicast(OutputStore),
|
||||
LeakyRelu(&'t Tensor),
|
||||
QScale(isize, RoundingPolicy, i32),
|
||||
RoundingShiftRight(usize, RoundingPolicy),
|
||||
ShiftLeft(usize),
|
||||
Store(OutputStore),
|
||||
AddMatMul {
|
||||
a: AsInputValue<'t>,
|
||||
b: AsInputValue<'t>,
|
||||
packing: usize,
|
||||
},
|
||||
}
|
||||
|
||||
impl FusedSpec<'_> {
|
||||
pub fn prefer_col_outer(&self) -> Option<bool> {
|
||||
if let FusedSpec::AddMatMul { a, b, .. } = self {
|
||||
let a_is_eager = a.format().is::<PackedFormat>();
|
||||
let b_is_eager = b.format().is::<PackedFormat>();
|
||||
if a_is_eager == b_is_eager {
|
||||
None
|
||||
} else {
|
||||
Some(a_is_eager)
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
//
|
||||
LoadTile(*const TI, *const TI), // jump_to:load_tile
|
||||
|
||||
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
|
||||
|
||||
LeakyRelu(TI), // jump_to:leaky_relu
|
||||
|
||||
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, packing: usize },
|
||||
}
|
||||
|
||||
unsafe impl<TI: Copy> Send for FusedKerSpec<TI> {}
|
||||
unsafe impl<TI: Copy> Sync for FusedKerSpec<TI> {}
|
||||
|
||||
#[cfg(test)]
|
||||
#[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::<FusedKerSpec<f32>>(),
|
||||
std::mem::size_of::<usize>() + std::mem::size_of::<OutputStoreKer>()
|
||||
);
|
||||
assert_eq!(
|
||||
std::mem::size_of::<FusedKerSpec<f32>>(),
|
||||
5 * std::mem::size_of::<usize>()
|
||||
);
|
||||
}
|
||||
Vendored
+157
@@ -0,0 +1,157 @@
|
||||
use downcast_rs::{Downcast, impl_downcast};
|
||||
use dyn_clone::DynClone;
|
||||
use dyn_eq::DynEq;
|
||||
use dyn_hash::DynHash;
|
||||
use std::alloc::Layout;
|
||||
use std::fmt::{Debug, Display};
|
||||
use std::hash::Hash;
|
||||
use std::sync::Arc;
|
||||
use tract_data::internal::*;
|
||||
|
||||
use crate::WeightType;
|
||||
|
||||
pub trait MMMInputFormat:
|
||||
Downcast + Debug + DynHash + dyn_eq::DynEq + DynClone + Send + Sync + Display
|
||||
{
|
||||
fn prepare_tensor(&self, t: &Tensor, k_axis: usize, mn_axis: usize) -> TractResult<Tensor>;
|
||||
fn prepare_one(
|
||||
&self,
|
||||
t: &Tensor,
|
||||
k_axis: usize,
|
||||
mn_axis: usize,
|
||||
) -> TractResult<Box<dyn MMMInputValue>>;
|
||||
fn precursor(&self) -> WeightType;
|
||||
fn r(&self) -> usize;
|
||||
fn k_alignment(&self) -> usize;
|
||||
fn merge_with<'o, 'a: 'o, 'b: 'o>(
|
||||
&'a self,
|
||||
other: &'b dyn MMMInputFormat,
|
||||
) -> Option<&'o dyn MMMInputFormat> {
|
||||
if self.dyn_eq(other) {
|
||||
Some(other)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
fn mem_size(&self, k: TDim, mn: TDim) -> TDim;
|
||||
fn extract_at_mn_f16(
|
||||
&self,
|
||||
data: &EagerPackedInput,
|
||||
mn: usize,
|
||||
slice: &mut [f16],
|
||||
) -> TractResult<()>;
|
||||
fn extract_at_mn_f32(
|
||||
&self,
|
||||
data: &EagerPackedInput,
|
||||
mn: usize,
|
||||
slice: &mut [f32],
|
||||
) -> TractResult<()>;
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(MMMInputFormat);
|
||||
impl_downcast!(MMMInputFormat);
|
||||
dyn_hash::hash_trait_object!(MMMInputFormat);
|
||||
dyn_eq::eq_trait_object!(MMMInputFormat);
|
||||
|
||||
pub trait MMMInputValue:
|
||||
DynClone + Debug + DynHash + dyn_eq::DynEq + Send + Sync + Display + Downcast
|
||||
{
|
||||
fn format(&self) -> &dyn MMMInputFormat;
|
||||
fn scratch_panel_buffer_layout(&self) -> Option<Layout>;
|
||||
fn panel_bytes(&self, i: usize, buffer: Option<*mut u8>) -> TractResult<*const u8>;
|
||||
fn panels_count(&self) -> usize {
|
||||
self.mn().divceil(self.format().r())
|
||||
}
|
||||
fn mn(&self) -> usize;
|
||||
fn k(&self) -> usize;
|
||||
fn exotic_fact(&self) -> &dyn ExoticFact;
|
||||
|
||||
fn extract_at_mn_f16(&self, mn: usize, slice: &mut [f16]) -> TractResult<()>;
|
||||
fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()>;
|
||||
}
|
||||
dyn_clone::clone_trait_object!(MMMInputValue);
|
||||
impl_downcast!(MMMInputValue);
|
||||
dyn_hash::hash_trait_object!(MMMInputValue);
|
||||
dyn_eq::eq_trait_object!(MMMInputValue);
|
||||
|
||||
#[allow(clippy::derived_hash_with_manual_eq)]
|
||||
#[derive(Clone, Hash, Debug)]
|
||||
pub struct PackedExoticFact {
|
||||
pub format: Box<dyn MMMInputFormat>,
|
||||
pub mn: TDim,
|
||||
pub k: usize,
|
||||
}
|
||||
|
||||
impl Display for PackedExoticFact {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"Eager {} tensor (mn={} k={})",
|
||||
self.format, self.mn, self.k
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ExoticFact for PackedExoticFact {
|
||||
fn buffer_sizes(&self) -> TVec<TDim> {
|
||||
tvec!(self.format.mem_size(self.k.to_dim(), self.mn.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for PackedExoticFact {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.format == other.format && self.mn == other.mn && self.k == other.k
|
||||
}
|
||||
}
|
||||
impl Eq for PackedExoticFact {}
|
||||
|
||||
#[derive(Clone, Hash, PartialEq, Eq)]
|
||||
pub struct EagerPackedInput {
|
||||
pub fact: PackedExoticFact,
|
||||
pub packed: Arc<Blob>,
|
||||
pub panel_bytes: usize,
|
||||
pub mn: usize,
|
||||
}
|
||||
|
||||
impl MMMInputValue for EagerPackedInput {
|
||||
fn scratch_panel_buffer_layout(&self) -> Option<Layout> {
|
||||
None
|
||||
}
|
||||
fn panel_bytes(&self, i: usize, _buffer: Option<*mut u8>) -> TractResult<*const u8> {
|
||||
unsafe { Ok(self.packed.as_ptr().add(i * self.panel_bytes)) }
|
||||
}
|
||||
fn k(&self) -> usize {
|
||||
self.fact.k
|
||||
}
|
||||
fn mn(&self) -> usize {
|
||||
self.mn
|
||||
}
|
||||
fn format(&self) -> &dyn MMMInputFormat {
|
||||
&*self.fact.format
|
||||
}
|
||||
fn exotic_fact(&self) -> &dyn ExoticFact {
|
||||
&self.fact
|
||||
}
|
||||
fn extract_at_mn_f16(&self, mn: usize, slice: &mut [f16]) -> TractResult<()> {
|
||||
ensure!(slice.len() == self.k());
|
||||
ensure!(mn < self.mn());
|
||||
self.fact.format.extract_at_mn_f16(self, mn, slice)
|
||||
}
|
||||
fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()> {
|
||||
ensure!(slice.len() == self.k());
|
||||
ensure!(mn < self.mn());
|
||||
self.fact.format.extract_at_mn_f32(self, mn, slice)
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for EagerPackedInput {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
(&self.fact as &dyn Display).fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
impl Debug for EagerPackedInput {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
<Self as Display>::fmt(self, f)
|
||||
}
|
||||
}
|
||||
+167
@@ -0,0 +1,167 @@
|
||||
use crate::frame::pack::PackedFormat;
|
||||
|
||||
use super::*;
|
||||
use std::borrow::Cow;
|
||||
use std::fmt::Debug;
|
||||
|
||||
use crate::LADatum;
|
||||
|
||||
pub trait MatMatMulKer: Clone + Debug + Send + Sync + 'static {
|
||||
type Acc: LADatum;
|
||||
fn name(&self) -> &str;
|
||||
fn kernel(&self, op: &[FusedKerSpec<Self::Acc>]) -> isize;
|
||||
fn mr(&self) -> usize;
|
||||
fn nr(&self) -> usize;
|
||||
|
||||
fn quality(&self) -> ImplementationQuality;
|
||||
fn dynamic_boost(&self) -> isize;
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)];
|
||||
fn stores(&self) -> Cow<'_, [DatumType]>;
|
||||
|
||||
#[allow(unused_variables)]
|
||||
fn can_fuse(&self, spec: &FusedSpec) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
#[allow(unused_variables)]
|
||||
fn is_supported_here(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
type Kernel<Acc> = unsafe fn(&[FusedKerSpec<Acc>]) -> isize;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct DynKernel<const MR: usize, const NR: usize, Acc: LADatum> {
|
||||
pub name: String,
|
||||
pub kernel: Kernel<Acc>,
|
||||
pub quality: ImplementationQuality,
|
||||
pub packings: Vec<(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)>,
|
||||
pub stores: Vec<DatumType>,
|
||||
pub supported_predicate: fn() -> bool,
|
||||
pub boost: fn() -> isize,
|
||||
pub can_fuse: fn(&FusedSpec) -> bool,
|
||||
}
|
||||
|
||||
impl<const MR: usize, const NR: usize, Acc: LADatum> DynKernel<MR, NR, Acc> {
|
||||
pub fn new(
|
||||
name: &str,
|
||||
kernel: Kernel<Acc>,
|
||||
packing_a: PackedFormat,
|
||||
packing_b: PackedFormat,
|
||||
quality: ImplementationQuality,
|
||||
) -> Self {
|
||||
let kernel = DynKernel {
|
||||
name: name.to_string(),
|
||||
kernel,
|
||||
quality,
|
||||
packings: vec![],
|
||||
stores: vec![Acc::datum_type()],
|
||||
supported_predicate: || true,
|
||||
boost: || 0,
|
||||
can_fuse: |_| true,
|
||||
};
|
||||
kernel.with_packing(packing_a, packing_b)
|
||||
}
|
||||
|
||||
pub fn with_platform_condition(mut self, f: fn() -> bool) -> Self {
|
||||
self.supported_predicate = f;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_boost(mut self, f: fn() -> isize) -> Self {
|
||||
self.boost = f;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_packing(mut self, a: impl MMMInputFormat, b: impl MMMInputFormat) -> Self {
|
||||
self.packings.push((Box::new(a), Box::new(b)));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_packing_a(self, a: impl MMMInputFormat) -> Self {
|
||||
let b = self.regular_pack_b();
|
||||
self.with_packing(a, b)
|
||||
}
|
||||
|
||||
pub fn regular_pack_a(&self) -> PackedFormat {
|
||||
*self.packings[0]
|
||||
.0
|
||||
.clone()
|
||||
.downcast::<PackedFormat>()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub fn regular_pack_b(&self) -> PackedFormat {
|
||||
*self.packings[0]
|
||||
.1
|
||||
.clone()
|
||||
.downcast::<PackedFormat>()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub fn with_can_fuse(self, can_fuse: fn(&FusedSpec) -> bool) -> Self {
|
||||
Self { can_fuse, ..self }
|
||||
}
|
||||
|
||||
pub fn with_store<D: LADatum>(mut self) -> Self {
|
||||
self.stores.push(D::datum_type());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn mmm(&self) -> Box<dyn MatMatMul> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<const MR: usize, const NR: usize, Acc: LADatum> Debug for DynKernel<MR, NR, Acc> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.name)
|
||||
}
|
||||
}
|
||||
|
||||
impl<const MR: usize, const NR: usize, Acc: LADatum> MatMatMulKer for DynKernel<MR, NR, Acc> {
|
||||
type Acc = Acc;
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn mr(&self) -> usize {
|
||||
MR
|
||||
}
|
||||
|
||||
fn nr(&self) -> usize {
|
||||
NR
|
||||
}
|
||||
|
||||
fn quality(&self) -> ImplementationQuality {
|
||||
self.quality
|
||||
}
|
||||
|
||||
fn is_supported_here(&self) -> bool {
|
||||
(self.supported_predicate)()
|
||||
}
|
||||
|
||||
fn can_fuse(&self, spec: &FusedSpec) -> bool {
|
||||
(self.can_fuse)(spec)
|
||||
}
|
||||
|
||||
fn kernel(&self, op: &[FusedKerSpec<Self::Acc>]) -> isize {
|
||||
unsafe { (self.kernel)(op) }
|
||||
}
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)] {
|
||||
&self.packings
|
||||
}
|
||||
|
||||
fn stores(&self) -> Cow<'_, [DatumType]> {
|
||||
Cow::Borrowed(&self.stores)
|
||||
}
|
||||
|
||||
fn dynamic_boost(&self) -> isize {
|
||||
(self.boost)()
|
||||
}
|
||||
}
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
macro_rules! MMMExternKernel {
|
||||
(
|
||||
$func:ident<$ti:ident>($mr: expr, $nr: expr)
|
||||
$(@($align_a:expr, $align_b:expr))?
|
||||
$(where($where:expr))?
|
||||
$(can_fuse($can_fuse:expr))?
|
||||
$(packing[$pnum:literal] = $pid:ident => $packing:expr;)*
|
||||
$(quality($quality:expr))?
|
||||
$(boost($boost:expr))?
|
||||
$(store($($store:ty),*))?
|
||||
) => {
|
||||
paste! {
|
||||
mod [<sys_ $func>] {
|
||||
#[allow(unused_imports)]
|
||||
use super::*;
|
||||
#[allow(unused_imports)]
|
||||
use crate::frame::mmm::*;
|
||||
extern_kernel!(fn $func(op: *const FusedKerSpec<$ti>) -> isize);
|
||||
|
||||
#[inline]
|
||||
pub unsafe fn rusty(op: &[FusedKerSpec<$ti>]) -> isize {
|
||||
unsafe { $func(op.as_ptr()) }
|
||||
}
|
||||
}
|
||||
|
||||
MMMKernel!([<sys_$func>]::rusty as $func<$ti>($mr, $nr)
|
||||
$(@($align_a, $align_b))?
|
||||
$(where($where))?
|
||||
$(can_fuse($can_fuse))?
|
||||
$(packing[$pnum] = $pid => $packing;)*
|
||||
$(quality($quality))?
|
||||
$(boost($boost))?
|
||||
$(store($($store),*))?
|
||||
);
|
||||
}
|
||||
};
|
||||
}
|
||||
macro_rules! MMMRustKernel {
|
||||
( $func: path =>
|
||||
$id:ident<$ti:ident>($mr: expr, $nr: expr)
|
||||
$(@($align_a:expr, $align_b:expr))?
|
||||
$(where($where:expr))?
|
||||
$(can_fuse($can_fuse:expr))?
|
||||
$(packing[$pnum:literal] = $pid:ident => $packing:expr;)*
|
||||
$(quality($quality:expr))?
|
||||
$(store($($store:ty),*))?
|
||||
) => {
|
||||
paste! {
|
||||
mod [<sys_ $id>] {
|
||||
#[allow(unused_imports)]
|
||||
use crate::frame::mmm::*;
|
||||
use super::*;
|
||||
#[inline]
|
||||
pub unsafe fn rusty(op: &[FusedKerSpec<$ti>]) -> isize {
|
||||
unsafe { $func(op.as_ptr()) }
|
||||
}
|
||||
}
|
||||
MMMKernel!([<sys_$id>]::rusty as $id<$ti>($mr, $nr)
|
||||
$(@($align_a, $align_b))?
|
||||
generic(true)
|
||||
$(where($where))?
|
||||
$(can_fuse($can_fuse))?
|
||||
$(packing[$pnum] = $pid => $packing;)*
|
||||
$(quality($quality))?
|
||||
$(store($($store),*))?
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! MMMKernel {
|
||||
(
|
||||
$func: path as
|
||||
$id:ident<$ti:ident>($mr: expr, $nr: expr)
|
||||
$(@($align_a:expr, $align_b:expr))?
|
||||
$(generic($generic:expr))?
|
||||
$(where($where:expr))?
|
||||
$(can_fuse($can_fuse:expr))?
|
||||
$(packing[$pnum:literal] = $pid:ident => $packing:expr;)*
|
||||
$(quality($quality:expr))?
|
||||
$(boost($boost:expr))?
|
||||
$(store($($store:ty),*))?
|
||||
) => {
|
||||
paste! {
|
||||
lazy_static::lazy_static! {
|
||||
pub static ref $id: $crate::mmm::DynKernel<$mr, $nr, $ti> = {
|
||||
use $crate::mmm::DynKernel;
|
||||
#[allow(unused_imports)]
|
||||
use tract_data::prelude::*;
|
||||
use $crate::pack::Packing;
|
||||
#[allow(unused_mut)]
|
||||
let (mut packing_a, mut packing_b) = ($ti::packing($mr), $ti::packing($nr));
|
||||
$(
|
||||
packing_a = packing_a.align($align_a);
|
||||
packing_b = packing_b.align($align_b);
|
||||
)?
|
||||
#[allow(unused_mut)]
|
||||
let mut k = DynKernel::<$mr, $nr, $ti>::new(stringify!($id), $func, packing_a, packing_b, $crate::frame::mmm::ImplementationQuality::Dreadful);
|
||||
$(k = k.with_platform_condition($where);)?
|
||||
$(
|
||||
assert!(k.packings.len() == $pnum);
|
||||
let f: fn(DynKernel<$mr, $nr, $ti>) -> DynKernel<$mr, $nr, $ti> = $packing;
|
||||
k = f(k);
|
||||
)*
|
||||
$($(
|
||||
k.stores.push(<$store>::datum_type());
|
||||
)*)?
|
||||
$(k.can_fuse = $can_fuse;)?
|
||||
$(k.quality = $quality;)?
|
||||
$(k = k.with_boost($boost);)?
|
||||
k
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod [<test_$id>] {
|
||||
use super::$id;
|
||||
test_mmm_kernel!($ti, &*super::$id);
|
||||
$(mmm_packed_packed_tests!(&*super::$id, $pid : $pnum);)*
|
||||
$($(mmm_store_test!(&*super::$id, $store);)*)?
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
+634
@@ -0,0 +1,634 @@
|
||||
#[macro_use]
|
||||
mod macros;
|
||||
|
||||
pub mod cost_model;
|
||||
#[macro_use]
|
||||
pub(crate) mod fuse;
|
||||
pub(crate) mod input_store;
|
||||
pub(crate) mod kernel;
|
||||
#[macro_use]
|
||||
pub(crate) mod panel_extract;
|
||||
mod scratch;
|
||||
mod storage;
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod tests;
|
||||
|
||||
use crate::multithread::Executor;
|
||||
use std::borrow::Cow;
|
||||
use std::cmp::Ordering;
|
||||
use std::fmt::Debug;
|
||||
use tract_data::internal::*;
|
||||
|
||||
pub use cost_model::*;
|
||||
pub use fuse::*;
|
||||
pub use input_store::*;
|
||||
pub use kernel::*;
|
||||
pub use panel_extract::*;
|
||||
pub use scratch::*;
|
||||
pub use storage::*;
|
||||
|
||||
pub fn no_prefetch(_ptr: *const u8, _len: usize) {}
|
||||
|
||||
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
|
||||
pub enum ImplementationQuality {
|
||||
/// Individual operations are emulated by individual conversion (f16->f32->f16)
|
||||
Dreadful,
|
||||
/// Rust scalar operation (with whatever optimisation the compiler manages)
|
||||
Generic,
|
||||
/// Implicit vectorization (e.g. Rust code, some unrolled loops, explicit template instantiations for small constant)
|
||||
RustOptimized,
|
||||
/// Explicit vectorization (e.g. intrinsics vector code)
|
||||
TargetOptimized,
|
||||
/// Hand optimized (assembly)
|
||||
ManuallyOptimized,
|
||||
}
|
||||
|
||||
impl ImplementationQuality {
|
||||
pub fn best_to_worst() -> &'static [ImplementationQuality] {
|
||||
use ImplementationQuality::*;
|
||||
&[
|
||||
ManuallyOptimized,
|
||||
TargetOptimized,
|
||||
RustOptimized,
|
||||
Generic,
|
||||
Dreadful,
|
||||
]
|
||||
}
|
||||
|
||||
pub fn cost(&self) -> usize {
|
||||
ImplementationQuality::best_to_worst()
|
||||
.iter()
|
||||
.position(|x| x == self)
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialOrd for ImplementationQuality {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
||||
Some(usize::from(*self).cmp(&usize::from(*other)))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ImplementationQuality> for usize {
|
||||
fn from(value: ImplementationQuality) -> Self {
|
||||
value.cost()
|
||||
}
|
||||
}
|
||||
|
||||
pub trait MatMatMul: Debug + dyn_clone::DynClone + Send + Sync + std::any::Any {
|
||||
fn name(&self) -> &str;
|
||||
fn mr(&self) -> usize;
|
||||
fn nr(&self) -> usize;
|
||||
|
||||
fn quality(&self) -> ImplementationQuality;
|
||||
fn dynamic_boost(&self) -> isize;
|
||||
|
||||
/// Whether this kernel is runnable on the current CPU (platform feature
|
||||
/// gate, e.g. FEAT_DotProd for the SDOT i8 kernel).
|
||||
fn is_supported_here(&self) -> bool;
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)];
|
||||
|
||||
fn internal_type(&self) -> DatumType;
|
||||
|
||||
unsafe fn c_view(&self, m_axis: Option<usize>, n_axis: Option<usize>) -> OutputStoreSpec;
|
||||
unsafe fn c_from_data_and_strides(
|
||||
&self,
|
||||
item_size: usize,
|
||||
row_stride: isize,
|
||||
col_stride: isize,
|
||||
) -> OutputStoreSpec;
|
||||
|
||||
fn can_fuse(&self, spec: &FusedSpec) -> bool;
|
||||
|
||||
fn stores(&self) -> Cow<'_, [DatumType]>;
|
||||
|
||||
unsafe fn run(&self, m: usize, n: usize, non_linear: &[FusedSpec]) -> TractResult<()> {
|
||||
unsafe {
|
||||
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],
|
||||
) -> TractResult<()>;
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(MatMatMul);
|
||||
|
||||
impl PartialEq for Box<dyn MatMatMul> {
|
||||
fn eq(&self, other: &Box<dyn MatMatMul>) -> bool {
|
||||
self.name() == other.name()
|
||||
}
|
||||
}
|
||||
impl Eq for Box<dyn MatMatMul> {}
|
||||
|
||||
impl std::hash::Hash for Box<dyn MatMatMul> {
|
||||
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
|
||||
self.name().hash(state)
|
||||
}
|
||||
}
|
||||
|
||||
impl<K: MatMatMulKer> MatMatMul for K {
|
||||
fn name(&self) -> &str {
|
||||
self.name()
|
||||
}
|
||||
fn mr(&self) -> usize {
|
||||
self.mr()
|
||||
}
|
||||
fn nr(&self) -> usize {
|
||||
self.nr()
|
||||
}
|
||||
|
||||
fn quality(&self) -> ImplementationQuality {
|
||||
MatMatMulKer::quality(self)
|
||||
}
|
||||
|
||||
fn dynamic_boost(&self) -> isize {
|
||||
MatMatMulKer::dynamic_boost(self)
|
||||
}
|
||||
|
||||
fn is_supported_here(&self) -> bool {
|
||||
MatMatMulKer::is_supported_here(self)
|
||||
}
|
||||
|
||||
fn packings(&self) -> &[(Box<dyn MMMInputFormat>, Box<dyn MMMInputFormat>)] {
|
||||
self.packings()
|
||||
}
|
||||
|
||||
fn internal_type(&self) -> DatumType {
|
||||
K::Acc::datum_type()
|
||||
}
|
||||
|
||||
fn can_fuse(&self, spec: &FusedSpec) -> bool {
|
||||
self.can_fuse(spec)
|
||||
}
|
||||
|
||||
unsafe fn c_view(&self, m_axis: Option<usize>, n_axis: Option<usize>) -> OutputStoreSpec {
|
||||
OutputStoreSpec::View {
|
||||
m_axis,
|
||||
n_axis,
|
||||
mr: self.mr(),
|
||||
nr: self.nr(),
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn c_from_data_and_strides(
|
||||
&self,
|
||||
item_size: 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: self.mr(),
|
||||
nr: self.nr(),
|
||||
}
|
||||
}
|
||||
|
||||
fn stores(&self) -> Cow<'_, [DatumType]> {
|
||||
self.stores()
|
||||
}
|
||||
|
||||
unsafe fn allocate_scratch_space(&self) -> Box<dyn ScratchSpace> {
|
||||
Box::<ScratchSpaceImpl<K::Acc>>::default()
|
||||
}
|
||||
|
||||
unsafe fn can_use_scratch_space(&self, scratch: &dyn ScratchSpace) -> bool {
|
||||
scratch.downcast_ref::<ScratchSpaceImpl<K::Acc>>().is_some()
|
||||
}
|
||||
|
||||
unsafe fn run_with_scratch_space(
|
||||
&self,
|
||||
m: usize,
|
||||
n: usize,
|
||||
scratch: &mut dyn ScratchSpace,
|
||||
non_linear: &[FusedSpec],
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
let scratch = scratch
|
||||
.downcast_mut::<ScratchSpaceImpl<K::Acc>>()
|
||||
.context("Wrong scratch space type")?;
|
||||
scratch.prepare(self, m, n, non_linear)?;
|
||||
if n == 1 && self.nr() == 1 {
|
||||
run_with_scratch_space_vec(self, m, scratch, non_linear)
|
||||
} else {
|
||||
let (mut prefer_col, mut prefer_row) = (0, 0);
|
||||
for uop in non_linear.iter() {
|
||||
if let Some(col) = uop.prefer_col_outer() {
|
||||
prefer_col = col as usize;
|
||||
prefer_row = (!col) as usize;
|
||||
}
|
||||
}
|
||||
// k drives the single-thread cache-block size; read it from the
|
||||
// first AddMatMul's packed input (0 if none → max block).
|
||||
let k = non_linear
|
||||
.iter()
|
||||
.find_map(|f| match f {
|
||||
FusedSpec::AddMatMul { a, .. } => Some(a.k()),
|
||||
_ => None,
|
||||
})
|
||||
.unwrap_or(0);
|
||||
if prefer_col > prefer_row {
|
||||
run_with_scratch_space_col_outer(self, m, n, k, scratch, non_linear)
|
||||
} else {
|
||||
run_with_scratch_space_row_outer(self, m, n, k, scratch, non_linear)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn run_with_scratch_space_vec<K: MatMatMulKer>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
scratch: &mut ScratchSpaceImpl<K::Acc>,
|
||||
non_linear: &[FusedSpec],
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
match crate::multithread::current_tract_executor() {
|
||||
Executor::SingleThread => scratch.run_in_tls_scope(|scratch, tls| {
|
||||
for ia in 0..m.divceil(ker.mr()) {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, 0)?;
|
||||
}
|
||||
TractResult::Ok(())
|
||||
}),
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
Executor::MultiThread(pool) => chunked_dispatch_rayon(
|
||||
Some(&pool),
|
||||
m.divceil(ker.mr()),
|
||||
1,
|
||||
|ia_start, ia_end, _, _| {
|
||||
scratch.run_in_tls_scope(|scratch, tls| {
|
||||
for ia in ia_start..ia_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, 0)?;
|
||||
}
|
||||
TractResult::Ok(())
|
||||
})
|
||||
},
|
||||
),
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
Executor::RayonGlobal => {
|
||||
chunked_dispatch_rayon(None, m.divceil(ker.mr()), 1, |ia_start, ia_end, _, _| {
|
||||
scratch.run_in_tls_scope(|scratch, tls| {
|
||||
for ia in ia_start..ia_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, 0)?;
|
||||
}
|
||||
TractResult::Ok(())
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Upper bound on the single-thread panel-block edge (matches the multithread
|
||||
/// `chunk_grid` default).
|
||||
const ST_BLK_MAX: usize = 16;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn parse_cache_size(s: &str) -> usize {
|
||||
let s = s.trim();
|
||||
let (num, mult) = if let Some(n) = s.strip_suffix(['K', 'k']) {
|
||||
(n, 1024)
|
||||
} else if let Some(n) = s.strip_suffix(['M', 'm']) {
|
||||
(n, 1024 * 1024)
|
||||
} else {
|
||||
(s, 1)
|
||||
};
|
||||
num.trim().parse::<usize>().unwrap_or(0) * mult
|
||||
}
|
||||
|
||||
/// Best-effort L2 data-cache size in bytes (per perf-core / cluster); 0 if
|
||||
/// unknown. Cached. Used to size the single-thread cache-block budget so it is
|
||||
/// correct across hardware instead of a hard-coded constant.
|
||||
fn detect_l2_bytes() -> usize {
|
||||
static L2: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
|
||||
*L2.get_or_init(|| {
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
let sysctl = |k: &str| -> Option<usize> {
|
||||
let o = std::process::Command::new("sysctl")
|
||||
.arg("-n")
|
||||
.arg(k)
|
||||
.output()
|
||||
.ok()?;
|
||||
if !o.status.success() {
|
||||
return None;
|
||||
}
|
||||
String::from_utf8_lossy(&o.stdout).trim().parse().ok()
|
||||
};
|
||||
// Prefer the performance-core L2 on hybrid Apple Silicon.
|
||||
sysctl("hw.perflevel0.l2cachesize")
|
||||
.or_else(|| sysctl("hw.l2cachesize"))
|
||||
.unwrap_or(0)
|
||||
}
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
// index2/index3 is typically the unified L2 (index0/1 are L1 d/i).
|
||||
for idx in [2usize, 3] {
|
||||
if let Ok(s) = std::fs::read_to_string(format!(
|
||||
"/sys/devices/system/cpu/cpu0/cache/index{idx}/size"
|
||||
)) {
|
||||
let b = parse_cache_size(s.trim());
|
||||
if b > 0 {
|
||||
return b;
|
||||
}
|
||||
}
|
||||
}
|
||||
0
|
||||
}
|
||||
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
|
||||
{
|
||||
0
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Working-set budget (bytes) for the single-thread cache-block: ~a third of L2
|
||||
/// (leaving room for the C accumulator tile + packing metadata). Conservative
|
||||
/// 256 KiB fallback when L2 is unknown (WASM/Windows/BSD) ⇒ small blocks ≈ the
|
||||
/// naive loop, so it can never over-block a cache it can't see.
|
||||
fn block_budget_bytes() -> usize {
|
||||
let l2 = detect_l2_bytes();
|
||||
if l2 == 0 {
|
||||
256 * 1024
|
||||
} else {
|
||||
(l2 / 3).clamp(64 * 1024, 8 * 1024 * 1024)
|
||||
}
|
||||
}
|
||||
|
||||
/// Cache-adaptive panel-block edge: large enough to amortise streaming, small
|
||||
/// enough that the block's A+B sub-panels (`~blk·(mr+nr)·k·elem_bytes`) stay
|
||||
/// L2-resident at the given `k`. Capped at [`ST_BLK_MAX`]; the floor of 1
|
||||
/// degrades exactly to the naive loop, so an unknown/small cache can never
|
||||
/// over-block (regression-safe). The budget is **cache-size derived** (not a
|
||||
/// hard-coded constant), so it is correct across hardware.
|
||||
#[inline]
|
||||
fn st_block_edge(mr: usize, nr: usize, k: usize, elem_bytes: usize) -> usize {
|
||||
if k == 0 {
|
||||
return ST_BLK_MAX;
|
||||
}
|
||||
let per_blk = ((mr + nr) * k * elem_bytes.max(1)).max(1);
|
||||
(block_budget_bytes() / per_blk).clamp(1, ST_BLK_MAX)
|
||||
}
|
||||
|
||||
/// Single-thread tile walk over the `m_panels × n_panels` grid, blocked into
|
||||
/// cache-sized panel blocks for locality (the naive nested loop re-streams the
|
||||
/// whole inner operand per outer panel at large k; the multithread path already
|
||||
/// blocks this way via `chunk_grid`). `col_outer` selects the within-block inner
|
||||
/// order (B-reuse vs A-reuse). Reordering independent tiles changes no result —
|
||||
/// bit-exact with the naive loop.
|
||||
#[inline]
|
||||
unsafe fn run_single_thread_blocked<K: MatMatMulKer>(
|
||||
ker: &K,
|
||||
m_panels: usize,
|
||||
n_panels: usize,
|
||||
k: usize,
|
||||
col_outer: bool,
|
||||
scratch: &mut ScratchSpaceImpl<K::Acc>,
|
||||
non_linear: &[FusedSpec],
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
let blk = st_block_edge(ker.mr(), ker.nr(), k, K::Acc::datum_type().size_of());
|
||||
scratch.run_in_tls_scope(|scratch, tls| {
|
||||
let mut jb = 0;
|
||||
while jb < n_panels {
|
||||
let jb_end = (jb + blk).min(n_panels);
|
||||
let mut ja = 0;
|
||||
while ja < m_panels {
|
||||
let ja_end = (ja + blk).min(m_panels);
|
||||
if col_outer {
|
||||
for ib in jb..jb_end {
|
||||
for ia in ja..ja_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, ib)?;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for ia in ja..ja_end {
|
||||
for ib in jb..jb_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, ib)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
ja = ja_end;
|
||||
}
|
||||
jb = jb_end;
|
||||
}
|
||||
TractResult::Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn run_with_scratch_space_col_outer<K: MatMatMulKer>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
k: usize,
|
||||
scratch: &mut ScratchSpaceImpl<K::Acc>,
|
||||
non_linear: &[FusedSpec],
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
match crate::multithread::current_tract_executor() {
|
||||
Executor::SingleThread => run_single_thread_blocked(
|
||||
ker,
|
||||
m.divceil(ker.mr()),
|
||||
n.divceil(ker.nr()),
|
||||
k,
|
||||
true,
|
||||
scratch,
|
||||
non_linear,
|
||||
),
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
Executor::MultiThread(pool) => chunked_dispatch_rayon(
|
||||
Some(&pool),
|
||||
m.divceil(ker.mr()),
|
||||
n.divceil(ker.nr()),
|
||||
|ia_start, ia_end, ib_start, ib_end| {
|
||||
scratch.run_in_tls_scope(|scratch, tls| {
|
||||
for ib in ib_start..ib_end {
|
||||
for ia in ia_start..ia_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, ib)?;
|
||||
}
|
||||
}
|
||||
TractResult::Ok(())
|
||||
})
|
||||
},
|
||||
),
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
Executor::RayonGlobal => chunked_dispatch_rayon(
|
||||
None,
|
||||
m.divceil(ker.mr()),
|
||||
n.divceil(ker.nr()),
|
||||
|ia_start, ia_end, ib_start, ib_end| {
|
||||
scratch.run_in_tls_scope(|scratch, tls| {
|
||||
for ib in ib_start..ib_end {
|
||||
for ia in ia_start..ia_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, ib)?;
|
||||
}
|
||||
}
|
||||
TractResult::Ok(())
|
||||
})
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn run_with_scratch_space_row_outer<K: MatMatMulKer>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
k: usize,
|
||||
scratch: &mut ScratchSpaceImpl<K::Acc>,
|
||||
non_linear: &[FusedSpec],
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
match crate::multithread::current_tract_executor() {
|
||||
Executor::SingleThread => run_single_thread_blocked(
|
||||
ker,
|
||||
m.divceil(ker.mr()),
|
||||
n.divceil(ker.nr()),
|
||||
k,
|
||||
false,
|
||||
scratch,
|
||||
non_linear,
|
||||
),
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
Executor::MultiThread(pool) => chunked_dispatch_rayon(
|
||||
Some(&pool),
|
||||
m.divceil(ker.mr()),
|
||||
n.divceil(ker.nr()),
|
||||
|ia_start, ia_end, ib_start, ib_end| {
|
||||
scratch.run_in_tls_scope(|scratch, tls| {
|
||||
for ia in ia_start..ia_end {
|
||||
for ib in ib_start..ib_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, ib)?;
|
||||
}
|
||||
}
|
||||
TractResult::Ok(())
|
||||
})
|
||||
},
|
||||
),
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
Executor::RayonGlobal => chunked_dispatch_rayon(
|
||||
None,
|
||||
m.divceil(ker.mr()),
|
||||
n.divceil(ker.nr()),
|
||||
|ia_start, ia_end, ib_start, ib_end| {
|
||||
scratch.run_in_tls_scope(|scratch, tls| {
|
||||
for ia in ia_start..ia_end {
|
||||
for ib in ib_start..ib_end {
|
||||
scratch.run_one_tile(ker, non_linear, tls, ia, ib)?;
|
||||
}
|
||||
}
|
||||
TractResult::Ok(())
|
||||
})
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Chunk grid for the 2D dispatch.
|
||||
///
|
||||
/// Mirrors ggml's `mul_mat` heuristic (`ggml/src/ggml-cpu/ggml-cpu.c:1378-1398`):
|
||||
/// * 16-tile panel chunks by default;
|
||||
/// * 64-tile chunks when one dimension is 1 (vec / vec-mat);
|
||||
/// * fallback to "block-per-thread along the longer axis" when the natural
|
||||
/// grid would have fewer than `4·nth` chunks.
|
||||
///
|
||||
/// Returns `(nchunks_m, nchunks_n, dr_m, dr_n)`.
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
fn chunk_grid(n_panels_m: usize, n_panels_n: usize, nth: usize) -> (usize, usize, usize, usize) {
|
||||
let chunk_size = if n_panels_m == 1 || n_panels_n == 1 {
|
||||
64
|
||||
} else {
|
||||
16
|
||||
};
|
||||
let mut nchunks_m = n_panels_m.div_ceil(chunk_size);
|
||||
let mut nchunks_n = n_panels_n.div_ceil(chunk_size);
|
||||
if nchunks_m * nchunks_n < 4 * nth {
|
||||
if n_panels_m > n_panels_n {
|
||||
nchunks_m = nth;
|
||||
nchunks_n = 1;
|
||||
} else {
|
||||
nchunks_m = 1;
|
||||
nchunks_n = nth;
|
||||
}
|
||||
}
|
||||
let dr_m = n_panels_m.div_ceil(nchunks_m).max(1);
|
||||
let dr_n = n_panels_n.div_ceil(nchunks_n).max(1);
|
||||
(nchunks_m, nchunks_n, dr_m, dr_n)
|
||||
}
|
||||
|
||||
/// 2D chunked dispatcher across the (m_panels × n_panels) grid for the
|
||||
/// rayon path. Replaces a 1D `into_par_iter` over a single panel axis.
|
||||
/// Better-utilises threads on small/skewed shapes where one dimension has
|
||||
/// fewer panels than there are workers.
|
||||
///
|
||||
/// The closure receives **chunk bounds** (`ia_start, ia_end, ib_start, ib_end`),
|
||||
/// not per-tile indices. This lets the caller amortise per-worker setup
|
||||
/// (e.g. `ScratchSpaceImpl::run_in_tls_scope`) across all tiles in the
|
||||
/// chunk, mirroring #2206 for the multi-threaded path. The closure is
|
||||
/// invoked exactly once per rayon work item (and once total when the
|
||||
/// small-graph fallback path is taken).
|
||||
///
|
||||
/// `pool`:
|
||||
/// * `Some(p)` with `p.current_num_threads() > 1` → scoped via `p.install`
|
||||
/// (native, custom pool path).
|
||||
/// * `Some(p)` with single-thread pool, or `None` → dispatched via
|
||||
/// `into_par_iter` directly, which uses rayon's GLOBAL pool. This is
|
||||
/// the only working path on `wasm32-unknown-unknown` via
|
||||
/// `wasm_bindgen_rayon::init_thread_pool`.
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
unsafe fn chunked_dispatch_rayon<F>(
|
||||
pool: Option<&rayon::ThreadPool>,
|
||||
n_panels_m: usize,
|
||||
n_panels_n: usize,
|
||||
run_chunk: F,
|
||||
) -> TractResult<()>
|
||||
where
|
||||
F: Fn(usize, usize, usize, usize) -> TractResult<()> + Sync,
|
||||
{
|
||||
use rayon::prelude::*;
|
||||
if n_panels_m == 0 || n_panels_n == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
if n_panels_m * n_panels_n < crate::multithread::current_threading_panel_threshold() {
|
||||
// Below the threading threshold: run the whole grid as a single chunk
|
||||
// on the calling thread. Closure handles its own TLS scope.
|
||||
return run_chunk(0, n_panels_m, 0, n_panels_n);
|
||||
}
|
||||
let use_global = pool.is_none_or(|p| p.current_num_threads() <= 1);
|
||||
let body = || {
|
||||
let nth = rayon::current_num_threads();
|
||||
let (nchunks_m, nchunks_n, dr_m, dr_n) = chunk_grid(n_panels_m, n_panels_n, nth);
|
||||
let total = nchunks_m * nchunks_n;
|
||||
(0..total).into_par_iter().try_for_each(|idx| {
|
||||
let im = idx % nchunks_m;
|
||||
let in_ = idx / nchunks_m;
|
||||
let ia_start = im * dr_m;
|
||||
let ia_end = (ia_start + dr_m).min(n_panels_m);
|
||||
let ib_start = in_ * dr_n;
|
||||
let ib_end = (ib_start + dr_n).min(n_panels_n);
|
||||
run_chunk(ia_start, ia_end, ib_start, ib_end)
|
||||
})
|
||||
};
|
||||
if use_global {
|
||||
body()
|
||||
} else {
|
||||
pool.unwrap().install(body)
|
||||
}
|
||||
}
|
||||
Vendored
+332
@@ -0,0 +1,332 @@
|
||||
use std::fmt::{Debug, Display};
|
||||
use tract_data::internal::*;
|
||||
|
||||
use super::{EagerPackedInput, MMMInputFormat, MMMInputValue};
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
type Kernel = unsafe fn(input: *const u8, output: *mut u8, k: usize);
|
||||
|
||||
#[allow(clippy::derived_hash_with_manual_eq)]
|
||||
#[derive(Hash, Clone)]
|
||||
pub struct PanelExtractor {
|
||||
pub name: String,
|
||||
pub from: Box<dyn MMMInputFormat>,
|
||||
pub to: PackedFormat,
|
||||
pub kernel: Kernel,
|
||||
pub supported_predicate: fn() -> bool,
|
||||
}
|
||||
|
||||
impl Debug for PanelExtractor {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{} ({:?} -> {:?})", self.name, self.from, self.to)
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for PanelExtractor {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.name)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for PanelExtractor {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.name == other.name && *self.from == *other.from && self.to == other.to
|
||||
}
|
||||
}
|
||||
impl Eq for PanelExtractor {}
|
||||
|
||||
impl PanelExtractor {
|
||||
#[allow(unused_variables)]
|
||||
pub fn is_supported_here(&self) -> bool {
|
||||
(self.supported_predicate)()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Hash, PartialEq, Eq)]
|
||||
pub struct PanelExtractInput {
|
||||
pub format: PanelExtractor,
|
||||
pub data: EagerPackedInput,
|
||||
}
|
||||
|
||||
impl MMMInputValue for PanelExtractInput {
|
||||
fn scratch_panel_buffer_layout(&self) -> Option<std::alloc::Layout> {
|
||||
Some(
|
||||
self.format
|
||||
.to
|
||||
.single_panel_layout(self.data.k(), self.format.to.dt.size_of()),
|
||||
)
|
||||
}
|
||||
fn panel_bytes(&self, i: usize, buffer: Option<*mut u8>) -> TractResult<*const u8> {
|
||||
let scratch = buffer.unwrap();
|
||||
unsafe {
|
||||
let source = self.data.packed.as_ptr().add(self.data.panel_bytes * i);
|
||||
(self.format.kernel)(source, scratch, self.data.k());
|
||||
}
|
||||
Ok(scratch)
|
||||
}
|
||||
fn mn(&self) -> usize {
|
||||
self.data.mn()
|
||||
}
|
||||
fn k(&self) -> usize {
|
||||
self.data.k()
|
||||
}
|
||||
fn format(&self) -> &dyn MMMInputFormat {
|
||||
&self.format.to
|
||||
}
|
||||
fn exotic_fact(&self) -> &dyn ExoticFact {
|
||||
self.data.exotic_fact()
|
||||
}
|
||||
fn extract_at_mn_f16(&self, mn: usize, slice: &mut [f16]) -> TractResult<()> {
|
||||
self.data.extract_at_mn_f16(mn, slice)
|
||||
}
|
||||
fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()> {
|
||||
self.data.extract_at_mn_f32(mn, slice)
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for PanelExtractInput {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "PanelExtract({})", self.data)
|
||||
}
|
||||
}
|
||||
|
||||
impl Debug for PanelExtractInput {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "PanelExtract({})", self.data)
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! panel_extractor {
|
||||
( $func:path as $id:ident($from:expr, $to: expr)
|
||||
$(where($where:expr))?
|
||||
) => {
|
||||
paste! {
|
||||
lazy_static::lazy_static! {
|
||||
pub static ref $id: $crate::mmm::PanelExtractor = {
|
||||
use $crate::mmm::MMMInputFormat;
|
||||
let (from, to) = ($from, $to);
|
||||
assert!(from.r() == to.r());
|
||||
#[allow(unused_mut)]
|
||||
let mut it = $crate::mmm::PanelExtractor {
|
||||
name: stringify!($id).to_string(),
|
||||
from,
|
||||
to,
|
||||
kernel: $func,
|
||||
supported_predicate: || true
|
||||
};
|
||||
$(
|
||||
it.supported_predicate = $where;
|
||||
)?
|
||||
it
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod [<test_$id>] {
|
||||
use super::$id;
|
||||
#[test]
|
||||
fn repack_0block_1panel() {
|
||||
$crate::frame::mmm::panel_extract::test::test_packing(&$id, 0, 1).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repack_1block_0panel() {
|
||||
$crate::frame::mmm::panel_extract::test::test_packing(&$id, 1, 0).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repack_1block_1panel() {
|
||||
$crate::frame::mmm::panel_extract::test::test_packing(&$id, 1, 1).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repack_2block_1panel() {
|
||||
$crate::frame::mmm::panel_extract::test::test_packing(&$id, 2, 1).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repack_1block_2panel() {
|
||||
$crate::frame::mmm::panel_extract::test::test_packing(&$id, 1, 2).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repack_2block_2panel() {
|
||||
$crate::frame::mmm::panel_extract::test::test_packing(&$id, 2, 2).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test {
|
||||
use crate::frame::block_quant::PackedBlockQuantFormat;
|
||||
use crate::mmm::PackedMatrixStorage;
|
||||
use tract_data::internal::*;
|
||||
use tract_ndarray::Array2;
|
||||
|
||||
use super::*;
|
||||
|
||||
pub fn test_packing(
|
||||
extractor: &PanelExtractor,
|
||||
blocks: usize,
|
||||
panels: usize,
|
||||
) -> TractResult<()> {
|
||||
if !extractor.is_supported_here() {
|
||||
return Ok(());
|
||||
}
|
||||
assert!(extractor.from.r() == extractor.to.r());
|
||||
assert!(extractor.to.dt == f32::datum_type() || extractor.to.dt == f16::datum_type());
|
||||
if let Some(from) = extractor.from.downcast_ref::<PackedBlockQuantFormat>() {
|
||||
test_packing_bq(extractor, from, blocks, panels)
|
||||
} else if let Some(from) = extractor.from.downcast_ref() {
|
||||
test_packing_plain(extractor, from, blocks, panels)
|
||||
} else {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn test_packing_plain(
|
||||
extractor: &PanelExtractor,
|
||||
from: &PackedFormat,
|
||||
blocks: usize,
|
||||
panels: usize,
|
||||
) -> TractResult<()> {
|
||||
let m = from.r * panels;
|
||||
let k = 8 * blocks; // 8 is arbitrary
|
||||
let to = &extractor.to;
|
||||
let weights_orig =
|
||||
Array2::from_shape_fn((m, k), |(m, k)| ((m * 31 + k * 17) % 20) as f32 - 10.)
|
||||
.into_tensor()
|
||||
.cast_to_dt(from.dt)?
|
||||
.into_owned();
|
||||
let packed_orig = from.prepare_tensor(&weights_orig, 1, 0)?;
|
||||
let packed_orig_storage = packed_orig.try_storage_as::<PackedMatrixStorage>()?;
|
||||
let packed_orig = packed_orig_storage
|
||||
.value()
|
||||
.downcast_ref::<EagerPackedInput>()
|
||||
.unwrap();
|
||||
|
||||
for panel in 0..panels {
|
||||
let orig_panel = &packed_orig.packed[packed_orig.panel_bytes * panel..]
|
||||
[..k * from.r * from.dt.size_of()];
|
||||
let mut reference_panel = Tensor::zero_dt(from.dt, &[k, from.r])?;
|
||||
reference_panel.as_bytes_mut().copy_from_slice(orig_panel);
|
||||
reference_panel = reference_panel.cast_to_dt(to.dt)?.into_owned();
|
||||
|
||||
let mut tested_panel = Tensor::zero_dt(to.dt, &[k, from.r])?;
|
||||
unsafe {
|
||||
(extractor.kernel)(
|
||||
orig_panel.as_ptr(),
|
||||
tested_panel.as_bytes_mut().as_mut_ptr(),
|
||||
k,
|
||||
);
|
||||
}
|
||||
compare_panels(&tested_panel, &reference_panel, from.r, k);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn test_packing_bq(
|
||||
extractor: &PanelExtractor,
|
||||
from: &PackedBlockQuantFormat,
|
||||
blocks: usize,
|
||||
panels: usize,
|
||||
) -> TractResult<()> {
|
||||
let m = from.r * panels;
|
||||
let k = from.bq.block_len() * blocks;
|
||||
let to = &extractor.to;
|
||||
let weights_orig =
|
||||
Array2::from_shape_fn((m, k), |(m, k)| ((m * 31 + k * 17) % 20) as f32 - 10.)
|
||||
.into_tensor()
|
||||
.cast_to_dt(to.dt)?
|
||||
.into_owned();
|
||||
let weights = if to.dt == f32::datum_type() {
|
||||
from.bq
|
||||
.dequant_f32(
|
||||
&from
|
||||
.bq
|
||||
.quant_f32(weights_orig.try_as_plain()?.as_slice::<f32>()?)?,
|
||||
)?
|
||||
.into_shape(&[m, k])?
|
||||
} else {
|
||||
from.bq
|
||||
.dequant_f16(
|
||||
&from
|
||||
.bq
|
||||
.quant_f16(weights_orig.try_as_plain()?.as_slice::<f16>()?)?,
|
||||
)?
|
||||
.into_shape(&[m, k])?
|
||||
};
|
||||
let block_quant = if to.dt == f32::datum_type() {
|
||||
from.bq
|
||||
.quant_f32(weights.try_as_plain()?.as_slice::<f32>()?)?
|
||||
} else {
|
||||
from.bq
|
||||
.quant_f16(weights.try_as_plain()?.as_slice::<f16>()?)?
|
||||
};
|
||||
let packed_block_quant =
|
||||
from.bq
|
||||
.pack(&block_quant, k, from.r, from.zip, from.scales_at_end)?;
|
||||
|
||||
let mut reference_panel = Tensor::zero_dt(to.dt, &[k, from.r])?;
|
||||
let mut tested_panel = Tensor::zero_dt(to.dt, &[k, from.r])?;
|
||||
|
||||
for panel in 0..packed_block_quant.panels_count() {
|
||||
unsafe {
|
||||
from.bq.extract_packed_panel(
|
||||
&packed_block_quant,
|
||||
to,
|
||||
panel,
|
||||
reference_panel.as_bytes_mut().as_mut_ptr(),
|
||||
)?;
|
||||
|
||||
let source = packed_block_quant
|
||||
.packed
|
||||
.as_ptr()
|
||||
.add(packed_block_quant.panel_bytes * panel);
|
||||
(extractor.kernel)(source, tested_panel.as_bytes_mut().as_mut_ptr(), k);
|
||||
}
|
||||
compare_panels(&tested_panel, &reference_panel, from.r, k);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn compare_panels(tested_panel: &Tensor, reference_panel: &Tensor, r: usize, k: usize) {
|
||||
if tested_panel != reference_panel {
|
||||
if reference_panel.datum_type() == f32::datum_type() {
|
||||
crate::frame::mmm::tests::display_error(
|
||||
tested_panel
|
||||
.try_as_plain()
|
||||
.unwrap()
|
||||
.as_slice::<f32>()
|
||||
.unwrap(),
|
||||
reference_panel
|
||||
.try_as_plain()
|
||||
.unwrap()
|
||||
.as_slice::<f32>()
|
||||
.unwrap(),
|
||||
r,
|
||||
k,
|
||||
);
|
||||
} else {
|
||||
crate::frame::mmm::tests::display_error(
|
||||
tested_panel
|
||||
.try_as_plain()
|
||||
.unwrap()
|
||||
.as_slice::<f16>()
|
||||
.unwrap(),
|
||||
reference_panel
|
||||
.try_as_plain()
|
||||
.unwrap()
|
||||
.as_slice::<f16>()
|
||||
.unwrap(),
|
||||
r,
|
||||
k,
|
||||
);
|
||||
}
|
||||
}
|
||||
assert_eq!(tested_panel, reference_panel);
|
||||
}
|
||||
}
|
||||
+590
@@ -0,0 +1,590 @@
|
||||
use super::{FusedKerSpec, FusedSpec, MatMatMulKer, OutputStoreKer};
|
||||
use crate::{BinOp, LADatum};
|
||||
use downcast_rs::{Downcast, impl_downcast};
|
||||
use std::cell::RefCell;
|
||||
use std::fmt::Debug;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use tract_data::internal::num_integer::Integer;
|
||||
use tract_data::internal::*;
|
||||
|
||||
static GENERATION: AtomicUsize = AtomicUsize::new(1);
|
||||
|
||||
thread_local! {
|
||||
static TLS: RefCell<TLSScratch> = Default::default();
|
||||
}
|
||||
|
||||
#[derive(Default, Debug)]
|
||||
pub(crate) struct TLSScratch {
|
||||
generation: usize,
|
||||
blob: Blob,
|
||||
ker_specs_16: Vec<FusedKerSpec<f16>>,
|
||||
ker_specs_32: Vec<FusedKerSpec<f32>>,
|
||||
ker_specs_64: Vec<FusedKerSpec<f64>>,
|
||||
}
|
||||
|
||||
impl TLSScratch {
|
||||
#[allow(unknown_lints, clippy::missing_transmute_annotations)]
|
||||
fn ker_specs<TI: LADatum>(&mut self) -> &mut Vec<FusedKerSpec<TI>> {
|
||||
unsafe {
|
||||
if TI::datum_type() == f32::datum_type() || TI::datum_type() == i32::datum_type() {
|
||||
std::mem::transmute(&mut self.ker_specs_32)
|
||||
} else if TI::datum_type() == f16::datum_type() {
|
||||
std::mem::transmute(&mut self.ker_specs_16)
|
||||
} else if TI::datum_type() == f64::datum_type() {
|
||||
std::mem::transmute(&mut self.ker_specs_64)
|
||||
} else {
|
||||
todo!();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn sync<TI: LADatum>(&mut self, scratch: &ScratchSpaceImpl<TI>) {
|
||||
if self.generation == scratch.generation {
|
||||
return;
|
||||
}
|
||||
let ker_specs = self.ker_specs::<TI>();
|
||||
ker_specs.clear();
|
||||
ker_specs.extend_from_slice(&scratch.ker_specs);
|
||||
|
||||
unsafe {
|
||||
self.blob
|
||||
.ensure_size_and_align(scratch.blob_size, scratch.blob_align);
|
||||
|
||||
for LocDependant { loc, ker_spec, .. } in &scratch.loc_dependant {
|
||||
#[allow(clippy::single_match)]
|
||||
if matches!(scratch.ker_specs[*ker_spec], FusedKerSpec::AddMatMul { .. }) {
|
||||
let scratch = &mut *(self.blob.as_ptr().add(*loc) as *mut AddMatMulTemp);
|
||||
scratch.panel_a_id = usize::MAX;
|
||||
scratch.panel_b_id = usize::MAX;
|
||||
};
|
||||
}
|
||||
}
|
||||
self.generation = scratch.generation;
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ScratchSpace: Downcast + Send {}
|
||||
impl_downcast!(ScratchSpace);
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ScratchSpaceImpl<TI: LADatum> {
|
||||
generation: usize,
|
||||
blob_size: usize,
|
||||
blob_align: usize,
|
||||
ker_specs: Vec<FusedKerSpec<TI>>,
|
||||
loc_dependant: TVec<LocDependant>,
|
||||
valid_down_tiles: usize,
|
||||
remnant_down: usize,
|
||||
valid_right_tiles: usize,
|
||||
remnant_right: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, new)]
|
||||
struct LocDependant {
|
||||
spec: usize,
|
||||
ker_spec: usize,
|
||||
// offset for the location dependant structure
|
||||
loc: usize,
|
||||
// offset of its associated dynamic-size buffers
|
||||
buffer_a: Option<usize>,
|
||||
buffer_b: Option<usize>,
|
||||
}
|
||||
|
||||
impl<TI: LADatum> ScratchSpace for ScratchSpaceImpl<TI> {}
|
||||
unsafe impl<TI: LADatum> Send for ScratchSpaceImpl<TI> {}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct AddMatMulTemp {
|
||||
ptr_a: *const u8,
|
||||
panel_a_id: usize,
|
||||
ptr_b: *const u8,
|
||||
panel_b_id: usize,
|
||||
}
|
||||
|
||||
impl<TI: LADatum> ScratchSpaceImpl<TI> {
|
||||
pub unsafe fn prepare(
|
||||
&mut self,
|
||||
ker: &impl MatMatMulKer<Acc = TI>,
|
||||
m: usize,
|
||||
n: usize,
|
||||
specs: &[FusedSpec],
|
||||
) -> TractResult<()> {
|
||||
use FusedKerSpec as FKS;
|
||||
use FusedSpec as FS;
|
||||
self.ker_specs.clear();
|
||||
self.loc_dependant.clear();
|
||||
self.ker_specs.reserve(specs.len() + 2);
|
||||
self.ker_specs.push(FusedKerSpec::Clear);
|
||||
self.valid_down_tiles = m / ker.mr();
|
||||
self.remnant_down = m % ker.mr();
|
||||
self.valid_right_tiles = n / ker.nr();
|
||||
self.remnant_right = n % ker.nr();
|
||||
let mut offset = 0;
|
||||
let mut align = std::mem::size_of::<*const ()>();
|
||||
fn ld(spec: usize, uspec: usize, loc: usize) -> LocDependant {
|
||||
LocDependant {
|
||||
spec,
|
||||
ker_spec: uspec,
|
||||
loc,
|
||||
buffer_a: None,
|
||||
buffer_b: None,
|
||||
}
|
||||
}
|
||||
for (ix, spec) in specs.iter().enumerate() {
|
||||
offset = offset.next_multiple_of(&align);
|
||||
let ker_spec = match spec {
|
||||
FS::BinScalar(t, op) => match op {
|
||||
BinOp::Min => FKS::ScalarMin(*t.try_as_plain()?.to_scalar()?),
|
||||
BinOp::Max => FKS::ScalarMax(*t.try_as_plain()?.to_scalar()?),
|
||||
BinOp::Mul => FKS::ScalarMul(*t.try_as_plain()?.to_scalar()?),
|
||||
BinOp::Add => FKS::ScalarAdd(*t.try_as_plain()?.to_scalar()?),
|
||||
BinOp::Sub => FKS::ScalarSub(*t.try_as_plain()?.to_scalar()?),
|
||||
BinOp::SubF => FKS::ScalarSubF(*t.try_as_plain()?.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.ker_specs.len(), offset));
|
||||
offset += TI::datum_type().size_of() * ker.mr();
|
||||
FusedKerSpec::Done
|
||||
}
|
||||
FS::BinPerCol(_, _) => {
|
||||
self.loc_dependant
|
||||
.push(ld(ix, self.ker_specs.len(), offset));
|
||||
offset += TI::datum_type().size_of() * ker.nr();
|
||||
FusedKerSpec::Done
|
||||
}
|
||||
FS::AddRowColProducts(_, _) => {
|
||||
self.loc_dependant
|
||||
.push(ld(ix, self.ker_specs.len(), offset));
|
||||
offset += TI::datum_type().size_of() * (ker.mr() + ker.nr());
|
||||
FusedKerSpec::Done
|
||||
}
|
||||
FS::AddUnicast(_) => {
|
||||
self.loc_dependant
|
||||
.push(ld(ix, self.ker_specs.len(), offset));
|
||||
offset += TI::datum_type().size_of() * ker.mr() * ker.nr();
|
||||
FusedKerSpec::Done
|
||||
}
|
||||
FS::Store(store) => {
|
||||
self.loc_dependant
|
||||
.push(ld(ix, self.ker_specs.len(), offset));
|
||||
offset += store.item_size * ker.mr() * ker.nr();
|
||||
FusedKerSpec::Done
|
||||
}
|
||||
FS::LeakyRelu(t) => FKS::LeakyRelu(*t.try_as_plain()?.to_scalar()?),
|
||||
FS::AddMatMul { a, b, packing } => {
|
||||
let mut ld = ld(ix, self.ker_specs.len(), offset);
|
||||
offset += std::mem::size_of::<AddMatMulTemp>();
|
||||
if let Some(tmp) = a.scratch_panel_buffer_layout() {
|
||||
align = tmp.align().lcm(&align);
|
||||
offset = Integer::next_multiple_of(&offset, &tmp.align());
|
||||
ld.buffer_a = Some(offset);
|
||||
offset += tmp.size();
|
||||
}
|
||||
if let Some(tmp) = b.scratch_panel_buffer_layout() {
|
||||
align = tmp.align().lcm(&align);
|
||||
offset = Integer::next_multiple_of(&offset, &tmp.align());
|
||||
ld.buffer_b = Some(offset);
|
||||
offset += tmp.size();
|
||||
}
|
||||
self.loc_dependant.push(ld);
|
||||
FusedKerSpec::AddMatMul {
|
||||
k: 0,
|
||||
pa: std::ptr::null(),
|
||||
pb: std::ptr::null(),
|
||||
packing: *packing,
|
||||
}
|
||||
}
|
||||
};
|
||||
self.ker_specs.push(ker_spec);
|
||||
}
|
||||
self.ker_specs.push(FKS::Done);
|
||||
self.blob_size = offset;
|
||||
self.blob_align = align;
|
||||
|
||||
self.generation = GENERATION.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub unsafe fn run(
|
||||
&self,
|
||||
ker: &impl MatMatMulKer<Acc = TI>,
|
||||
specs: &[FusedSpec],
|
||||
down: usize,
|
||||
right: usize,
|
||||
) -> TractResult<()> {
|
||||
// Per-tile entry: enter the TLS scope (does sync once) then run a single
|
||||
// tile. Single-threaded callers should prefer `run_in_tls_scope`+
|
||||
// `run_one_tile` to amortise the TLS borrow + sync over many tiles.
|
||||
unsafe {
|
||||
self.run_in_tls_scope(|this, tls| this.run_one_tile(ker, specs, tls, down, right))
|
||||
}
|
||||
}
|
||||
|
||||
/// Borrow the per-thread scratch blob for a single MMM call and `sync` it
|
||||
/// once. The closure is invoked once with a mutable reference to the TLS
|
||||
/// scratch and to `self`. Used by single-threaded matmul drivers to avoid
|
||||
/// re-entering TLS / re-running `sync` per tile.
|
||||
pub(crate) unsafe fn run_in_tls_scope<F, R>(&self, f: F) -> R
|
||||
where
|
||||
F: FnOnce(&Self, &mut TLSScratch) -> R,
|
||||
{
|
||||
TLS.with_borrow_mut(|tls| {
|
||||
tls.sync(self);
|
||||
f(self, tls)
|
||||
})
|
||||
}
|
||||
|
||||
/// Run a single tile against an already-borrowed TLS scratch. Caller is
|
||||
/// responsible for entering `run_in_tls_scope` first (so `sync` has run).
|
||||
#[inline(always)]
|
||||
pub(crate) unsafe fn run_one_tile(
|
||||
&self,
|
||||
ker: &impl MatMatMulKer<Acc = TI>,
|
||||
specs: &[FusedSpec],
|
||||
tls: &mut TLSScratch,
|
||||
down: usize,
|
||||
right: usize,
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
if down < self.valid_down_tiles && right < self.valid_right_tiles {
|
||||
self.for_valid_tile(ker, specs, tls, down, right)?;
|
||||
let err = ker.kernel(tls.ker_specs());
|
||||
debug_assert_eq!(err, 0, "Kernel return error {err}");
|
||||
} else {
|
||||
let remnant_down = if down < self.valid_down_tiles {
|
||||
ker.mr()
|
||||
} else {
|
||||
self.remnant_down
|
||||
};
|
||||
let remnant_right = if right < self.valid_right_tiles {
|
||||
ker.nr()
|
||||
} else {
|
||||
self.remnant_right
|
||||
};
|
||||
self.for_border_tile(ker, specs, tls, down, right, remnant_down, remnant_right)?;
|
||||
let err = ker.kernel(tls.ker_specs());
|
||||
debug_assert_eq!(err, 0, "Kernel return error {err}");
|
||||
self.postprocess_tile(specs, tls, down, right, remnant_down, remnant_right)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
unsafe fn for_valid_tile(
|
||||
&self,
|
||||
ker: &impl MatMatMulKer<Acc = TI>,
|
||||
specs: &[FusedSpec],
|
||||
tls: &mut TLSScratch,
|
||||
down: usize,
|
||||
right: usize,
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
use FusedKerSpec as FKS;
|
||||
use FusedSpec as FS;
|
||||
let ScratchSpaceImpl {
|
||||
ker_specs,
|
||||
loc_dependant,
|
||||
..
|
||||
} = self;
|
||||
debug_assert!(specs.len() + 2 == ker_specs.len());
|
||||
for LocDependant {
|
||||
spec,
|
||||
ker_spec,
|
||||
loc,
|
||||
buffer_a,
|
||||
buffer_b,
|
||||
} in loc_dependant
|
||||
{
|
||||
let spec = specs.get_unchecked(*spec);
|
||||
let it = match spec {
|
||||
FS::BinPerRow(v, op) => {
|
||||
let v = v.as_ptr_unchecked::<TI>().add(down * ker.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 * ker.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 * ker.mr());
|
||||
let col_ptr = cols.as_ptr_unchecked::<TI>().add(right * ker.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 { a, b, packing } => {
|
||||
let scratch = (tls.blob.as_mut_ptr().add(*loc) as *mut AddMatMulTemp)
|
||||
.as_mut()
|
||||
.unwrap();
|
||||
if scratch.panel_a_id != down {
|
||||
scratch.ptr_a = a.panel_bytes(
|
||||
down,
|
||||
buffer_a.map(|o| tls.blob.as_mut_ptr().add(o)),
|
||||
)?;
|
||||
scratch.panel_a_id = down;
|
||||
}
|
||||
if scratch.panel_b_id != right {
|
||||
scratch.ptr_b = b.panel_bytes(
|
||||
right,
|
||||
buffer_b.map(|o| tls.blob.as_mut_ptr().add(o)),
|
||||
)?;
|
||||
scratch.panel_b_id = right;
|
||||
}
|
||||
FKS::AddMatMul {
|
||||
k: b.k(),
|
||||
pa: scratch.ptr_a,
|
||||
pb: scratch.ptr_b,
|
||||
packing: *packing,
|
||||
}
|
||||
}
|
||||
_ => std::hint::unreachable_unchecked(),
|
||||
};
|
||||
*tls.ker_specs().get_unchecked_mut(*ker_spec) = it;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
unsafe fn for_border_tile(
|
||||
&self,
|
||||
ker: &impl MatMatMulKer<Acc = TI>,
|
||||
specs: &[FusedSpec],
|
||||
tls: &mut TLSScratch,
|
||||
down: usize,
|
||||
right: usize,
|
||||
m_remnant: usize,
|
||||
n_remnant: usize,
|
||||
) -> TractResult<()> {
|
||||
unsafe {
|
||||
use FusedKerSpec as FKS;
|
||||
use FusedSpec as FS;
|
||||
for LocDependant {
|
||||
spec,
|
||||
ker_spec: uspec,
|
||||
loc,
|
||||
buffer_a,
|
||||
buffer_b,
|
||||
} in &self.loc_dependant
|
||||
{
|
||||
let loc = tls.blob.as_mut_ptr().add(*loc);
|
||||
let spec = specs.get_unchecked(*spec);
|
||||
let it = match spec {
|
||||
FS::BinPerRow(v, op) => {
|
||||
let buf = std::slice::from_raw_parts_mut(loc as *mut TI, ker.mr());
|
||||
let ptr = if m_remnant < ker.mr() {
|
||||
if m_remnant > 0 {
|
||||
buf.get_unchecked_mut(..m_remnant).copy_from_slice(
|
||||
v.as_slice_unchecked()
|
||||
.get_unchecked(down * ker.mr()..)
|
||||
.get_unchecked(..m_remnant),
|
||||
);
|
||||
}
|
||||
if cfg!(debug_assertions) {
|
||||
buf.get_unchecked_mut(m_remnant..)
|
||||
.iter_mut()
|
||||
.for_each(|x| *x = TI::zero());
|
||||
}
|
||||
buf.as_ptr()
|
||||
} else {
|
||||
v.as_ptr_unchecked::<TI>().add(down * ker.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, ker.nr());
|
||||
let ptr = if n_remnant < ker.nr() {
|
||||
if n_remnant > 0 {
|
||||
buf.get_unchecked_mut(..n_remnant).copy_from_slice(
|
||||
v.as_slice_unchecked()
|
||||
.get_unchecked(right * ker.nr()..)
|
||||
.get_unchecked(..n_remnant),
|
||||
);
|
||||
}
|
||||
if cfg!(debug_assertions) {
|
||||
buf.get_unchecked_mut(n_remnant..)
|
||||
.iter_mut()
|
||||
.for_each(|x| *x = TI::zero());
|
||||
}
|
||||
buf.as_ptr()
|
||||
} else {
|
||||
v.as_ptr_unchecked::<TI>().add(right * ker.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, ker.mr());
|
||||
let row_ptr = if m_remnant < ker.mr() {
|
||||
r.get_unchecked_mut(..m_remnant).copy_from_slice(
|
||||
rows.as_slice_unchecked()
|
||||
.get_unchecked(down * ker.mr()..)
|
||||
.get_unchecked(..m_remnant),
|
||||
);
|
||||
if cfg!(debug_assertions) {
|
||||
r.get_unchecked_mut(m_remnant..)
|
||||
.iter_mut()
|
||||
.for_each(|x| *x = TI::zero());
|
||||
}
|
||||
r.as_ptr()
|
||||
} else {
|
||||
rows.as_ptr_unchecked::<TI>().add(down * ker.mr())
|
||||
};
|
||||
let c = std::slice::from_raw_parts_mut(
|
||||
(loc as *mut TI).add(ker.mr()),
|
||||
ker.nr(),
|
||||
);
|
||||
let col_ptr = if n_remnant < ker.nr() {
|
||||
c.get_unchecked_mut(..n_remnant).copy_from_slice(
|
||||
cols.as_slice_unchecked()
|
||||
.get_unchecked(right * ker.nr()..)
|
||||
.get_unchecked(..n_remnant),
|
||||
);
|
||||
if cfg!(debug_assertions) {
|
||||
r.get_unchecked_mut(n_remnant..)
|
||||
.iter_mut()
|
||||
.for_each(|x| *x = TI::zero());
|
||||
}
|
||||
c.as_ptr()
|
||||
} else {
|
||||
cols.as_ptr_unchecked::<TI>().add(right * ker.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 * ker.mr() as isize
|
||||
+ col_byte_stride * right as isize * ker.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, ker.mr() * ker.nr());
|
||||
if cfg!(debug_assertions) {
|
||||
tmp_d_tile.iter_mut().for_each(|t| *t = TI::zero());
|
||||
}
|
||||
for r in 0..m_remnant as isize {
|
||||
for c in 0..n_remnant 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 * ker.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>() * ker.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 * ker.mr()) as isize,
|
||||
};
|
||||
FKS::Store(tmpc)
|
||||
}
|
||||
FS::AddMatMul { a, b, packing } => {
|
||||
let scratch = (loc as *mut AddMatMulTemp).as_mut().unwrap();
|
||||
if scratch.panel_a_id != down {
|
||||
scratch.ptr_a = a.panel_bytes(
|
||||
down,
|
||||
buffer_a.map(|o| tls.blob.as_mut_ptr().add(o)),
|
||||
)?;
|
||||
scratch.panel_a_id = down;
|
||||
}
|
||||
if scratch.panel_b_id != right {
|
||||
scratch.ptr_b = b.panel_bytes(
|
||||
right,
|
||||
buffer_b.map(|o| tls.blob.as_mut_ptr().add(o)),
|
||||
)?;
|
||||
scratch.panel_b_id = right;
|
||||
}
|
||||
FKS::AddMatMul {
|
||||
k: b.k(),
|
||||
pa: scratch.ptr_a,
|
||||
pb: scratch.ptr_b,
|
||||
packing: *packing,
|
||||
}
|
||||
}
|
||||
_ => std::hint::unreachable_unchecked(),
|
||||
};
|
||||
*tls.ker_specs().get_unchecked_mut(*uspec) = it;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn uspecs(&self) -> &[FusedKerSpec<TI>] {
|
||||
&self.ker_specs
|
||||
}
|
||||
|
||||
unsafe fn postprocess_tile(
|
||||
&self,
|
||||
specs: &[FusedSpec],
|
||||
tls: &mut TLSScratch,
|
||||
down: usize,
|
||||
right: usize,
|
||||
m_remnant: usize,
|
||||
n_remnant: usize,
|
||||
) -> TractResult<()>
|
||||
where
|
||||
TI: LADatum,
|
||||
{
|
||||
unsafe {
|
||||
for LocDependant {
|
||||
spec,
|
||||
ker_spec: uspec,
|
||||
..
|
||||
} in self.loc_dependant.iter()
|
||||
{
|
||||
let spec = specs.get_unchecked(*spec);
|
||||
let ker_spec = tls.ker_specs::<TI>().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)
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
+326
@@ -0,0 +1,326 @@
|
||||
use std::fmt;
|
||||
use std::fmt::Debug;
|
||||
use tract_data::internal::*;
|
||||
|
||||
use super::MMMInputValue;
|
||||
|
||||
/// Non-plain tensor storage for packed matrices.
|
||||
///
|
||||
/// Holds one or more `Box<dyn MMMInputValue>` values with an optional batch
|
||||
/// shape, replacing the previous `Tensor` + double-downcast pattern.
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct PackedMatrixStorage {
|
||||
values: Vec<Box<dyn MMMInputValue>>,
|
||||
batch_shape: TVec<usize>,
|
||||
batch_strides: TVec<isize>,
|
||||
}
|
||||
|
||||
impl PackedMatrixStorage {
|
||||
/// Scalar storage (one value, empty shape).
|
||||
pub fn new(value: Box<dyn MMMInputValue>) -> Self {
|
||||
PackedMatrixStorage {
|
||||
values: vec![value],
|
||||
batch_shape: tvec![],
|
||||
batch_strides: tvec![],
|
||||
}
|
||||
}
|
||||
|
||||
/// Batched storage (shape like `[batch, group]`).
|
||||
pub fn new_batched(shape: &[usize], values: Vec<Box<dyn MMMInputValue>>) -> Self {
|
||||
let expected: usize = shape.iter().product();
|
||||
assert_eq!(
|
||||
values.len(),
|
||||
expected,
|
||||
"values length must match shape product"
|
||||
);
|
||||
let strides = Self::compute_strides(shape);
|
||||
PackedMatrixStorage {
|
||||
values,
|
||||
batch_shape: shape.into(),
|
||||
batch_strides: strides,
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_strides(shape: &[usize]) -> TVec<isize> {
|
||||
let mut strides: TVec<isize> = tvec![0; shape.len()];
|
||||
if !shape.is_empty() {
|
||||
strides[shape.len() - 1] = 1;
|
||||
for i in (0..shape.len() - 1).rev() {
|
||||
strides[i] = strides[i + 1] * shape[i + 1] as isize;
|
||||
}
|
||||
}
|
||||
strides
|
||||
}
|
||||
|
||||
/// Scalar access (asserts single value).
|
||||
#[inline]
|
||||
pub fn value(&self) -> &dyn MMMInputValue {
|
||||
debug_assert_eq!(self.values.len(), 1);
|
||||
&*self.values[0]
|
||||
}
|
||||
|
||||
/// Batched access by coordinates.
|
||||
pub fn value_at(&self, coords: &[usize]) -> &dyn MMMInputValue {
|
||||
let idx = self.flat_index(coords);
|
||||
&*self.values[idx]
|
||||
}
|
||||
|
||||
/// Batched access by flat (pre-computed) index.
|
||||
#[inline]
|
||||
pub fn value_at_flat(&self, idx: usize) -> &dyn MMMInputValue {
|
||||
&*self.values[idx]
|
||||
}
|
||||
|
||||
pub fn values(&self) -> &[Box<dyn MMMInputValue>] {
|
||||
&self.values
|
||||
}
|
||||
|
||||
pub fn batch_shape(&self) -> &[usize] {
|
||||
&self.batch_shape
|
||||
}
|
||||
|
||||
pub fn batch_strides(&self) -> &[isize] {
|
||||
&self.batch_strides
|
||||
}
|
||||
|
||||
/// Convert to a Tensor with the given logical datum type.
|
||||
pub fn into_tensor(self, dt: DatumType) -> Tensor {
|
||||
let shape: TVec<usize> = self.batch_shape.clone();
|
||||
Tensor::from_storage(dt, &shape, self)
|
||||
}
|
||||
|
||||
fn flat_index(&self, coords: &[usize]) -> usize {
|
||||
coords
|
||||
.iter()
|
||||
.zip(self.batch_strides.iter())
|
||||
.map(|(c, s)| *c as isize * s)
|
||||
.sum::<isize>() as usize
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for PackedMatrixStorage {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"PackedMatrixStorage({} values, shape={:?})",
|
||||
self.values.len(),
|
||||
self.batch_shape
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for PackedMatrixStorage {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"PackedMatrixStorage({} values, shape={:?})",
|
||||
self.values.len(),
|
||||
self.batch_shape
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl TensorStorage for PackedMatrixStorage {
|
||||
fn byte_len(&self) -> usize {
|
||||
// Approximate: sum of individual value sizes isn't precise but gives a ballpark
|
||||
self.values.len() * std::mem::size_of::<Box<dyn MMMInputValue>>()
|
||||
}
|
||||
|
||||
fn is_empty(&self) -> bool {
|
||||
self.values.is_empty()
|
||||
}
|
||||
|
||||
fn deep_clone(&self) -> Box<dyn TensorStorage> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
|
||||
fn as_plain(&self) -> Option<&PlainStorage> {
|
||||
None
|
||||
}
|
||||
|
||||
fn as_plain_mut(&mut self) -> Option<&mut PlainStorage> {
|
||||
None
|
||||
}
|
||||
|
||||
fn into_plain(self: Box<Self>) -> Option<PlainStorage> {
|
||||
None
|
||||
}
|
||||
|
||||
fn dyn_hash(&self, state: &mut dyn std::hash::Hasher) {
|
||||
for v in &self.values {
|
||||
v.dyn_hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
fn exotic_fact(&self, _shape: &[usize]) -> TractResult<Option<Box<dyn ExoticFact>>> {
|
||||
if self.values.len() == 1 {
|
||||
Ok(Some(dyn_clone::clone_box(self.values[0].exotic_fact())))
|
||||
} else {
|
||||
let facts: TVec<Box<dyn ExoticFact>> = self
|
||||
.values
|
||||
.iter()
|
||||
.map(|v| dyn_clone::clone_box(v.exotic_fact()))
|
||||
.collect();
|
||||
Ok(Some(Box::new(facts)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
|
||||
pub enum OutputStoreSpec {
|
||||
View {
|
||||
m_axis: Option<usize>,
|
||||
n_axis: Option<usize>,
|
||||
mr: usize,
|
||||
nr: usize,
|
||||
},
|
||||
Strides {
|
||||
row_byte_stride: isize,
|
||||
col_byte_stride: isize,
|
||||
mr: usize,
|
||||
nr: 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,
|
||||
}
|
||||
|
||||
unsafe impl Send for OutputStore {}
|
||||
unsafe impl Sync for OutputStore {}
|
||||
|
||||
impl OutputStoreSpec {
|
||||
#[inline]
|
||||
pub unsafe fn wrap(&self, tensor: &TensorView) -> OutputStore {
|
||||
let (mr, nr, row_byte_stride, col_byte_stride) = unsafe { self.compute_strides(tensor) };
|
||||
OutputStore {
|
||||
ptr: unsafe { 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(),
|
||||
}
|
||||
}
|
||||
|
||||
#[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 = m_axis
|
||||
.map(|ax| *unsafe { tensor_strides.get_unchecked(ax) })
|
||||
.unwrap_or(0);
|
||||
let col_item_stride = n_axis
|
||||
.map(|ax| *unsafe { tensor_strides.get_unchecked(ax) })
|
||||
.unwrap_or(0);
|
||||
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 {
|
||||
unsafe {
|
||||
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,
|
||||
) {
|
||||
unsafe {
|
||||
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,
|
||||
) {
|
||||
unsafe {
|
||||
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,
|
||||
}
|
||||
Vendored
+304
@@ -0,0 +1,304 @@
|
||||
use crate::frame::mmm::*;
|
||||
use crate::{BinOp, LADatum};
|
||||
use num_traits::AsPrimitive;
|
||||
use std::ops::Neg;
|
||||
use tests::display_error;
|
||||
use tract_data::internal::*;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! mmm_frame_tests {
|
||||
($ker:expr, $ta:ty, $tb:ty, $tc:ty, $ti:ty) => {
|
||||
mod frame {
|
||||
use tract_data::internal::*;
|
||||
#[allow(unused_imports)]
|
||||
use $crate::frame::mmm::tests::frame::*;
|
||||
|
||||
#[test]
|
||||
fn row_mul_2_1_3() -> TractResult<()> {
|
||||
unsafe { row_mul::<_, $ta, $tb, $tc, $ti>($ker, 2, 3)? }
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn row_add_2_1_3() -> TractResult<()> {
|
||||
unsafe { row_add::<_, $ta, $tb, $tc, $ti>($ker, 2, 3)? }
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn col_mul_2_1_3() -> TractResult<()> {
|
||||
unsafe { col_mul::<_, $ta, $tb, $tc, $ti>($ker, 2, 3)? }
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn col_add_2_1_3() -> TractResult<()> {
|
||||
unsafe { col_add::<_, $ta, $tb, $tc, $ti>($ker, 2, 3)? }
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_2_1_3() -> TractResult<()> {
|
||||
unsafe { max::<_, $ta, $tb, $tc, $ti>($ker, 2, 3)? }
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn min_2_1_3() -> TractResult<()> {
|
||||
unsafe { min::<_, $ta, $tb, $tc, $ti>($ker, 2, 3)? }
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn add_d_2_1_3() -> TractResult<()> {
|
||||
unsafe { add_d::<_, $ta, $tb, $tc, $ti>($ker, 2, 3)? }
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn add_d_big() -> TractResult<()> {
|
||||
unsafe { add_d::<_, $ta, $tb, $tc, $ti>($ker, 197, 1)? }
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub unsafe fn fused_ops<
|
||||
K: MatMatMulKer<Acc = TI> + 'static,
|
||||
TA,
|
||||
TB,
|
||||
TC,
|
||||
TI,
|
||||
F: Fn(usize, usize) -> TC,
|
||||
>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
spec: &[FusedSpec],
|
||||
expect: F,
|
||||
) -> TractResult<()>
|
||||
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>,
|
||||
{
|
||||
if !ker.is_supported_here() {
|
||||
return Ok(());
|
||||
};
|
||||
crate::setup_test_logger();
|
||||
|
||||
let mut found = Tensor::zero::<TC>(&[m, n])?;
|
||||
let c_store = unsafe {
|
||||
ker.c_from_data_and_strides(TC::datum_type().size_of(), n as isize, 1)
|
||||
.wrap(&found.view_mut())
|
||||
};
|
||||
let mut spec: TVec<FusedSpec> = spec.into();
|
||||
spec.push(FusedSpec::Store(c_store));
|
||||
|
||||
unsafe { ker.run(m, n, &spec) }?;
|
||||
let expected =
|
||||
tract_ndarray::prelude::Array2::from_shape_fn((m, n), |(r, c)| expect(r, c)).into_tensor();
|
||||
let err = found.close_enough(&expected, true);
|
||||
if err.is_err() {
|
||||
display_error(
|
||||
found.try_as_plain()?.as_slice::<TC>()?,
|
||||
expected.try_as_plain()?.as_slice::<TC>()?,
|
||||
m,
|
||||
n,
|
||||
);
|
||||
}
|
||||
err
|
||||
}
|
||||
|
||||
pub unsafe fn row_add<K: MatMatMulKer<Acc = TI> + 'static, TA, TB, TC, TI>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
) -> TractResult<()>
|
||||
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>>();
|
||||
unsafe {
|
||||
fused_ops::<K, TA, TB, TC, TI, _>(
|
||||
ker,
|
||||
m,
|
||||
n,
|
||||
&[FusedSpec::BinPerRow(tensor1(&bias).view(), BinOp::Add)],
|
||||
|r, _| bias[r].as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub unsafe fn row_mul<K: MatMatMulKer<Acc = TI> + 'static, TA, TB, TC, TI>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
) -> TractResult<()>
|
||||
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>>();
|
||||
unsafe {
|
||||
fused_ops::<K, TA, TB, TC, TI, _>(
|
||||
ker,
|
||||
m,
|
||||
n,
|
||||
&[
|
||||
FusedSpec::BinScalar(&tensor0(1i32.as_()), BinOp::Add),
|
||||
FusedSpec::BinPerRow(tensor1(&bias).view(), BinOp::Mul),
|
||||
],
|
||||
|r, _| bias[r].as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub unsafe fn col_add<K: MatMatMulKer<Acc = TI> + 'static, TA, TB, TC, TI>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
) -> TractResult<()>
|
||||
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>>();
|
||||
unsafe {
|
||||
fused_ops::<K, TA, TB, TC, TI, _>(
|
||||
ker,
|
||||
m,
|
||||
n,
|
||||
&[FusedSpec::BinPerCol(tensor1(&bias).view(), BinOp::Add)],
|
||||
|_, c| bias[c].as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub unsafe fn col_mul<K: MatMatMulKer<Acc = TI> + 'static, TA, TB, TC, TI>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
) -> TractResult<()>
|
||||
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>>();
|
||||
unsafe {
|
||||
fused_ops::<K, TA, TB, TC, TI, _>(
|
||||
ker,
|
||||
m,
|
||||
n,
|
||||
&[
|
||||
FusedSpec::BinScalar(&tensor0(1i32.as_()), BinOp::Add),
|
||||
FusedSpec::BinPerCol(tensor1(&bias).view(), BinOp::Mul),
|
||||
],
|
||||
|_, c| bias[c].as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub unsafe fn add_d<K: MatMatMulKer<Acc = TI> + 'static, TA, TB, TC, TI>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
) -> TractResult<()>
|
||||
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])?;
|
||||
let store_spec = OutputStoreSpec::View {
|
||||
m_axis: Some(0),
|
||||
n_axis: Some(1),
|
||||
mr: ker.mr(),
|
||||
nr: ker.nr(),
|
||||
};
|
||||
let view_d = d.to_plain_array_view::<TI>()?.into_dimensionality()?;
|
||||
unsafe {
|
||||
fused_ops::<K, TA, TB, TC, TI, _>(
|
||||
ker,
|
||||
m,
|
||||
n,
|
||||
&[FusedSpec::AddUnicast(store_spec.wrap(&d.view()))],
|
||||
|r, c| view_d[(r, c)].as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub unsafe fn max<K: MatMatMulKer<Acc = TI>, TA, TB, TC, TI>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
) -> TractResult<()>
|
||||
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_();
|
||||
unsafe {
|
||||
fused_ops::<K, TA, TB, TC, TI, _>(
|
||||
ker,
|
||||
m,
|
||||
n,
|
||||
&[FusedSpec::BinScalar(&tensor0(five), BinOp::Max)],
|
||||
|_, _| five.as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub unsafe fn min<K: MatMatMulKer<Acc = TI>, TA, TB, TC, TI>(
|
||||
ker: &K,
|
||||
m: usize,
|
||||
n: usize,
|
||||
) -> TractResult<()>
|
||||
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_();
|
||||
unsafe {
|
||||
fused_ops::<K, TA, TB, TC, TI, _>(
|
||||
ker,
|
||||
m,
|
||||
n,
|
||||
&[FusedSpec::BinScalar(&tensor0(five), BinOp::Min)],
|
||||
|_, _| TC::zero(),
|
||||
)
|
||||
}
|
||||
}
|
||||
+351
@@ -0,0 +1,351 @@
|
||||
use crate::frame::mmm::fuse::FusedKerSpec;
|
||||
use crate::frame::mmm::storage::*;
|
||||
use crate::frame::mmm::tests::display_error;
|
||||
use crate::frame::mmm::tests::store::mmm_stride_storage;
|
||||
use crate::frame::mmm::*;
|
||||
use num_traits::{AsPrimitive, Bounded};
|
||||
use proptest::prelude::*;
|
||||
use tract_data::internal::*;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! mmm_kernel_fuse_tests {
|
||||
($ker:expr, $tc:ty, $ti: ty) => {
|
||||
mod fuse {
|
||||
use num_traits::Zero;
|
||||
#[allow(unused_imports)]
|
||||
use tract_data::prelude::f16;
|
||||
use tract_data::prelude::tensor0;
|
||||
use $crate::frame::mmm::MatMatMulKer;
|
||||
use $crate::frame::mmm::tests::fuse as test;
|
||||
#[allow(unused_imports)]
|
||||
use $crate::frame::mmm::tests::fuse::*;
|
||||
|
||||
#[test]
|
||||
fn return_zeros() {
|
||||
test::return_zeros::<_, $tc, $ti>($ker)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn store_non_contiguous() {
|
||||
test::store_non_contiguous::<_, $tc, $ti>($ker)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn add_unicast_non_contiguous() {
|
||||
test::add_unicast_non_contiguous::<_, $ti>($ker)
|
||||
}
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn return_c_prop(c in tile::<_, $ti>($ker)) {
|
||||
test::return_c::<_, $ti>($ker, &c)
|
||||
}
|
||||
}
|
||||
|
||||
fn fmin<T: PartialOrd>(a: T, b: T) -> T {
|
||||
if a < b { a } else { b }
|
||||
}
|
||||
|
||||
fn fmax<T: PartialOrd>(a: T, b: T) -> T {
|
||||
if a > b { a } else { b }
|
||||
}
|
||||
|
||||
macro_rules! bin {
|
||||
($FKS:ident, $geo:expr, $f:expr, $extra_cond:expr) => {
|
||||
paste! {
|
||||
#[test]
|
||||
fn [<$FKS:snake>]() {
|
||||
if ($ker).is_supported_here() && $extra_cond {
|
||||
test::$geo::<_, $ti>($ker, $crate::mmm::FusedKerSpec::$FKS, $f);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
bin!(PerColMin, per_col, fmin, true);
|
||||
bin!(PerColMax, per_col, fmax, true);
|
||||
bin!(PerColAdd, per_col, |a, b| a + b, true);
|
||||
bin!(PerColMul, per_col, |a, b| a * b, true);
|
||||
bin!(PerColSub, per_col, |a, b| a - b, true);
|
||||
bin!(PerColSubF, per_col, |a, b| b - a, true);
|
||||
|
||||
bin!(PerRowMin, per_row, fmin, true);
|
||||
bin!(PerRowMax, per_row, fmax, true);
|
||||
bin!(PerRowAdd, per_row, |a, b| a + b, true);
|
||||
bin!(PerRowMul, per_row, |a, b| a * b, true);
|
||||
bin!(PerRowSub, per_row, |a, b| a - b, true);
|
||||
bin!(PerRowSubF, per_row, |a, b| b - a, true);
|
||||
|
||||
bin!(ScalarMin, scalar, fmin, true);
|
||||
bin!(ScalarMax, scalar, fmax, true);
|
||||
bin!(ScalarAdd, scalar, |a, b| a + b, true);
|
||||
bin!(ScalarMul, scalar, |a, b| a * b, true);
|
||||
bin!(ScalarSub, scalar, |a, b| a - b, true);
|
||||
bin!(ScalarSubF, scalar, |a, b| b - a, true);
|
||||
|
||||
bin!(
|
||||
LeakyRelu,
|
||||
scalar,
|
||||
|a, b| if b > <$ti>::zero() { b } else { a * b },
|
||||
($ker).can_fuse(&$crate::mmm::FusedSpec::LeakyRelu(&tensor0(<$ti>::from(
|
||||
1_u8
|
||||
))))
|
||||
);
|
||||
|
||||
#[test]
|
||||
fn return_c_add_row_col_product() {
|
||||
test::return_c_add_row_col_product::<_, $ti>($ker)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn return_c_plus_d() {
|
||||
test::return_c_plus_d::<_, $ti, $ti>($ker)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn return_c_clear() {
|
||||
test::return_c_clear::<_, $ti>($ker)
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
use crate::LADatum;
|
||||
pub fn return_zeros<K, TC, TI>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TC: LADatum,
|
||||
TI: LADatum + Bounded + PartialEq,
|
||||
{
|
||||
if !ker.is_supported_here() {
|
||||
return;
|
||||
}
|
||||
let v = vec![TC::max_value(); ker.mr() * ker.nr()];
|
||||
let c = mmm_stride_storage(&v, ker.nr());
|
||||
let non_linear = tvec![
|
||||
FusedKerSpec::Clear,
|
||||
FusedKerSpec::Store(c),
|
||||
FusedKerSpec::Done
|
||||
];
|
||||
let err = ker.kernel(&non_linear);
|
||||
assert_eq!(err, 0);
|
||||
let expected = vec![TC::zero(); v.len()];
|
||||
display_error(&v, &expected, ker.mr(), ker.nr());
|
||||
assert_eq!(v, expected);
|
||||
}
|
||||
|
||||
pub fn store_non_contiguous<K, TC, TI>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TC: LADatum,
|
||||
TI: LADatum + Bounded + PartialEq,
|
||||
{
|
||||
if !ker.is_supported_here() {
|
||||
return;
|
||||
}
|
||||
let v = vec![TC::max_value(); ker.mr() * 5 * ker.nr() * 3];
|
||||
let c = OutputStoreKer {
|
||||
ptr: v.as_ptr() as _,
|
||||
row_byte_stride: (std::mem::size_of::<TC>() * 3 * ker.nr() * 5) as isize,
|
||||
col_byte_stride: std::mem::size_of::<TC>() as isize * 3,
|
||||
item_size: std::mem::size_of::<TC>(),
|
||||
};
|
||||
let non_linear = tvec![
|
||||
FusedKerSpec::Clear,
|
||||
FusedKerSpec::Store(c),
|
||||
FusedKerSpec::Done
|
||||
];
|
||||
let err = ker.kernel(&non_linear);
|
||||
assert_eq!(err, 0);
|
||||
let mut expected = vec![TC::max_value(); v.len()];
|
||||
for c in 0..ker.nr() {
|
||||
for r in 0..ker.mr() {
|
||||
expected[c * 3 + r * 3 * 5 * ker.nr()] = TC::zero();
|
||||
}
|
||||
}
|
||||
assert_eq!(v, expected);
|
||||
}
|
||||
|
||||
/// `Clear` + `AddUnicast(strided)` + `Store(contiguous)` and check the
|
||||
/// source pattern reaches the destination. Counterpart of
|
||||
/// `store_non_contiguous` on the read side; `return_c_plus_d` uses
|
||||
/// `mmm_stride_storage` (tightly packed) and so doesn't exercise this.
|
||||
pub fn add_unicast_non_contiguous<K, TI>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum + AsPrimitive<TI>,
|
||||
usize: AsPrimitive<TI>,
|
||||
{
|
||||
if !ker.is_supported_here() {
|
||||
return;
|
||||
}
|
||||
let item = std::mem::size_of::<TI>();
|
||||
let row_stride_items = 3 * ker.nr() * 5;
|
||||
let col_stride_items = 3;
|
||||
// Source: a non-contiguous buffer with distinct values at the used
|
||||
// (r, c) cells and sentinel garbage everywhere else.
|
||||
let mut src: Vec<TI> = vec![TI::max_value(); ker.mr() * row_stride_items];
|
||||
for r in 0..ker.mr() {
|
||||
for c in 0..ker.nr() {
|
||||
src[r * row_stride_items + c * col_stride_items] = (1 + c + r * ker.nr()).as_();
|
||||
}
|
||||
}
|
||||
let src_store = OutputStoreKer {
|
||||
ptr: src.as_ptr() as _,
|
||||
row_byte_stride: (item * row_stride_items) as isize,
|
||||
col_byte_stride: (item * col_stride_items) as isize,
|
||||
item_size: item,
|
||||
};
|
||||
// Destination: tightly-packed output for easy comparison.
|
||||
let mut dst: Vec<TI> = vec![TI::min_value(); ker.mr() * ker.nr()];
|
||||
let dst_store = OutputStoreKer {
|
||||
ptr: dst.as_ptr() as _,
|
||||
row_byte_stride: (item * ker.nr()) as isize,
|
||||
col_byte_stride: item as isize,
|
||||
item_size: item,
|
||||
};
|
||||
let non_linear = tvec![
|
||||
FusedKerSpec::Clear,
|
||||
FusedKerSpec::AddUnicast(src_store),
|
||||
FusedKerSpec::Store(dst_store),
|
||||
FusedKerSpec::Done,
|
||||
];
|
||||
let err = ker.kernel(&non_linear);
|
||||
assert_eq!(err, 0);
|
||||
let expected: Vec<TI> = (0..ker.mr() * ker.nr()).map(|i| (1 + i).as_()).collect();
|
||||
display_error(&dst, &expected, ker.mr(), ker.nr());
|
||||
assert_eq!(dst, expected);
|
||||
}
|
||||
|
||||
pub fn fused_ops<K, TI, E>(ker: &K, c: &[TI], ops: &[FusedKerSpec<TI>], expect: E)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
E: Fn(usize, usize, TI) -> TI,
|
||||
{
|
||||
if !ker.is_supported_here() {
|
||||
return;
|
||||
}
|
||||
assert!(c.len() == ker.mr() * ker.nr());
|
||||
let v = c.to_vec();
|
||||
let c = mmm_stride_storage(&v, ker.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 / ker.nr(), ix % ker.nr(), v[ix]))
|
||||
.collect::<Vec<TI>>();
|
||||
let err = ker.kernel(&ops);
|
||||
assert_eq!(err, 0);
|
||||
display_error(&v, &expected, ker.mr(), ker.nr());
|
||||
assert_eq!(v, expected);
|
||||
}
|
||||
|
||||
pub fn return_c<K, TI>(ker: &K, v: &[TI])
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
usize: AsPrimitive<TI>,
|
||||
{
|
||||
fused_ops::<K, TI, _>(ker, v, &[], |_, _, c| c + 1.as_() - 1.as_())
|
||||
}
|
||||
|
||||
pub fn return_c_plus_d<K, TI, TD>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
TD: LADatum + AsPrimitive<TI>,
|
||||
usize: AsPrimitive<TI> + AsPrimitive<TD>,
|
||||
{
|
||||
let len = ker.mr() * ker.nr();
|
||||
let v: Vec<TI> = (0..len).map(|f| f.as_()).collect();
|
||||
let d: Vec<TD> = (0..len).map(|f| ((3 * f) % 7).as_()).collect();
|
||||
fused_ops::<K, TI, _>(
|
||||
ker,
|
||||
&v,
|
||||
&[FusedKerSpec::AddUnicast(mmm_stride_storage(&d, ker.nr()))],
|
||||
|row, col, c| c + d[row * ker.nr() + col].as_(),
|
||||
);
|
||||
}
|
||||
|
||||
pub fn per_col<K, TI>(ker: &K, op: impl Fn(*const TI) -> FusedKerSpec<TI>, f: impl Fn(TI, TI) -> TI)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
usize: AsPrimitive<TI>,
|
||||
{
|
||||
let len = ker.mr() * ker.nr();
|
||||
let v: Vec<TI> = (0..len).map(|f| f.as_()).collect();
|
||||
let bias: Vec<TI> = (0..ker.nr()).map(|f| (f + 1).as_()).collect();
|
||||
fused_ops::<K, TI, _>(ker, &v, &[op(bias.as_ptr())], |_, col, c| f(bias[col], c))
|
||||
}
|
||||
|
||||
pub fn per_row<K, TI>(ker: &K, op: impl Fn(*const TI) -> FusedKerSpec<TI>, f: impl Fn(TI, TI) -> TI)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
usize: AsPrimitive<TI>,
|
||||
{
|
||||
let len = ker.mr() * ker.nr();
|
||||
let v: Vec<TI> = (0..len).map(|f| f.as_()).collect();
|
||||
let bias: Vec<TI> = (0..ker.mr()).map(|f| (f + 1).as_()).collect();
|
||||
fused_ops::<K, TI, _>(ker, &v, &[op(bias.as_ptr())], |row, _, c| f(bias[row], c))
|
||||
}
|
||||
|
||||
pub fn scalar<K, TI>(ker: &K, op: impl Fn(TI) -> FusedKerSpec<TI>, f: impl Fn(TI, TI) -> TI)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
isize: AsPrimitive<TI>,
|
||||
{
|
||||
let len = ker.mr() * ker.nr();
|
||||
let v: Vec<TI> = (0..len as isize)
|
||||
.map(|f| (f - len as isize / 2).as_())
|
||||
.collect();
|
||||
let five: TI = 5.as_();
|
||||
fused_ops::<K, TI, _>(ker, &v, &[op(five)], |_, _, c| f(five, c))
|
||||
}
|
||||
|
||||
pub fn return_c_add_row_col_product<K, TI>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
usize: AsPrimitive<TI>,
|
||||
{
|
||||
let len = ker.mr() * ker.nr();
|
||||
let v: Vec<TI> = (0..len).map(|f| (f + 1).as_()).collect();
|
||||
let rows: Vec<TI> = (0..ker.mr()).map(|f| (f + 3).as_()).collect();
|
||||
let cols: Vec<TI> = (0..ker.nr()).map(|f| (f + 2).as_()).collect();
|
||||
fused_ops::<K, TI, _>(
|
||||
ker,
|
||||
&v,
|
||||
&[FusedKerSpec::AddRowColProducts(
|
||||
rows.as_ptr(),
|
||||
cols.as_ptr(),
|
||||
)],
|
||||
|row, col, c| c + cols[col] * rows[row],
|
||||
)
|
||||
}
|
||||
|
||||
pub fn return_c_clear<K, TI>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
usize: AsPrimitive<TI>,
|
||||
{
|
||||
let len = ker.mr() * ker.nr();
|
||||
let v: Vec<TI> = (0..len).map(|f| f.as_()).collect();
|
||||
fused_ops::<K, TI, _>(ker, &v, &[FusedKerSpec::Clear], |_, _, _| 0.as_())
|
||||
}
|
||||
|
||||
pub fn tile<K, TI>(ker: &K) -> BoxedStrategy<Vec<TI>>
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TI: LADatum,
|
||||
i8: AsPrimitive<TI>,
|
||||
{
|
||||
let len = ker.mr() * ker.nr();
|
||||
proptest::collection::vec(any::<i8>().prop_map(|c| c.as_()), len..=len).boxed()
|
||||
}
|
||||
+89
@@ -0,0 +1,89 @@
|
||||
use crate::LADatum;
|
||||
|
||||
#[macro_use]
|
||||
pub mod fuse;
|
||||
#[macro_use]
|
||||
pub mod frame;
|
||||
#[macro_use]
|
||||
pub mod packed_packed;
|
||||
#[macro_use]
|
||||
pub mod q_scale;
|
||||
#[macro_use]
|
||||
pub mod store;
|
||||
|
||||
#[cfg(test)]
|
||||
macro_rules! test_mmm_kernel {
|
||||
(f16, $ker:expr) => {
|
||||
test_mmm_kernel_f16!($ker);
|
||||
};
|
||||
(f32, $ker:expr) => {
|
||||
test_mmm_kernel_f32!($ker);
|
||||
};
|
||||
(f64, $ker:expr) => {
|
||||
test_mmm_kernel_f64!($ker);
|
||||
};
|
||||
(i32, $ker:expr) => {
|
||||
test_mmm_kernel_i32!($ker);
|
||||
};
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! test_mmm_kernel_f16 {
|
||||
($ker: expr) => {
|
||||
mmm_packed_packed_tests!(&*$ker, f16f16:0);
|
||||
mmm_frame_tests!(&*$ker, f16, f16, f16, f16);
|
||||
mmm_kernel_fuse_tests!(&*$ker, f16, f16);
|
||||
mmm_store_test!(&*$ker, f16);
|
||||
};
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! test_mmm_kernel_f32 {
|
||||
($ker: expr) => {
|
||||
mmm_packed_packed_tests!(&*$ker, f32f32:0);
|
||||
mmm_frame_tests!(&*$ker, f32, f32, f32, f32);
|
||||
mmm_kernel_fuse_tests!(&*$ker, f32, f32);
|
||||
mmm_store_test!(&*$ker, f32);
|
||||
};
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! test_mmm_kernel_f64 {
|
||||
($ker:expr) => {
|
||||
mmm_packed_packed_tests!(&*$ker, f64f64:0);
|
||||
mmm_frame_tests!(&*$ker, f64, f64, f64, f64);
|
||||
mmm_kernel_fuse_tests!(&*$ker, f64, f64);
|
||||
mmm_store_test!(&*$ker, f64);
|
||||
};
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! test_mmm_kernel_i32 {
|
||||
($ker: expr) => {
|
||||
mmm_packed_packed_tests!(&*$ker, i32i32:0);
|
||||
mmm_kernel_fuse_tests!(&*$ker, i32, i32);
|
||||
mmm_frame_tests!(&*$ker, i32, i32, i32, i32);
|
||||
mmm_q_scale_tests!(&*$ker);
|
||||
mmm_store_test!(&*$ker, i32);
|
||||
};
|
||||
}
|
||||
|
||||
pub fn display_error<TC: LADatum>(v: &[TC], expected: &[TC], m: usize, n: usize) {
|
||||
if v != expected {
|
||||
for ixm in 0..m {
|
||||
print!("|");
|
||||
for ixn in 0..n {
|
||||
use nu_ansi_term::Color::*;
|
||||
let f = v[ixm * n + ixn];
|
||||
let e = expected[ixm * n + ixn];
|
||||
let color = if f != e { Red.bold() } else { Green.into() };
|
||||
print!("{}|", color.paint(format!("{f:5}")));
|
||||
}
|
||||
print!(" # ");
|
||||
for ixn in 0..n {
|
||||
print!("{:5} ", expected[ixm * n + ixn]);
|
||||
}
|
||||
println!();
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+442
@@ -0,0 +1,442 @@
|
||||
use crate::WeightType;
|
||||
use crate::block_quant::PackedBlockQuantFormat;
|
||||
use crate::mmm::tests::display_error;
|
||||
use crate::mmm::{AsInputValue, FusedKerSpec, FusedSpec, MatMatMul, MatMatMulKer, OutputStoreKer};
|
||||
use proptest::collection::vec;
|
||||
use proptest::prelude::*;
|
||||
use std::fmt::Debug;
|
||||
use tract_data::internal::*;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! mmm_packed_packed_tests {
|
||||
($ker:expr, $packing_id:ident : $packing: expr) => {
|
||||
mod $packing_id {
|
||||
use super::*;
|
||||
#[allow(unused_imports)]
|
||||
use proptest::prelude::*;
|
||||
#[allow(unused_imports)]
|
||||
use tract_data::prelude::f16;
|
||||
use tract_data::prelude::*;
|
||||
use tract_itertools::Itertools;
|
||||
use $crate::frame::mmm::kernel::MatMatMulKer;
|
||||
#[allow(unused_imports)]
|
||||
use $crate::frame::mmm::tests::packed_packed::*;
|
||||
|
||||
mod fuse {
|
||||
use super::*;
|
||||
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(pb in arbitrary_problem(false, $ker, $packing)) {
|
||||
pb.check().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
fn t(a: impl Into<Vec<f32>>, b: impl Into<Vec<f32>>) -> TractResult<()> {
|
||||
PackedPackedProblem::kernel($ker, $packing, a, b).check()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_1() -> TractResult<()> {
|
||||
t(vec![1f32; $ker.mr()], vec![1f32; $ker.nr()])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_2() -> TractResult<()> {
|
||||
t(vec![1f32; $ker.mr() * 2], vec![1f32; $ker.nr() * 2])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_13() -> TractResult<()> {
|
||||
t(vec![1f32; $ker.mr() * 13], vec![1f32; $ker.nr() * 13])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_a_scale() -> TractResult<()> {
|
||||
t(
|
||||
(1..=$ker.mr() as i64).map(|x| x as f32).collect_vec(),
|
||||
vec![1f32; $ker.nr()],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_a_scale_times_2() -> TractResult<()> {
|
||||
t(
|
||||
(1..=2 * $ker.mr() as i64).map(|x| x as f32).collect_vec(),
|
||||
vec![1f32; $ker.nr() * 2],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_empty() -> TractResult<()> {
|
||||
t(vec![0f32; 0], vec![0f32; 0])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_bug_1() -> TractResult<()> {
|
||||
t(vec![0f32; $ker.mr()], vec![0f32; $ker.nr()])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_bug_2() -> TractResult<()> {
|
||||
let mut a = vec![0f32; $ker.mr()];
|
||||
a[0] = 1.;
|
||||
let mut b = vec![0f32; $ker.nr()];
|
||||
b[0] = 1.;
|
||||
t(a, b)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_bug_3() -> TractResult<()> {
|
||||
if $ker.mr() >= 4 {
|
||||
let mut a = vec![0f32; 2 * $ker.mr()];
|
||||
let mut b = vec![0f32; 2 * $ker.nr()];
|
||||
a[2] = -0.7548828f32;
|
||||
a[3] = 0.23547363f32;
|
||||
b[2 * $ker.nr() - 1] = 0.93603516;
|
||||
t(a, b)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_bug_4() -> TractResult<()> {
|
||||
if $ker.mr() > 16 {
|
||||
let mut a = vec![0f32; $ker.mr()];
|
||||
let mut b = vec![0f32; $ker.nr()];
|
||||
a[16] = 1.;
|
||||
b[0] = 1.;
|
||||
t(a, b)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
mod frame {
|
||||
use super::*;
|
||||
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(pb in arbitrary_problem(true, $ker, $packing)) {
|
||||
pb.check().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
fn t(
|
||||
m: usize,
|
||||
n: usize,
|
||||
a: impl Into<Vec<f32>>,
|
||||
b: impl Into<Vec<f32>>,
|
||||
) -> TractResult<()> {
|
||||
PackedPackedProblem::frame($ker, $packing, m, n, a, b).check()
|
||||
}
|
||||
|
||||
fn ti(
|
||||
m: usize,
|
||||
n: usize,
|
||||
a: impl Into<Vec<i32>>,
|
||||
b: impl Into<Vec<i32>>,
|
||||
) -> TractResult<()> {
|
||||
let a = a.into().into_iter().map(|i| i as f32).collect_vec();
|
||||
let b = b.into().into_iter().map(|i| i as f32).collect_vec();
|
||||
t(m, n, a, b)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trivial_1x2() -> TractResult<()> {
|
||||
ti(1, 2, [0], [0, 0])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_empty() -> TractResult<()> {
|
||||
t($ker.mr(), $ker.nr(), [], [])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn packed_packed_empty_2() -> TractResult<()> {
|
||||
t(2 * $ker.mr(), 2 * $ker.nr(), [], [])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mat_mul_1() -> TractResult<()> {
|
||||
ti(
|
||||
3,
|
||||
2,
|
||||
[-3, 3, 5, -5, 6, 0, -6, -5, 0, 0, 9, 7],
|
||||
[-8, 5, 5, -3, 5, 7, -8, -1],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mat_mul_2() -> TractResult<()> {
|
||||
ti(1, 3, [122, 82], [0, 0, 37, 0, 0, 57])
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[derive(Debug, new)]
|
||||
pub struct PackedPackedProblem<K>
|
||||
where
|
||||
K: MatMatMulKer,
|
||||
{
|
||||
pub frame_test: Option<(usize, usize)>,
|
||||
pub ker: K,
|
||||
pub packing: usize,
|
||||
pub a: Vec<f32>,
|
||||
pub b: Vec<f32>,
|
||||
}
|
||||
|
||||
pub fn arbitrary_problem<K: MatMatMulKer>(
|
||||
frame_test: bool,
|
||||
ker: &K,
|
||||
packing: usize,
|
||||
) -> BoxedStrategy<PackedPackedProblem<K>> {
|
||||
let (mr, nr) = (ker.mr(), ker.nr());
|
||||
let item_range = if ker.internal_type().is_integer() {
|
||||
(-5f32)..5f32
|
||||
} else {
|
||||
(-1f32)..1f32
|
||||
};
|
||||
let (m_range, n_range) = if frame_test {
|
||||
(1usize..3 * mr, 1usize..3 * nr)
|
||||
} else {
|
||||
(mr..mr + 1, nr..nr + 1)
|
||||
};
|
||||
let ker = ker.clone();
|
||||
(m_range, 0usize..40, n_range)
|
||||
.prop_flat_map(move |(m, k, n)| {
|
||||
(
|
||||
vec(item_range.clone(), k * m..=k * m),
|
||||
vec(item_range.clone(), k * n..=k * n),
|
||||
Just((m, n)),
|
||||
)
|
||||
})
|
||||
.prop_map(move |(mut a, mut b, mn)| {
|
||||
a.reverse();
|
||||
b.reverse();
|
||||
PackedPackedProblem {
|
||||
frame_test: Some(mn).filter(|_| frame_test),
|
||||
ker: ker.clone(),
|
||||
packing,
|
||||
a,
|
||||
b,
|
||||
}
|
||||
})
|
||||
.boxed()
|
||||
}
|
||||
|
||||
impl<K: MatMatMulKer> PackedPackedProblem<K> {
|
||||
pub fn kernel(
|
||||
ker: &K,
|
||||
packing: usize,
|
||||
a: impl Into<Vec<f32>>,
|
||||
b: impl Into<Vec<f32>>,
|
||||
) -> PackedPackedProblem<K> {
|
||||
PackedPackedProblem {
|
||||
frame_test: None,
|
||||
ker: ker.clone(),
|
||||
packing,
|
||||
a: a.into(),
|
||||
b: b.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn frame(
|
||||
ker: &K,
|
||||
packing: usize,
|
||||
m: usize,
|
||||
n: usize,
|
||||
a: impl Into<Vec<f32>>,
|
||||
b: impl Into<Vec<f32>>,
|
||||
) -> PackedPackedProblem<K> {
|
||||
PackedPackedProblem {
|
||||
frame_test: Some((m, n)),
|
||||
ker: ker.clone(),
|
||||
packing,
|
||||
a: a.into(),
|
||||
b: b.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn mkn(&self) -> (usize, usize, usize) {
|
||||
let (m, n) = self.frame_test.unwrap_or((self.ker.mr(), self.ker.nr()));
|
||||
assert!(m != 0 && n != 0);
|
||||
let k = self.a.len() / m;
|
||||
assert_eq!(self.b.len() / n, k);
|
||||
(m, k, n)
|
||||
}
|
||||
|
||||
pub fn padded_inputs(&self) -> TractResult<(Tensor, Tensor)> {
|
||||
let (pack_a, pack_b) = &self.ker.packings()[self.packing];
|
||||
let (m, k, n) = self.mkn();
|
||||
let k_aligned = k.next_multiple_of(pack_a.k_alignment().max(pack_b.k_alignment()));
|
||||
|
||||
let mut a = Tensor::zero::<f32>(&[m, k_aligned])?;
|
||||
for row in 0..m {
|
||||
for col in 0..k {
|
||||
a.try_as_plain_mut()?.to_array_view_mut()?[[row, col]] = self.a[col + k * row];
|
||||
}
|
||||
}
|
||||
if let WeightType::Plain(dt) = pack_a.precursor() {
|
||||
a = a.cast_to_dt(dt)?.into_owned();
|
||||
}
|
||||
let mut b = Tensor::zero::<f32>(&[k_aligned, n])?;
|
||||
for row in 0..k {
|
||||
for col in 0..n {
|
||||
b.try_as_plain_mut()?.to_array_view_mut()?[[row, col]] = self.b[col + n * row];
|
||||
}
|
||||
}
|
||||
if let WeightType::Plain(dt) = pack_b.precursor() {
|
||||
b = b.cast_to_dt(dt)?.into_owned();
|
||||
}
|
||||
|
||||
Ok((a, b))
|
||||
}
|
||||
|
||||
pub fn reference(&self) -> TractResult<Tensor> {
|
||||
let (m, k, n) = self.mkn();
|
||||
let (pack_a, pack_b) = &self.ker.packings()[self.packing];
|
||||
let (mut a, b) = self.padded_inputs()?;
|
||||
let k_aligned = k.next_multiple_of(pack_a.k_alignment().max(pack_b.k_alignment()));
|
||||
if let Some(pbqf) = pack_a.downcast_ref::<PackedBlockQuantFormat>() {
|
||||
a = pbqf.simulate_precision_loss(a, 1)?;
|
||||
};
|
||||
let mut c = Tensor::zero::<K::Acc>(&[m, n])?;
|
||||
|
||||
let a = a.cast_to::<K::Acc>()?;
|
||||
let a = a.try_as_plain()?.as_slice::<K::Acc>()?;
|
||||
let b = b.cast_to::<K::Acc>()?;
|
||||
let b = b.try_as_plain()?.as_slice::<K::Acc>()?;
|
||||
let mut c_plain = c.try_as_plain_mut()?;
|
||||
let mut view = c_plain
|
||||
.to_array_view_mut::<K::Acc>()?
|
||||
.into_dimensionality()?;
|
||||
for ix_m in 0..m {
|
||||
for ix_n in 0..n {
|
||||
for ix_k in 0..k {
|
||||
let a = a[ix_k + k_aligned * ix_m];
|
||||
let b = b[ix_n + n * ix_k];
|
||||
view[(ix_m, ix_n)] += a * b;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(c)
|
||||
}
|
||||
|
||||
pub fn run(&self) -> TractResult<Tensor> {
|
||||
let (m, k, n) = self.mkn();
|
||||
let (pack_a, pack_b) = &self.ker.packings()[self.packing];
|
||||
let k_aligned = k.next_multiple_of(pack_a.k_alignment().max(pack_b.k_alignment()));
|
||||
|
||||
let (a, b) = self.padded_inputs()?;
|
||||
let pa = pack_a.prepare_one(&a, 1, 0)?;
|
||||
let pb = pack_b.prepare_one(&b, 0, 1)?;
|
||||
|
||||
let mut v = unsafe { Tensor::uninitialized_dt(self.ker.internal_type(), &[m, n])? };
|
||||
let item_size = self.ker.internal_type().size_of();
|
||||
|
||||
if self.frame_test.is_some() {
|
||||
unsafe {
|
||||
let c = self.ker.c_view(Some(0), Some(1)).wrap(&v.view_mut());
|
||||
let ops = tvec!(
|
||||
FusedSpec::AddMatMul {
|
||||
a: AsInputValue::Borrowed(&*pa),
|
||||
b: AsInputValue::Borrowed(&*pb),
|
||||
packing: self.packing
|
||||
},
|
||||
FusedSpec::Store(c)
|
||||
);
|
||||
self.ker.run(m, n, &ops)?;
|
||||
}
|
||||
} else {
|
||||
let c = OutputStoreKer {
|
||||
ptr: v.as_bytes_mut().as_mut_ptr(),
|
||||
row_byte_stride: (item_size * self.ker.nr()) as isize,
|
||||
col_byte_stride: item_size as isize,
|
||||
item_size,
|
||||
};
|
||||
|
||||
let non_linear_ops = tvec!(
|
||||
FusedKerSpec::Clear,
|
||||
FusedKerSpec::AddMatMul {
|
||||
k: k_aligned,
|
||||
pa: pa.panel_bytes(0, None)?,
|
||||
pb: pb.panel_bytes(0, None)?,
|
||||
packing: self.packing
|
||||
},
|
||||
FusedKerSpec::Store(c),
|
||||
FusedKerSpec::Done
|
||||
);
|
||||
let err = self.ker.kernel(&non_linear_ops);
|
||||
assert_eq!(err, 0);
|
||||
}
|
||||
Ok(v)
|
||||
}
|
||||
|
||||
pub fn check(&self) -> TractResult<()> {
|
||||
if !self.ker.is_supported_here() {
|
||||
return Ok(());
|
||||
}
|
||||
let expected = self.reference()?;
|
||||
let found = self.run()?;
|
||||
let app = if K::Acc::datum_type() == f16::datum_type() {
|
||||
Approximation::SuperApproximate
|
||||
} else {
|
||||
Approximation::Approximate
|
||||
};
|
||||
let result = found.close_enough(&expected, app);
|
||||
if result.is_err() {
|
||||
let exp = expected.try_as_plain()?.as_slice::<K::Acc>()?;
|
||||
let found = found.try_as_plain()?.as_slice::<K::Acc>()?;
|
||||
let (m, _, n) = self.mkn();
|
||||
display_error(found, exp, m, n);
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
// Large-shape frame tests that exercise the single-thread 2D-blocked tile walk
|
||||
// (`run_single_thread_blocked`): the existing `arbitrary_problem` frame proptests
|
||||
// only reach 3 panels per dim (m,n < 3·mr), below the ST_BLK=16 blocking
|
||||
// threshold, so the blocked path was otherwise uncovered. generic_f32_4x4 has
|
||||
// mr=nr=4, so m,n=80 → 20×20 panels → multiple blocks. Compares the frame
|
||||
// output against the naive reference (must be bit/approx-exact).
|
||||
#[cfg(test)]
|
||||
mod single_thread_blocking {
|
||||
use super::PackedPackedProblem;
|
||||
use crate::generic::mmm::generic_f32_4x4;
|
||||
use tract_data::internal::TractResult;
|
||||
|
||||
fn check_large(m: usize, n: usize, k: usize) -> TractResult<()> {
|
||||
let a: Vec<f32> = (0..m * k)
|
||||
.map(|i| ((i * 7 + 3) % 13) as f32 - 6.0)
|
||||
.collect();
|
||||
let b: Vec<f32> = (0..k * n)
|
||||
.map(|i| ((i * 5 + 1) % 11) as f32 - 5.0)
|
||||
.collect();
|
||||
PackedPackedProblem::frame(&*generic_f32_4x4, 0, m, n, a, b).check()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocked_80x80() -> TractResult<()> {
|
||||
check_large(80, 80, 24) // 20×20 panels, multiple ST_BLK blocks
|
||||
}
|
||||
#[test]
|
||||
fn blocked_skew_200x40() -> TractResult<()> {
|
||||
check_large(200, 40, 8) // 50×10 panels (m-axis chunked)
|
||||
}
|
||||
#[test]
|
||||
fn blocked_40x200() -> TractResult<()> {
|
||||
check_large(40, 200, 8) // 10×50 panels (n-axis chunked)
|
||||
}
|
||||
#[test]
|
||||
fn blocked_64x64_exact() -> TractResult<()> {
|
||||
check_large(64, 64, 16) // exactly 16×16 panels (block boundary)
|
||||
}
|
||||
#[test]
|
||||
fn blocked_68x68_offset() -> TractResult<()> {
|
||||
check_large(68, 68, 10) // 17×17 panels (one full block + a 1-panel remainder)
|
||||
}
|
||||
}
|
||||
Vendored
+178
@@ -0,0 +1,178 @@
|
||||
use crate::Scaler;
|
||||
use crate::frame::mmm::MatMatMulKer;
|
||||
use crate::frame::mmm::fuse::RoundingPolicy;
|
||||
use crate::generic::rounding::ScaleShiftAndRound;
|
||||
use crate::mmm::{FusedKerSpec, FusedSpec};
|
||||
use proptest::prelude::*;
|
||||
|
||||
use super::fuse::fused_ops;
|
||||
|
||||
#[derive(Debug, new)]
|
||||
pub struct QScaleProblem<K>
|
||||
where
|
||||
K: MatMatMulKer<Acc = i32>,
|
||||
{
|
||||
pub ker: K,
|
||||
pub c: Vec<i32>,
|
||||
pub scaler: Scaler,
|
||||
pub boo: std::marker::PhantomData<K>,
|
||||
}
|
||||
|
||||
pub fn arbitrary_qscale_problem<K: MatMatMulKer<Acc = i32>>(
|
||||
ker: &K,
|
||||
) -> BoxedStrategy<QScaleProblem<K>> {
|
||||
use RoundingPolicy::*;
|
||||
let ker = ker.clone();
|
||||
let len = ker.mr() * ker.nr();
|
||||
(
|
||||
proptest::collection::vec(-20i32..20, 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(move |(c, scale_pot, scale_mult, policy)| QScaleProblem {
|
||||
ker: ker.clone(),
|
||||
c,
|
||||
scaler: Scaler::new(scale_mult * 2f32.powi(scale_pot), policy),
|
||||
boo: std::marker::PhantomData,
|
||||
})
|
||||
.boxed()
|
||||
}
|
||||
|
||||
impl<K> QScaleProblem<K>
|
||||
where
|
||||
K: MatMatMulKer<Acc = i32>,
|
||||
{
|
||||
pub fn run(&self) {
|
||||
if !self.ker.is_supported_here() {
|
||||
return;
|
||||
}
|
||||
if let FusedSpec::QScale(shift, policy, mult) = self.scaler.as_fused_spec() {
|
||||
fused_ops::<K, i32, _>(
|
||||
&self.ker,
|
||||
&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, i32, _>(
|
||||
&self.ker,
|
||||
&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, i32, _>(
|
||||
&self.ker,
|
||||
&self.c,
|
||||
&[FusedKerSpec::ShiftLeft(shift)],
|
||||
|_, _, c| c.q_shl(shift),
|
||||
)
|
||||
} else {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn return_c_scale_bigpot<K>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = i32>,
|
||||
{
|
||||
let ker = ker.clone();
|
||||
let len = ker.mr() * ker.nr();
|
||||
let v: Vec<i32> = (-(len as i32) / 2..).take(len).collect();
|
||||
fused_ops::<K, i32, _>(&ker, &v, &[FusedKerSpec::ShiftLeft(1)], |_, _, c| {
|
||||
c.q_shl(1)
|
||||
})
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! mmm_q_scale_tests {
|
||||
($ker:expr) => {
|
||||
use $crate::frame::mmm::fuse::RoundingPolicy;
|
||||
use $crate::frame::mmm::tests::q_scale::arbitrary_qscale_problem;
|
||||
use $crate::frame::mmm::tests::q_scale::QScaleProblem;
|
||||
use $crate::frame::mmm::MatMatMulKer;
|
||||
use $crate::generic::Scaler;
|
||||
// FIXME: Scaler should be arbitrary
|
||||
macro_rules! test_q_scale {
|
||||
($policy: ident) => {
|
||||
paste! {
|
||||
#[test]
|
||||
fn [<return_q_scale_halfpos_ $policy:lower>]() {
|
||||
let ker = $ker;
|
||||
let len = (ker.mr() * ker.nr()) as i64;
|
||||
let v = (0..len).map(|i| (i - len / 2) as i32).collect();
|
||||
QScaleProblem::new(ker.clone(), v, Scaler::new(0.5f32, RoundingPolicy::$policy)).run()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn [<return_q_scale_halfneg_ $policy:lower>]() {
|
||||
let ker = $ker;
|
||||
let len = (ker.mr() * ker.nr()) as i64;
|
||||
let v = (0..len).map(|i| (i - len / 2) as i32).collect();
|
||||
QScaleProblem::new(ker.clone(), v, Scaler::new(-0.5f32, RoundingPolicy::$policy)).run()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn [<return_q_scale_pot_ $policy:lower>]() {
|
||||
let ker = $ker;
|
||||
let len = (ker.mr() * ker.nr()) as i64;
|
||||
let v = (0..len).map(|i| (i - len / 2) as i32).collect();
|
||||
QScaleProblem::new(ker.clone(), v, Scaler::new(0.25f32, RoundingPolicy::$policy)).run()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn [<return_q_scale_nonpot_ $policy:lower>]() {
|
||||
let ker = $ker;
|
||||
let len = (ker.mr() * ker.nr()) as i64;
|
||||
let v = (0..len).map(|i| (i - len / 2) as i32).collect();
|
||||
QScaleProblem::new(ker.clone(), v, Scaler::new(1f32 / 5., RoundingPolicy::$policy)).run()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn [<return_q_scale_bigpot_ $policy:lower>]() {
|
||||
let ker = $ker;
|
||||
let len = (ker.mr() * ker.nr()) as i64;
|
||||
let v = (0..len).map(|i| (i - len / 2) as i32).collect();
|
||||
QScaleProblem::new(ker.clone(), v, Scaler::new(4f32, RoundingPolicy::$policy)).run()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn [<return_q_scale_bignonpot_ $policy:lower>]() {
|
||||
let ker = $ker;
|
||||
let len = (ker.mr() * ker.nr()) as i64;
|
||||
let v = (0..len).map(|i| (i - len / 2) as i32).collect();
|
||||
QScaleProblem::new(ker.clone(), 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 arbitrary_qscale_problem($ker)) {
|
||||
pb.run()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn return_c_scale_bigpot() {
|
||||
$crate::frame::mmm::tests::q_scale::return_c_scale_bigpot::<_>($ker)
|
||||
}
|
||||
};
|
||||
}
|
||||
Vendored
+143
@@ -0,0 +1,143 @@
|
||||
use crate::LADatum;
|
||||
use crate::frame::mmm::fuse::FusedKerSpec;
|
||||
use crate::frame::mmm::storage::*;
|
||||
use crate::frame::mmm::tests::display_error;
|
||||
use crate::frame::mmm::*;
|
||||
use num_traits::Bounded;
|
||||
use tract_data::internal::*;
|
||||
use tract_itertools::Itertools;
|
||||
use tract_ndarray::Axis;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! mmm_store_test {
|
||||
($ker:expr, $tc:ident) => {
|
||||
paste! {
|
||||
mod [<store_$tc>] {
|
||||
#[allow(unused_imports)]
|
||||
use tract_data::prelude::f16;
|
||||
use $crate::frame::mmm::tests::store::StoreLayout;
|
||||
|
||||
#[test] fn store_zeros() {
|
||||
$crate::frame::mmm::tests::store::store_zeros::<_,$tc,_>($ker);
|
||||
}
|
||||
|
||||
#[test] fn store_col_major() {
|
||||
$crate::frame::mmm::tests::store::store_pattern::<_,$tc,_>($ker, StoreLayout::ColMajor);
|
||||
}
|
||||
|
||||
#[test] fn store_row_major() {
|
||||
$crate::frame::mmm::tests::store::store_pattern::<_,$tc,_>($ker, StoreLayout::RowMajor);
|
||||
}
|
||||
|
||||
#[test] fn store_arbitrary() {
|
||||
$crate::frame::mmm::tests::store::store_pattern::<_,$tc,_>($ker, StoreLayout::Arbitrary);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
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>(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn store_zeros<K, TC, TI>(ker: &K)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TC: LADatum,
|
||||
TI: LADatum + Bounded + PartialEq,
|
||||
{
|
||||
if !ker.is_supported_here() {
|
||||
return;
|
||||
}
|
||||
let v = vec![TC::max_value(); ker.mr() * ker.nr()];
|
||||
let c = mmm_stride_storage(&v, ker.nr());
|
||||
let non_linear = tvec![
|
||||
FusedKerSpec::Clear,
|
||||
FusedKerSpec::Store(c),
|
||||
FusedKerSpec::Done
|
||||
];
|
||||
let err = ker.kernel(&non_linear);
|
||||
assert_eq!(err, 0);
|
||||
let expected = vec![TC::zero(); v.len()];
|
||||
display_error(&v, &expected, ker.mr(), ker.nr());
|
||||
assert_eq!(v, expected);
|
||||
}
|
||||
|
||||
pub enum StoreLayout {
|
||||
ColMajor,
|
||||
RowMajor,
|
||||
Arbitrary,
|
||||
}
|
||||
|
||||
pub fn store_pattern<K, TC, TI>(ker: &K, layout: StoreLayout)
|
||||
where
|
||||
K: MatMatMulKer<Acc = TI>,
|
||||
TC: LADatum,
|
||||
TI: LADatum + Bounded + PartialEq,
|
||||
{
|
||||
if !ker.is_supported_here() {
|
||||
return;
|
||||
}
|
||||
let (mr, nr) = (ker.mr(), ker.nr());
|
||||
let pattern = tensor1(&(0..).take(mr * nr).collect_vec())
|
||||
.cast_to::<TI>()
|
||||
.unwrap()
|
||||
.into_owned()
|
||||
.into_shape(&[mr, nr])
|
||||
.unwrap();
|
||||
let pattern_aligned = Blob::from_bytes_alignment(pattern.as_bytes(), 128).unwrap();
|
||||
let pattern_col_major = pattern.clone().permute_axes(&[1, 0]).unwrap();
|
||||
let pattern_col_major_aligned =
|
||||
Blob::from_bytes_alignment(pattern_col_major.as_bytes(), 128).unwrap();
|
||||
let size_of_tc = std::mem::size_of::<TC>();
|
||||
let (row_stride, col_stride, result_size) = match layout {
|
||||
StoreLayout::RowMajor => (nr, 1, mr * nr),
|
||||
StoreLayout::ColMajor => (1, mr, mr * nr),
|
||||
// like row major, but storing every other third column
|
||||
StoreLayout::Arbitrary => (nr * 3, 3, mr * nr * 3),
|
||||
};
|
||||
let mut result = tensor0(TC::max_value())
|
||||
.broadcast_to_shape(&[result_size])
|
||||
.unwrap();
|
||||
let non_linear = tvec![
|
||||
FusedKerSpec::LoadTile(
|
||||
pattern_col_major_aligned.as_ptr() as *const TI,
|
||||
pattern_aligned.as_ptr() as *const TI,
|
||||
),
|
||||
FusedKerSpec::Store(OutputStoreKer {
|
||||
ptr: result.as_bytes_mut().as_mut_ptr(),
|
||||
row_byte_stride: (size_of_tc * row_stride) as isize,
|
||||
col_byte_stride: (size_of_tc * col_stride) as isize,
|
||||
item_size: size_of_tc,
|
||||
}),
|
||||
FusedKerSpec::Done
|
||||
];
|
||||
let err = ker.kernel(&non_linear);
|
||||
assert_eq!(err, 0);
|
||||
let expected = pattern.cast_to::<TC>().unwrap().into_owned();
|
||||
let result = match layout {
|
||||
StoreLayout::RowMajor => result,
|
||||
StoreLayout::ColMajor => result
|
||||
.into_shape(&[ker.nr(), ker.mr()])
|
||||
.unwrap()
|
||||
.permute_axes(&[1, 0])
|
||||
.unwrap(),
|
||||
StoreLayout::Arbitrary => result
|
||||
.into_plain_array::<TC>()
|
||||
.unwrap()
|
||||
.into_shape_with_order((mr, nr, 3))
|
||||
.unwrap()
|
||||
.index_axis_move(Axis(2), 0)
|
||||
.into_tensor(),
|
||||
};
|
||||
let expected = expected.try_as_plain().unwrap().as_slice::<TC>().unwrap();
|
||||
let result = result.try_as_plain().unwrap().as_slice::<TC>().unwrap();
|
||||
display_error(result, expected, ker.mr(), ker.nr());
|
||||
assert_eq!(result, expected);
|
||||
}
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
#[macro_use]
|
||||
pub mod block_quant;
|
||||
#[macro_use]
|
||||
pub mod element_wise;
|
||||
pub mod element_wise_helper;
|
||||
#[macro_use]
|
||||
pub mod unicast;
|
||||
#[macro_use]
|
||||
pub mod by_scalar;
|
||||
#[macro_use]
|
||||
pub mod erf;
|
||||
#[macro_use]
|
||||
pub mod gelu;
|
||||
#[macro_use]
|
||||
pub mod hardswish;
|
||||
#[macro_use]
|
||||
pub mod leaky_relu;
|
||||
#[macro_use]
|
||||
pub mod lut;
|
||||
#[macro_use]
|
||||
pub mod mmm;
|
||||
#[macro_use]
|
||||
pub mod pack;
|
||||
#[macro_use]
|
||||
pub mod reduce;
|
||||
#[macro_use]
|
||||
pub mod sigmoid;
|
||||
#[macro_use]
|
||||
pub mod silu;
|
||||
#[macro_use]
|
||||
pub mod tanh;
|
||||
#[macro_use]
|
||||
pub mod weights;
|
||||
+1561
File diff suppressed because it is too large
Load Diff
+42
@@ -0,0 +1,42 @@
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::reduce::ReduceKer;
|
||||
use num_traits::{AsPrimitive, Float};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! max_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-25f32..25.0, 0..100)) {
|
||||
if $cond {
|
||||
$crate::frame::reduce::max::test::test_max::<$ker, $t>(&*xs).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::max::test::test_max::<$ker, $t>(&[]).unwrap()
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_max<K: ReduceKer<T>, T: LADatum + Float>(values: &[f32]) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
let values: Vec<T> = values.iter().copied().map(|x| x.as_()).collect();
|
||||
crate::frame::reduce::test::test_reduce::<K, _>(
|
||||
&values,
|
||||
<T as Float>::min_value(),
|
||||
|a, b| a.max(b),
|
||||
)
|
||||
}
|
||||
}
|
||||
+303
@@ -0,0 +1,303 @@
|
||||
pub mod max;
|
||||
pub mod softmax;
|
||||
pub mod sum;
|
||||
|
||||
use std::fmt::Debug;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use tract_data::TractResult;
|
||||
|
||||
use crate::LADatum;
|
||||
|
||||
use super::element_wise_helper::{map_reduce_slice_with_alignment, reduce_slice_with_alignment};
|
||||
|
||||
macro_rules! reduce_impl_wrap {
|
||||
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $neutral: expr, $run: item, $reduce_two: item) => {
|
||||
paste! {
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
#[allow(non_camel_case_types)]
|
||||
pub struct $func;
|
||||
|
||||
impl crate::frame::reduce::ReduceKer<$ti, $params> 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(always)]
|
||||
fn neutral() -> $ti {
|
||||
$neutral
|
||||
}
|
||||
$run
|
||||
$reduce_two
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub trait Reduce<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
T: Copy + Debug + PartialEq + Send + Sync,
|
||||
{
|
||||
fn name(&self) -> &'static str;
|
||||
fn run(&self, vec: &[T]) -> TractResult<T> {
|
||||
self.run_with_params(vec, Params::default())
|
||||
}
|
||||
fn run_with_params(&self, vec: &[T], params: Params) -> TractResult<T>;
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(<T, Params> Reduce<T, Params> where T: Copy, Params: Copy);
|
||||
|
||||
#[derive(Debug, Clone, new)]
|
||||
pub struct ReduceImpl<K, T, Params = ()>
|
||||
where
|
||||
T: LADatum,
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
K: ReduceKer<T, Params> + Clone,
|
||||
{
|
||||
phantom: PhantomData<(K, T, Params)>,
|
||||
}
|
||||
|
||||
impl<K, T, Params> Reduce<T, Params> for ReduceImpl<K, T, Params>
|
||||
where
|
||||
T: LADatum,
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
K: ReduceKer<T, Params> + Clone,
|
||||
{
|
||||
fn name(&self) -> &'static str {
|
||||
K::name()
|
||||
}
|
||||
|
||||
fn run_with_params(&self, vec: &[T], params: Params) -> TractResult<T> {
|
||||
reduce_slice_with_alignment(
|
||||
vec,
|
||||
|data| K::run(data, params),
|
||||
K::nr(),
|
||||
K::alignment_bytes(),
|
||||
K::neutral(),
|
||||
K::reduce_two,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ReduceKer<T, Params = ()>:
|
||||
Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
T: LADatum,
|
||||
{
|
||||
fn name() -> &'static str;
|
||||
fn alignment_bytes() -> usize {
|
||||
Self::alignment_items() * T::datum_type().size_of()
|
||||
}
|
||||
fn alignment_items() -> usize;
|
||||
fn nr() -> usize;
|
||||
fn neutral() -> T;
|
||||
fn reduce_two(a: T, b: T) -> T;
|
||||
fn run(vec: &[T], params: Params) -> T;
|
||||
fn red() -> Box<dyn Reduce<T, Params>> {
|
||||
Box::new(ReduceImpl::<Self, T, Params>::new())
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused_macros)]
|
||||
macro_rules! map_reduce_impl_wrap {
|
||||
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $map_neutral: expr, $reduce_neutral: expr, $run: item, $reduce_two: item) => {
|
||||
paste! {
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
#[allow(non_camel_case_types)]
|
||||
pub struct $func;
|
||||
|
||||
impl crate::frame::reduce::MapReduceKer<$ti, $params> 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(always)]
|
||||
fn map_neutral() -> $ti {
|
||||
$map_neutral
|
||||
}
|
||||
#[inline(always)]
|
||||
fn reduce_neutral() -> $ti {
|
||||
$reduce_neutral
|
||||
}
|
||||
$run
|
||||
$reduce_two
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub trait MapReduce<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
T: Copy + Debug + PartialEq + Send + Sync,
|
||||
{
|
||||
fn name(&self) -> &'static str;
|
||||
fn run(&self, vec: &mut [T]) -> TractResult<T> {
|
||||
self.run_with_params(vec, Params::default())
|
||||
}
|
||||
fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<T>;
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(<T, Params> MapReduce<T, Params> where T: Copy, Params: Copy);
|
||||
|
||||
#[derive(Debug, Clone, new)]
|
||||
pub struct MapReduceImpl<K, T, Params = ()>
|
||||
where
|
||||
T: LADatum,
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
K: MapReduceKer<T, Params> + Clone,
|
||||
{
|
||||
phantom: PhantomData<(K, T, Params)>,
|
||||
}
|
||||
|
||||
impl<K, T, Params> MapReduce<T, Params> for MapReduceImpl<K, T, Params>
|
||||
where
|
||||
T: LADatum,
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
K: MapReduceKer<T, Params> + Clone,
|
||||
{
|
||||
fn name(&self) -> &'static str {
|
||||
K::name()
|
||||
}
|
||||
fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<T> {
|
||||
map_reduce_slice_with_alignment(
|
||||
vec,
|
||||
|data| K::run(data, params),
|
||||
K::nr(),
|
||||
K::alignment_bytes(),
|
||||
K::map_neutral(),
|
||||
K::reduce_neutral(),
|
||||
K::reduce_two,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub trait MapReduceKer<T, Params = ()>:
|
||||
Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
T: LADatum,
|
||||
{
|
||||
fn name() -> &'static str;
|
||||
fn alignment_bytes() -> usize {
|
||||
Self::alignment_items() * T::datum_type().size_of()
|
||||
}
|
||||
fn alignment_items() -> usize;
|
||||
fn nr() -> usize;
|
||||
fn map_neutral() -> T;
|
||||
fn reduce_neutral() -> T;
|
||||
fn reduce_two(a: T, b: T) -> T;
|
||||
fn run(vec: &mut [T], params: Params) -> T;
|
||||
fn red() -> Box<dyn MapReduce<T, Params>> {
|
||||
Box::new(MapReduceImpl::<Self, T, Params>::new())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test {
|
||||
use super::*;
|
||||
use proptest::test_runner::{TestCaseError, TestCaseResult};
|
||||
use tract_data::internal::*;
|
||||
use tract_data::itertools::Itertools;
|
||||
|
||||
pub fn test_reduce<K: ReduceKer<T, ()>, T: LADatum>(
|
||||
values: &[T],
|
||||
neutral: T,
|
||||
reference_reduce: impl Fn(T, T) -> T,
|
||||
) -> TestCaseResult {
|
||||
test_reduce_params::<K, T, ()>(values, neutral, reference_reduce, ())
|
||||
}
|
||||
|
||||
pub fn test_reduce_params<K: ReduceKer<T, Params>, T: LADatum, Params>(
|
||||
values: &[T],
|
||||
neutral: T,
|
||||
reference_reducer: impl Fn(T, T) -> T,
|
||||
params: Params,
|
||||
) -> TestCaseResult
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
let op = K::red();
|
||||
let expected = values
|
||||
.iter()
|
||||
.fold(neutral, |acc, i| reference_reducer(acc, *i));
|
||||
let found = values;
|
||||
let red = op.run_with_params(found, params).unwrap();
|
||||
tensor0(red)
|
||||
.close_enough(&tensor0(expected), true)
|
||||
.map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn test_map_reduce<K: MapReduceKer<T, ()>, T: LADatum>(
|
||||
values: &[T],
|
||||
map_neutral: T,
|
||||
neutral: T,
|
||||
reference_map: impl Fn(T) -> T,
|
||||
reference_reduce: impl Fn(T, T) -> T,
|
||||
) -> TestCaseResult {
|
||||
test_map_reduce_params::<K, T, ()>(
|
||||
values,
|
||||
map_neutral,
|
||||
neutral,
|
||||
reference_map,
|
||||
reference_reduce,
|
||||
(),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn test_map_reduce_params<K: MapReduceKer<T, Params>, T: LADatum, Params>(
|
||||
values: &[T],
|
||||
_neutral: T,
|
||||
map_neutral: T,
|
||||
reference_map: impl Fn(T) -> T,
|
||||
reference_reducer: impl Fn(T, T) -> T,
|
||||
params: Params,
|
||||
) -> TestCaseResult
|
||||
where
|
||||
Params: Copy + Send + Sync + Debug + 'static + Default,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
let op = K::red();
|
||||
let mut found = values.to_vec();
|
||||
let expected_values = values.iter().copied().map(reference_map).collect_vec();
|
||||
let expected_reduced = expected_values
|
||||
.iter()
|
||||
.fold(map_neutral, |acc, i| reference_reducer(acc, *i));
|
||||
let red = op.run_with_params(&mut found, params).unwrap();
|
||||
tensor1(&found)
|
||||
.close_enough(&tensor1(&expected_values), Approximation::SuperApproximate)
|
||||
.map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
|
||||
tensor0(red)
|
||||
.close_enough(&tensor0(expected_reduced), Approximation::SuperApproximate)
|
||||
.map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::reduce::MapReduceKer;
|
||||
use num_traits::{AsPrimitive, Float};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! softmax_l2_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-25f32..25.0, 1..100)) {
|
||||
if $cond {
|
||||
$crate::frame::reduce::softmax::test::test_softmax_l2::<$ker, $t>(&*xs).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::softmax::test::test_softmax_l2::<$ker, $t>(&[0.0]).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_zeros() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::softmax::test::test_softmax_l2::<$ker, $t>(&[0.0, 0.0]).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_0() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::softmax::test::test_softmax_l2::<$ker, $t>(&[
|
||||
16.62555, 21.950674,
|
||||
])
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_1() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::softmax::test::test_softmax_l2::<$ker, $t>(&[0.0f32, 0.38132212])
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_missing_max() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::softmax::test::test_softmax_l2::<$ker, $t>(&[
|
||||
-46.15512, 42.875168,
|
||||
])
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_softmax_l2<K: MapReduceKer<T, T>, T>(values: &[f32]) -> TestCaseResult
|
||||
where
|
||||
T: LADatum + Float + AsPrimitive<f32>,
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
use crate::generic::reduce::softmax_l2::fast_compact_exp_f32;
|
||||
crate::setup_test_logger();
|
||||
let max = values.iter().max_by(|a, b| a.total_cmp(b)).unwrap();
|
||||
let values: Vec<T> = values.iter().copied().map(|x| x.as_()).collect();
|
||||
crate::frame::reduce::test::test_map_reduce_params::<K, T, T>(
|
||||
&values,
|
||||
<T as Float>::min_value(),
|
||||
T::zero(),
|
||||
// |x| (x - max.as_()).exp(),
|
||||
|x| fast_compact_exp_f32(x.as_() - max).as_(),
|
||||
|a, b| a + b,
|
||||
max.as_(),
|
||||
)
|
||||
}
|
||||
}
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::reduce::ReduceKer;
|
||||
use num_traits::{AsPrimitive, Float, Zero};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! sum_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-25_isize..25, 0..100)) {
|
||||
if $cond {
|
||||
let xs_float = xs.into_iter().map(|it| it as f32).collect::<Vec<_>>();
|
||||
$crate::frame::reduce::sum::test::test_sum::<$ker, $t>(&*xs_float).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::sum::test::test_sum::<$ker, $t>(&[]).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn simple() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::sum::test::test_sum::<$ker, $t>(&[1.0, 2.0]).unwrap()
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn multiple_tile() {
|
||||
if $cond {
|
||||
$crate::frame::reduce::sum::test::test_sum::<$ker, $t>(&[1.0; 35]).unwrap()
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_sum<K, T>(values: &[f32]) -> TestCaseResult
|
||||
where
|
||||
K: ReduceKer<T>,
|
||||
f32: AsPrimitive<T>,
|
||||
T: LADatum + Float + Zero + AsPrimitive<f32>,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
let values: Vec<T> = values.iter().copied().map(|x| x.as_()).collect();
|
||||
crate::frame::reduce::test::test_reduce::<K, _>(&values, <T as Zero>::zero(), |a, b| a + b)
|
||||
}
|
||||
}
|
||||
+96
@@ -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>,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
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())
|
||||
})
|
||||
}
|
||||
}
|
||||
+58
@@ -0,0 +1,58 @@
|
||||
#[allow(unused_macros)]
|
||||
macro_rules! silu_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::*;
|
||||
silu_frame_tests!($cond, $ti, $func);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use crate::LADatum;
|
||||
use crate::frame::element_wise::*;
|
||||
use num_traits::{AsPrimitive, Float};
|
||||
use proptest::test_runner::TestCaseResult;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! silu_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty) => {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn prop(xs in proptest::collection::vec(-10f32..10.0, 0..100)) {
|
||||
if $cond {
|
||||
$crate::frame::silu::test::test_silu::<$ker, $t>(&*xs).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn trivial() {
|
||||
if $cond {
|
||||
$crate::frame::silu::test::test_silu::<$ker, $t>(&[-5f32, -1.0, 0.0, 1.0, 5.0])
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub fn test_silu<K: ElementWiseKer<T>, T: LADatum + Float>(values: &[f32]) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
let data = tract_data::prelude::tensor1(values);
|
||||
let data = data.cast_to::<T>().unwrap();
|
||||
let data = data.try_as_plain().unwrap().as_slice::<T>().unwrap();
|
||||
crate::frame::element_wise::test::test_element_wise::<K, T, _>(data, |x: T| {
|
||||
let one: T = 1f32.as_();
|
||||
let neg_x = T::zero() - x;
|
||||
let sigmoid = one / (one + neg_x.exp());
|
||||
x * sigmoid
|
||||
})
|
||||
}
|
||||
}
|
||||
+101
@@ -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>,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
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())
|
||||
}
|
||||
}
|
||||
+237
@@ -0,0 +1,237 @@
|
||||
use std::fmt::Debug;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
use tract_data::TractResult;
|
||||
use tract_data::internal::TensorView;
|
||||
|
||||
use crate::frame::element_wise_helper::TempBuffer;
|
||||
use crate::{LADatum, LinalgFn};
|
||||
|
||||
macro_rules! unicast_impl_wrap {
|
||||
($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $run: item) => {
|
||||
paste! {
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
#[allow(non_camel_case_types)]
|
||||
pub struct $func;
|
||||
|
||||
impl crate::frame::unicast::UnicastKer<$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
|
||||
}
|
||||
$run
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub trait Unicast<T>: Send + Sync + Debug + dyn_clone::DynClone
|
||||
where
|
||||
T: Copy + Debug + PartialEq + Send + Sync,
|
||||
{
|
||||
fn name(&self) -> &'static str;
|
||||
fn run(&self, a: &mut [T], b: &[T]) -> TractResult<()>;
|
||||
}
|
||||
|
||||
dyn_clone::clone_trait_object!(<T> Unicast<T> where T: Copy);
|
||||
|
||||
#[derive(Debug, Clone, new)]
|
||||
pub struct UnicastImpl<K, T>
|
||||
where
|
||||
T: LADatum,
|
||||
K: UnicastKer<T> + Clone,
|
||||
{
|
||||
phantom: PhantomData<(K, T)>,
|
||||
}
|
||||
|
||||
impl<K, T> UnicastImpl<K, T>
|
||||
where
|
||||
T: LADatum,
|
||||
K: UnicastKer<T> + Clone,
|
||||
{
|
||||
}
|
||||
impl<K, T> Unicast<T> for UnicastImpl<K, T>
|
||||
where
|
||||
T: LADatum,
|
||||
K: UnicastKer<T> + Clone,
|
||||
{
|
||||
fn name(&self) -> &'static str {
|
||||
K::name()
|
||||
}
|
||||
fn run(&self, a: &mut [T], b: &[T]) -> TractResult<()> {
|
||||
unicast_with_alignment(a, b, |a, b| K::run(a, b), K::nr(), K::alignment_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait UnicastKer<T>: Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
|
||||
where
|
||||
T: LADatum,
|
||||
{
|
||||
fn name() -> &'static str;
|
||||
fn alignment_bytes() -> usize {
|
||||
Self::alignment_items() * T::datum_type().size_of()
|
||||
}
|
||||
fn alignment_items() -> usize;
|
||||
fn nr() -> usize;
|
||||
fn run(a: &mut [T], b: &[T]);
|
||||
fn bin() -> Box<LinalgFn> {
|
||||
Box::new(|a: &mut TensorView, b: &TensorView| {
|
||||
let a_slice = a.as_slice_mut()?;
|
||||
let b_slice = b.as_slice()?;
|
||||
UnicastImpl::<Self, T>::new().run(a_slice, b_slice)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
std::thread_local! {
|
||||
static TMP: std::cell::RefCell<(TempBuffer, TempBuffer)> = std::cell::RefCell::new((TempBuffer::default(), TempBuffer::default()));
|
||||
}
|
||||
|
||||
pub(crate) fn unicast_with_alignment<T>(
|
||||
a: &mut [T],
|
||||
b: &[T],
|
||||
f: impl Fn(&mut [T], &[T]),
|
||||
nr: usize,
|
||||
alignment_bytes: usize,
|
||||
) -> TractResult<()>
|
||||
where
|
||||
T: LADatum,
|
||||
{
|
||||
if a.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
unsafe {
|
||||
TMP.with(|buffers| {
|
||||
let mut buffers = buffers.borrow_mut();
|
||||
buffers.0.ensure(nr * T::datum_type().size_of(), alignment_bytes);
|
||||
buffers.1.ensure(nr * T::datum_type().size_of(), alignment_bytes);
|
||||
let tmp_a = std::slice::from_raw_parts_mut(buffers.0.buffer as *mut T, nr);
|
||||
let tmp_b = std::slice::from_raw_parts_mut(buffers.1.buffer as *mut T, nr);
|
||||
let mut compute_via_temp_buffer = |a: &mut [T], b: &[T]| {
|
||||
tmp_a[..a.len()].copy_from_slice(a);
|
||||
tmp_b[..b.len()].copy_from_slice(b);
|
||||
f(tmp_a, tmp_b);
|
||||
a.copy_from_slice(&tmp_a[..a.len()])
|
||||
};
|
||||
|
||||
let mut num_element_processed = 0;
|
||||
let a_prefix_len = a.as_ptr().align_offset(alignment_bytes).min(a.len());
|
||||
let b_prefix_len = b.as_ptr().align_offset(alignment_bytes).min(b.len());
|
||||
assert!(
|
||||
a_prefix_len == b_prefix_len,
|
||||
"Both inputs should be of the same alignement, got {a_prefix_len:?}, {b_prefix_len:?}"
|
||||
);
|
||||
let mut applied_prefix_len = 0;
|
||||
if a_prefix_len > 0 {
|
||||
// Incomplete tile needs to be created to process unaligned data.
|
||||
let sub_a = &mut a[..a_prefix_len];
|
||||
let sub_b = &b[..a_prefix_len];
|
||||
compute_via_temp_buffer(sub_a, sub_b);
|
||||
num_element_processed += a_prefix_len;
|
||||
applied_prefix_len = a_prefix_len;
|
||||
}
|
||||
|
||||
let num_complete_tiles = (a.len() - applied_prefix_len) / nr;
|
||||
if num_complete_tiles > 0 {
|
||||
// Process all tiles that are complete.
|
||||
let sub_a = &mut a[applied_prefix_len..][..(num_complete_tiles * nr)];
|
||||
let sub_b = &b[applied_prefix_len..][..(num_complete_tiles * nr)];
|
||||
f(sub_a, sub_b);
|
||||
num_element_processed += num_complete_tiles * nr;
|
||||
}
|
||||
|
||||
if num_element_processed < a.len() {
|
||||
// Incomplete tile needs to be created to process remaining elements.
|
||||
compute_via_temp_buffer(
|
||||
&mut a[num_element_processed..],
|
||||
&b[num_element_processed..],
|
||||
);
|
||||
}
|
||||
})
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod test {
|
||||
use super::*;
|
||||
use crate::LADatum;
|
||||
use proptest::test_runner::{TestCaseError, TestCaseResult};
|
||||
use tract_data::internal::*;
|
||||
use tract_num_traits::{AsPrimitive, Float};
|
||||
|
||||
pub fn test_unicast<K: UnicastKer<T>, T: LADatum>(
|
||||
a: &mut [T],
|
||||
b: &[T],
|
||||
reference: impl Fn(T, T) -> T,
|
||||
) -> TestCaseResult {
|
||||
crate::setup_test_logger();
|
||||
let op = UnicastImpl::<K, T>::new();
|
||||
let expected = a
|
||||
.iter()
|
||||
.zip(b.iter())
|
||||
.map(|(a, b)| (reference)(*a, *b))
|
||||
.collect::<Vec<_>>();
|
||||
op.run(a, b).unwrap();
|
||||
tensor1(a)
|
||||
.close_enough(&tensor1(&expected), true)
|
||||
.map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn test_unicast_t<K: UnicastKer<T>, T: LADatum + Float>(
|
||||
a: &[f32],
|
||||
b: &[f32],
|
||||
func: impl Fn(T, T) -> T,
|
||||
) -> TestCaseResult
|
||||
where
|
||||
f32: AsPrimitive<T>,
|
||||
{
|
||||
crate::setup_test_logger();
|
||||
let vec_a: Vec<T> = a.iter().copied().map(|x| x.as_()).collect();
|
||||
// We allocate a tensor to ensure allocation is done with alignement
|
||||
let mut a = unsafe { Tensor::from_slice_align(vec_a.as_slice(), vector_size()).unwrap() };
|
||||
let vec_b: Vec<T> = b.iter().copied().map(|x| x.as_()).collect();
|
||||
// We allocate a tensor to ensure allocation is done with alignement
|
||||
let b = unsafe { Tensor::from_slice_align(vec_b.as_slice(), vector_size()).unwrap() };
|
||||
crate::frame::unicast::test::test_unicast::<K, _>(
|
||||
a.try_as_plain_mut().unwrap().as_slice_mut::<T>().unwrap(),
|
||||
b.try_as_plain().unwrap().as_slice::<T>().unwrap(),
|
||||
func,
|
||||
)
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! unicast_frame_tests {
|
||||
($cond:expr, $t: ty, $ker:ty, $func:expr) => {
|
||||
pastey::paste! {
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn [<prop_ $ker:snake>](
|
||||
(a, b) in (0..100_usize).prop_flat_map(|len| (vec![-25f32..25.0; len], vec![-25f32..25.0; len]))
|
||||
) {
|
||||
if $cond {
|
||||
$crate::frame::unicast::test::test_unicast_t::<$ker, $t>(&*a, &*b, $func).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn [<empty_ $ker:snake>]() {
|
||||
if $cond {
|
||||
$crate::frame::unicast::test::test_unicast_t::<$ker, $t>(&[], &[], $func).unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
use std::fmt::Debug;
|
||||
use tract_data::prelude::DatumType;
|
||||
|
||||
use crate::block_quant::{BlockQuant, PackedBlockQuantFormat};
|
||||
|
||||
use crate::mmm::MMMInputFormat;
|
||||
use crate::pack::PackedFormat;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum WeightType {
|
||||
Plain(DatumType),
|
||||
BlockQuant(Box<dyn BlockQuant>),
|
||||
}
|
||||
|
||||
impl From<DatumType> for WeightType {
|
||||
fn from(value: DatumType) -> Self {
|
||||
match value {
|
||||
DatumType::F16 => WeightType::Plain(DatumType::F16),
|
||||
DatumType::F32 => WeightType::Plain(DatumType::F32),
|
||||
DatumType::F64 => WeightType::Plain(DatumType::F64),
|
||||
DatumType::I32 => WeightType::Plain(DatumType::I32),
|
||||
DatumType::I8 | DatumType::QI8(_) => WeightType::Plain(DatumType::I8),
|
||||
DatumType::U8 | DatumType::QU8(_) => WeightType::Plain(DatumType::U8),
|
||||
_ => panic!("Can't build a WeightType from {value:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Box<dyn MMMInputFormat>> for WeightType {
|
||||
fn from(value: Box<dyn MMMInputFormat>) -> Self {
|
||||
(&*value).into()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&dyn MMMInputFormat> for WeightType {
|
||||
fn from(value: &dyn MMMInputFormat) -> Self {
|
||||
if let Some(pf) = value.downcast_ref::<PackedFormat>() {
|
||||
WeightType::Plain(pf.dt)
|
||||
} else if let Some(pbqf) = value.downcast_ref::<PackedBlockQuantFormat>() {
|
||||
WeightType::BlockQuant(dyn_clone::clone_box(&*pbqf.bq))
|
||||
} else {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for WeightType {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
use WeightType::*;
|
||||
match (self, other) {
|
||||
(Plain(a), Plain(b)) => a == b,
|
||||
(BlockQuant(a), BlockQuant(b)) => a == b,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<BQ: BlockQuant> From<BQ> for WeightType {
|
||||
fn from(value: BQ) -> Self {
|
||||
WeightType::BlockQuant(dyn_clone::clone_box(&value))
|
||||
}
|
||||
}
|
||||
|
||||
impl Debug for WeightType {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Plain(p) => write!(f, "{p:?}"),
|
||||
Self::BlockQuant(bq) => write!(f, "{bq:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl WeightType {
|
||||
pub fn as_dt(&self) -> Option<DatumType> {
|
||||
match self {
|
||||
WeightType::Plain(dt) => Some(*dt),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
pub mod by_scalar;
|
||||
pub mod erf;
|
||||
pub mod gelu;
|
||||
pub mod hardswish;
|
||||
pub mod leaky_relu;
|
||||
pub mod lut;
|
||||
pub mod mmm;
|
||||
pub mod reduce;
|
||||
pub mod rms_norm;
|
||||
pub mod rounding;
|
||||
pub mod sigmoid;
|
||||
pub mod silu;
|
||||
pub mod tanh;
|
||||
pub mod unicast;
|
||||
|
||||
use tract_data::prelude::DatumType;
|
||||
|
||||
use crate::by_scalar::ByScalarKer;
|
||||
use crate::unicast::UnicastKer;
|
||||
use crate::{BinOp, LinalgRegistry};
|
||||
|
||||
pub use self::by_scalar::{HMulByScalar8, SMulByScalar4};
|
||||
pub use self::erf::SErf4;
|
||||
pub use self::gelu::{HGelu8, SGelu4};
|
||||
pub use self::hardswish::{HHardSwish8, SHardSwish4};
|
||||
pub use self::leaky_relu::{HLeakyRelu8, SLeakyRelu4};
|
||||
pub use self::lut::GenericLut8;
|
||||
pub use self::reduce::softmax_l2::SSoftMaxL2;
|
||||
pub use self::rounding::{ScaleShiftAndRound, Scaler};
|
||||
pub use self::sigmoid::{HSigmoid8, SSigmoid4};
|
||||
pub use self::silu::{HSiLU8, SSiLU4};
|
||||
pub use self::tanh::{HTanh8, STanh4};
|
||||
|
||||
pub(crate) fn register_all_unicast(registry: &mut LinalgRegistry) {
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F32),
|
||||
Box::new(|| unicast::SUnicastMul4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F16),
|
||||
Box::new(|| unicast::HUnicastMul8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F32),
|
||||
Box::new(|| unicast::SUnicastAdd4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F16),
|
||||
Box::new(|| unicast::HUnicastAdd8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F32),
|
||||
Box::new(|| unicast::SUnicastSub4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F16),
|
||||
Box::new(|| unicast::HUnicastSub8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F32),
|
||||
Box::new(|| unicast::SUnicastSubF4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F16),
|
||||
Box::new(|| unicast::HUnicastSubF8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F32),
|
||||
Box::new(|| unicast::SUnicastMin4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F16),
|
||||
Box::new(|| unicast::HUnicastMin8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F32),
|
||||
Box::new(|| unicast::SUnicastMax4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F16),
|
||||
Box::new(|| unicast::HUnicastMax8::bin()),
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn register_all_by_scalar(registry: &mut LinalgRegistry) {
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F32),
|
||||
Box::new(|| by_scalar::SMulByScalar4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Mul, DatumType::F16),
|
||||
Box::new(|| by_scalar::HMulByScalar8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F32),
|
||||
Box::new(|| by_scalar::SAddByScalar4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Add, DatumType::F16),
|
||||
Box::new(|| by_scalar::HAddByScalar8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F32),
|
||||
Box::new(|| by_scalar::SSubByScalar4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Sub, DatumType::F16),
|
||||
Box::new(|| by_scalar::HSubByScalar8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F32),
|
||||
Box::new(|| by_scalar::SSubFByScalar4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::SubF, DatumType::F16),
|
||||
Box::new(|| by_scalar::HSubFByScalar8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F32),
|
||||
Box::new(|| by_scalar::SMinByScalar4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Min, DatumType::F16),
|
||||
Box::new(|| by_scalar::HMinByScalar8::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F32),
|
||||
Box::new(|| by_scalar::SMaxByScalar4::bin()),
|
||||
);
|
||||
registry.insert(
|
||||
(BinOp::Max, DatumType::F16),
|
||||
Box::new(|| by_scalar::HMaxByScalar8::bin()),
|
||||
);
|
||||
}
|
||||
+181
@@ -0,0 +1,181 @@
|
||||
use tract_data::internal::f16;
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
SMulByScalar4,
|
||||
4,
|
||||
4,
|
||||
f32,
|
||||
fn run(x: &mut [f32], s: 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 *= s)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
SAddByScalar4,
|
||||
4,
|
||||
4,
|
||||
f32,
|
||||
fn run(x: &mut [f32], s: 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 += s)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
SSubByScalar4,
|
||||
4,
|
||||
4,
|
||||
f32,
|
||||
fn run(x: &mut [f32], s: 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 -= s)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
SSubFByScalar4,
|
||||
4,
|
||||
4,
|
||||
f32,
|
||||
fn run(x: &mut [f32], s: 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 = s - *px)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
SMinByScalar4,
|
||||
4,
|
||||
4,
|
||||
f32,
|
||||
fn run(x: &mut [f32], s: 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 = px.min(s))
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f32,
|
||||
SMaxByScalar4,
|
||||
4,
|
||||
4,
|
||||
f32,
|
||||
fn run(x: &mut [f32], s: 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 = px.max(s))
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod mul_by_scalar_f32 {
|
||||
use super::*;
|
||||
by_scalar_frame_tests!(true, f32, SMulByScalar4, |a, b| a * b);
|
||||
by_scalar_frame_tests!(true, f32, SAddByScalar4, |a, b| a + b);
|
||||
by_scalar_frame_tests!(true, f32, SSubByScalar4, |a, b| a - b);
|
||||
by_scalar_frame_tests!(true, f32, SSubFByScalar4, |a, b| b - a);
|
||||
by_scalar_frame_tests!(true, f32, SMinByScalar4, |a, b| a.min(b));
|
||||
by_scalar_frame_tests!(true, f32, SMaxByScalar4, |a, b| a.max(b));
|
||||
}
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
HMulByScalar8,
|
||||
8,
|
||||
8,
|
||||
f16,
|
||||
fn run(x: &mut [f16], s: 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 *= s)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
HAddByScalar8,
|
||||
8,
|
||||
8,
|
||||
f16,
|
||||
fn run(x: &mut [f16], s: 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 += s)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
HSubByScalar8,
|
||||
8,
|
||||
8,
|
||||
f16,
|
||||
fn run(x: &mut [f16], s: 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 -= s)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
HSubFByScalar8,
|
||||
8,
|
||||
8,
|
||||
f16,
|
||||
fn run(x: &mut [f16], s: 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 = s - *px)
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
HMinByScalar8,
|
||||
8,
|
||||
8,
|
||||
f16,
|
||||
fn run(x: &mut [f16], s: 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 = px.min(s))
|
||||
}
|
||||
);
|
||||
|
||||
by_scalar_impl_wrap!(
|
||||
f16,
|
||||
HMaxByScalar8,
|
||||
8,
|
||||
8,
|
||||
f16,
|
||||
fn run(x: &mut [f16], s: 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 = px.max(s))
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod mul_by_scalar_f16 {
|
||||
use super::*;
|
||||
by_scalar_frame_tests!(true, f16, HMulByScalar8, |a, b| a * b);
|
||||
by_scalar_frame_tests!(true, f16, HAddByScalar8, |a, b| a + b);
|
||||
by_scalar_frame_tests!(true, f16, HSubByScalar8, |a, b| a - b);
|
||||
by_scalar_frame_tests!(true, f16, HSubFByScalar8, |a, b| b - a);
|
||||
by_scalar_frame_tests!(true, f16, HMinByScalar8, |a, b| a.min(b));
|
||||
by_scalar_frame_tests!(true, f16, HMaxByScalar8, |a, b| a.max(b));
|
||||
}
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_serf4 {
|
||||
use super::*;
|
||||
crate::erf_frame_tests!(true, f32, SErf4);
|
||||
}
|
||||
+92
@@ -0,0 +1,92 @@
|
||||
#![allow(clippy::excessive_precision)]
|
||||
use crate::frame::element_wise::ElementWiseKer;
|
||||
use tract_data::internal::*;
|
||||
|
||||
// Tanh-form GELU approximation matching tract's GeluApproximate (pow=3, the
|
||||
// canonical Hendrycks-Gimpel/Open-AI form):
|
||||
//
|
||||
// gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
|
||||
//
|
||||
// The fast variant (pow=2) is not exposed here; the graph op falls back to
|
||||
// scalar when fast_impl=true.
|
||||
|
||||
const SQRT_2_OVER_PI: f32 = 0.7978845608028654;
|
||||
const COEF: f32 = 0.044715;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SGelu4;
|
||||
|
||||
impl ElementWiseKer<f32> for SGelu4 {
|
||||
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| {
|
||||
let v = *px;
|
||||
let inner = SQRT_2_OVER_PI * (v + COEF * v * v * v);
|
||||
*px = 0.5 * v * (1.0 + inner.tanh());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct HGelu8;
|
||||
|
||||
impl ElementWiseKer<f16> for HGelu8 {
|
||||
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| {
|
||||
let v = px.to_f32();
|
||||
let inner = SQRT_2_OVER_PI * (v + COEF * v * v * v);
|
||||
*px = f16::from_f32(0.5 * v * (1.0 + inner.tanh()));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
gelu_frame_tests!(true, f32, crate::generic::gelu::SGelu4);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
gelu_frame_tests!(
|
||||
true,
|
||||
tract_data::internal::f16,
|
||||
crate::generic::gelu::HGelu8
|
||||
);
|
||||
}
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
#![allow(clippy::excessive_precision)]
|
||||
use crate::frame::element_wise::ElementWiseKer;
|
||||
use tract_data::internal::*;
|
||||
use tract_num_traits::Zero;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SHardSwish4;
|
||||
|
||||
impl ElementWiseKer<f32> for SHardSwish4 {
|
||||
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);
|
||||
const INV6: f32 = 1.0 / 6.0;
|
||||
x.iter_mut().for_each(|px| {
|
||||
let relu6 = ((*px + 3.0).min(6.0)).max(0.0);
|
||||
*px = *px * relu6 * INV6;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct HHardSwish8;
|
||||
|
||||
impl ElementWiseKer<f16> for HHardSwish8 {
|
||||
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);
|
||||
let three = f16::from_f32(3.0);
|
||||
let six = f16::from_f32(6.0);
|
||||
let inv6 = f16::from_f32(1.0 / 6.0);
|
||||
x.iter_mut().for_each(|px| {
|
||||
let relu6 = ((*px + three).min(six)).max(f16::zero());
|
||||
*px = *px * relu6 * inv6;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
hardswish_frame_tests!(true, f32, crate::generic::hardswish::SHardSwish4);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
hardswish_frame_tests!(
|
||||
true,
|
||||
tract_data::internal::f16,
|
||||
crate::generic::hardswish::HHardSwish8
|
||||
);
|
||||
}
|
||||
+76
@@ -0,0 +1,76 @@
|
||||
#![allow(clippy::excessive_precision)]
|
||||
use crate::frame::element_wise::ElementWiseKer;
|
||||
use tract_data::internal::*;
|
||||
use tract_num_traits::Zero;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SLeakyRelu4;
|
||||
|
||||
impl ElementWiseKer<f32, f32> for SLeakyRelu4 {
|
||||
fn name() -> &'static str {
|
||||
"generic"
|
||||
}
|
||||
|
||||
fn alignment_bytes() -> usize {
|
||||
16
|
||||
}
|
||||
|
||||
fn alignment_items() -> usize {
|
||||
4
|
||||
}
|
||||
|
||||
fn nr() -> usize {
|
||||
4
|
||||
}
|
||||
|
||||
fn run(x: &mut [f32], alpha: 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 = if *px < 0. { *px * alpha } else { *px });
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct HLeakyRelu8;
|
||||
|
||||
impl ElementWiseKer<f16, f16> for HLeakyRelu8 {
|
||||
fn name() -> &'static str {
|
||||
"generic"
|
||||
}
|
||||
|
||||
fn alignment_bytes() -> usize {
|
||||
16
|
||||
}
|
||||
|
||||
fn alignment_items() -> usize {
|
||||
4
|
||||
}
|
||||
|
||||
fn nr() -> usize {
|
||||
8
|
||||
}
|
||||
|
||||
fn run(x: &mut [f16], alpha: 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 = if *px < f16::zero() { *px * alpha } else { *px })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
leaky_relu_frame_tests!(true, f32, crate::generic::leaky_relu::SLeakyRelu4);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
leaky_relu_frame_tests!(
|
||||
true,
|
||||
tract_data::internal::f16,
|
||||
crate::generic::leaky_relu::HLeakyRelu8
|
||||
);
|
||||
}
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
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) {
|
||||
unsafe {
|
||||
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);
|
||||
}
|
||||
+455
@@ -0,0 +1,455 @@
|
||||
#![allow(clippy::needless_range_loop)]
|
||||
use num_traits::AsPrimitive;
|
||||
|
||||
use tract_data::prelude::f16;
|
||||
use tract_data::prelude::*;
|
||||
|
||||
use super::*;
|
||||
use crate::frame::block_quant::{BlockQuant, NibbleReader, PackedBlockQuantFormat, Q4_0};
|
||||
use crate::frame::mmm::*;
|
||||
use crate::{LADatum, Ops, has_fp16};
|
||||
|
||||
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])
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
unsafe fn add_mat_mul<const MR: usize, const NR: usize, TI, TA, TB>(
|
||||
pa: *const u8,
|
||||
pb: *const u8,
|
||||
k: usize,
|
||||
ab: &mut [[TI; NR]; MR],
|
||||
) where
|
||||
TA: LADatum + AsPrimitive<TI>,
|
||||
TB: LADatum + AsPrimitive<TI>,
|
||||
TI: LADatum,
|
||||
{
|
||||
unsafe {
|
||||
let a = pa as *const TA;
|
||||
let b = pb as *const TB;
|
||||
for ik in 0..k {
|
||||
let a = std::slice::from_raw_parts(a.add(MR * ik), MR);
|
||||
let b = std::slice::from_raw_parts(b.add(NR * ik), NR);
|
||||
for i in 0..MR {
|
||||
for j in 0..NR {
|
||||
ab[i][j] += a[i].as_() * b[j].as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn add_mat_mul_pq40<const MR: usize, const NR: usize, TB, TI>(
|
||||
pa: *const u8,
|
||||
pb: *const u8,
|
||||
k: usize,
|
||||
ab: &mut [[TI; NR]; MR],
|
||||
) where
|
||||
TI: LADatum,
|
||||
f16: AsPrimitive<TI>,
|
||||
TB: AsPrimitive<TI>,
|
||||
i8: AsPrimitive<TI>,
|
||||
{
|
||||
unsafe {
|
||||
assert!(k % Q4_0.block_len() == 0);
|
||||
let len = (k * MR) / Q4_0.block_len() * Q4_0.block_bytes();
|
||||
let mut pa = NibbleReader::for_slice(std::slice::from_raw_parts(pa, len));
|
||||
let b = pb as *const TB;
|
||||
for bk in 0..k / 32 {
|
||||
let mut scales: [TI; MR] = [TI::zero(); MR];
|
||||
scales.iter_mut().for_each(|x| *x = pa.read_f16().as_());
|
||||
for ik in 0..32 {
|
||||
let mut a: [TI; MR] = [TI::zero(); MR];
|
||||
a.iter_mut()
|
||||
.zip(&scales)
|
||||
.for_each(|(x, s)| *x = *s * (pa.read_i4() - 8).as_());
|
||||
let b = std::slice::from_raw_parts(b.add(NR * (ik + 32 * bk)), NR);
|
||||
for i in 0..MR {
|
||||
for j in 0..NR {
|
||||
ab[i][j] += a[i] * b[j].as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn add_mat_mul_pq40_scales_at_end<const MR: usize, const NR: usize, TB, TI>(
|
||||
pa: *const u8,
|
||||
pb: *const u8,
|
||||
k: usize,
|
||||
ab: &mut [[TI; NR]; MR],
|
||||
) where
|
||||
TI: LADatum,
|
||||
f16: AsPrimitive<TI>,
|
||||
TB: AsPrimitive<TI>,
|
||||
i8: AsPrimitive<TI>,
|
||||
{
|
||||
unsafe {
|
||||
assert!(k % Q4_0.block_len() == 0);
|
||||
let len = (k * MR) / Q4_0.block_len() * Q4_0.block_bytes();
|
||||
let mut pa = NibbleReader::for_slice(std::slice::from_raw_parts(pa, len));
|
||||
let b = pb as *const TB;
|
||||
for bk in 0..k / 32 {
|
||||
let mut temp = [[TI::zero(); NR]; MR];
|
||||
for ik in 0..32 {
|
||||
let mut a: [TI; MR] = [TI::zero(); MR];
|
||||
a.iter_mut().for_each(|x| *x = (pa.read_i4() - 8).as_());
|
||||
let b = std::slice::from_raw_parts(b.add(NR * (ik + 32 * bk)), NR);
|
||||
for i in 0..MR {
|
||||
for j in 0..NR {
|
||||
temp[i][j] += a[i] * b[j].as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
for i in 0..MR {
|
||||
let scale = pa.read_f16().as_();
|
||||
for j in 0..NR {
|
||||
ab[i][j] += temp[i][j] * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn add_unicast<const MR: usize, const NR: usize, TI, TO>(
|
||||
ab: &mut [[TI; NR]; MR],
|
||||
other: &OutputStoreKer,
|
||||
) where
|
||||
TI: LADatum,
|
||||
TO: LADatum + AsPrimitive<TI>,
|
||||
{
|
||||
unsafe {
|
||||
for i in 0usize..MR {
|
||||
for j in 0usize..NR {
|
||||
let value: *const TO = other
|
||||
.ptr
|
||||
.offset(other.row_byte_stride * i as isize + other.col_byte_stride * j as isize)
|
||||
as _;
|
||||
ab[i].as_mut()[j] += (*value).as_();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn store_t<const MR: usize, const NR: usize, TC, TI>(
|
||||
tile: &OutputStoreKer,
|
||||
ab: &[[TI; NR]; MR],
|
||||
) where
|
||||
TC: Copy,
|
||||
{
|
||||
unsafe {
|
||||
for i in 0usize..MR {
|
||||
for j in 0usize..NR {
|
||||
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_float_t<const MR: usize, const NR: usize, TC, TI>(
|
||||
tile: &OutputStoreKer,
|
||||
ab: &[[TI; NR]; MR],
|
||||
) where
|
||||
TC: Copy + 'static,
|
||||
TI: Copy + 'static + AsPrimitive<TC>,
|
||||
{
|
||||
unsafe {
|
||||
for i in 0usize..MR {
|
||||
for j in 0usize..NR {
|
||||
let loc: *mut TC = tile
|
||||
.ptr
|
||||
.offset(tile.row_byte_stride * i as isize + tile.col_byte_stride * j as isize)
|
||||
as _;
|
||||
let val = ab[i].as_ref()[j].as_();
|
||||
*loc = val
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
unsafe fn kernel<TI, const MR: usize, const NR: usize>(mut pnl: *const FusedKerSpec<TI>) -> isize
|
||||
where
|
||||
TI: LADatum + ScaleShiftAndRound + AsPrimitive<TI>,
|
||||
TI: AsPrimitive<f16> + AsPrimitive<f32> + AsPrimitive<f64>,
|
||||
usize: AsPrimitive<TI>,
|
||||
f16: AsPrimitive<TI>,
|
||||
f32: AsPrimitive<TI>,
|
||||
f64: AsPrimitive<TI>,
|
||||
i8: AsPrimitive<TI>,
|
||||
i32: AsPrimitive<TI>,
|
||||
{
|
||||
unsafe {
|
||||
let mut ab = [[TI::zero(); NR]; MR];
|
||||
loop {
|
||||
if pnl.is_null() {
|
||||
break;
|
||||
}
|
||||
match *pnl {
|
||||
FusedKerSpec::Done => break,
|
||||
FusedKerSpec::Clear => ab = std::mem::zeroed(),
|
||||
FusedKerSpec::LoadTile(col_major, _row_major) => {
|
||||
for row in 0..MR {
|
||||
for col in 0..NR {
|
||||
ab[row][col] = *col_major.add(col * MR + row);
|
||||
}
|
||||
}
|
||||
}
|
||||
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::LeakyRelu(m) => {
|
||||
scalar!(ab, m, |a, b| if b > TI::zero() { b } else { 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::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..MR {
|
||||
for j in 0..NR {
|
||||
ab[i][j] += *rows.add(i) * *cols.add(j);
|
||||
}
|
||||
}
|
||||
}
|
||||
FusedKerSpec::AddUnicast(other) => {
|
||||
if TI::datum_type().is_float() && other.item_size == 2 {
|
||||
add_unicast::<MR, NR, TI, f16>(&mut ab, &other)
|
||||
} else if TI::datum_type().is_float() && other.item_size == 4 {
|
||||
add_unicast::<MR, NR, TI, f32>(&mut ab, &other)
|
||||
} else if TI::datum_type().is_float() && other.item_size == 8 {
|
||||
add_unicast::<MR, NR, TI, f64>(&mut ab, &other)
|
||||
} else if TI::datum_type() == i32::datum_type() && other.item_size == 1 {
|
||||
add_unicast::<MR, NR, TI, i8>(&mut ab, &other)
|
||||
} else if TI::datum_type() == i32::datum_type() && other.item_size == 4 {
|
||||
add_unicast::<MR, NR, TI, i32>(&mut ab, &other)
|
||||
} else {
|
||||
unimplemented!("Missing AddUnicast type");
|
||||
}
|
||||
}
|
||||
FusedKerSpec::ShiftLeft(shift) => {
|
||||
for i in 0..MR {
|
||||
for j in 0..NR {
|
||||
ab[i][j] = ab[i][j].q_shl(shift);
|
||||
}
|
||||
}
|
||||
}
|
||||
FusedKerSpec::RoundingShiftRight(shift, rp) => {
|
||||
for i in 0..MR {
|
||||
for j in 0..NR {
|
||||
ab[i][j] = ab[i][j].q_shr(shift, rp);
|
||||
}
|
||||
}
|
||||
}
|
||||
FusedKerSpec::QScale(shift, rp, mult) => {
|
||||
for i in 0..MR {
|
||||
for j in 0..NR {
|
||||
ab[i][j] = ab[i][j].q_scale(Scaler::from_fuse_params(shift, rp, mult));
|
||||
}
|
||||
}
|
||||
}
|
||||
FusedKerSpec::AddMatMul { k, pa, pb, packing } => {
|
||||
use std::mem::transmute;
|
||||
if TI::datum_type().is_float() {
|
||||
match packing {
|
||||
0 => add_mat_mul::<MR, NR, TI, TI, TI>(pa, pb, k, &mut ab),
|
||||
1 => add_mat_mul::<MR, NR, TI, f16, f16>(pa, pb, k, &mut ab),
|
||||
2 => add_mat_mul::<MR, NR, TI, f32, f32>(pa, pb, k, &mut ab),
|
||||
3 => add_mat_mul::<MR, NR, TI, f16, f32>(pa, pb, k, &mut ab),
|
||||
4 => add_mat_mul::<MR, NR, TI, f32, f16>(pa, pb, k, &mut ab),
|
||||
5 => add_mat_mul_pq40::<MR, NR, f16, TI>(pa, pb, k, &mut ab),
|
||||
6 => add_mat_mul_pq40_scales_at_end::<MR, NR, f16, TI>(
|
||||
pa, pb, k, &mut ab,
|
||||
),
|
||||
7 => add_mat_mul_pq40::<MR, NR, f32, TI>(pa, pb, k, &mut ab),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
} else if TI::datum_type() == i32::datum_type() {
|
||||
// transmute to allow using explicitly i3 in add_mat_mul generic params
|
||||
let ab = transmute::<&mut [[TI; NR]; MR], &mut [[i32; NR]; MR]>(&mut ab);
|
||||
if packing == 0 {
|
||||
add_mat_mul::<MR, NR, i32, i32, i32>(pa, pb, k, ab)
|
||||
} else if packing == 1 {
|
||||
add_mat_mul::<MR, NR, i32, i8, i8>(pa, pb, k, ab)
|
||||
} else {
|
||||
return 1;
|
||||
}
|
||||
} else {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
FusedKerSpec::Store(tile) => {
|
||||
if TI::datum_type().is_float() {
|
||||
match tile.item_size {
|
||||
2 => store_float_t::<MR, NR, f16, _>(&tile, &ab),
|
||||
4 => store_float_t::<MR, NR, f32, _>(&tile, &ab),
|
||||
8 => store_float_t::<MR, NR, f64, _>(&tile, &ab),
|
||||
_ => unimplemented!(),
|
||||
}
|
||||
} else {
|
||||
match tile.item_size {
|
||||
1 => store_t::<MR, NR, u8, _>(&tile, &ab),
|
||||
2 => store_t::<MR, NR, u16, _>(&tile, &ab),
|
||||
4 => store_t::<MR, NR, u32, _>(&tile, &ab),
|
||||
8 => store_t::<MR, NR, u64, _>(&tile, &ab),
|
||||
_ => unimplemented!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
pnl = pnl.add(1);
|
||||
}
|
||||
}
|
||||
0
|
||||
}
|
||||
|
||||
fn pq40_r4() -> PackedBlockQuantFormat {
|
||||
PackedBlockQuantFormat::new(&Q4_0, 4, 0, false)
|
||||
}
|
||||
|
||||
fn pq40_r4_se() -> PackedBlockQuantFormat {
|
||||
PackedBlockQuantFormat::new(&Q4_0, 4, 0, true)
|
||||
}
|
||||
|
||||
// f16 kernels
|
||||
MMMRustKernel!(kernel::<f16, 4, 4> => generic_f16_4x4<f16>(4,4)
|
||||
packing[1] = f16f16bis => |k| k.with_packing(f16::packing(4), f16::packing(4));
|
||||
packing[2] = f32f32 => |k| k.with_packing(f32::packing(4), f32::packing(4));
|
||||
packing[3] = f16f32 => |k| k.with_packing(f16::packing(4), f32::packing(4));
|
||||
packing[4] = f32f16 => |k| k.with_packing(f32::packing(4), f16::packing(4));
|
||||
packing[5] = q40f16 => |k| k.with_packing(pq40_r4(), f16::packing(4));
|
||||
packing[6] = q40f16se => |k| k.with_packing(pq40_r4_se(), f16::packing(4));
|
||||
packing[7] = q40f32 => |k| k.with_packing(pq40_r4(), f32::packing(4));
|
||||
quality(if has_fp16() { ImplementationQuality::Generic } else { ImplementationQuality::Dreadful })
|
||||
store(f32, f64)
|
||||
);
|
||||
|
||||
MMMRustKernel! {kernel::<f16, 4, 1> => generic_f16_4x1<f16>(4,1)
|
||||
packing[1] = f16f16bis => |k| k.with_packing(f16::packing(4), f16::packing(1));
|
||||
packing[2] = f32f32 => |k| k.with_packing(f32::packing(4), f32::packing(1));
|
||||
packing[3] = f16f32 => |k| k.with_packing(f16::packing(4), f32::packing(1));
|
||||
packing[4] = f32f16 => |k| k.with_packing(f32::packing(4), f16::packing(1));
|
||||
packing[5] = q40f16 => |k| k.with_packing(pq40_r4(), f16::packing(1));
|
||||
packing[6] = q40f16se => |k| k.with_packing(pq40_r4_se(), f16::packing(1));
|
||||
packing[7] = q40f32 => |k| k.with_packing(pq40_r4(), f32::packing(1));
|
||||
quality(if has_fp16() { ImplementationQuality::Generic } else { ImplementationQuality::Dreadful })
|
||||
store(f32, f64)
|
||||
}
|
||||
|
||||
// f32 kernels
|
||||
MMMRustKernel!(kernel::<f32, 4, 4> => generic_f32_4x4<f32>(4,4)
|
||||
packing[1] = f16f16 => |k| k.with_packing(f16::packing(4), f16::packing(4));
|
||||
packing[2] = f32f32bis => |k| k.with_packing(f32::packing(4), f32::packing(4));
|
||||
packing[3] = f16f32 => |k| k.with_packing(f16::packing(4), f32::packing(4));
|
||||
packing[4] = f32f16 => |k| k.with_packing(f32::packing(4), f16::packing(4));
|
||||
packing[5] = q40f16 => |k| k.with_packing(pq40_r4(), f16::packing(4));
|
||||
packing[6] = q40f16se => |k| k.with_packing(pq40_r4_se(), f16::packing(4));
|
||||
packing[7] = q40f32 => |k| k.with_packing(pq40_r4(), f32::packing(4));
|
||||
quality(ImplementationQuality::Generic)
|
||||
store(f16, f64)
|
||||
);
|
||||
MMMRustKernel! {kernel::<f32, 4, 1> => generic_f32_4x1<f32>(4,1)
|
||||
packing[1] = f16f16 => |k| k.with_packing(f16::packing(4), f16::packing(1));
|
||||
packing[2] = f32f32bis => |k| k.with_packing(f32::packing(4), f32::packing(1));
|
||||
packing[3] = f16f32 => |k| k.with_packing(f16::packing(4), f32::packing(1));
|
||||
packing[4] = f32f16 => |k| k.with_packing(f32::packing(4), f16::packing(1));
|
||||
packing[5] = q40f16 => |k| k.with_packing(pq40_r4(), f16::packing(1));
|
||||
packing[6] = q40f16se => |k| k.with_packing(pq40_r4_se(), f16::packing(1));
|
||||
packing[7] = q40f32 => |k| k.with_packing(pq40_r4(), f32::packing(1));
|
||||
quality(ImplementationQuality::Generic)
|
||||
store(f16, f64)
|
||||
}
|
||||
|
||||
// f64 kernels
|
||||
MMMRustKernel!(kernel::<f64, 4, 4> => generic_f64_4x4<f64>(4,4)
|
||||
quality(ImplementationQuality::Generic)
|
||||
store(f16, f32));
|
||||
MMMRustKernel!(kernel::<f64, 4, 1> => generic_f64_4x1<f64>(4,1)
|
||||
quality(ImplementationQuality::Generic)
|
||||
store(f16, f32));
|
||||
|
||||
// I32 kernels
|
||||
MMMRustKernel! {kernel::<i32, 4, 4> => generic_i32_4x4<i32>(4,4)
|
||||
packing[1] = i8i8 => |k| k.with_packing(i8::packing(4), i8::packing(4));
|
||||
quality(ImplementationQuality::Generic)
|
||||
store(i8)
|
||||
}
|
||||
|
||||
MMMRustKernel! {kernel::<i32, 4, 1> => generic_i32_4x1<i32>(4,1)
|
||||
packing[1] = i8i8 => |k| k.with_packing(i8::packing(4), i8::packing(1));
|
||||
quality(ImplementationQuality::Generic)
|
||||
store(i8)
|
||||
}
|
||||
|
||||
// extra tests kernels
|
||||
#[cfg(test)]
|
||||
MMMRustKernel!(kernel::<f32, 3, 2> => generic_f32_3x2<f32>(3,2) store(f16, f64));
|
||||
|
||||
#[cfg(test)]
|
||||
MMMRustKernel! {kernel::<i32, 3, 2> => generic_i32_3x2<i32>(3,2)
|
||||
packing[1] = i8i8 => |k| k.with_packing(i8::packing(3), i8::packing(2));
|
||||
store(i8)
|
||||
}
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
ops.mmm_impls.push(generic_f16_4x4.mmm());
|
||||
ops.mmm_impls.push(generic_f16_4x1.mmm());
|
||||
ops.mmm_impls.push(generic_f32_4x4.mmm());
|
||||
ops.mmm_impls.push(generic_f32_4x1.mmm());
|
||||
ops.mmm_impls.push(generic_f64_4x4.mmm());
|
||||
ops.mmm_impls.push(generic_f64_4x1.mmm());
|
||||
ops.mmm_impls.push(generic_i32_4x4.mmm());
|
||||
ops.mmm_impls.push(generic_i32_4x1.mmm());
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
|
||||
#[test]
|
||||
fn kits() {
|
||||
let mut ops = crate::generic();
|
||||
super::plug(&mut ops);
|
||||
}
|
||||
}
|
||||
+187
@@ -0,0 +1,187 @@
|
||||
// Reduce<max> generic implementation
|
||||
pub mod max {
|
||||
pub use tract_data::internal::f16;
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f32,
|
||||
SMax4,
|
||||
4,
|
||||
4,
|
||||
(),
|
||||
f32::MIN,
|
||||
fn run(x: &[f32], _: ()) -> f32 {
|
||||
debug_assert!(x.len() % Self::nr() == 0);
|
||||
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
*x.iter().max_by(|a, b| a.total_cmp(b)).unwrap()
|
||||
},
|
||||
fn reduce_two(a: f32, b: f32) -> f32 {
|
||||
a.max(b)
|
||||
}
|
||||
);
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f16,
|
||||
HMax8,
|
||||
8,
|
||||
8,
|
||||
(),
|
||||
f16::MIN,
|
||||
fn run(x: &[f16], _: ()) -> f16 {
|
||||
debug_assert!(x.len() % Self::nr() == 0);
|
||||
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
*x.iter().max_by(|a, b| a.total_cmp(b)).unwrap()
|
||||
},
|
||||
fn reduce_two(a: f16, b: f16) -> f16 {
|
||||
a.max(b)
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
crate::max_frame_tests!(true, f32, crate::generic::reduce::max::SMax4);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
use super::*;
|
||||
crate::max_frame_tests!(true, f16, crate::generic::reduce::max::HMax8);
|
||||
}
|
||||
}
|
||||
|
||||
// Reduce<sum> generic implementation
|
||||
pub mod sum {
|
||||
use crate::num_traits::Zero;
|
||||
pub use tract_data::internal::f16;
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f32,
|
||||
SSum4,
|
||||
4,
|
||||
4,
|
||||
(),
|
||||
0.0,
|
||||
fn run(x: &[f32], _: ()) -> f32 {
|
||||
debug_assert!(x.len() % Self::nr() == 0);
|
||||
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
x.iter().sum::<f32>()
|
||||
},
|
||||
fn reduce_two(a: f32, b: f32) -> f32 {
|
||||
a + b
|
||||
}
|
||||
);
|
||||
|
||||
reduce_impl_wrap!(
|
||||
f16,
|
||||
HSum8,
|
||||
8,
|
||||
8,
|
||||
(),
|
||||
f16::zero(),
|
||||
fn run(x: &[f16], _: ()) -> f16 {
|
||||
debug_assert!(x.len() % Self::nr() == 0);
|
||||
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
x.iter().sum::<f16>()
|
||||
},
|
||||
fn reduce_two(a: f16, b: f16) -> f16 {
|
||||
a + b
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
crate::sum_frame_tests!(true, f32, crate::generic::reduce::sum::SSum4);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
use super::*;
|
||||
crate::sum_frame_tests!(true, f16, crate::generic::reduce::sum::HSum8);
|
||||
}
|
||||
}
|
||||
|
||||
// Softmax generic implementation
|
||||
pub mod softmax_l2 {
|
||||
use crate::num_traits::Zero;
|
||||
use tract_data::internal::f16;
|
||||
|
||||
map_reduce_impl_wrap!(
|
||||
f32,
|
||||
SSoftMaxL2,
|
||||
4,
|
||||
4,
|
||||
f32,
|
||||
f32::MIN,
|
||||
0.0,
|
||||
fn run(x: &mut [f32], max: f32) -> f32 {
|
||||
debug_assert!(x.len() % Self::nr() == 0);
|
||||
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
let mut sum = 0.;
|
||||
for v in x.iter_mut() {
|
||||
let y = *v - max;
|
||||
let y = fast_compact_exp_f32(y);
|
||||
*v = y;
|
||||
sum += y;
|
||||
}
|
||||
sum
|
||||
},
|
||||
fn reduce_two(a: f32, b: f32) -> f32 {
|
||||
a + b
|
||||
}
|
||||
);
|
||||
|
||||
map_reduce_impl_wrap!(
|
||||
f16,
|
||||
HSoftMaxL2,
|
||||
8,
|
||||
8,
|
||||
f16,
|
||||
f16::MIN,
|
||||
f16::zero(),
|
||||
fn run(x: &mut [f16], max: f16) -> f16 {
|
||||
debug_assert!(x.len() % Self::nr() == 0);
|
||||
debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
let mut sum = f16::zero();
|
||||
for v in x.iter_mut() {
|
||||
let y = *v - max;
|
||||
let y = f16::from_f32(fast_compact_exp_f32(y.to_f32()));
|
||||
*v = y;
|
||||
sum += y;
|
||||
}
|
||||
sum
|
||||
},
|
||||
fn reduce_two(a: f16, b: f16) -> f16 {
|
||||
a + b
|
||||
}
|
||||
);
|
||||
|
||||
// ported from https://github.com/gnuradio/volk/blob/master/kernels/volk/volk_32f_expfast_32f.h
|
||||
// probably inspired from https://nic.schraudolph.org/pubs/Schraudolph99.pdf
|
||||
// not that the cast to u32 deals with negative right, while implem in volk code are wrong in some
|
||||
// corner cases (need a max(0,x) before the u32 conversion)
|
||||
pub fn fast_compact_exp_f32(v: f32) -> f32 {
|
||||
const MLN2: f32 = 0.6931471805f32;
|
||||
const A: f32 = 8388608.0f32;
|
||||
const B: f32 = 1065353216.0f32;
|
||||
const C: f32 = 60801.0f32;
|
||||
const SLOPE: f32 = A / MLN2;
|
||||
const OFFSET: f32 = B - C;
|
||||
f32::from_bits(((SLOPE * v) + OFFSET) as u32)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
crate::softmax_l2_frame_tests!(true, f32, super::SSoftMaxL2);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
use super::*;
|
||||
crate::softmax_l2_frame_tests!(true, f16, HSoftMaxL2);
|
||||
}
|
||||
}
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
/// Generic scalar reference implementation of fused row-wise RmsNorm.
|
||||
/// out_i = x_i * rsqrt(mean(x_i²) + eps)
|
||||
///
|
||||
/// Replaces tract-core's 4-call composition (`Reducer::MeanOfSquares` + `Add` +
|
||||
/// `Rsqrt` + `Mul`) with a single 2-pass kernel. Overridden by AVX-512 on
|
||||
/// x86_64; non-x86 / non-AVX512 hosts keep this scalar version.
|
||||
pub fn rms_norm_f32(buf: &mut [f32], eps: f32) {
|
||||
if buf.is_empty() {
|
||||
return;
|
||||
}
|
||||
let n = buf.len() as f32;
|
||||
let sum_sq: f32 = buf.iter().map(|x| x * x).sum();
|
||||
let mean_sq = sum_sq / n;
|
||||
let inv_std = (mean_sq + eps).sqrt().recip();
|
||||
for x in buf.iter_mut() {
|
||||
*x *= inv_std;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn close_enough(got: f32, want: f32) -> bool {
|
||||
(got - want).abs() < 1e-5
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rms_norm_constant() {
|
||||
// RmsNorm of all-ones with eps=0: mean(1²)=1, rsqrt=1, output=1.
|
||||
let mut buf = [1.0; 16];
|
||||
rms_norm_f32(&mut buf, 0.0);
|
||||
for v in buf {
|
||||
assert!(close_enough(v, 1.0), "got {v}, want 1.0");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rms_norm_one_two_three_four() {
|
||||
// mean(1+4+9+16)/4 = 7.5, rsqrt(7.5) ≈ 0.3651
|
||||
let mut buf = [1.0_f32, 2.0, 3.0, 4.0];
|
||||
rms_norm_f32(&mut buf, 0.0);
|
||||
let inv = (7.5_f32).sqrt().recip();
|
||||
for (i, v) in buf.iter().enumerate() {
|
||||
let want = (i + 1) as f32 * inv;
|
||||
assert!(close_enough(*v, want), "i={i}: got {v}, want {want}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rms_norm_eps_added_under_root() {
|
||||
// eps inside the sqrt, not added afterward.
|
||||
let mut buf = [0.0_f32; 4];
|
||||
rms_norm_f32(&mut buf, 1e-5);
|
||||
for v in buf {
|
||||
// 0 * anything = 0; just verify no NaN/inf.
|
||||
assert!(v.is_finite());
|
||||
assert_eq!(v, 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rms_norm_empty() {
|
||||
let mut buf: [f32; 0] = [];
|
||||
rms_norm_f32(&mut buf, 1e-5);
|
||||
}
|
||||
}
|
||||
+534
@@ -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);
|
||||
}
|
||||
}
|
||||
+142
@@ -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
|
||||
);
|
||||
}
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
#![allow(clippy::excessive_precision)]
|
||||
use crate::frame::element_wise::ElementWiseKer;
|
||||
use tract_data::internal::*;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SSiLU4;
|
||||
|
||||
impl ElementWiseKer<f32> for SSiLU4 {
|
||||
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| {
|
||||
let sigmoid = 1.0 / (1.0 + (-*px).exp());
|
||||
*px = *px * sigmoid;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct HSiLU8;
|
||||
|
||||
impl ElementWiseKer<f16> for HSiLU8 {
|
||||
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| {
|
||||
let x_f32 = px.to_f32();
|
||||
let sigmoid = 1.0 / (1.0 + (-x_f32).exp());
|
||||
*px = f16::from_f32(x_f32 * sigmoid);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
silu_frame_tests!(true, f32, crate::generic::silu::SSiLU4);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
silu_frame_tests!(
|
||||
true,
|
||||
tract_data::internal::f16,
|
||||
crate::generic::silu::HSiLU8
|
||||
);
|
||||
}
|
||||
+137
@@ -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
|
||||
);
|
||||
}
|
||||
+194
@@ -0,0 +1,194 @@
|
||||
pub use tract_data::internal::f16;
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
SUnicastMul4,
|
||||
4,
|
||||
4,
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a *= b)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
HUnicastMul8,
|
||||
8,
|
||||
8,
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a *= b)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
SUnicastAdd4,
|
||||
4,
|
||||
4,
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a += b)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
HUnicastAdd8,
|
||||
8,
|
||||
8,
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a += b)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
SUnicastSub4,
|
||||
4,
|
||||
4,
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a -= b)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
HUnicastSub8,
|
||||
8,
|
||||
8,
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a -= b)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
SUnicastSubF4,
|
||||
4,
|
||||
4,
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a = *b - *a)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
HUnicastSubF8,
|
||||
8,
|
||||
8,
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a = *b - *a)
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
SUnicastMin4,
|
||||
4,
|
||||
4,
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a = a.min(*b))
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
HUnicastMin8,
|
||||
8,
|
||||
8,
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a = a.min(*b))
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f32,
|
||||
SUnicastMax4,
|
||||
4,
|
||||
4,
|
||||
fn run(a: &mut [f32], b: &[f32]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a = a.max(*b))
|
||||
}
|
||||
);
|
||||
|
||||
unicast_impl_wrap!(
|
||||
f16,
|
||||
HUnicastMax8,
|
||||
8,
|
||||
8,
|
||||
fn run(a: &mut [f16], b: &[f16]) {
|
||||
debug_assert!(a.len() == b.len());
|
||||
debug_assert!(a.len() % Self::nr() == 0);
|
||||
debug_assert!(a.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
debug_assert!(b.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
a.iter_mut().zip(b.iter()).for_each(|(a, b)| *a = a.max(*b))
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod s {
|
||||
use super::*;
|
||||
use proptest::strategy::Strategy;
|
||||
crate::unicast_frame_tests!(true, f32, SUnicastMul4, |a, b| a * b);
|
||||
crate::unicast_frame_tests!(true, f32, SUnicastAdd4, |a, b| a + b);
|
||||
crate::unicast_frame_tests!(true, f32, SUnicastSub4, |a, b| a - b);
|
||||
crate::unicast_frame_tests!(true, f32, SUnicastSubF4, |a, b| b - a);
|
||||
crate::unicast_frame_tests!(true, f32, SUnicastMin4, |a, b| a.min(b));
|
||||
crate::unicast_frame_tests!(true, f32, SUnicastMax4, |a, b| a.max(b));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[macro_use]
|
||||
pub mod h {
|
||||
use super::*;
|
||||
use proptest::strategy::Strategy;
|
||||
crate::unicast_frame_tests!(true, f16, HUnicastMul8, |a, b| a * b);
|
||||
crate::unicast_frame_tests!(true, f16, HUnicastAdd8, |a, b| a + b);
|
||||
crate::unicast_frame_tests!(true, f16, HUnicastSub8, |a, b| a - b);
|
||||
crate::unicast_frame_tests!(true, f16, HUnicastSubF8, |a, b| b - a);
|
||||
crate::unicast_frame_tests!(true, f16, HUnicastMin8, |a, b| a.min(b));
|
||||
crate::unicast_frame_tests!(true, f16, HUnicastMax8, |a, b| a.max(b));
|
||||
}
|
||||
+159
@@ -0,0 +1,159 @@
|
||||
use tract_data::itertools::Itertools;
|
||||
use tract_data::prelude::Blob;
|
||||
|
||||
use super::runner;
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
static mut HAS_AVX512: bool = false;
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
#[inline(never)]
|
||||
fn load_a_slice(slice: &[u8], loops: usize) {
|
||||
unsafe {
|
||||
if HAS_AVX512 {
|
||||
for _ in 0..loops {
|
||||
let mut ptr = slice.as_ptr();
|
||||
let end = ptr.add(slice.len());
|
||||
while ptr < end {
|
||||
std::arch::asm!("
|
||||
vmovaps zmm0, [rsi]
|
||||
vmovaps zmm1, [rsi + 64]
|
||||
vmovaps zmm2, [rsi + 128]
|
||||
vmovaps zmm3, [rsi + 192]
|
||||
vmovaps zmm4, [rsi + 256]
|
||||
vmovaps zmm5, [rsi + 320]
|
||||
vmovaps zmm6, [rsi + 384]
|
||||
vmovaps zmm7, [rsi + 448]
|
||||
", inout("rsi") ptr,
|
||||
out("zmm0") _,
|
||||
out("zmm1") _,
|
||||
);
|
||||
ptr = ptr.add(512);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let mut ptr = slice.as_ptr();
|
||||
let end = ptr.add(slice.len());
|
||||
for _ in 0..loops {
|
||||
while ptr < end {
|
||||
std::arch::asm!("
|
||||
vmovaps ymm0, [rsi]
|
||||
vmovaps ymm1, [rsi + 32]
|
||||
vmovaps ymm2, [rsi + 64]
|
||||
vmovaps ymm3, [rsi + 96]
|
||||
", inout("rsi") ptr,
|
||||
out("ymm0") _,
|
||||
out("ymm1") _,
|
||||
out("ymm2") _,
|
||||
out("ymm3") _,
|
||||
);
|
||||
ptr = ptr.add(128);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
#[inline]
|
||||
fn load_a_slice(slice: &[u8], loops: usize) {
|
||||
unsafe {
|
||||
for _ in 0..loops {
|
||||
let mut ptr = slice.as_ptr();
|
||||
let end = ptr.add(slice.len());
|
||||
while ptr < end {
|
||||
std::arch::asm!("
|
||||
ld1 {{v0.16b-v3.16b}}, [x0], #64
|
||||
ld1 {{v4.16b-v7.16b}}, [x0], #64
|
||||
", inout("x0") ptr,
|
||||
out("v0") _,
|
||||
out("v1") _,
|
||||
out("v2") _,
|
||||
out("v3") _,
|
||||
out("v4") _,
|
||||
out("v5") _,
|
||||
out("v6") _,
|
||||
out("v7") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "arm")]
|
||||
#[inline(never)]
|
||||
fn load_a_slice(slice: &[u8], loops: usize) {
|
||||
unsafe {
|
||||
for _ in 0..loops {
|
||||
let mut ptr = slice.as_ptr();
|
||||
let end = ptr.add(slice.len());
|
||||
while ptr < end {
|
||||
std::arch::asm!("
|
||||
vldmia r1!, {{q0-q3}}
|
||||
vldmia r1!, {{q4-q7}}
|
||||
", inout("r1") ptr,
|
||||
out("d0") _, out("d1") _, out("d2") _, out("d3") _,
|
||||
out("d4") _, out("d5") _, out("d6") _, out("d7") _,
|
||||
out("d8") _, out("d9") _, out("d10") _, out("d11") _,
|
||||
out("d12") _, out("d13") _, out("d14") _, out("d15") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn bandwidth_seq(slice_len: usize, threads: usize) -> f64 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
unsafe {
|
||||
HAS_AVX512 = std::is_x86_feature_detected!("avx512f");
|
||||
}
|
||||
std::thread::scope(|s| {
|
||||
let gards = (0..threads)
|
||||
.map(|_| {
|
||||
s.spawn(|| {
|
||||
let buffer = unsafe { Blob::new_for_size_and_align(slice_len, 1024) };
|
||||
runner::run_bench(|loops| load_a_slice(&buffer, loops))
|
||||
})
|
||||
})
|
||||
.collect_vec();
|
||||
let time = gards.into_iter().map(|t| t.join().unwrap()).sum::<f64>() / threads as f64;
|
||||
(slice_len * threads) as f64 / time
|
||||
})
|
||||
}
|
||||
|
||||
pub fn what_is_big() -> usize {
|
||||
1024 * 1024 * if cfg!(target_arch = "arm") { 64 } else { 256 }
|
||||
}
|
||||
|
||||
pub fn l1_bandwidth_seq(threads: usize) -> f64 {
|
||||
// [1024, 2048, 4096, 8192, 16384, 32768, 65536]
|
||||
[1024]
|
||||
.into_iter()
|
||||
.map(|slice_len| bandwidth_seq(slice_len, threads))
|
||||
.max_by_key(|x| *x as i64)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub fn main_memory_bandwith_seq(threads: usize) -> f64 {
|
||||
bandwidth_seq(what_is_big(), threads)
|
||||
}
|
||||
|
||||
#[ignore]
|
||||
#[test]
|
||||
fn b() {
|
||||
let max = what_is_big();
|
||||
for threads in [1, 2, 3, 4] {
|
||||
println!("Threads: {}", threads);
|
||||
for size in (0..)
|
||||
.flat_map(|po2| (0..2).map(move |f| (1024 + 512 * f) * (1 << po2)))
|
||||
.take_while(|&s| s < max)
|
||||
{
|
||||
let bw = bandwidth_seq(size, threads);
|
||||
println!(
|
||||
"threads: {threads} slice: {} KiB bandwidth: {} GiB/s",
|
||||
size as f64 / 1024.,
|
||||
(bw / (1024. * 1024. * 1024.)) as usize
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
pub mod runner;
|
||||
|
||||
#[cfg(feature = "hwbench")]
|
||||
pub mod bandwidth;
|
||||
+127
@@ -0,0 +1,127 @@
|
||||
#![allow(unused_macros)]
|
||||
|
||||
use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! r1 { ($($stat:stmt)*) => { $( $stat )* } }
|
||||
#[macro_export]
|
||||
macro_rules! r2 { ($($stat:stmt)*) => { $( $stat )* $( $stat )* } }
|
||||
#[macro_export]
|
||||
macro_rules! r4 { ($($stat:stmt)*) => { r2!(r2!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r8 { ($($stat:stmt)*) => { r2!(r4!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r16 { ($($stat:stmt)*) => { r2!(r8!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r32 { ($($stat:stmt)*) => { r2!(r16!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r64 { ($($stat:stmt)*) => { r2!(r32!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r128 { ($($stat:stmt)*) => { r2!(r64!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r256 { ($($stat:stmt)*) => { r2!(r128!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r512 { ($($stat:stmt)*) => { r2!(r256!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r1024 { ($($stat:stmt)*) => { r2!(r512!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r2048 { ($($stat:stmt)*) => { r2!(r1024!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r4096 { ($($stat:stmt)*) => { r2!(r2048!($($stat)*)) }}
|
||||
#[macro_export]
|
||||
macro_rules! r8192 { ($($stat:stmt)*) => { r2!(r4096!($($stat)*)) }}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! b1 { ($($stat:stmt)*) => { nano::run_bench(|| { r1!($($stat)*); }) / 1.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b2 { ($($stat:stmt)*) => { nano::run_bench(|| { r2!($($stat)*); }) / 2.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b4 { ($($stat:stmt)*) => { nano::run_bench(|| { r4!($($stat)*); }) / 4.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b8 { ($($stat:stmt)*) => { nano::run_bench(|| { r8!($($stat)*); }) / 8.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b16 { ($($stat:stmt)*) => { nano::run_bench(|| { r16!($($stat)*); }) / 16.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b32 { ($($stat:stmt)*) => { nano::run_bench(|| { r32!($($stat)*); }) / 32.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b64 { ($($stat:stmt)*) => { nano::run_bench(|| { r64!($($stat)*); }) / 64.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b128 { ($($stat:stmt)*) => { nano::run_bench(|| { r128!($($stat)*); }) / 128.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b256 { ($($stat:stmt)*) => { nano::run_bench(|| { r256!($($stat)*); }) / 256.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b512 { ($($stat:stmt)*) => { nano::run_bench(|| { r512!($($stat)*); }) / 512.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b1024 { ($($stat:stmt)*) => { nano::run_bench(|| { r1024!($($stat)*); }) / 1024.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b2048 { ($($stat:stmt)*) => { nano::run_bench(|| { r2048!($($stat)*); }) / 2048.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b4096 { ($($stat:stmt)*) => { nano::run_bench(|| { r4096!($($stat)*); }) / 4096.0 } }
|
||||
#[macro_export]
|
||||
macro_rules! b8192 { ($($stat:stmt)*) => { nano::run_bench(|| { r8192!($($stat)*); }) / 8192.0 } }
|
||||
|
||||
#[inline]
|
||||
fn black_box<T>(dummy: T) -> T {
|
||||
unsafe {
|
||||
let ret = std::ptr::read_volatile(&dummy);
|
||||
std::mem::forget(dummy);
|
||||
ret
|
||||
}
|
||||
}
|
||||
|
||||
pub fn run_bench<T, F: FnMut(usize) -> T + Copy>(f: F) -> f64 {
|
||||
let start = Instant::now();
|
||||
let mut f = black_box(f);
|
||||
black_box(f(1));
|
||||
let once = start.elapsed();
|
||||
let evaled = if once < Duration::from_millis(1) {
|
||||
let start = Instant::now();
|
||||
black_box(f)(1000);
|
||||
start.elapsed().as_secs_f64() / 1000.
|
||||
} else {
|
||||
once.as_secs_f64()
|
||||
};
|
||||
// raw evaluation is over a second. stop right there
|
||||
if evaled > 1.0 {
|
||||
return evaled;
|
||||
}
|
||||
|
||||
// we want each individual sample to run for no less than
|
||||
let minimum_sampling_time_s = 0.01;
|
||||
let minimum_samples = 25;
|
||||
let desired_bench_time = 1.0;
|
||||
|
||||
let inner_loops = (minimum_sampling_time_s / evaled).max(1.0) as usize;
|
||||
|
||||
let samples =
|
||||
((desired_bench_time / (inner_loops as f64 * evaled)) as usize).max(minimum_samples);
|
||||
let warmup = (1.0 / evaled) as usize;
|
||||
|
||||
// println!(
|
||||
// "evaled: {:?} samples:{samples} inner_loops:{inner_loops} time:{}",
|
||||
// Duration::from_secs_f64(evaled),
|
||||
// (samples * inner_loops) as f64 * evaled
|
||||
// );
|
||||
let mut measures = vec![0.0; samples];
|
||||
|
||||
black_box(f(warmup));
|
||||
for m in &mut measures {
|
||||
let start = Instant::now();
|
||||
black_box(black_box(f))(inner_loops);
|
||||
let time = start.elapsed().as_secs_f64();
|
||||
*m = time / inner_loops as f64
|
||||
}
|
||||
measures.sort_by(|a, b| {
|
||||
if a < b {
|
||||
std::cmp::Ordering::Less
|
||||
} else {
|
||||
std::cmp::Ordering::Greater
|
||||
}
|
||||
});
|
||||
let q1 = measures[samples / 4];
|
||||
let q3 = measures[samples - samples / 4];
|
||||
let iq = q3 - q1;
|
||||
measures.retain(|&x| x >= q1 - 3. * iq && x <= q3 + 3. * iq);
|
||||
measures.iter().copied().sum::<f64>() / measures.len() as f64
|
||||
}
|
||||
@@ -0,0 +1,438 @@
|
||||
#![allow(clippy::missing_safety_doc)]
|
||||
#![allow(clippy::redundant_closure_call)]
|
||||
#![allow(clippy::len_zero)]
|
||||
#![allow(clippy::excessive_precision)]
|
||||
#![allow(clippy::approx_constant)]
|
||||
#![allow(clippy::manual_is_multiple_of)]
|
||||
#![allow(unexpected_cfgs)]
|
||||
#![allow(unused_macros)]
|
||||
#[macro_use]
|
||||
extern crate derive_new;
|
||||
extern crate lazy_static;
|
||||
extern crate log;
|
||||
extern crate num_traits;
|
||||
#[macro_use]
|
||||
extern crate pastey;
|
||||
#[cfg(test)]
|
||||
extern crate proptest;
|
||||
|
||||
include!(concat!(env!("OUT_DIR"), "/extern_kernel_macro.rs"));
|
||||
|
||||
#[macro_use]
|
||||
mod frame;
|
||||
pub mod generic;
|
||||
pub mod multithread;
|
||||
pub use frame::weights::WeightType;
|
||||
pub use generic::{ScaleShiftAndRound, Scaler};
|
||||
use lazy_static::lazy_static;
|
||||
use mmm::{MMMInputFormat, MatMatMul, PanelExtractor};
|
||||
use tract_data::internal::TensorView;
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
pub mod x86_64_fma;
|
||||
|
||||
pub mod hwbench;
|
||||
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
pub mod arm64;
|
||||
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
pub use arm64::has_fp16;
|
||||
use tract_itertools::Itertools;
|
||||
|
||||
#[cfg(not(target_arch = "aarch64"))]
|
||||
pub fn has_fp16() -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
#[cfg(any(target_arch = "arm", target_arch = "armv7", target_arch = "arm"))]
|
||||
pub mod arm32;
|
||||
|
||||
#[cfg(all(target_family = "wasm", target_feature = "simd128"))]
|
||||
pub mod wasm;
|
||||
|
||||
pub use self::frame::*;
|
||||
|
||||
use tract_data::prelude::*;
|
||||
|
||||
pub 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_impls: Vec<Box<dyn mmm::MatMatMul>>,
|
||||
panel_extractors: Vec<mmm::PanelExtractor>,
|
||||
|
||||
mmm_f64: MMMImpl,
|
||||
mmv_f64: MMVImpl,
|
||||
|
||||
mmm_f32: MMMImpl,
|
||||
mmv_f32: MMVImpl,
|
||||
|
||||
mmm_f16: MMMImpl,
|
||||
mmv_f16: MMVImpl,
|
||||
|
||||
qmmm_i32: MMMImpl,
|
||||
qmmv_i32: MMVImpl,
|
||||
|
||||
pub leaky_relu_f16: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f16, f16>> + Send + Sync>,
|
||||
pub leaky_relu_f32: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f32, f32>> + Send + Sync>,
|
||||
pub mul_by_scalar_f32:
|
||||
Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f32, f32>> + Send + Sync>,
|
||||
pub mul_by_scalar_f16:
|
||||
Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f16, f16>> + Send + Sync>,
|
||||
|
||||
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 hardswish_f16: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f16>> + Send + Sync>,
|
||||
pub hardswish_f32: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f32>> + Send + Sync>,
|
||||
pub silu_f16: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f16>> + Send + Sync>,
|
||||
pub silu_f32: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f32>> + Send + Sync>,
|
||||
pub gelu_f16: Box<dyn Fn() -> Box<dyn element_wise::ElementWise<f16>> + Send + Sync>,
|
||||
pub gelu_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>,
|
||||
|
||||
pub max_f16: Box<dyn Fn() -> Box<dyn reduce::Reduce<f16>> + Send + Sync>,
|
||||
pub max_f32: Box<dyn Fn() -> Box<dyn reduce::Reduce<f32>> + Send + Sync>,
|
||||
|
||||
pub sum_f16: Box<dyn Fn() -> Box<dyn reduce::Reduce<f16>> + Send + Sync>,
|
||||
pub sum_f32: Box<dyn Fn() -> Box<dyn reduce::Reduce<f32>> + Send + Sync>,
|
||||
|
||||
pub softmax2_fastcompact_f16:
|
||||
Box<dyn Fn() -> Box<dyn reduce::MapReduce<f16, f16>> + Send + Sync>,
|
||||
pub softmax2_fastcompact_f32:
|
||||
Box<dyn Fn() -> Box<dyn reduce::MapReduce<f32, f32>> + Send + Sync>,
|
||||
|
||||
/// Fused row-wise RmsNorm: out_i = x_i * rsqrt(mean(x_i²) + eps).
|
||||
/// Replaces a 4-call composition (MeanOfSquares + Add + Rsqrt + Mul) with
|
||||
/// a single 2-pass kernel. Called once per row by `core::ops::nn::RmsNorm`
|
||||
/// when the input is f32 and the axis is the last (contiguous) one.
|
||||
pub rms_norm_f32: Box<dyn Fn(&mut [f32], f32) + Send + Sync>,
|
||||
}
|
||||
|
||||
impl Ops {
|
||||
pub fn mmm_impls(&self) -> &[Box<dyn mmm::MatMatMul>] {
|
||||
&self.mmm_impls
|
||||
}
|
||||
|
||||
pub fn all_possible_packing(
|
||||
&self,
|
||||
weight_type: impl Into<WeightType>,
|
||||
) -> impl Iterator<Item = &dyn MMMInputFormat> {
|
||||
let weight_type = weight_type.into();
|
||||
self.mmm_impls
|
||||
.iter()
|
||||
.flat_map(|m| m.packings())
|
||||
.map(|p| &*p.0)
|
||||
.flat_map(move |p| {
|
||||
let mut packs: Vec<&dyn MMMInputFormat> = vec![];
|
||||
if p.precursor() == weight_type {
|
||||
packs.push(p)
|
||||
};
|
||||
for pe in &self.panel_extractors {
|
||||
if pe.from.precursor() == weight_type && pe.to.dyn_eq(p) {
|
||||
packs.push(&*pe.from);
|
||||
}
|
||||
}
|
||||
packs.into_iter()
|
||||
})
|
||||
.sorted_by_key(|p| p.to_string())
|
||||
.dedup()
|
||||
}
|
||||
|
||||
pub fn filter_impls<'o>(
|
||||
&'o self,
|
||||
weight: &'o dyn MMMInputFormat,
|
||||
acc: &[DatumType],
|
||||
act: DatumType,
|
||||
store: DatumType,
|
||||
) -> impl Iterator<
|
||||
Item = (
|
||||
&'o dyn MatMatMul,
|
||||
usize,
|
||||
&'o dyn MMMInputFormat,
|
||||
Option<&'o PanelExtractor>,
|
||||
&'o dyn MMMInputFormat,
|
||||
),
|
||||
> {
|
||||
let acc = acc.to_vec();
|
||||
self.mmm_impls
|
||||
.iter()
|
||||
.filter(move |mmm| acc.contains(&mmm.internal_type()) && mmm.stores().contains(&store))
|
||||
.flat_map(|mmm| {
|
||||
mmm.packings()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(pack_ix, (a, b))| (&**mmm, pack_ix, &**a, &**b))
|
||||
})
|
||||
.filter_map(|(mmm, ix, a, b)| {
|
||||
if a.dyn_eq(weight) {
|
||||
Some((mmm, ix, a, None, b))
|
||||
} else {
|
||||
self.panel_extractors
|
||||
.iter()
|
||||
.find(|pe| pe.from.dyn_eq(weight) && pe.to.dyn_eq(a))
|
||||
.map(|pe| (mmm, ix, a, Some(pe), b))
|
||||
}
|
||||
})
|
||||
.filter(move |(_mmm, _ix, _a, _pe, b)| {
|
||||
b.precursor().as_dt().is_some_and(|dt| dt == act)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn panel_extractors(&self) -> &[mmm::panel_extract::PanelExtractor] {
|
||||
&self.panel_extractors
|
||||
}
|
||||
|
||||
pub fn mmm(
|
||||
&self,
|
||||
accumulator: DatumType,
|
||||
m: Option<usize>,
|
||||
k: Option<usize>,
|
||||
n: Option<usize>,
|
||||
) -> Option<Box<dyn mmm::MatMatMul>> {
|
||||
use DatumType::*;
|
||||
match accumulator {
|
||||
F64 => Some(if n == Some(1) {
|
||||
(self.mmv_f64)(m, k)
|
||||
} else {
|
||||
(self.mmm_f64)(m, k, n)
|
||||
}),
|
||||
F32 => Some(if n == Some(1) {
|
||||
(self.mmv_f32)(m, k)
|
||||
} else {
|
||||
(self.mmm_f32)(m, k, n)
|
||||
}),
|
||||
F16 => Some(if n == Some(1) {
|
||||
(self.mmv_f16)(m, k)
|
||||
} else {
|
||||
(self.mmm_f16)(m, k, n)
|
||||
}),
|
||||
I32 => Some(if n == Some(1) {
|
||||
(self.qmmv_i32)(m, k)
|
||||
} else {
|
||||
(self.qmmm_i32)(m, k, n)
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn generic() -> Ops {
|
||||
use crate::generic::mmm::*;
|
||||
use element_wise::ElementWiseKer;
|
||||
use reduce::{MapReduceKer, ReduceKer};
|
||||
let mut ops = Ops {
|
||||
mmm_impls: vec![],
|
||||
panel_extractors: vec![],
|
||||
mmm_f64: Box::new(|_, _, _| generic_f64_4x4.mmm()),
|
||||
mmv_f64: Box::new(|_, _| generic_f64_4x1.mmm()),
|
||||
mmm_f32: Box::new(|_, _, _| generic_f32_4x4.mmm()),
|
||||
mmv_f32: Box::new(|_, _| generic_f32_4x1.mmm()),
|
||||
mmm_f16: Box::new(|_, _, _| generic_f16_4x4.mmm()),
|
||||
mmv_f16: Box::new(|_, _| generic_f16_4x1.mmm()),
|
||||
qmmm_i32: Box::new(|_, _, _| generic_i32_4x4.mmm()),
|
||||
qmmv_i32: Box::new(|_, _| generic_i32_4x4.mmm()),
|
||||
leaky_relu_f16: Box::new(|| generic::HLeakyRelu8::ew()),
|
||||
leaky_relu_f32: Box::new(|| generic::SLeakyRelu4::ew()),
|
||||
mul_by_scalar_f16: Box::new(|| generic::HMulByScalar8::ew()),
|
||||
mul_by_scalar_f32: Box::new(|| generic::SMulByScalar4::ew()),
|
||||
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()),
|
||||
hardswish_f16: Box::new(|| generic::HHardSwish8::ew()),
|
||||
hardswish_f32: Box::new(|| generic::SHardSwish4::ew()),
|
||||
silu_f16: Box::new(|| generic::HSiLU8::ew()),
|
||||
silu_f32: Box::new(|| generic::SSiLU4::ew()),
|
||||
gelu_f16: Box::new(|| generic::HGelu8::ew()),
|
||||
gelu_f32: Box::new(|| generic::SGelu4::ew()),
|
||||
lut_u8: Box::new(|table: &[u8]| Box::new(lut::LutImpl::<generic::GenericLut8>::new(table))),
|
||||
max_f16: Box::new(|| generic::reduce::max::HMax8::red()),
|
||||
max_f32: Box::new(|| generic::reduce::max::SMax4::red()),
|
||||
sum_f16: Box::new(|| generic::reduce::sum::HSum8::red()),
|
||||
sum_f32: Box::new(|| generic::reduce::sum::SSum4::red()),
|
||||
/*
|
||||
activation_f32: Box::new(|microcode| generic::SActivation::new(microcode))
|
||||
*/
|
||||
softmax2_fastcompact_f16: Box::new(|| generic::reduce::softmax_l2::HSoftMaxL2::red()),
|
||||
softmax2_fastcompact_f32: Box::new(|| generic::reduce::softmax_l2::SSoftMaxL2::red()),
|
||||
rms_norm_f32: Box::new(generic::rms_norm::rms_norm_f32),
|
||||
};
|
||||
crate::generic::mmm::plug(&mut ops);
|
||||
ops
|
||||
}
|
||||
|
||||
#[allow(unreachable_code, unused_mut, unexpected_cfgs)]
|
||||
pub fn best() -> Ops {
|
||||
let mut ops = generic();
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
x86_64_fma::plug(&mut ops);
|
||||
#[cfg(any(target_arch = "arm", target_arch = "armv7"))]
|
||||
arm32::plug(&mut ops);
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
arm64::plug(&mut ops);
|
||||
#[cfg(all(target_family = "wasm", target_feature = "simd128"))]
|
||||
wasm::plug(&mut ops);
|
||||
|
||||
ops
|
||||
}
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref OPS: Ops = {
|
||||
best()
|
||||
};
|
||||
}
|
||||
|
||||
#[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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn register_all_unicast(registry: &mut LinalgRegistry) {
|
||||
generic::register_all_unicast(registry);
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
arm64::register_all_unicast(registry);
|
||||
}
|
||||
|
||||
fn register_all_by_scalar(registry: &mut LinalgRegistry) {
|
||||
generic::register_all_by_scalar(registry);
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
arm64::register_all_by_scalar(registry);
|
||||
}
|
||||
|
||||
pub type LinalgFn = dyn Fn(&mut TensorView, &TensorView) -> TractResult<()> + Send + Sync;
|
||||
type LinalgRegistry = HashMap<(BinOp, DatumType), Box<dyn Fn() -> Box<LinalgFn> + Send + Sync>>;
|
||||
lazy_static! {
|
||||
static ref BIN_UNICAST_OPS: Mutex<LinalgRegistry> = {
|
||||
let mut registry = HashMap::default();
|
||||
register_all_unicast(&mut registry);
|
||||
Mutex::new(registry)
|
||||
};
|
||||
static ref BIN_BY_SCALAR_OPS: Mutex<LinalgRegistry> = {
|
||||
let mut registry = HashMap::default();
|
||||
register_all_by_scalar(&mut registry);
|
||||
Mutex::new(registry)
|
||||
};
|
||||
}
|
||||
|
||||
pub fn bin_by_scalar(dt: DatumType, bin: BinOp) -> Option<Box<LinalgFn>> {
|
||||
let map = BIN_BY_SCALAR_OPS.lock().unwrap();
|
||||
if (dt == DatumType::F16) && !has_fp16() {
|
||||
return None;
|
||||
}
|
||||
map.get(&(bin, dt)).map(|it| (it)())
|
||||
}
|
||||
|
||||
pub fn bin_unicast(dt: DatumType, bin: BinOp) -> Option<Box<LinalgFn>> {
|
||||
let map = BIN_UNICAST_OPS.lock().unwrap();
|
||||
if (dt == DatumType::F16) && !has_fp16() {
|
||||
return None;
|
||||
}
|
||||
map.get(&(bin, dt)).map(|it| (it)())
|
||||
}
|
||||
|
||||
pub fn ops() -> &'static Ops {
|
||||
&OPS
|
||||
}
|
||||
|
||||
use dyn_eq::DynEq;
|
||||
use num_traits::*;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Debug;
|
||||
use std::ops::*;
|
||||
use std::sync::Mutex;
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(dead_code)]
|
||||
fn setup_test_logger() {
|
||||
let _ = env_logger::Builder::from_env("TRACT_LOG").try_init();
|
||||
}
|
||||
+93
@@ -0,0 +1,93 @@
|
||||
use std::cell::RefCell;
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
#[allow(unused_imports)]
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
use rayon::{ThreadPool, ThreadPoolBuilder};
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub enum Executor {
|
||||
#[default]
|
||||
SingleThread,
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
MultiThread(Arc<ThreadPool>),
|
||||
/// Use rayon's GLOBAL thread pool — the one set up by
|
||||
/// `wasm_bindgen_rayon::init_thread_pool` on `wasm32-unknown-unknown`,
|
||||
/// or rayon's auto-initialised default on native.
|
||||
///
|
||||
/// Exists because `Arc<rayon::ThreadPool>` cannot be constructed on
|
||||
/// `wasm32-unknown-unknown`: rayon's default `spawn_handler` calls
|
||||
/// `std::thread::spawn`, which is unsupported there. The only working
|
||||
/// route is rayon's global pool, accessed via `into_par_iter` directly.
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
RayonGlobal,
|
||||
}
|
||||
|
||||
impl Executor {
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
pub fn multithread(n: usize) -> Executor {
|
||||
Executor::multithread_with_name(n, "tract-default")
|
||||
}
|
||||
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
pub fn multithread_with_name(n: usize, name: &str) -> Executor {
|
||||
let name = name.to_string();
|
||||
let pool = ThreadPoolBuilder::new()
|
||||
.thread_name(move |n| format!("{name}-{n}"))
|
||||
.num_threads(n)
|
||||
.build()
|
||||
.unwrap();
|
||||
Executor::MultiThread(Arc::new(pool))
|
||||
}
|
||||
}
|
||||
|
||||
static DEFAULT_EXECUTOR: Mutex<Executor> = Mutex::new(Executor::SingleThread);
|
||||
|
||||
thread_local! {
|
||||
static TLS_EXECUTOR_OVERRIDE: RefCell<Option<Executor>> = Default::default();
|
||||
}
|
||||
|
||||
pub fn current_tract_executor() -> Executor {
|
||||
if let Some(over_ride) = TLS_EXECUTOR_OVERRIDE.with_borrow(|tls| tls.clone()) {
|
||||
over_ride
|
||||
} else {
|
||||
DEFAULT_EXECUTOR.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_default_executor(executor: Executor) {
|
||||
*DEFAULT_EXECUTOR.lock().unwrap() = executor;
|
||||
}
|
||||
|
||||
pub fn multithread_tract_scope<R, F: FnOnce() -> R>(pool: Executor, f: F) -> R {
|
||||
let previous = TLS_EXECUTOR_OVERRIDE.replace(Some(pool));
|
||||
let result = f();
|
||||
TLS_EXECUTOR_OVERRIDE.set(previous);
|
||||
result
|
||||
}
|
||||
|
||||
/// Threshold (in panels) below which the rayon MMM dispatcher skips
|
||||
/// parallelism and runs inline single-threaded. Below this size,
|
||||
/// per-call dispatch overhead (~5 µs native, ~50 µs wasm-bindgen-rayon
|
||||
/// worker) exceeds the parallel speedup.
|
||||
///
|
||||
/// Default `64`. Tune higher for many-small-MMM workloads (mobile vision,
|
||||
/// streaming RNN) or lower for transformer-class workloads where every MMM
|
||||
/// is large. `0` disables the gate entirely (always thread).
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
static THREADING_PANEL_THRESHOLD: AtomicUsize = AtomicUsize::new(64);
|
||||
|
||||
/// Read the current MMM panel-count threshold for the rayon path.
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
pub fn current_threading_panel_threshold() -> usize {
|
||||
THREADING_PANEL_THRESHOLD.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Set the MMM panel-count threshold for the rayon path. Default is `64`.
|
||||
/// Pass `0` to thread regardless of size.
|
||||
#[cfg(feature = "multithread-mm")]
|
||||
pub fn set_threading_panel_threshold(panels: usize) {
|
||||
THREADING_PANEL_THRESHOLD.store(panels, Ordering::Relaxed);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
+128
@@ -0,0 +1,128 @@
|
||||
use crate::Ops;
|
||||
use crate::frame::element_wise::ElementWiseKer;
|
||||
use crate::frame::reduce::{MapReduceKer, ReduceKer};
|
||||
use crate::x86_64_fma::softmax::x86_64_avx512_softmax2_fastcompact_f16_64n;
|
||||
use crate::x86_64_fma::softmax::x86_64_fma_softmax2_fastcompact_f32_32n;
|
||||
|
||||
pub mod mmm;
|
||||
|
||||
pub mod act;
|
||||
pub mod act_f16;
|
||||
pub mod act_f16_fp16;
|
||||
pub mod by_scalar;
|
||||
pub mod erf;
|
||||
mod intel;
|
||||
pub mod max;
|
||||
pub mod panel_extract;
|
||||
pub mod rms_norm;
|
||||
pub mod softmax;
|
||||
|
||||
const AVX2: fn() -> bool = || is_x86_feature_detected!("avx2");
|
||||
const FMA: fn() -> bool = || is_x86_feature_detected!("fma");
|
||||
const AVX512F: fn() -> bool = || is_x86_feature_detected!("avx512f");
|
||||
#[cfg(tract_avx512vnni)]
|
||||
const AVX512VNNI: fn() -> bool = || is_x86_feature_detected!("avx512vnni");
|
||||
|
||||
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"));
|
||||
|
||||
// AVX-512 (zmm, 16-wide) variants. The assembly lives in x86_64/avx512/; the
|
||||
// main loop handles 64 lanes (4 zmm) per iteration with a 16-lane tail, so
|
||||
// nr()=16 (any multiple of 16 is safe).
|
||||
tanh_impl!(
|
||||
f32,
|
||||
avx512_tanh_f32,
|
||||
16,
|
||||
16,
|
||||
is_x86_feature_detected!("avx512f")
|
||||
);
|
||||
sigmoid_impl!(
|
||||
f32,
|
||||
avx512_sigmoid_f32,
|
||||
16,
|
||||
16,
|
||||
is_x86_feature_detected!("avx512f")
|
||||
);
|
||||
|
||||
fn plug_avx2(_ops: &mut Ops) {}
|
||||
|
||||
fn plug_fma(ops: &mut Ops) {
|
||||
panel_extract::plug(ops);
|
||||
|
||||
ops.sigmoid_f32 = Box::new(|| fma_sigmoid_f32::ew());
|
||||
ops.tanh_f32 = Box::new(|| fma_tanh_f32::ew());
|
||||
|
||||
ops.mul_by_scalar_f32 = Box::new(|| by_scalar::x86_64_avx_f32_mul_by_scalar_32n::ew());
|
||||
ops.max_f32 = Box::new(|| max::x86_64_fma_max_f32_32n::red());
|
||||
ops.softmax2_fastcompact_f32 = Box::new(|| x86_64_fma_softmax2_fastcompact_f32_32n::red());
|
||||
|
||||
log::info!("sigmoid_f32, tanh_f32: x86_64/fma activated");
|
||||
}
|
||||
|
||||
/// On hosts that also support AVX-512_FP16 (Sapphire Rapids / Granite Rapids /
|
||||
/// later, and recent Xeon-D / consumer parts), upgrade the f16 element-wise
|
||||
/// kernels from the f32-roundtrip implementations in `act_f16.rs` to the
|
||||
/// native f16 implementations in `act_f16_fp16.rs` where the native path is
|
||||
/// actually faster on this uarch. We benched each op against its f32-roundtrip
|
||||
/// equivalent on Sapphire Rapids and only plug in the ones that win:
|
||||
///
|
||||
/// hardswish_f16: 8.71 → 31.6 Gelem/s (3.62× native) — plug in
|
||||
/// leaky_relu_f16: 9.44 → 5.85 Gelem/s (0.62× native — regression) — keep
|
||||
/// the f32-roundtrip version from act_f16.rs. The native
|
||||
/// kernel exists in act_f16_fp16.rs for future revisits but
|
||||
/// is not wired here.
|
||||
fn plug_avx512fp16(ops: &mut Ops) {
|
||||
ops.hardswish_f16 = Box::new(|| act_f16_fp16::x86_64_avx512fp16_hardswish_f16_128n::ew());
|
||||
|
||||
log::info!("hardswish_f16: x86_64/avx512fp16 native activated");
|
||||
}
|
||||
|
||||
fn plug_avx512f(ops: &mut Ops) {
|
||||
ops.sigmoid_f32 = Box::new(|| avx512_sigmoid_f32::ew());
|
||||
ops.tanh_f32 = Box::new(|| avx512_tanh_f32::ew());
|
||||
ops.hardswish_f32 = Box::new(|| act::x86_64_avx512_hardswish_f32_64n::ew());
|
||||
ops.leaky_relu_f32 = Box::new(|| act::x86_64_avx512_leaky_relu_f32_64n::ew());
|
||||
ops.silu_f32 = Box::new(|| act::x86_64_avx512_silu_f32_16n::ew());
|
||||
ops.gelu_f32 = Box::new(|| act::x86_64_avx512_gelu_f32_16n::ew());
|
||||
|
||||
ops.sigmoid_f16 = Box::new(|| act_f16::x86_64_avx512_sigmoid_f16_16n::ew());
|
||||
ops.tanh_f16 = Box::new(|| act_f16::x86_64_avx512_tanh_f16_16n::ew());
|
||||
ops.hardswish_f16 = Box::new(|| act_f16::x86_64_avx512_hardswish_f16_64n::ew());
|
||||
ops.leaky_relu_f16 = Box::new(|| act_f16::x86_64_avx512_leaky_relu_f16_64n::ew());
|
||||
ops.silu_f16 = Box::new(|| act_f16::x86_64_avx512_silu_f16_16n::ew());
|
||||
ops.gelu_f16 = Box::new(|| act_f16::x86_64_avx512_gelu_f16_16n::ew());
|
||||
|
||||
ops.max_f32 = Box::new(|| max::x86_64_avx512_max_f32_64n::red());
|
||||
ops.softmax2_fastcompact_f32 =
|
||||
Box::new(|| softmax::x86_64_avx512_softmax2_fastcompact_f32_64n::red());
|
||||
ops.softmax2_fastcompact_f16 = Box::new(|| x86_64_avx512_softmax2_fastcompact_f16_64n::red());
|
||||
|
||||
ops.erf_f32 = Box::new(|| erf::x86_64_avx512_erf_f32_64n::ew());
|
||||
|
||||
ops.rms_norm_f32 = Box::new(rms_norm::rms_norm_f32);
|
||||
|
||||
log::info!(
|
||||
"sigmoid_f32, tanh_f32, hardswish_f32, leaky_relu_f32, \
|
||||
silu_f32, gelu_f32, \
|
||||
sigmoid_f16, tanh_f16, hardswish_f16, leaky_relu_f16, \
|
||||
silu_f16, gelu_f16, \
|
||||
max_f32, softmax2_fastcompact_f32, softmax2_fastcompact_f16, erf_f32, \
|
||||
rms_norm_f32: x86_64/avx512f activated"
|
||||
);
|
||||
}
|
||||
|
||||
pub fn plug(ops: &mut Ops) {
|
||||
mmm::plug(ops);
|
||||
if is_x86_feature_detected!("avx2") {
|
||||
plug_avx2(ops);
|
||||
if is_x86_feature_detected!("fma") {
|
||||
plug_fma(ops);
|
||||
if is_x86_feature_detected!("avx512f") {
|
||||
plug_avx512f(ops);
|
||||
if is_x86_feature_detected!("avx512fp16") {
|
||||
plug_avx512fp16(ops);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+266
@@ -0,0 +1,266 @@
|
||||
// AVX-512 (zmm, 16-wide) element-wise activation kernels with no FMA
|
||||
// predecessor on x86: hardswish and leaky_relu. They mirror the aarch64 NEON
|
||||
// kernels (arm64simd_hardswish_f32_8n / arm64simd_leaky_relu_f32_8n) but use
|
||||
// 512-bit zmm registers, processing 64 f32 lanes per iteration. Validated
|
||||
// against the generic scalar reference via the *_frame_tests! macros.
|
||||
|
||||
// hardswish(x) = x * relu6(x + 3) / 6
|
||||
// = x * max(0, min(6, x + 3)) * (1/6)
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
x86_64_avx512_hardswish_f32_64n,
|
||||
64,
|
||||
16,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
debug_assert!(buf.len() % Self::nr() == 0);
|
||||
debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
if buf.is_empty() {
|
||||
return;
|
||||
}
|
||||
unsafe { x86_64_avx512_hardswish_f32_64n_run(buf) }
|
||||
}
|
||||
);
|
||||
|
||||
#[target_feature(enable = "avx512f")]
|
||||
unsafe fn x86_64_avx512_hardswish_f32_64n_run(buf: &mut [f32]) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
vbroadcastss zmm0, xmm0 // 3.0
|
||||
vbroadcastss zmm1, xmm1 // 6.0
|
||||
vbroadcastss zmm2, xmm2 // 1/6
|
||||
vpxord zmm3, zmm3, zmm3 // 0.0
|
||||
2:
|
||||
vmovaps zmm4, [{ptr}]
|
||||
vmovaps zmm5, [{ptr} + 64]
|
||||
vmovaps zmm6, [{ptr} + 128]
|
||||
vmovaps zmm7, [{ptr} + 192]
|
||||
|
||||
vaddps zmm8, zmm4, zmm0
|
||||
vaddps zmm9, zmm5, zmm0
|
||||
vaddps zmm10, zmm6, zmm0
|
||||
vaddps zmm11, zmm7, zmm0
|
||||
|
||||
vminps zmm8, zmm8, zmm1
|
||||
vminps zmm9, zmm9, zmm1
|
||||
vminps zmm10, zmm10, zmm1
|
||||
vminps zmm11, zmm11, zmm1
|
||||
|
||||
vmaxps zmm8, zmm8, zmm3
|
||||
vmaxps zmm9, zmm9, zmm3
|
||||
vmaxps zmm10, zmm10, zmm3
|
||||
vmaxps zmm11, zmm11, zmm3
|
||||
|
||||
vmulps zmm8, zmm8, zmm4
|
||||
vmulps zmm9, zmm9, zmm5
|
||||
vmulps zmm10, zmm10, zmm6
|
||||
vmulps zmm11, zmm11, zmm7
|
||||
|
||||
vmulps zmm8, zmm8, zmm2
|
||||
vmulps zmm9, zmm9, zmm2
|
||||
vmulps zmm10, zmm10, zmm2
|
||||
vmulps zmm11, zmm11, zmm2
|
||||
|
||||
vmovaps [{ptr}], zmm8
|
||||
vmovaps [{ptr} + 64], zmm9
|
||||
vmovaps [{ptr} + 128], zmm10
|
||||
vmovaps [{ptr} + 192], zmm11
|
||||
|
||||
add {ptr}, 256
|
||||
sub {len}, 64
|
||||
jnz 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
inout("xmm0") 3.0f32 => _,
|
||||
inout("xmm1") 6.0f32 => _,
|
||||
inout("xmm2") 1.0f32 / 6.0f32 => _,
|
||||
out("zmm3") _,
|
||||
out("zmm4") _, out("zmm5") _, out("zmm6") _, out("zmm7") _,
|
||||
out("zmm8") _, out("zmm9") _, out("zmm10") _, out("zmm11") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_x86_64_avx512_hardswish_f32_64n {
|
||||
use super::*;
|
||||
hardswish_frame_tests!(
|
||||
is_x86_feature_detected!("avx512f"),
|
||||
f32,
|
||||
x86_64_avx512_hardswish_f32_64n
|
||||
);
|
||||
}
|
||||
|
||||
// leaky_relu(x) = x > 0 ? x : alpha * x
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
x86_64_avx512_leaky_relu_f32_64n,
|
||||
64,
|
||||
16,
|
||||
f32,
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], alpha: f32) {
|
||||
debug_assert!(buf.len() % Self::nr() == 0);
|
||||
debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
if buf.is_empty() {
|
||||
return;
|
||||
}
|
||||
unsafe { x86_64_avx512_leaky_relu_f32_64n_run(buf, alpha) }
|
||||
}
|
||||
);
|
||||
|
||||
#[target_feature(enable = "avx512f")]
|
||||
unsafe fn x86_64_avx512_leaky_relu_f32_64n_run(buf: &mut [f32], alpha: f32) {
|
||||
unsafe {
|
||||
let len = buf.len();
|
||||
let ptr = buf.as_ptr();
|
||||
std::arch::asm!("
|
||||
vbroadcastss zmm0, xmm0 // alpha
|
||||
vpxord zmm1, zmm1, zmm1 // 0.0
|
||||
2:
|
||||
vmovaps zmm4, [{ptr}]
|
||||
vmovaps zmm5, [{ptr} + 64]
|
||||
vmovaps zmm6, [{ptr} + 128]
|
||||
vmovaps zmm7, [{ptr} + 192]
|
||||
|
||||
// alpha * x in zmm8..11
|
||||
vmulps zmm8, zmm4, zmm0
|
||||
vmulps zmm9, zmm5, zmm0
|
||||
vmulps zmm10, zmm6, zmm0
|
||||
vmulps zmm11, zmm7, zmm0
|
||||
|
||||
// mask = x > 0
|
||||
vcmpps k1, zmm4, zmm1, 14
|
||||
vcmpps k2, zmm5, zmm1, 14
|
||||
vcmpps k3, zmm6, zmm1, 14
|
||||
vcmpps k4, zmm7, zmm1, 14
|
||||
|
||||
// where x > 0, overwrite alpha*x with x
|
||||
vmovaps zmm8{{k1}}, zmm4
|
||||
vmovaps zmm9{{k2}}, zmm5
|
||||
vmovaps zmm10{{k3}}, zmm6
|
||||
vmovaps zmm11{{k4}}, zmm7
|
||||
|
||||
vmovaps [{ptr}], zmm8
|
||||
vmovaps [{ptr} + 64], zmm9
|
||||
vmovaps [{ptr} + 128], zmm10
|
||||
vmovaps [{ptr} + 192], zmm11
|
||||
|
||||
add {ptr}, 256
|
||||
sub {len}, 64
|
||||
jnz 2b
|
||||
",
|
||||
len = inout(reg) len => _,
|
||||
ptr = inout(reg) ptr => _,
|
||||
inout("xmm0") alpha => _,
|
||||
out("zmm1") _,
|
||||
out("zmm4") _, out("zmm5") _, out("zmm6") _, out("zmm7") _,
|
||||
out("zmm8") _, out("zmm9") _, out("zmm10") _, out("zmm11") _,
|
||||
out("k1") _, out("k2") _, out("k3") _, out("k4") _,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_x86_64_avx512_leaky_relu_f32_64n {
|
||||
use super::*;
|
||||
leaky_relu_frame_tests!(
|
||||
is_x86_feature_detected!("avx512f"),
|
||||
f32,
|
||||
x86_64_avx512_leaky_relu_f32_64n
|
||||
);
|
||||
}
|
||||
|
||||
// SiLU(x) = x * sigmoid(x). Composed at the kernel level (mirrors arm64): save
|
||||
// the input chunk, run the AVX-512 sigmoid kernel in place, then multiply back
|
||||
// by the saved original. nr() and CHUNK (256) are multiples of 16 so the
|
||||
// sigmoid kernel always receives a 64-byte-aligned slice whose length is a
|
||||
// multiple of 16.
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
x86_64_avx512_silu_f32_16n,
|
||||
16,
|
||||
16,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
debug_assert!(buf.len() % Self::nr() == 0);
|
||||
debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
const CHUNK: usize = 256;
|
||||
let mut scratch = [0f32; CHUNK];
|
||||
let mut start = 0;
|
||||
while start < buf.len() {
|
||||
let end = (start + CHUNK).min(buf.len());
|
||||
let chunk = &mut buf[start..end];
|
||||
let n = chunk.len();
|
||||
scratch[..n].copy_from_slice(chunk);
|
||||
super::avx512_sigmoid_f32::run(chunk, ());
|
||||
for i in 0..n {
|
||||
chunk[i] *= scratch[i];
|
||||
}
|
||||
start = end;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_x86_64_avx512_silu_f32_16n {
|
||||
use super::*;
|
||||
silu_frame_tests!(
|
||||
is_x86_feature_detected!("avx512f"),
|
||||
f32,
|
||||
x86_64_avx512_silu_f32_16n
|
||||
);
|
||||
}
|
||||
|
||||
// Tanh-form GELU (pow=3) matching tract's GeluApproximate:
|
||||
// gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
|
||||
// Composed at the kernel level (mirrors arm64): save the original x, compute
|
||||
// the tanh argument in place, run the AVX-512 tanh kernel, then finish with the
|
||||
// 0.5 * x * (1 + tanh) combine.
|
||||
ew_impl_wrap!(
|
||||
f32,
|
||||
x86_64_avx512_gelu_f32_16n,
|
||||
16,
|
||||
16,
|
||||
(),
|
||||
#[inline(never)]
|
||||
fn run(buf: &mut [f32], _: ()) {
|
||||
debug_assert!(buf.len() % Self::nr() == 0);
|
||||
debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0);
|
||||
const SQRT_2_OVER_PI: f32 = 0.7978845608028654;
|
||||
const COEF: f32 = 0.044715;
|
||||
const CHUNK: usize = 256;
|
||||
let mut scratch = [0f32; CHUNK];
|
||||
let mut start = 0;
|
||||
while start < buf.len() {
|
||||
let end = (start + CHUNK).min(buf.len());
|
||||
let chunk = &mut buf[start..end];
|
||||
let n = chunk.len();
|
||||
for i in 0..n {
|
||||
let x = chunk[i];
|
||||
scratch[i] = x;
|
||||
chunk[i] = SQRT_2_OVER_PI * (x + COEF * x * x * x);
|
||||
}
|
||||
super::avx512_tanh_f32::run(chunk, ());
|
||||
for i in 0..n {
|
||||
chunk[i] = 0.5 * scratch[i] * (1.0 + chunk[i]);
|
||||
}
|
||||
start = end;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test_x86_64_avx512_gelu_f32_16n {
|
||||
use super::*;
|
||||
gelu_frame_tests!(
|
||||
is_x86_feature_detected!("avx512f"),
|
||||
f32,
|
||||
x86_64_avx512_gelu_f32_16n
|
||||
);
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user