Create by_scalar & unicast registries in linalg

core/src/ops/math/mod.rs

@@ -112,11 +112,15 @@ bin_to_super_type!(mul, Mul,
  }
  },
  eval_by_scalar: |a: &mut TensorView, b: &TensorView | -> TractResult<bool> {
- let res = tract_linalg::bin_by_scalar(tract_linalg::BinOp::Mul)(a, b).is_ok();
+ let res = tract_linalg::bin_by_scalar(a.datum_type(), tract_linalg::BinOp::Mul)
+ .context("unimplemented mul by scalar")?(a, b)
+ .is_ok();
  Ok(res)
  },
  eval_unicast: |a: &mut TensorView, b: &TensorView | -> TractResult<bool> {
- let res = tract_linalg::bin_unicast(tract_linalg::BinOp::Mul)(a, b).is_ok();
+ let res = tract_linalg::bin_unicast(a.datum_type(), tract_linalg::BinOp::Mul)
+ .context("unimplemented mul unicast")?(a, b)
+ .is_ok();
  Ok(res)
  },
  neutral_element: 1,

linalg/src/arm64.rs

@@ -14,9 +14,10 @@ mod arm64fp16;
 pub use arm64fp16::*;
 
 use crate::f16;
-use crate::Ops;
+use crate::{Ops, LinalgRegistry, DatumType, BinOp};
 
 use crate::frame::unicast::UnicastKer;
+use crate::frame::by_scalar::ByScalarKer;
 use crate::frame::element_wise::ElementWiseKer;
 use crate::frame::reduce::{MapReduceKer, ReduceKer};
 
@@ -215,6 +216,16 @@ impl Kind {
  }
 }
 
+pub(crate)fn register_all_unicast(registry: &mut LinalgRegistry) {
+ registry.insert((BinOp::Mul, DatumType::F32),Box::new(|| arm64simd_unicast_mul_f32_16n::bin_1()));
+ registry.insert((BinOp::Mul, DatumType::F16),Box::new(|| arm64fp16_unicast_mul_f16_32n::bin_1()));
+}
+
+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_1()));
+ registry.insert((BinOp::Mul, DatumType::F16),Box::new(|| arm64fp16_mul_by_scalar_f16_32n::bin_1()));
+}
+
 pub fn plug(ops: &mut Ops) {
  ops.mmm_impls.extend(
  [

linalg/src/arm64/arm64fp16/by_scalar.rs

@@ -1,6 +1,6 @@
 use crate::f16;
 
-ew_impl_wrap!(
+by_scalar_impl_wrap!(
  f16,
  arm64fp16_mul_by_scalar_f16_32n,
  32,

linalg/src/arm64/arm64simd/by_scalar.rs

@@ -1,4 +1,4 @@
-ew_impl_wrap!(
+by_scalar_impl_wrap!(
  f32,
  arm64simd_mul_by_scalar_f32_16n,
  16,

linalg/src/frame/by_scalar.rs

@@ -1,3 +1,61 @@
+use std::{fmt::Debug, marker::PhantomData};
+
+use tract_data::{TractResult, internal::TensorView};
+
+use crate::{LADatum, element_wise::ElementWiseKer};
+
+use super::{ElementWise, element_wise_helper::map_slice_with_alignment};
+
+
+/// 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_1() -> Box<dyn Fn(&mut TensorView, &TensorView) -> TractResult<()>> {
+ 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 {

linalg/src/frame/unicast/mod.rs

@@ -4,6 +4,7 @@ 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;
@@ -53,6 +54,13 @@ where
  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,
@@ -80,6 +88,13 @@ where
  fn bin() -> Box<dyn Unicast<T>> {
  Box::new(UnicastImpl::<Self, T>::new())
  }
+ fn bin_1() -> Box<dyn Fn(&mut TensorView, &TensorView) -> TractResult<()>> {
+ Box::new(|a: &mut TensorView, b: &TensorView| {
+ let a_slice = a.as_slice_mut()?;
+ let b_slice = b.as_slice()?;
+ (Self::bin()).run(a_slice, b_slice)
+ })
+ }
 }
 
 std::thread_local! {

linalg/src/generic.rs

@@ -9,6 +9,10 @@ pub mod sigmoid;
 pub mod tanh;
 pub mod unicast;
 
+use tract_data::prelude::DatumType;
+
+use crate::{LinalgRegistry, BinOp, UnicastKer, ByScalarKer};
+
 pub use self::by_scalar::{HMulByScalar8, SMulByScalar4};
 pub use self::erf::SErf4;
 pub use self::leaky_relu::{HLeakyRelu8, SLeakyRelu4};
@@ -17,3 +21,13 @@ pub use self::rounding::{ScaleShiftAndRound, Scaler};
 pub use self::sigmoid::{HSigmoid8, SSigmoid4};
 pub use self::reduce::softmax_l2::SSoftMaxL2;
 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_1()));
+ registry.insert((BinOp::Mul, DatumType::F16),Box::new(|| unicast::HUnicastMul8::bin_1()));
+}
+
+pub(crate)fn register_all_by_scalar(registry: &mut LinalgRegistry) {
+ registry.insert((BinOp::Mul, DatumType::F32),Box::new(|| by_scalar::SMulByScalar4::bin_1()));
+ registry.insert((BinOp::Mul, DatumType::F16),Box::new(|| by_scalar::HMulByScalar8::bin_1()));
+}

linalg/src/generic/by_scalar.rs

@@ -1,58 +1,36 @@
 use tract_data::internal::f16;
 
-use crate::element_wise::ElementWiseKer;
-
-#[derive(Clone, Debug)]
-pub struct SMulByScalar4;
-
-impl ElementWiseKer<f32, f32> for SMulByScalar4 {
- fn name() -> &'static str {
- "generic"
- }
-
- fn alignment_items() -> usize {
- 4
- }
-
- fn nr() -> usize {
- 4
- }
-
+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)
  }
-}
+);
 
 #[cfg(test)]
 #[macro_use]
 pub mod mul_by_scalar_f32 {
  mul_by_scalar_frame_tests!(true, f32, crate::generic::by_scalar::SMulByScalar4);
 }
 
-#[derive(Clone, Debug)]
-pub struct HMulByScalar8;
-
-impl ElementWiseKer<f16, f16> for HMulByScalar8 {
- fn name() -> &'static str {
- "generic"
- }
-
- fn alignment_items() -> usize {
- 8
- }
-
- fn nr() -> usize {
- 8
- }
-
+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)
  }
-}
+);
 
 #[cfg(test)]
 #[macro_use]