forked from M-Labs/nac3
core/ndstrides: implement ndarray np_{zeros,ones,full,empty}
This commit is contained in:
parent
701e45364c
commit
1c317f9205
|
@ -12,12 +12,16 @@ use crate::{
|
||||||
call_ndarray_calc_size,
|
call_ndarray_calc_size,
|
||||||
},
|
},
|
||||||
llvm_intrinsics::{self, call_memcpy_generic},
|
llvm_intrinsics::{self, call_memcpy_generic},
|
||||||
|
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},
|
stmt::{gen_for_callback_incrementing, gen_for_range_callback, gen_if_else_expr_callback},
|
||||||
CodeGenContext, CodeGenerator,
|
CodeGenContext, CodeGenerator,
|
||||||
},
|
},
|
||||||
symbol_resolver::ValueEnum,
|
symbol_resolver::ValueEnum,
|
||||||
toplevel::{
|
toplevel::{
|
||||||
helper::PrimDef,
|
helper::{extract_ndims, PrimDef},
|
||||||
numpy::{make_ndarray_ty, unpack_ndarray_var_tys},
|
numpy::{make_ndarray_ty, unpack_ndarray_var_tys},
|
||||||
DefinitionId,
|
DefinitionId,
|
||||||
},
|
},
|
||||||
|
@ -1742,8 +1746,13 @@ pub fn gen_ndarray_empty<'ctx>(
|
||||||
let shape_ty = fun.0.args[0].ty;
|
let shape_ty = fun.0.args[0].ty;
|
||||||
let shape_arg = args[0].1.clone().to_basic_value_enum(context, generator, shape_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)
|
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
|
||||||
.map(NDArrayValue::into)
|
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`.
|
/// Generates LLVM IR for `ndarray.zeros`.
|
||||||
|
@ -1760,8 +1769,13 @@ pub fn gen_ndarray_zeros<'ctx>(
|
||||||
let shape_ty = fun.0.args[0].ty;
|
let shape_ty = fun.0.args[0].ty;
|
||||||
let shape_arg = args[0].1.clone().to_basic_value_enum(context, generator, shape_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)
|
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
|
||||||
.map(NDArrayValue::into)
|
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`.
|
/// Generates LLVM IR for `ndarray.ones`.
|
||||||
|
@ -1778,8 +1792,13 @@ pub fn gen_ndarray_ones<'ctx>(
|
||||||
let shape_ty = fun.0.args[0].ty;
|
let shape_ty = fun.0.args[0].ty;
|
||||||
let shape_arg = args[0].1.clone().to_basic_value_enum(context, generator, shape_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)
|
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
|
||||||
.map(NDArrayValue::into)
|
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`.
|
/// Generates LLVM IR for `ndarray.full`.
|
||||||
|
@ -1799,8 +1818,14 @@ pub fn gen_ndarray_full<'ctx>(
|
||||||
let fill_value_arg =
|
let fill_value_arg =
|
||||||
args[1].1.clone().to_basic_value_enum(context, generator, fill_value_ty)?;
|
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)
|
let (dtype, ndims) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
|
||||||
.map(NDArrayValue::into)
|
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>(
|
pub fn gen_ndarray_array<'ctx>(
|
||||||
|
|
|
@ -0,0 +1,126 @@
|
||||||
|
use inkwell::values::BasicValueEnum;
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
codegen::{
|
||||||
|
irrt::call_nac3_ndarray_util_assert_shape_no_negative, model::*, CodeGenContext,
|
||||||
|
CodeGenerator,
|
||||||
|
},
|
||||||
|
typecheck::typedef::Type,
|
||||||
|
};
|
||||||
|
|
||||||
|
use super::NDArrayObject;
|
||||||
|
|
||||||
|
/// 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);
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
|
@ -1,6 +1,13 @@
|
||||||
|
pub mod factory;
|
||||||
pub mod nditer;
|
pub mod nditer;
|
||||||
|
pub mod shape_util;
|
||||||
|
|
||||||
use inkwell::{context::Context, types::BasicType, values::PointerValue, AddressSpace};
|
use inkwell::{
|
||||||
|
context::Context,
|
||||||
|
types::BasicType,
|
||||||
|
values::{BasicValueEnum, PointerValue},
|
||||||
|
AddressSpace,
|
||||||
|
};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
codegen::{
|
codegen::{
|
||||||
|
@ -353,4 +360,21 @@ impl<'ctx> NDArrayObject<'ctx> {
|
||||||
assert!(ctx.unifier.unioned(self.dtype, src.dtype), "self and src dtype should match");
|
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);
|
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();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -0,0 +1,105 @@
|
||||||
|
use util::gen_for_model;
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
codegen::{
|
||||||
|
model::*,
|
||||||
|
object::{any::AnyObject, list::ListObject, tuple::TupleObject},
|
||||||
|
CodeGenContext, CodeGenerator,
|
||||||
|
},
|
||||||
|
typecheck::typedef::TypeEnum,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// 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, i as u64, 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)
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue