newell
/
nac3
forked from M-Labs/nac3
1
0
Fork 0
Code Pull Requests Activity

core/ndstrides: implement np_{zeros,ones,full,empty}

This commit is contained in:
lyken 2024-08-20 12:25:01 +08:00 committed by David Mak
parent 7f4b4597c5
commit c1913f11c6
4 changed files with 289 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -21,12 +21,16 @@ use crate::{
},
llvm_intrinsics::{self, call_memcpy_generic},
macros::codegen_unreachable,
object::{
any::AnyObject,
ndarray::{shape_util::parse_numpy_int_sequence, NDArrayObject},
},
stmt::{gen_for_callback_incrementing, gen_for_range_callback, gen_if_else_expr_callback},
CodeGenContext, CodeGenerator,
},
symbol_resolver::ValueEnum,
toplevel::{
helper::PrimDef,
helper::{extract_ndims, PrimDef},
numpy::{make_ndarray_ty, unpack_ndarray_var_tys},
DefinitionId,
},
@ -1743,8 +1747,13 @@ pub fn gen_ndarray_empty<'ctx>(
let shape_ty = fun.0.args[0].ty;
let shape_arg = args[0].1.clone().to_basic_value_enum(context, generator, shape_ty)?;
call_ndarray_empty_impl(generator, context, context.primitives.float, shape_arg)
.map(NDArrayValue::into)
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
let ndims = extract_ndims(&context.unifier, ndims);
let shape = AnyObject { value: shape_arg, ty: shape_ty };
let (_, shape) = parse_numpy_int_sequence(generator, context, shape);
let ndarray = NDArrayObject::make_np_empty(generator, context, dtype, ndims, shape);
Ok(ndarray.instance.value)
}
/// Generates LLVM IR for `ndarray.zeros`.
@ -1761,8 +1770,13 @@ pub fn gen_ndarray_zeros<'ctx>(
let shape_ty = fun.0.args[0].ty;
let shape_arg = args[0].1.clone().to_basic_value_enum(context, generator, shape_ty)?;
call_ndarray_zeros_impl(generator, context, context.primitives.float, shape_arg)
.map(NDArrayValue::into)
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
let ndims = extract_ndims(&context.unifier, ndims);
let shape = AnyObject { value: shape_arg, ty: shape_ty };
let (_, shape) = parse_numpy_int_sequence(generator, context, shape);
let ndarray = NDArrayObject::make_np_zeros(generator, context, dtype, ndims, shape);
Ok(ndarray.instance.value)
}
/// Generates LLVM IR for `ndarray.ones`.
@ -1779,8 +1793,13 @@ pub fn gen_ndarray_ones<'ctx>(
let shape_ty = fun.0.args[0].ty;
let shape_arg = args[0].1.clone().to_basic_value_enum(context, generator, shape_ty)?;
call_ndarray_ones_impl(generator, context, context.primitives.float, shape_arg)
.map(NDArrayValue::into)
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
let ndims = extract_ndims(&context.unifier, ndims);
let shape = AnyObject { value: shape_arg, ty: shape_ty };
let (_, shape) = parse_numpy_int_sequence(generator, context, shape);
let ndarray = NDArrayObject::make_np_ones(generator, context, dtype, ndims, shape);
Ok(ndarray.instance.value)
}
/// Generates LLVM IR for `ndarray.full`.
@ -1800,8 +1819,14 @@ pub fn gen_ndarray_full<'ctx>(
let fill_value_arg =
args[1].1.clone().to_basic_value_enum(context, generator, fill_value_ty)?;
call_ndarray_full_impl(generator, context, fill_value_ty, shape_arg, fill_value_arg)
.map(NDArrayValue::into)
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
let ndims = extract_ndims(&context.unifier, ndims);
let shape = AnyObject { value: shape_arg, ty: shape_ty };
let (_, shape) = parse_numpy_int_sequence(generator, context, shape);
let ndarray =
NDArrayObject::make_np_full(generator, context, dtype, ndims, shape, fill_value_arg);
Ok(ndarray.instance.value)
}
pub fn gen_ndarray_array<'ctx>(

125
nac3core/src/codegen/object/ndarray/factory.rs Normal file
View File

@ -0,0 +1,125 @@
use inkwell::values::BasicValueEnum;
use super::NDArrayObject;
use crate::{
codegen::{
irrt::call_nac3_ndarray_util_assert_shape_no_negative, model::*, CodeGenContext,
CodeGenerator,
},
typecheck::typedef::Type,
};
/// Get the zero value in `np.zeros()` of a `dtype`.
fn ndarray_zero_value<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
) -> BasicValueEnum<'ctx> {
if [ctx.primitives.int32, ctx.primitives.uint32]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
ctx.ctx.i32_type().const_zero().into()
} else if [ctx.primitives.int64, ctx.primitives.uint64]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
ctx.ctx.i64_type().const_zero().into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.float) {
ctx.ctx.f64_type().const_zero().into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.bool) {
ctx.ctx.bool_type().const_zero().into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.str) {
ctx.gen_string(generator, "").into()
} else {
panic!("unrecognized dtype: {}", ctx.unifier.stringify(dtype));
}
}
/// Get the one value in `np.ones()` of a `dtype`.
fn ndarray_one_value<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
) -> BasicValueEnum<'ctx> {
if [ctx.primitives.int32, ctx.primitives.uint32]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
let is_signed = ctx.unifier.unioned(dtype, ctx.primitives.int32);
ctx.ctx.i32_type().const_int(1, is_signed).into()
} else if [ctx.primitives.int64, ctx.primitives.uint64]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
let is_signed = ctx.unifier.unioned(dtype, ctx.primitives.int64);
ctx.ctx.i64_type().const_int(1, is_signed).into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.float) {
ctx.ctx.f64_type().const_float(1.0).into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.bool) {
ctx.ctx.bool_type().const_int(1, false).into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.str) {
ctx.gen_string(generator, "1").into()
} else {
panic!("unrecognized dtype: {}", ctx.unifier.stringify(dtype));
}
}
impl<'ctx> NDArrayObject<'ctx> {
/// Create an ndarray like `np.empty`.
pub fn make_np_empty<G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
ndims: u64,
shape: Instance<'ctx, Ptr<Int<SizeT>>>,
) -> Self {
// Validate `shape`
let ndims_llvm = Int(SizeT).const_int(generator, ctx.ctx, ndims, false);
call_nac3_ndarray_util_assert_shape_no_negative(generator, ctx, ndims_llvm, shape);
let ndarray = NDArrayObject::alloca(generator, ctx, dtype, ndims);
ndarray.copy_shape_from_array(generator, ctx, shape);
ndarray.create_data(generator, ctx);
ndarray
}
/// Create an ndarray like `np.full`.
pub fn make_np_full<G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
ndims: u64,
shape: Instance<'ctx, Ptr<Int<SizeT>>>,
fill_value: BasicValueEnum<'ctx>,
) -> Self {
let ndarray = NDArrayObject::make_np_empty(generator, ctx, dtype, ndims, shape);
ndarray.fill(generator, ctx, fill_value);
ndarray
}
/// Create an ndarray like `np.zero`.
pub fn make_np_zeros<G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
ndims: u64,
shape: Instance<'ctx, Ptr<Int<SizeT>>>,
) -> Self {
let fill_value = ndarray_zero_value(generator, ctx, dtype);
NDArrayObject::make_np_full(generator, ctx, dtype, ndims, shape, fill_value)
}
/// Create an ndarray like `np.ones`.
pub fn make_np_ones<G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
ndims: u64,
shape: Instance<'ctx, Ptr<Int<SizeT>>>,
) -> Self {
let fill_value = ndarray_one_value(generator, ctx, dtype);
NDArrayObject::make_np_full(generator, ctx, dtype, ndims, shape, fill_value)
}
}

29
nac3core/src/codegen/object/ndarray/mod.rs
View File

@ -1,5 +1,11 @@
use inkwell::{context::Context, types::BasicType, values::PointerValue, AddressSpace};
use inkwell::{
context::Context,
types::BasicType,
values::{BasicValueEnum, PointerValue},
AddressSpace,
};
use super::any::AnyObject;
use crate::{
codegen::{
irrt::{
@ -15,9 +21,9 @@ use crate::{
typecheck::typedef::Type,
};
use super::any::AnyObject;
pub mod factory;
pub mod nditer;
pub mod shape_util;
/// Fields of [`NDArray`]
pub struct NDArrayFields<'ctx, F: FieldTraversal<'ctx>> {
@ -345,4 +351,21 @@ impl<'ctx> NDArrayObject<'ctx> {
assert!(ctx.unifier.unioned(self.dtype, src.dtype), "self and src dtype should match");
call_nac3_ndarray_copy_data(generator, ctx, src.instance, self.instance);
}
/// Fill the ndarray with a scalar.
///
/// `fill_value` must have the same LLVM type as the `dtype` of this ndarray.
pub fn fill<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
value: BasicValueEnum<'ctx>,
) {
self.foreach(generator, ctx, |generator, ctx, _hooks, nditer| {
let p = nditer.get_pointer(generator, ctx);
ctx.builder.build_store(p, value).unwrap();
Ok(())
})
.unwrap();
}
}

104
nac3core/src/codegen/object/ndarray/shape_util.rs Normal file
View File

@ -0,0 +1,104 @@
use crate::{
codegen::{
model::*,
object::{any::AnyObject, list::ListObject, tuple::TupleObject},
CodeGenContext, CodeGenerator,
},
typecheck::typedef::TypeEnum,
};
use util::gen_for_model;
/// Parse a NumPy-like "int sequence" input and return the int sequence as an array and its length.
///
/// * `sequence` - The `sequence` parameter.
/// * `sequence_ty` - The typechecker type of `sequence`
///
/// The `sequence` argument type may only be one of the following:
/// 1. A list of `int32`; e.g., `np.empty([600, 800, 3])`
/// 2. A tuple of `int32`; e.g., `np.empty((600, 800, 3))`
/// 3. A scalar `int32`; e.g., `np.empty(3)`, this is functionally equivalent to `np.empty([3])`
///
/// All `int32` values will be sign-extended to `SizeT`.
pub fn parse_numpy_int_sequence<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
input_sequence: AnyObject<'ctx>,
) -> (Instance<'ctx, Int<SizeT>>, Instance<'ctx, Ptr<Int<SizeT>>>) {
let zero = Int(SizeT).const_0(generator, ctx.ctx);
let one = Int(SizeT).const_1(generator, ctx.ctx);
// The result `list` to return.
match &*ctx.unifier.get_ty(input_sequence.ty) {
TypeEnum::TObj { obj_id, .. }
if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
{
// 1. A list of `int32`; e.g., `np.empty([600, 800, 3])`
// Check `input_sequence`
let input_sequence = ListObject::from_object(generator, ctx, input_sequence);
let len = input_sequence.instance.get(generator, ctx, |f| f.len);
let result = Int(SizeT).array_alloca(generator, ctx, len.value);
// Load all the `int32`s from the input_sequence, cast them to `SizeT`, and store them into `result`
gen_for_model(generator, ctx, zero, len, one, |generator, ctx, _hooks, i| {
// Load the i-th int32 in the input sequence
let int = input_sequence
.instance
.get(generator, ctx, |f| f.items)
.get_index(generator, ctx, i.value)
.value
.into_int_value();
// Cast to SizeT
let int = Int(SizeT).s_extend_or_bit_cast(generator, ctx, int);
// Store
result.set_index(ctx, i.value, int);
Ok(())
})
.unwrap();
(len, result)
}
TypeEnum::TTuple { .. } => {
// 2. A tuple of ints; e.g., `np.empty((600, 800, 3))`
let input_sequence = TupleObject::from_object(ctx, input_sequence);
let len = input_sequence.len(generator, ctx);
let result = Int(SizeT).array_alloca(generator, ctx, len.value);
for i in 0..input_sequence.num_elements() {
// Get the i-th element off of the tuple and load it into `result`.
let int = input_sequence.index(ctx, i).value.into_int_value();
let int = Int(SizeT).s_extend_or_bit_cast(generator, ctx, int);
result.set_index_const(ctx, i64::try_from(i).unwrap(), int);
}
(len, result)
}
TypeEnum::TObj { obj_id, .. }
if *obj_id == ctx.primitives.int32.obj_id(&ctx.unifier).unwrap() =>
{
// 3. A scalar int; e.g., `np.empty(3)`, this is functionally equivalent to `np.empty([3])`
let input_int = input_sequence.value.into_int_value();
let len = Int(SizeT).const_1(generator, ctx.ctx);
let result = Int(SizeT).array_alloca(generator, ctx, len.value);
let int = Int(SizeT).s_extend_or_bit_cast(generator, ctx, input_int);
// Storing into result[0]
result.store(ctx, int);
(len, result)
}
_ => panic!(
"encountered unknown sequence type: {}",
ctx.unifier.stringify(input_sequence.ty)
),
}
}