From 18dcbf5bbc350bec88bbfbc8f2df54c2f398e7b0 Mon Sep 17 00:00:00 2001 From: lyken Date: Wed, 14 Aug 2024 11:33:56 +0800 Subject: [PATCH] WIP: core/ndstrides: {make,from}_simple_ndarray --- nac3core/src/codegen/model/core.rs | 1 + nac3core/src/codegen/model/structure.rs | 2 +- .../src/codegen/object/ndarray/broadcast.rs | 2 +- .../src/codegen/object/ndarray/mapping.rs | 2 +- nac3core/src/codegen/object/ndarray/mod.rs | 123 ++++++++++++++++-- .../src/codegen/object/ndarray/nalgebra.rs | 0 .../src/codegen/object/ndarray/product.rs | 4 +- nac3core/src/codegen/stmt.rs | 2 +- nac3core/src/codegen/structure.rs | 43 +++++- 9 files changed, 160 insertions(+), 19 deletions(-) create mode 100644 nac3core/src/codegen/object/ndarray/nalgebra.rs diff --git a/nac3core/src/codegen/model/core.rs b/nac3core/src/codegen/model/core.rs index 1afe80f2..08250a64 100644 --- a/nac3core/src/codegen/model/core.rs +++ b/nac3core/src/codegen/model/core.rs @@ -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 } } diff --git a/nac3core/src/codegen/model/structure.rs b/nac3core/src/codegen/model/structure.rs index c2c9adf2..94f438ff 100644 --- a/nac3core/src/codegen/model/structure.rs +++ b/nac3core/src/codegen/model/structure.rs @@ -22,7 +22,7 @@ pub trait FieldTraversal<'ctx> { fn add>(&mut self, name: &'static str, model: M) -> Self::Out; - /// Like [`FieldTraversal::visit`] but [`Model`] is automatically inferred. + /// Like [`FieldTraversal::visit`] but [`Model`] is automatically inferred from [`Default`] trait. fn add_auto + Default>(&mut self, name: &'static str) -> Self::Out { self.add(name, M::default()) } diff --git a/nac3core/src/codegen/object/ndarray/broadcast.rs b/nac3core/src/codegen/object/ndarray/broadcast.rs index 68098d7b..89ad7a5c 100644 --- a/nac3core/src/codegen/object/ndarray/broadcast.rs +++ b/nac3core/src/codegen/object/ndarray/broadcast.rs @@ -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( + pub fn broadcast( generator: &mut G, ctx: &mut CodeGenContext<'ctx, '_>, ndarrays: &[Self], diff --git a/nac3core/src/codegen/object/ndarray/mapping.rs b/nac3core/src/codegen/object/ndarray/mapping.rs index 4486d941..234a56e7 100644 --- a/nac3core/src/codegen/object/ndarray/mapping.rs +++ b/nac3core/src/codegen/object/ndarray/mapping.rs @@ -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 } => { diff --git a/nac3core/src/codegen/object/ndarray/mod.rs b/nac3core/src/codegen/object/ndarray/mod.rs index fc0bfbcd..2549c343 100644 --- a/nac3core/src/codegen/object/ndarray/mod.rs +++ b/nac3core/src/codegen/object/ndarray/mod.rs @@ -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>>, - pub itemsize: F::Out>, - pub ndims: F::Out>, - pub shape: F::Out>>, - pub strides: F::Out>>, -} - /// 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>( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + item_model: Item, + ) -> Ptr<'ctx, StructModel>> { + // 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>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + simple_ndarray: Ptr<'ctx, StructModel>>, + 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( &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>, ) -> 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"); diff --git a/nac3core/src/codegen/object/ndarray/nalgebra.rs b/nac3core/src/codegen/object/ndarray/nalgebra.rs new file mode 100644 index 00000000..e69de29b diff --git a/nac3core/src/codegen/object/ndarray/product.rs b/nac3core/src/codegen/object/ndarray/product.rs index cb0f99a1..0903494c 100644 --- a/nac3core/src/codegen/object/ndarray/product.rs +++ b/nac3core/src/codegen/object/ndarray/product.rs @@ -13,7 +13,7 @@ use super::{NDArrayObject, NDArrayOut}; impl<'ctx> NDArrayObject<'ctx> { /// TODO: Document me - fn np_matmul_helper( + fn matmul_helper( 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); diff --git a/nac3core/src/codegen/stmt.rs b/nac3core/src/codegen/stmt.rs index 0eaaf86d..554599d1 100644 --- a/nac3core/src/codegen/stmt.rs +++ b/nac3core/src/codegen/stmt.rs @@ -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]; diff --git a/nac3core/src/codegen/structure.rs b/nac3core/src/codegen/structure.rs index 1c0c1531..049f1197 100644 --- a/nac3core/src/codegen/structure.rs +++ b/nac3core/src/codegen/structure.rs @@ -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 { } } +/// Fields of [`NDArray`] +pub struct NDArrayFields<'ctx, F: FieldTraversal<'ctx>> { + pub data: F::Out>>, + pub itemsize: F::Out>, + pub ndims: F::Out>, + pub shape: F::Out>>, + pub strides: F::Out>>, +} + +/// 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> = NpArrayFields<'ctx, F>; + type Fields> = NDArrayFields<'ctx, F>; fn traverse_fields>(&self, traversal: &mut F) -> Self::Fields { 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>, + pub shape: F::Out>>, + pub data: F::Out>, +} + +/// An ndarray without strides and non-opaque `data` field in NAC3. +#[derive(Debug, Clone, Copy)] +pub struct SimpleNDArray { + /// [`Model`] of the items. + pub item: M, +} + +impl<'ctx, Item: Model<'ctx>> StructKind<'ctx> for SimpleNDArray { + type Fields> = SimpleNDArrayFields<'ctx, F, Item>; + + fn traverse_fields>(&self, traversal: &mut F) -> Self::Fields { + Self::Fields { + ndims: traversal.add_auto("ndims"), + shape: traversal.add_auto("shape"), + data: traversal.add("data", PtrModel(self.item)), + } + } +}