forked from M-Labs/nac3
WIP: core/ndstrides: {make,from}_simple_ndarray
This commit is contained in:
parent
bb1687f8a4
commit
18dcbf5bbc
|
@ -34,6 +34,7 @@ pub trait Model<'ctx>: fmt::Debug + Clone + Copy {
|
|||
/// Create an instance from a value with [`Instance::model`] being this model.
|
||||
///
|
||||
/// Caller must make sure the type of `value` and the type of this `model` are equivalent.
|
||||
#[must_use]
|
||||
fn believe_value(&self, value: Self::Value) -> Instance<'ctx, Self> {
|
||||
Instance { model: *self, value }
|
||||
}
|
||||
|
|
|
@ -22,7 +22,7 @@ pub trait FieldTraversal<'ctx> {
|
|||
|
||||
fn add<M: Model<'ctx>>(&mut self, name: &'static str, model: M) -> Self::Out<M>;
|
||||
|
||||
/// Like [`FieldTraversal::visit`] but [`Model`] is automatically inferred.
|
||||
/// Like [`FieldTraversal::visit`] but [`Model`] is automatically inferred from [`Default`] trait.
|
||||
fn add_auto<M: Model<'ctx> + Default>(&mut self, name: &'static str) -> Self::Out<M> {
|
||||
self.add(name, M::default())
|
||||
}
|
||||
|
|
|
@ -107,7 +107,7 @@ pub fn broadcast_shapes<'ctx, G: CodeGenerator + ?Sized>(
|
|||
|
||||
impl<'ctx> NDArrayObject<'ctx> {
|
||||
// TODO: DOCUMENT: Behaves like `np.broadcast()`, except returns results differently.
|
||||
pub fn broadcast_all<G: CodeGenerator + ?Sized>(
|
||||
pub fn broadcast<G: CodeGenerator + ?Sized>(
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
ndarrays: &[Self],
|
||||
|
|
|
@ -35,7 +35,7 @@ impl<'ctx> NDArrayObject<'ctx> {
|
|||
let sizet_model = IntModel(SizeT);
|
||||
|
||||
// Broadcast inputs
|
||||
let broadcast_result = NDArrayObject::broadcast_all(generator, ctx, ndarrays);
|
||||
let broadcast_result = NDArrayObject::broadcast(generator, ctx, ndarrays);
|
||||
|
||||
let out_ndarray = match out {
|
||||
NDArrayOut::NewNDArray { dtype } => {
|
||||
|
|
|
@ -3,6 +3,7 @@ pub mod broadcast;
|
|||
pub mod functions;
|
||||
pub mod indexing;
|
||||
pub mod mapping;
|
||||
pub mod nalgebra;
|
||||
pub mod product;
|
||||
pub mod scalar;
|
||||
pub mod shape_util;
|
||||
|
@ -18,7 +19,7 @@ use crate::{
|
|||
},
|
||||
model::*,
|
||||
stmt::BreakContinueHooks,
|
||||
structure::NDArray,
|
||||
structure::{NDArray, SimpleNDArray},
|
||||
CodeGenContext, CodeGenerator,
|
||||
},
|
||||
toplevel::{helper::extract_ndims, numpy::unpack_ndarray_var_tys},
|
||||
|
@ -34,14 +35,6 @@ use inkwell::{
|
|||
use scalar::{ScalarObject, ScalarOrNDArray};
|
||||
use util::{call_memcpy_model, gen_for_model_auto};
|
||||
|
||||
pub struct NpArrayFields<'ctx, F: FieldTraversal<'ctx>> {
|
||||
pub data: F::Out<PtrModel<IntModel<Byte>>>,
|
||||
pub itemsize: F::Out<IntModel<SizeT>>,
|
||||
pub ndims: F::Out<IntModel<SizeT>>,
|
||||
pub shape: F::Out<PtrModel<IntModel<SizeT>>>,
|
||||
pub strides: F::Out<PtrModel<IntModel<SizeT>>>,
|
||||
}
|
||||
|
||||
/// A NAC3 Python ndarray object.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct NDArrayObject<'ctx> {
|
||||
|
@ -77,6 +70,113 @@ impl<'ctx> NDArrayObject<'ctx> {
|
|||
NDArrayObject { dtype, ndims, value }
|
||||
}
|
||||
|
||||
/// Create a [`SimpleNDArray`] from the contents of this ndarray.
|
||||
///
|
||||
/// This function may or may not be expensive depending on if this ndarray has contiguous data.
|
||||
///
|
||||
/// If this ndarray is not C-contiguous, this function will allocate memory on the stack for the `data` field of
|
||||
/// the returned [`SimpleNDArray`] and copy contents of this ndarray to there.
|
||||
///
|
||||
/// If this ndarray is C-contiguous, contents of this ndarray will not be copied. The created [`SimpleNDArray`]
|
||||
/// will have the same `data` field as this ndarray.
|
||||
///
|
||||
/// The `item_model` sets the [`Model`] of the returned [`SimpleNDArray`]'s `Item` model, and should match the
|
||||
/// `ctx.get_llvm_type()` of this ndarray's `dtype`. Otherwise this function panics.
|
||||
pub fn make_simple_ndarray<G: CodeGenerator + ?Sized, Item: Model<'ctx>>(
|
||||
&self,
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
item_model: Item,
|
||||
) -> Ptr<'ctx, StructModel<SimpleNDArray<Item>>> {
|
||||
// Sanity check on `self.dtype` and `item_model`.
|
||||
let dtype_llvm = ctx.get_llvm_type(generator, self.dtype);
|
||||
item_model.check_type(generator, ctx.ctx, dtype_llvm).unwrap();
|
||||
|
||||
let simple_ndarray_model = StructModel(SimpleNDArray { item: item_model });
|
||||
|
||||
let current_bb = ctx.builder.get_insert_block().unwrap();
|
||||
let then_bb = ctx.ctx.insert_basic_block_after(current_bb, "then_bb");
|
||||
let else_bb = ctx.ctx.insert_basic_block_after(then_bb, "else_bb");
|
||||
let end_bb = ctx.ctx.insert_basic_block_after(else_bb, "end_bb");
|
||||
|
||||
// Allocate and setup the resulting [`SimpleNDArray`].
|
||||
let result = simple_ndarray_model.alloca(generator, ctx, "simple_ndarray");
|
||||
|
||||
// Set ndims and shape.
|
||||
let ndims = self.get_ndims(generator, ctx.ctx);
|
||||
result.set(ctx, |f| f.ndims, ndims);
|
||||
|
||||
let shape = self.value.get(generator, ctx, |f| f.shape, "shape");
|
||||
result.set(ctx, |f| f.shape, shape);
|
||||
|
||||
// Set data, but we do things differently if this ndarray is contiguous.
|
||||
let is_contiguous = self.is_c_contiguous(generator, ctx);
|
||||
ctx.builder.build_conditional_branch(is_contiguous.value, then_bb, else_bb).unwrap();
|
||||
|
||||
// Inserting into then_bb; This ndarray is contiguous.
|
||||
let data = self.value.get(generator, ctx, |f| f.data, "");
|
||||
let data = data.pointer_cast(generator, ctx, item_model, "");
|
||||
result.set(ctx, |f| f.data, data);
|
||||
ctx.builder.build_unconditional_branch(end_bb).unwrap();
|
||||
|
||||
// Inserting into else_bb; This ndarray is not contiguous. Do a full-copy on `data`.
|
||||
// TODO: Reimplement this? This method does give us the contiguous `data`, but
|
||||
// this creates a few extra bytes of useless information because an entire NDArray
|
||||
// is allocated. Though this is super convenient.
|
||||
let data = self.make_clone(generator, ctx, "").value.get(generator, ctx, |f| f.data, "");
|
||||
let data = data.pointer_cast(generator, ctx, item_model, "");
|
||||
result.set(ctx, |f| f.data, data);
|
||||
ctx.builder.build_unconditional_branch(end_bb).unwrap();
|
||||
|
||||
// Reposition to end_bb for continuation
|
||||
ctx.builder.position_at_end(end_bb);
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Create an [`NDArrayObject`] from a [`SimpleNDArray`].
|
||||
///
|
||||
/// This operation is super cheap. The newly created [`NDArray`] will not copy contents
|
||||
/// from `simple_ndarray`, but only having its `data` and `shape` pointing to `simple_array`.
|
||||
pub fn from_simple_ndarray<G: CodeGenerator + ?Sized, Item: Model<'ctx>>(
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
simple_ndarray: Ptr<'ctx, StructModel<SimpleNDArray<Item>>>,
|
||||
dtype: Type,
|
||||
ndims: u64,
|
||||
) -> Self {
|
||||
// Sanity check on `dtype` and `simple_array`'s `Item` model.
|
||||
let dtype_llvm = ctx.get_llvm_type(generator, dtype);
|
||||
simple_ndarray.model.0 .0.item.check_type(generator, ctx.ctx, dtype_llvm).unwrap();
|
||||
|
||||
let byte_model = IntModel(Byte);
|
||||
|
||||
// TODO: Check if `ndims` is consistent with that in `simple_array`?
|
||||
|
||||
// Allocate the resulting ndarray.
|
||||
let ndarray = NDArrayObject::alloca_uninitialized(
|
||||
generator,
|
||||
ctx,
|
||||
dtype,
|
||||
ndims,
|
||||
"from_simple_ndarray",
|
||||
);
|
||||
|
||||
// Set data, shape by simply copying addresses.
|
||||
let data = simple_ndarray
|
||||
.get(generator, ctx, |f| f.data, "")
|
||||
.pointer_cast(generator, ctx, byte_model, "data");
|
||||
ndarray.value.set(ctx, |f| f.data, data);
|
||||
|
||||
let shape = simple_ndarray.get(generator, ctx, |f| f.shape, "shape");
|
||||
ndarray.value.set(ctx, |f| f.shape, shape);
|
||||
|
||||
// Set strides. We know `data` is contiguous.
|
||||
ndarray.update_strides_by_shape(generator, ctx);
|
||||
|
||||
ndarray
|
||||
}
|
||||
|
||||
/// Get the `np.size()` of this ndarray.
|
||||
pub fn size<G: CodeGenerator + ?Sized>(
|
||||
&self,
|
||||
|
@ -175,7 +275,7 @@ impl<'ctx> NDArrayObject<'ctx> {
|
|||
/// Allocate an ndarray on the stack given its `ndims` and `dtype`.
|
||||
///
|
||||
/// `shape` and `strides` will be automatically allocated on the stack.
|
||||
///
|
||||
//e
|
||||
/// The returned ndarray's content will be:
|
||||
/// - `data`: set to `nullptr`.
|
||||
/// - `itemsize`: set to the `sizeof()` of `dtype`.
|
||||
|
@ -479,7 +579,8 @@ impl<'ctx> NDArrayObject<'ctx> {
|
|||
new_shape: Ptr<'ctx, IntModel<SizeT>>,
|
||||
) -> Self {
|
||||
// TODO: The current criterion for whether to do a full copy or not is by checking `is_c_contiguous`,
|
||||
// but this is not optimal. Look into how numpy does it.
|
||||
// but this is not optimal - there are cases when the ndarray is not contiguous but could be reshaped
|
||||
// without copying data. Look into how numpy does it.
|
||||
|
||||
let current_bb = ctx.builder.get_insert_block().unwrap();
|
||||
let then_bb = ctx.ctx.insert_basic_block_after(current_bb, "then_bb");
|
||||
|
|
|
@ -13,7 +13,7 @@ use super::{NDArrayObject, NDArrayOut};
|
|||
|
||||
impl<'ctx> NDArrayObject<'ctx> {
|
||||
/// TODO: Document me
|
||||
fn np_matmul_helper<G: CodeGenerator + ?Sized>(
|
||||
fn matmul_helper<G: CodeGenerator + ?Sized>(
|
||||
generator: &mut G,
|
||||
ctx: &mut CodeGenContext<'ctx, '_>,
|
||||
a: Self,
|
||||
|
@ -113,7 +113,7 @@ impl<'ctx> NDArrayObject<'ctx> {
|
|||
|
||||
// NOTE: `result` will always be a newly allocated ndarray.
|
||||
// Current implementation cannot do in-place matrix muliplication.
|
||||
let mut result = NDArrayObject::np_matmul_helper(generator, ctx, new_a, new_b);
|
||||
let mut result = NDArrayObject::matmul_helper(generator, ctx, new_a, new_b);
|
||||
|
||||
let i32_model = IntModel(Int32); // TODO: Upgrade to SizeT
|
||||
let zero = i32_model.const_0(generator, ctx.ctx);
|
||||
|
|
|
@ -437,7 +437,7 @@ pub fn gen_setitem<'ctx, G: CodeGenerator>(
|
|||
let value =
|
||||
split_scalar_or_ndarray(generator, ctx, value, value_ty).as_ndarray(generator, ctx);
|
||||
|
||||
let broadcast_result = NDArrayObject::broadcast_all(generator, ctx, &[target, value]);
|
||||
let broadcast_result = NDArrayObject::broadcast(generator, ctx, &[target, value]);
|
||||
|
||||
let target = broadcast_result.ndarrays[0];
|
||||
let value = broadcast_result.ndarrays[1];
|
||||
|
|
|
@ -2,7 +2,7 @@ use inkwell::context::Context;
|
|||
|
||||
use crate::codegen::model::*;
|
||||
|
||||
use super::{object::ndarray::NpArrayFields, CodeGenerator};
|
||||
use super::CodeGenerator;
|
||||
|
||||
/// Fields of [`CSlice`]
|
||||
pub struct CSliceFields<'ctx, F: FieldTraversal<'ctx>> {
|
||||
|
@ -130,11 +130,23 @@ impl<'ctx, Item: Model<'ctx>> StructKind<'ctx> for List<Item> {
|
|||
}
|
||||
}
|
||||
|
||||
/// Fields of [`NDArray`]
|
||||
pub struct NDArrayFields<'ctx, F: FieldTraversal<'ctx>> {
|
||||
pub data: F::Out<PtrModel<IntModel<Byte>>>,
|
||||
pub itemsize: F::Out<IntModel<SizeT>>,
|
||||
pub ndims: F::Out<IntModel<SizeT>>,
|
||||
pub shape: F::Out<PtrModel<IntModel<SizeT>>>,
|
||||
pub strides: F::Out<PtrModel<IntModel<SizeT>>>,
|
||||
}
|
||||
|
||||
/// A strided ndarray in NAC3.
|
||||
///
|
||||
/// See IRRT implementation for details about its fields.
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct NDArray;
|
||||
|
||||
impl<'ctx> StructKind<'ctx> for NDArray {
|
||||
type Fields<F: FieldTraversal<'ctx>> = NpArrayFields<'ctx, F>;
|
||||
type Fields<F: FieldTraversal<'ctx>> = NDArrayFields<'ctx, F>;
|
||||
|
||||
fn traverse_fields<F: FieldTraversal<'ctx>>(&self, traversal: &mut F) -> Self::Fields<F> {
|
||||
Self::Fields {
|
||||
|
@ -146,3 +158,30 @@ impl<'ctx> StructKind<'ctx> for NDArray {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fields of [`SimpleNDArray`]
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct SimpleNDArrayFields<'ctx, F: FieldTraversal<'ctx>, Item: Model<'ctx>> {
|
||||
pub ndims: F::Out<IntModel<SizeT>>,
|
||||
pub shape: F::Out<PtrModel<IntModel<SizeT>>>,
|
||||
pub data: F::Out<PtrModel<Item>>,
|
||||
}
|
||||
|
||||
/// An ndarray without strides and non-opaque `data` field in NAC3.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct SimpleNDArray<M> {
|
||||
/// [`Model`] of the items.
|
||||
pub item: M,
|
||||
}
|
||||
|
||||
impl<'ctx, Item: Model<'ctx>> StructKind<'ctx> for SimpleNDArray<Item> {
|
||||
type Fields<F: FieldTraversal<'ctx>> = SimpleNDArrayFields<'ctx, F, Item>;
|
||||
|
||||
fn traverse_fields<F: FieldTraversal<'ctx>>(&self, traversal: &mut F) -> Self::Fields<F> {
|
||||
Self::Fields {
|
||||
ndims: traversal.add_auto("ndims"),
|
||||
shape: traversal.add_auto("shape"),
|
||||
data: traversal.add("data", PtrModel(self.item)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue