WIP: core/ndstrides: {make,from}_simple_ndarray

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 }
     }

nac3core/src/codegen/model/structure.rs

@@ -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())
     }

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<G: CodeGenerator + ?Sized>(
+    pub fn broadcast<G: CodeGenerator + ?Sized>(
         generator: &mut G,
         ctx: &mut CodeGenContext<'ctx, '_>,
         ndarrays: &[Self],

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 } => {

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<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");

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<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);

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];

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<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)),
+        }
+    }
+}