Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Comp ops #1508

Merged
merged 7 commits into from
Aug 29, 2024
Merged

Comp ops #1508

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
acaf25d
refactor comparison ops
kali Aug 26, 2024
cdabc51
fix TDim::max
kali Aug 28, 2024
a0bf24a
missing file
kali Aug 28, 2024
b045827
proper error for gather input ranks
kali Aug 28, 2024
f62a5a3
clip
kali Aug 28, 2024
48ecf90
cleanup
kali Aug 29, 2024
b647a66
clip
kali Aug 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions core/src/ops/array/gather.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ impl Gather {
input_shape: &[D],
indices_shape: &[D],
) -> TractResult<TVec<D>> {
ensure!(input_shape.len() > self.axis);
let mut output_shape: TVec<D> = input_shape[..self.axis].into();
output_shape.extend(indices_shape.iter().cloned());
output_shape.extend(input_shape[self.axis + 1..].iter().cloned());
Expand Down
118 changes: 0 additions & 118 deletions core/src/ops/binary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -594,124 +594,6 @@ macro_rules! bin_to_super_type {
};
}

macro_rules! bin_to_bool {
($func:ident, $Op:ident,
$( codegen: $codegen:expr, )?
$( cost: $cost:expr, )?
$( declutter: $declutter:expr, )?
$( operating_datum_type: $operating_datum_type:expr, )?
$( [$($typ:ident),*] => $cab:expr),*) => {
#[derive(Debug, Clone, Hash)]
pub struct $Op;
impl $crate::ops::binary::BinMiniOp for $Op {
fn name(&self) -> &'static str {
stringify!($Op)
}

fn eval_uniform_in_place(&self, a: &Tensor, b: &mut Tensor) -> TractResult<()> {
$(
$(if a.datum_type() == $typ::datum_type() {
let cab: fn(&mut bool, &bool, &bool) -> () = $cab;
let a = &a.as_slice::<bool>()?[0];
let b = b.as_slice_mut::<bool>()?;
unsafe {
for i in 0..b.len() {
let mut c = bool::default();
cab(&mut c, a, b.get_unchecked(i));
*b.get_unchecked_mut(i) = c;
}
}
return Ok(())
}
)*
)*
bail!("{} does not support {:?} (inplace uniform)", self.name(), a.datum_type());
}

#[allow(unreachable_code)]
fn eval_unicast_in_place(&self, a: &Tensor, b: &mut Tensor) -> TractResult<()> {
$(
$(if a.datum_type() == $typ::datum_type() {
let cab: fn(&mut bool, &bool, &bool) -> () = $cab;
let a = a.as_slice::<bool>()?;
let b = b.as_slice_mut::<bool>()?;
unsafe {
for i in 0..a.len() {
let mut c = bool::default();
cab(&mut c, a.get_unchecked(i), b.get_unchecked(i));
*b.get_unchecked_mut(i) = c;
}
}
return Ok(())
}
)*
)*
bail!("{} does not support {:?}", self.name(), a.datum_type());
}

fn eval_out_of_place(&self, c: &mut Tensor, a: &Tensor, b: &Tensor) -> TractResult<()> {
$(
$(if a.datum_type() == $typ::datum_type() {
let cab: fn(&mut bool, &$typ, &$typ) -> () = $cab;
let a = a.to_array_view::<$typ>()?;
let b = b.to_array_view::<$typ>()?;
let mut c = c.to_array_view_mut::<bool>()?;
ndarray::Zip::from(&mut c).and_broadcast(a).and_broadcast(b).for_each(cab);
return Ok(())
}
)*
)*
bail!("{} does not support {:?}", self.name(), a.datum_type());
}

fn eval_in_a(&self, a: &mut Tensor, _b: &Tensor) -> TractResult<()> {
bail!("{} does not support {:?}", self.name(), a.datum_type());
}

fn result_datum_type(&self, _a: DatumType, _b: DatumType) -> TractResult<DatumType> {
Ok(bool::datum_type())
}

$(
fn codegen(
&self,
model: &TypedModel,
node: &TypedNode,
) -> TractResult<Option<TypedModelPatch>> {
($codegen)(self, model, node)
}
)?


$(
fn declutter(
&self,
model: &TypedModel,
node: &TypedNode,
) -> TractResult<Option<TypedModelPatch>> {
($declutter)(self, model, node)
}
)?

$(
fn cost_per_element(&self, dt: DatumType) -> TVec<(Cost, usize)> {
($cost)(dt)
}
)?

$(
fn operating_datum_type(&self, a: DatumType, b: DatumType) -> TractResult<DatumType> {
($operating_datum_type)(a, b)
})?

}

pub fn $func() -> $crate::ops::binary::TypedBinOp {
$crate::ops::binary::TypedBinOp(Box::new($Op), None)
}
};
}

#[derive(Debug)]
pub(crate) struct OneUniformInput {
pub uni: Arc<Tensor>,
Expand Down
94 changes: 2 additions & 92 deletions core/src/ops/logic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,111 +1,21 @@
#![allow(clippy::bool_comparison)]
#![allow(clippy::unnecessary_cast)]

mod comparison;
mod ite;
pub use ite::IfThenElse;
pub use comparison::Comp;

use ndarray::*;

use crate::broadcast::multi_broadcast;
use crate::internal::*;

use super::binary::BinMiniOp;
use super::element_wise::ElementWiseOp;

bin_to_super_type!(and, And,
[bool, u8, u16, u32, u64, i8, i16, i32, i64] => |c, &a, &b| *c = (a as i64 != 0 && b as i64 != 0) as _);
bin_to_super_type!(or, Or,
[bool, u8, u16, u32, u64, i8, i16, i32, i64] => |c, &a, &b| *c = (a as i64 != 0 || b as i64 != 0) as _);
bin_to_super_type!(xor, Xor, /*flip: commute, */ [bool] => |c, &a, &b| *c = a ^ b);
bin_to_bool!(equals, Equals,
[bool, u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64, TDim] => |c, a, b | *c = a == b
);
bin_to_bool!(not_equals, NotEquals, /* flip: commute, */
[bool, u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64, TDim] => |c, a, b | *c = a != b
);

bin_to_bool!(less, Less,
codegen: codegen_compare_to_zero,
operating_datum_type: operating_datum_type_for_cmp,
[bool, u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64] => |c, &a, &b | *c = a < b);
bin_to_bool!(less_equal, LessEqual,
codegen: codegen_compare_to_zero,
operating_datum_type: operating_datum_type_for_cmp,
[bool, u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64] => |c, &a, &b | *c = a <= b);
bin_to_bool!(greater, Greater,
codegen: codegen_compare_to_zero,
operating_datum_type: operating_datum_type_for_cmp,
[bool, u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64] => |c, &a, &b | *c = a > b);
bin_to_bool!(greater_equal, GreaterEqual,
codegen: codegen_compare_to_zero,
operating_datum_type: operating_datum_type_for_cmp,
[bool, u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64] => |c, &a, &b | *c = a >= b);

pub fn operating_datum_type_for_cmp(a: DatumType, b: DatumType) -> TractResult<DatumType> {
let dt = a
.common_super_type(b)
.with_context(|| format_err!("No super type for {:?} and {:?}", a, b))?;
if dt == DatumType::TDim {
Ok(DatumType::I64)
} else {
Ok(dt)
}
}

fn codegen_compare_to_zero(
op: &dyn BinMiniOp,
model: &TypedModel,
node: &TypedNode,
) -> TractResult<Option<TypedModelPatch>> {
let facts = model.node_input_facts(node.id)?;
if let Some(uniform) = crate::ops::binary::one_input_is_uniform(model, node)? {
let dt = facts[0].datum_type;
if (dt.is_signed() || dt.is_float()) && *uniform.uni == Tensor::zero_scalar_dt(dt)? {
let reversed = uniform.left_is_uniform;
let mapped = || -> Box<dyn ElementWiseMiniOp> {
macro_rules! m {
($bin: ty, $same: expr, $other: expr) => {
if op.is::<$bin>() {
return if reversed { Box::new($other) } else { Box::new($same) };
};
};
}
m!(Less, LessThanZero {}, GreaterEqualThanZero {});
m!(LessEqual, LessEqualThanZero {}, GreaterThanZero {});
m!(Greater, GreaterThanZero {}, LessEqualThanZero {});
m!(GreaterEqual, GreaterEqualThanZero {}, LessThanZero {});
unreachable!();
};
return Ok(Some(TypedModelPatch::replace_single_op(
model,
node,
&[uniform.var],
ElementWiseOp(mapped(), None),
)?));
}
}
Ok(None)
}

element_wise_oop!(less_than_zero, LessThanZero, [f16, f32, f64, i8, i16, i32, i64] => bool |_op, xs, ys| {
xs.iter().zip(ys.iter_mut()).for_each(|(x,y)| *y = *x < num_traits::Zero::zero());
Ok(())
});

element_wise_oop!(less_equal_than_zero, LessEqualThanZero, [f16, f32, f64, i8, i16, i32, i64] => bool |_op, xs, ys| {
xs.iter().zip(ys.iter_mut()).for_each(|(x,y)| *y = *x <= num_traits::Zero::zero());
Ok(())
});

element_wise_oop!(greater_than_zero, GreaterThanZero, [f16, f32, f64, i8, i16, i32, i64] => bool |_op, xs, ys| {
xs.iter().zip(ys.iter_mut()).for_each(|(x,y)| *y = *x > num_traits::Zero::zero());
Ok(())
});

element_wise_oop!(greater_equal_than_zero, GreaterEqualThanZero, [f16, f32, f64, i8, i16, i32, i64] => bool |_op, xs, ys| {
xs.iter().zip(ys.iter_mut()).for_each(|(x,y)| *y = *x >= num_traits::Zero::zero());
Ok(())
});

element_wise!(not, Not, [bool] => |_, vs| {
vs.iter_mut().for_each(|a| *a = !*a);
Expand Down
Loading
Loading