forked from M-Labs/nac3
core: Implement elementwise unary operators
This commit is contained in:
parent
6af13a8261
commit
727a1886b3
|
@ -1296,7 +1296,7 @@ pub fn gen_binop_expr<'ctx, G: CodeGenerator>(
|
|||
/// Generates LLVM IR for a unary operator expression using the [`Type`] and
|
||||
/// [LLVM value][`BasicValueEnum`] of the operands.
|
||||
pub fn gen_unaryop_expr_with_values<'ctx, G: CodeGenerator>(
|
||||
_generator: &mut G,
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
op: &ast::Unaryop,
|
||||
operand: (&Option<Type>, BasicValueEnum<'ctx>),
|
||||
|
@ -1336,6 +1336,33 @@ pub fn gen_unaryop_expr_with_values<'ctx, G: CodeGenerator>(
|
|||
.unwrap(),
|
||||
_ => val.into(),
|
||||
}
|
||||
} else if ty.obj_id(&ctx.unifier).is_some_and(|id| id == PRIMITIVE_DEF_IDS.ndarray) {
|
||||
let llvm_usize = generator.get_size_type(ctx.ctx);
|
||||
let (ndarray_dtype, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty);
|
||||
|
||||
let val = NDArrayValue::from_ptr_val(
|
||||
val.into_pointer_value(),
|
||||
llvm_usize,
|
||||
None,
|
||||
);
|
||||
|
||||
let res = numpy::ndarray_elementwise_unaryop_impl(
|
||||
generator,
|
||||
ctx,
|
||||
ndarray_dtype,
|
||||
None,
|
||||
val,
|
||||
|generator, ctx, val| {
|
||||
gen_unaryop_expr_with_values(
|
||||
generator,
|
||||
ctx,
|
||||
op,
|
||||
(&Some(ndarray_dtype), val)
|
||||
)?.unwrap().to_basic_value_enum(ctx, generator, ndarray_dtype)
|
||||
},
|
||||
)?;
|
||||
|
||||
res.as_ptr_value().into()
|
||||
} else {
|
||||
unimplemented!()
|
||||
}))
|
||||
|
|
|
@ -340,6 +340,31 @@ fn ndarray_fill_indexed<'ctx, G, ValueFn>(
|
|||
)
|
||||
}
|
||||
|
||||
fn ndarray_fill_mapping<'ctx, G, MapFn>(
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
src: NDArrayValue<'ctx>,
|
||||
dest: NDArrayValue<'ctx>,
|
||||
map_fn: MapFn,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
G: CodeGenerator + ?Sized,
|
||||
MapFn: Fn(&mut G, &mut CodeGenContext<'ctx, '_>, BasicValueEnum<'ctx>) -> Result<BasicValueEnum<'ctx>, String>,
|
||||
{
|
||||
ndarray_fill_flattened(
|
||||
generator,
|
||||
ctx,
|
||||
dest,
|
||||
|generator, ctx, i| {
|
||||
let elem = unsafe {
|
||||
src.data().get_unchecked(ctx, generator, i, None)
|
||||
};
|
||||
|
||||
map_fn(generator, ctx, elem)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Generates the LLVM IR for checking whether the source `ndarray` can be broadcast to the shape of
|
||||
/// the target `ndarray`.
|
||||
fn ndarray_assert_is_broadcastable<'ctx, G: CodeGenerator + ?Sized>(
|
||||
|
@ -656,6 +681,48 @@ fn ndarray_copy_impl<'ctx, G: CodeGenerator + ?Sized>(
|
|||
Ok(ndarray)
|
||||
}
|
||||
|
||||
pub fn ndarray_elementwise_unaryop_impl<'ctx, G, MapFn>(
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
elem_ty: Type,
|
||||
res: Option<NDArrayValue<'ctx>>,
|
||||
operand: NDArrayValue<'ctx>,
|
||||
map_fn: MapFn,
|
||||
) -> Result<NDArrayValue<'ctx>, String>
|
||||
where
|
||||
G: CodeGenerator,
|
||||
MapFn: Fn(&mut G, &mut CodeGenContext<'ctx, '_>, BasicValueEnum<'ctx>) -> Result<BasicValueEnum<'ctx>, String>,
|
||||
{
|
||||
let res = res.unwrap_or_else(|| {
|
||||
create_ndarray_dyn_shape(
|
||||
generator,
|
||||
ctx,
|
||||
elem_ty,
|
||||
&operand,
|
||||
|_, ctx, v| {
|
||||
Ok(v.load_ndims(ctx))
|
||||
},
|
||||
|generator, ctx, v, idx| {
|
||||
unsafe {
|
||||
Ok(v.dim_sizes().get_typed_unchecked(ctx, generator, idx, None))
|
||||
}
|
||||
},
|
||||
).unwrap()
|
||||
});
|
||||
|
||||
ndarray_fill_mapping(
|
||||
generator,
|
||||
ctx,
|
||||
operand,
|
||||
res,
|
||||
|generator, ctx, elem| {
|
||||
map_fn(generator, ctx, elem)
|
||||
}
|
||||
)?;
|
||||
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
/// LLVM-typed implementation for computing elementwise binary operations on two input operands.
|
||||
///
|
||||
/// If the operand is a `ndarray`, the broadcast index corresponding to each element in the output
|
||||
|
|
|
@ -5,7 +5,7 @@ expression: res_vec
|
|||
[
|
||||
"Class {\nname: \"Generic_A\",\nancestors: [\"Generic_A[V]\", \"B\"],\nfields: [\"aa\", \"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"foo\", \"fn[[b:T], none]\"), (\"fun\", \"fn[[a:int32], V]\")],\ntype_vars: [\"V\"]\n}\n",
|
||||
"Function {\nname: \"Generic_A.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"Generic_A.fun\",\nsig: \"fn[[a:int32], V]\",\nvar_id: [124]\n}\n",
|
||||
"Function {\nname: \"Generic_A.fun\",\nsig: \"fn[[a:int32], V]\",\nvar_id: [127]\n}\n",
|
||||
"Class {\nname: \"B\",\nancestors: [\"B\"],\nfields: [\"aa\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"foo\", \"fn[[b:T], none]\")],\ntype_vars: []\n}\n",
|
||||
"Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"B.foo\",\nsig: \"fn[[b:T], none]\",\nvar_id: []\n}\n",
|
||||
|
|
|
@ -7,7 +7,7 @@ expression: res_vec
|
|||
"Function {\nname: \"A.__init__\",\nsig: \"fn[[t:T], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"A.fun\",\nsig: \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"A.foo\",\nsig: \"fn[[c:C], none]\",\nvar_id: []\n}\n",
|
||||
"Class {\nname: \"B\",\nancestors: [\"B[typevar113]\", \"A[float]\"],\nfields: [\"a\", \"b\", \"c\", \"d\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\"), (\"foo\", \"fn[[c:C], none]\")],\ntype_vars: [\"typevar113\"]\n}\n",
|
||||
"Class {\nname: \"B\",\nancestors: [\"B[typevar116]\", \"A[float]\"],\nfields: [\"a\", \"b\", \"c\", \"d\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\"), (\"foo\", \"fn[[c:C], none]\")],\ntype_vars: [\"typevar116\"]\n}\n",
|
||||
"Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"B.fun\",\nsig: \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\",\nvar_id: []\n}\n",
|
||||
"Class {\nname: \"C\",\nancestors: [\"C\", \"B[bool]\", \"A[float]\"],\nfields: [\"a\", \"b\", \"c\", \"d\", \"e\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\"), (\"foo\", \"fn[[c:C], none]\")],\ntype_vars: []\n}\n",
|
||||
|
|
|
@ -5,8 +5,8 @@ expression: res_vec
|
|||
[
|
||||
"Function {\nname: \"foo\",\nsig: \"fn[[a:list[int32], b:tuple[T, float]], A[B, bool]]\",\nvar_id: []\n}\n",
|
||||
"Class {\nname: \"A\",\nancestors: [\"A[T, V]\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[v:V], none]\"), (\"fun\", \"fn[[a:T], V]\")],\ntype_vars: [\"T\", \"V\"]\n}\n",
|
||||
"Function {\nname: \"A.__init__\",\nsig: \"fn[[v:V], none]\",\nvar_id: [126]\n}\n",
|
||||
"Function {\nname: \"A.fun\",\nsig: \"fn[[a:T], V]\",\nvar_id: [131]\n}\n",
|
||||
"Function {\nname: \"A.__init__\",\nsig: \"fn[[v:V], none]\",\nvar_id: [129]\n}\n",
|
||||
"Function {\nname: \"A.fun\",\nsig: \"fn[[a:T], V]\",\nvar_id: [134]\n}\n",
|
||||
"Function {\nname: \"gfun\",\nsig: \"fn[[a:A[list[float], int32]], none]\",\nvar_id: []\n}\n",
|
||||
"Class {\nname: \"B\",\nancestors: [\"B\"],\nfields: [],\nmethods: [(\"__init__\", \"fn[[], none]\")],\ntype_vars: []\n}\n",
|
||||
"Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
|
||||
|
|
|
@ -3,7 +3,7 @@ source: nac3core/src/toplevel/test.rs
|
|||
expression: res_vec
|
||||
---
|
||||
[
|
||||
"Class {\nname: \"A\",\nancestors: [\"A[typevar112, typevar113]\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[a:A[float, bool], b:B], none]\"), (\"fun\", \"fn[[a:A[float, bool]], A[bool, int32]]\")],\ntype_vars: [\"typevar112\", \"typevar113\"]\n}\n",
|
||||
"Class {\nname: \"A\",\nancestors: [\"A[typevar115, typevar116]\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[a:A[float, bool], b:B], none]\"), (\"fun\", \"fn[[a:A[float, bool]], A[bool, int32]]\")],\ntype_vars: [\"typevar115\", \"typevar116\"]\n}\n",
|
||||
"Function {\nname: \"A.__init__\",\nsig: \"fn[[a:A[float, bool], b:B], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"A.fun\",\nsig: \"fn[[a:A[float, bool]], A[bool, int32]]\",\nvar_id: []\n}\n",
|
||||
"Class {\nname: \"B\",\nancestors: [\"B\", \"A[int64, bool]\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:A[float, bool]], A[bool, int32]]\"), (\"foo\", \"fn[[b:B], B]\"), (\"bar\", \"fn[[a:A[list[B], int32]], tuple[A[virtual[A[B, int32]], bool], B]]\")],\ntype_vars: []\n}\n",
|
||||
|
|
|
@ -6,12 +6,12 @@ expression: res_vec
|
|||
"Class {\nname: \"A\",\nancestors: [\"A\"],\nfields: [\"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[b:B], none]\"), (\"foo\", \"fn[[a:T, b:V], none]\")],\ntype_vars: []\n}\n",
|
||||
"Function {\nname: \"A.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"A.fun\",\nsig: \"fn[[b:B], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"A.foo\",\nsig: \"fn[[a:T, b:V], none]\",\nvar_id: [132]\n}\n",
|
||||
"Function {\nname: \"A.foo\",\nsig: \"fn[[a:T, b:V], none]\",\nvar_id: [135]\n}\n",
|
||||
"Class {\nname: \"B\",\nancestors: [\"B\", \"C\", \"A\"],\nfields: [\"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[b:B], none]\"), (\"foo\", \"fn[[a:T, b:V], none]\")],\ntype_vars: []\n}\n",
|
||||
"Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
|
||||
"Class {\nname: \"C\",\nancestors: [\"C\", \"A\"],\nfields: [\"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[b:B], none]\"), (\"foo\", \"fn[[a:T, b:V], none]\")],\ntype_vars: []\n}\n",
|
||||
"Function {\nname: \"C.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"C.fun\",\nsig: \"fn[[b:B], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"foo\",\nsig: \"fn[[a:A], none]\",\nvar_id: []\n}\n",
|
||||
"Function {\nname: \"ff\",\nsig: \"fn[[a:T], V]\",\nvar_id: [140]\n}\n",
|
||||
"Function {\nname: \"ff\",\nsig: \"fn[[a:T], V]\",\nvar_id: [143]\n}\n",
|
||||
]
|
||||
|
|
|
@ -466,6 +466,23 @@ pub fn typeof_binop(
|
|||
}))
|
||||
}
|
||||
|
||||
pub fn typeof_unaryop(
|
||||
unifier: &mut Unifier,
|
||||
primitives: &PrimitiveStore,
|
||||
op: &Unaryop,
|
||||
operand: Type,
|
||||
) -> Result<Option<Type>, String> {
|
||||
if *op == Unaryop::Not && operand.obj_id(unifier).is_some_and(|id| id == primitives.ndarray.obj_id(unifier).unwrap()) {
|
||||
return Err("The truth value of an array with more than one element is ambiguous".to_string())
|
||||
}
|
||||
|
||||
Ok(if operand.obj_id(unifier).is_some_and(|id| PRIMITIVE_DEF_IDS.iter().any(|prim_id| id == prim_id)) {
|
||||
Some(operand)
|
||||
} else {
|
||||
None
|
||||
})
|
||||
}
|
||||
|
||||
pub fn set_primitives_magic_methods(store: &PrimitiveStore, unifier: &mut Unifier) {
|
||||
let PrimitiveStore {
|
||||
int32: int32_t,
|
||||
|
@ -525,4 +542,6 @@ pub fn set_primitives_magic_methods(store: &PrimitiveStore, unifier: &mut Unifie
|
|||
impl_div(unifier, store, ndarray_t, &[ndarray_t, ndarray_dtype_t], None);
|
||||
impl_floordiv(unifier, store, ndarray_t, &[ndarray_unsized_t, ndarray_unsized_dtype_t], None);
|
||||
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));
|
||||
}
|
||||
|
|
|
@ -549,7 +549,9 @@ impl<'a> Fold<()> for Inferencer<'a> {
|
|||
ExprKind::BinOp { left, op, right } => {
|
||||
Some(self.infer_bin_ops(expr.location, left, op, right, false)?)
|
||||
}
|
||||
ExprKind::UnaryOp { op, operand } => Some(self.infer_unary_ops(op, operand)?),
|
||||
ExprKind::UnaryOp { op, operand } => {
|
||||
Some(self.infer_unary_ops(expr.location, op, operand)?)
|
||||
}
|
||||
ExprKind::Compare { left, ops, comparators } => {
|
||||
Some(self.infer_compare(left, ops, comparators)?)
|
||||
}
|
||||
|
@ -1247,11 +1249,20 @@ impl<'a> Inferencer<'a> {
|
|||
|
||||
fn infer_unary_ops(
|
||||
&mut self,
|
||||
location: Location,
|
||||
op: &ast::Unaryop,
|
||||
operand: &ast::Expr<Option<Type>>,
|
||||
) -> InferenceResult {
|
||||
let method = unaryop_name(op).into();
|
||||
self.build_method_call(operand.location, method, operand.custom.unwrap(), vec![], None)
|
||||
|
||||
let ret = typeof_unaryop(
|
||||
self.unifier,
|
||||
self.primitives,
|
||||
op,
|
||||
operand.custom.unwrap(),
|
||||
).map_err(|e| HashSet::from([format!("{e} (at {location})")]))?;
|
||||
|
||||
self.build_method_call(operand.location, method, operand.custom.unwrap(), vec![], ret)
|
||||
}
|
||||
|
||||
fn infer_compare(
|
||||
|
|
|
@ -1,3 +1,7 @@
|
|||
@extern
|
||||
def output_bool(x: bool):
|
||||
...
|
||||
|
||||
@extern
|
||||
def output_int32(x: int32):
|
||||
...
|
||||
|
@ -6,10 +10,20 @@ def output_int32(x: int32):
|
|||
def output_float64(x: float):
|
||||
...
|
||||
|
||||
def output_ndarray_bool_2(n: ndarray[bool, Literal[2]]):
|
||||
for r in range(len(n)):
|
||||
for c in range(len(n[r])):
|
||||
output_bool(n[r][c])
|
||||
|
||||
def output_ndarray_int32_1(n: ndarray[int32, Literal[1]]):
|
||||
for i in range(len(n)):
|
||||
output_int32(n[i])
|
||||
|
||||
def output_ndarray_int32_2(n: ndarray[int32, Literal[2]]):
|
||||
for r in range(len(n)):
|
||||
for c in range(len(n[r])):
|
||||
output_int32(n[r][c])
|
||||
|
||||
def output_ndarray_float_1(n: ndarray[float, Literal[1]]):
|
||||
for i in range(len(n)):
|
||||
output_float64(n[i])
|
||||
|
@ -408,6 +422,39 @@ def test_ndarray_ipow_broadcast_scalar():
|
|||
|
||||
output_ndarray_float_2(x)
|
||||
|
||||
def test_ndarray_pos():
|
||||
x_int32 = np_full([2, 2], -2)
|
||||
y_int32 = +x_int32
|
||||
|
||||
output_ndarray_int32_2(x_int32)
|
||||
output_ndarray_int32_2(y_int32)
|
||||
|
||||
x_float = np_full([2, 2], -2.0)
|
||||
y_float = +x_float
|
||||
|
||||
output_ndarray_float_2(x_float)
|
||||
output_ndarray_float_2(y_float)
|
||||
|
||||
def test_ndarray_neg():
|
||||
x_int32 = np_full([2, 2], -2)
|
||||
y_int32 = -x_int32
|
||||
|
||||
output_ndarray_int32_2(x_int32)
|
||||
output_ndarray_int32_2(y_int32)
|
||||
|
||||
x_float = np_full([2, 2], 2.0)
|
||||
y_float = -x_float
|
||||
|
||||
output_ndarray_float_2(x_float)
|
||||
output_ndarray_float_2(y_float)
|
||||
|
||||
def test_ndarray_inv():
|
||||
x_int32 = np_full([2, 2], -2)
|
||||
y_int32 = ~x_int32
|
||||
|
||||
output_ndarray_int32_2(x_int32)
|
||||
output_ndarray_int32_2(y_int32)
|
||||
|
||||
def run() -> int32:
|
||||
test_ndarray_ctor()
|
||||
test_ndarray_empty()
|
||||
|
@ -467,5 +514,8 @@ def run() -> int32:
|
|||
test_ndarray_ipow()
|
||||
test_ndarray_ipow_broadcast()
|
||||
test_ndarray_ipow_broadcast_scalar()
|
||||
test_ndarray_pos()
|
||||
test_ndarray_neg()
|
||||
test_ndarray_inv()
|
||||
|
||||
return 0
|
||||
|
|
Loading…
Reference in New Issue