M-Labs
/
nac3
8
0
Fork 5
Code Issues 87 Pull Requests 3 Releases Wiki Activity

core: Implement elementwise comparison operators

David Mak 2024-03-27 12:57:11 +08:00
parent d457e6c986
commit 29ae48faad
5 changed files with 360 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

84
nac3core/src/codegen/expr.rs
View File

@ -39,6 +39,7 @@ use inkwell::{
types::{AnyType, BasicType, BasicTypeEnum},
values::{BasicValueEnum, CallSiteValue, FunctionValue, IntValue, PointerValue}
};
use inkwell::values::BasicValue;
use itertools::{chain, izip, Itertools, Either};
use nac3parser::ast::{
self, Boolop, Comprehension, Constant, Expr, ExprKind, Location, Operator, StrRef,
@ -1380,12 +1381,90 @@ pub fn gen_unaryop_expr<'ctx, G: CodeGenerator>(
}
pub fn gen_cmpop_expr_with_values<'ctx, G: CodeGenerator>(
_generator: &mut G,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
left: (Option<Type>, BasicValueEnum<'ctx>),
ops: &[ast::Cmpop],
comparators: &[(Option<Type>, BasicValueEnum<'ctx>)],
) -> Result<Option<ValueEnum<'ctx>>, String> {
debug_assert_eq!(comparators.len(), ops.len());
if comparators.len() == 1 {
let left_ty = ctx.unifier.get_representative(left.0.unwrap());
let right_ty = ctx.unifier.get_representative(comparators[0].0.unwrap());
if left_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PRIMITIVE_DEF_IDS.ndarray) || right_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PRIMITIVE_DEF_IDS.ndarray) {
let llvm_usize = generator.get_size_type(ctx.ctx);
let (Some(left_ty), lhs) = left else { unreachable!() };
let (Some(right_ty), rhs) = comparators[0] else { unreachable!() };
let op = ops[0].clone();
let is_ndarray1 = left_ty.get_obj_id(&ctx.unifier) == PRIMITIVE_DEF_IDS.ndarray;
let is_ndarray2 = right_ty.get_obj_id(&ctx.unifier) == PRIMITIVE_DEF_IDS.ndarray;
return if is_ndarray1 && is_ndarray2 {
let (ndarray_dtype1, _) = unpack_ndarray_var_tys(&mut ctx.unifier, left_ty);
let (ndarray_dtype2, _) = unpack_ndarray_var_tys(&mut ctx.unifier, right_ty);
assert!(ctx.unifier.unioned(ndarray_dtype1, ndarray_dtype2));
let left_val = NDArrayValue::from_ptr_val(
lhs.into_pointer_value(),
llvm_usize,
None
);
let res = numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ctx.primitives.bool,
None,
(left_val.as_ptr_value().into(), false),
(rhs, false),
|generator, ctx, elem_ty, (lhs, rhs)| {
let val = gen_cmpop_expr_with_values(
generator,
ctx,
(Some(ndarray_dtype1), lhs),
&[op.clone()],
&[(Some(ndarray_dtype2), rhs)],
)?.unwrap().to_basic_value_enum(ctx, generator, elem_ty)?;
Ok(generator.bool_to_i8(ctx, val.into_int_value()).into())
},
)?;
Ok(Some(res.as_ptr_value().into()))
} else {
let (ndarray_dtype, _) = unpack_ndarray_var_tys(
&mut ctx.unifier,
if is_ndarray1 { left_ty } else { right_ty },
);
let res = numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ctx.primitives.bool,
None,
(lhs, !is_ndarray1),
(rhs, !is_ndarray2),
|generator, ctx, elem_ty, (lhs, rhs)| {
let val = gen_cmpop_expr_with_values(
generator,
ctx,
(Some(ndarray_dtype), lhs),
&[op.clone()],
&[(Some(ndarray_dtype), rhs)],
)?.unwrap().to_basic_value_enum(ctx, generator, elem_ty)?;
Ok(generator.bool_to_i8(ctx, val.into_int_value()).into())
},
)?;
Ok(Some(res.as_ptr_value().into()))
}
}
}
let cmp_val = izip!(chain(once(&left), comparators.iter()), comparators.iter(), ops.iter(),)
.fold(Ok(None), |prev: Result<Option<_>, String>, (lhs, rhs, op)| {
let (left_ty, lhs) = lhs;
@ -1441,7 +1520,7 @@ pub fn gen_cmpop_expr_with_values<'ctx, G: CodeGenerator>(
let lhs = lhs.into_float_value();
let rhs = rhs.into_float_value();
let op = match op {
ast::Cmpop::Eq | ast::Cmpop::Is => inkwell::FloatPredicate::OEQ,
ast::Cmpop::NotEq => inkwell::FloatPredicate::ONE,
@ -1455,6 +1534,7 @@ pub fn gen_cmpop_expr_with_values<'ctx, G: CodeGenerator>(
} else {
unimplemented!()
};
Ok(prev?.map(|v| ctx.builder.build_and(v, current, "cmp").unwrap()).or(Some(current)))
})?;

1
nac3core/src/codegen/numpy.rs
View File

@ -728,6 +728,7 @@ pub fn ndarray_elementwise_unaryop_impl<'ctx, G, MapFn>(
/// # Panic
///
/// This function will panic if neither input operands (`lhs` or `rhs`) is a `ndarray`.
// TODO: Remove elem_ty from value_fn
pub fn ndarray_elementwise_binop_impl<'ctx, G, ValueFn>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,

71
nac3core/src/typecheck/magic_methods.rs
View File

@ -170,19 +170,35 @@ pub fn impl_unaryop(unifier: &mut Unifier, ty: Type, ret_ty: Option<Type>, ops:
pub fn impl_cmpop(
unifier: &mut Unifier,
store: &PrimitiveStore,
_store: &PrimitiveStore,
ty: Type,
other_ty: Type,
other_ty: &[Type],
ops: &[Cmpop],
ret_ty: Option<Type>,
) {
with_fields(unifier, ty, |unifier, fields| {
let (other_ty, other_var_id) = if other_ty.len() == 1 {
(other_ty[0], None)
} else {
let (ty, var_id) = unifier.get_fresh_var_with_range(other_ty, Some("N".into()), None);
(ty, Some(var_id))
};
let function_vars = if let Some(var_id) = other_var_id {
vec![(var_id, other_ty)].into_iter().collect::<VarMap>()
} else {
VarMap::new()
};
let ret_ty = ret_ty.unwrap_or_else(|| unifier.get_fresh_var(None, None).0);
for op in ops {
fields.insert(
comparison_name(op).unwrap().into(),
(
unifier.add_ty(TypeEnum::TFunc(FunSignature {
ret: store.bool,
vars: VarMap::new(),
ret: ret_ty,
vars: function_vars.clone(),
args: vec![FuncArg {
ty: other_ty,
default_value: None,
@ -291,19 +307,20 @@ pub fn impl_not(unifier: &mut Unifier, store: &PrimitiveStore, ty: Type, ret_ty:
}
/// `Lt`, `LtE`, `Gt`, `GtE`
pub fn impl_comparison(unifier: &mut Unifier, store: &PrimitiveStore, ty: Type, other_ty: Type) {
pub fn impl_comparison(unifier: &mut Unifier, store: &PrimitiveStore, ty: Type, other_ty: &[Type], ret_ty: Option<Type>) {
impl_cmpop(
unifier,
store,
ty,
other_ty,
&[Cmpop::Lt, Cmpop::Gt, Cmpop::LtE, Cmpop::GtE],
ret_ty,
);
}
/// `Eq`, `NotEq`
pub fn impl_eq(unifier: &mut Unifier, store: &PrimitiveStore, ty: Type) {
impl_cmpop(unifier, store, ty, ty, &[Cmpop::Eq, Cmpop::NotEq]);
pub fn impl_eq(unifier: &mut Unifier, store: &PrimitiveStore, ty: Type, other_ty: &[Type], ret_ty: Option<Type>) {
impl_cmpop(unifier, store, ty, other_ty, &[Cmpop::Eq, Cmpop::NotEq], ret_ty);
}
/// Returns the expected return type of binary operations with at least one `ndarray` operand.
@ -458,6 +475,33 @@ pub fn typeof_unaryop(
})
}
/// Returns the return type given a comparison operator and its primitive operands.
pub fn typeof_cmpop(
unifier: &mut Unifier,
primitives: &PrimitiveStore,
_op: &Cmpop,
lhs: Type,
rhs: Type,
) -> Result<Option<Type>, String> {
let is_left_ndarray = lhs
.obj_id(unifier)
.is_some_and(|id| id == PRIMITIVE_DEF_IDS.ndarray);
let is_right_ndarray = rhs
.obj_id(unifier)
.is_some_and(|id| id == PRIMITIVE_DEF_IDS.ndarray);
Ok(Some(if is_left_ndarray || is_right_ndarray {
let brd = typeof_ndarray_broadcast(unifier, primitives, lhs, rhs)?;
let (_, ndims) = unpack_ndarray_var_tys(unifier, brd);
make_ndarray_ty(unifier, primitives, Some(primitives.bool), Some(ndims))
} else if unifier.unioned(lhs, rhs) {
primitives.bool
} else {
return Ok(None)
}))
}
pub fn set_primitives_magic_methods(store: &PrimitiveStore, unifier: &mut Unifier) {
let PrimitiveStore {
int32: int32_t,
@ -483,8 +527,8 @@ pub fn set_primitives_magic_methods(store: &PrimitiveStore, unifier: &mut Unifie
impl_mod(unifier, store, t, &[t, ndarray_int_t], None);
impl_invert(unifier, store, t, Some(t));
impl_not(unifier, store, t, Some(bool_t));
impl_comparison(unifier, store, t, t);
impl_eq(unifier, store, t);
impl_comparison(unifier, store, t, &[t, ndarray_int_t], None);
impl_eq(unifier, store, t, &[t, ndarray_int_t], None);
}
for t in [int32_t, int64_t] {
impl_sign(unifier, store, t, Some(t));
@ -500,12 +544,13 @@ pub fn set_primitives_magic_methods(store: &PrimitiveStore, unifier: &mut Unifie
impl_mod(unifier, store, float_t, &[float_t, ndarray_float_t], None);
impl_sign(unifier, store, float_t, Some(float_t));
impl_not(unifier, store, float_t, Some(bool_t));
impl_comparison(unifier, store, float_t, float_t);
impl_eq(unifier, store, float_t);
impl_comparison(unifier, store, float_t, &[float_t, ndarray_float_t], None);
impl_eq(unifier, store, float_t, &[float_t, ndarray_float_t], None);
/* bool ======== */
let ndarray_bool_t = make_ndarray_ty(unifier, store, Some(bool_t), None);
impl_not(unifier, store, bool_t, Some(bool_t));
impl_eq(unifier, store, bool_t);
impl_eq(unifier, store, bool_t, &[bool_t, ndarray_bool_t], None);
/* ndarray ===== */
let ndarray_usized_ndims_tvar = unifier.get_fresh_const_generic_var(size_t, Some("ndarray_ndims".into()), None);
@ -519,4 +564,6 @@ pub fn set_primitives_magic_methods(store: &PrimitiveStore, unifier: &mut Unifie
impl_mod(unifier, store, ndarray_t, &[ndarray_unsized_t, ndarray_unsized_dtype_t], None);
impl_sign(unifier, store, ndarray_t, Some(ndarray_t));
impl_invert(unifier, store, ndarray_t, Some(ndarray_t));
impl_eq(unifier, store, ndarray_t, &[ndarray_unsized_t, ndarray_unsized_dtype_t], None);
impl_comparison(unifier, store, ndarray_t, &[ndarray_unsized_t, ndarray_unsized_dtype_t], None);
}

29
nac3core/src/typecheck/type_inferencer/mod.rs
View File

@ -1271,22 +1271,45 @@ impl<'a> Inferencer<'a> {
ops: &[ast::Cmpop],
comparators: &[ast::Expr<Option<Type>>],
) -> InferenceResult {
let boolean = self.primitives.bool;
if ops.len() > 1 && once(left).chain(comparators).any(|expr| expr.custom.unwrap().obj_id(self.unifier).is_some_and(|id| id == PRIMITIVE_DEF_IDS.ndarray)) {
return Err(HashSet::from([String::from("Comparator chaining with ndarray types not supported")]))
}
for (a, b, c) in izip!(once(left).chain(comparators), comparators, ops) {
let method = comparison_name(c)
.ok_or_else(|| HashSet::from([
"unsupported comparator".to_string()
]))?
.into();
let ret = typeof_cmpop(
self.unifier,
self.primitives,
c,
a.custom.unwrap(),
b.custom.unwrap(),
).map_err(|e| HashSet::from([format!("{e} (at {})", b.location)]))?;
self.build_method_call(
a.location,
method,
a.custom.unwrap(),
vec![b.custom.unwrap()],
Some(boolean),
ret,
)?;
}
Ok(boolean)
let res_lhs = comparators.iter().rev().nth(1).unwrap_or(left);
let res_rhs = comparators.iter().rev().nth(0).unwrap();
let res_op = ops.iter().rev().nth(0).unwrap();
Ok(typeof_cmpop(
self.unifier,
self.primitives,
res_op,
res_lhs.custom.unwrap(),
res_rhs.custom.unwrap(),
).unwrap().unwrap())
}
/// Infers the type of a subscript expression on an `ndarray`.

192
nac3standalone/demo/src/ndarray.py
View File

@ -455,6 +455,174 @@ def test_ndarray_inv():
output_ndarray_int32_2(x_int32)
output_ndarray_int32_2(y_int32)
def test_ndarray_eq():
x = np_identity(2)
y = x == np_full([2, 2], 0.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_eq_broadcast():
x = np_identity(2)
y = x == np_full([2], 0.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_eq_broadcast_lhs_scalar():
x = np_identity(2)
y = 0.0 == x
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_eq_broadcast_rhs_scalar():
x = np_identity(2)
y = x == 0.0
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ne():
x = np_identity(2)
y = x != np_full([2, 2], 0.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ne_broadcast():
x = np_identity(2)
y = x != np_full([2], 0.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ne_broadcast_lhs_scalar():
x = np_identity(2)
y = 0.0 != x
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ne_broadcast_rhs_scalar():
x = np_identity(2)
y = x != 0.0
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_lt():
x = np_identity(2)
y = x < np_full([2, 2], 1.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_lt_broadcast():
x = np_identity(2)
y = x < np_full([2], 1.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_lt_broadcast_lhs_scalar():
x = np_identity(2)
y = 1.0 < x
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_lt_broadcast_rhs_scalar():
x = np_identity(2)
y = x < 1.0
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_le():
x = np_identity(2)
y = x <= np_full([2, 2], 0.5)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_le_broadcast():
x = np_identity(2)
y = x <= np_full([2], 0.5)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_le_broadcast_lhs_scalar():
x = np_identity(2)
y = 0.5 <= x
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_le_broadcast_rhs_scalar():
x = np_identity(2)
y = x <= 0.5
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_gt():
x = np_identity(2)
y = x > np_full([2, 2], 0.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_gt_broadcast():
x = np_identity(2)
y = x > np_full([2], 0.0)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_gt_broadcast_lhs_scalar():
x = np_identity(2)
y = 0.0 > x
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_gt_broadcast_rhs_scalar():
x = np_identity(2)
y = x > 0.0
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ge():
x = np_identity(2)
y = x >= np_full([2, 2], 0.5)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ge_broadcast():
x = np_identity(2)
y = x >= np_full([2], 0.5)
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ge_broadcast_lhs_scalar():
x = np_identity(2)
y = 0.5 >= x
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def test_ndarray_ge_broadcast_rhs_scalar():
x = np_identity(2)
y = x >= 0.5
output_ndarray_float_2(x)
output_ndarray_bool_2(y)
def run() -> int32:
test_ndarray_ctor()
test_ndarray_empty()
@ -517,5 +685,29 @@ def run() -> int32:
test_ndarray_pos()
test_ndarray_neg()
test_ndarray_inv()
test_ndarray_eq()
test_ndarray_eq_broadcast()
test_ndarray_eq_broadcast_lhs_scalar()
test_ndarray_eq_broadcast_rhs_scalar()
test_ndarray_ne()
test_ndarray_ne_broadcast()
test_ndarray_ne_broadcast_lhs_scalar()
test_ndarray_ne_broadcast_rhs_scalar()
test_ndarray_lt()
test_ndarray_lt_broadcast()
test_ndarray_lt_broadcast_lhs_scalar()
test_ndarray_lt_broadcast_rhs_scalar()
test_ndarray_lt()
test_ndarray_le_broadcast()
test_ndarray_le_broadcast_lhs_scalar()
test_ndarray_le_broadcast_rhs_scalar()
test_ndarray_gt()
test_ndarray_gt_broadcast()
test_ndarray_gt_broadcast_lhs_scalar()
test_ndarray_gt_broadcast_rhs_scalar()
test_ndarray_gt()
test_ndarray_ge_broadcast()
test_ndarray_ge_broadcast_lhs_scalar()
test_ndarray_ge_broadcast_rhs_scalar()
return 0