[core] codegen: Implement NDArray functions from a0a1f35b

nac3artiq/src/codegen.rs

@@ -732,7 +732,7 @@ fn format_rpc_ret<'ctx>(
                 );
             }
 
-            ndarray.create_data(generator, ctx, llvm_elem_ty, num_elements);
+            unsafe { ndarray.create_data(generator, ctx) };
 
             let ndarray_data = ndarray.data().base_ptr(ctx, generator);
             let ndarray_data_i8 =

nac3core/src/codegen/expr.rs

@@ -2560,7 +2560,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
     generator: &mut G,
     ctx: &mut CodeGenContext<'ctx, '_>,
     ty: Type,
-    ndims: Type,
+    ndims_ty: Type,
     v: NDArrayValue<'ctx>,
     slice: &Expr<Option<Type>>,
 ) -> Result<Option<ValueEnum<'ctx>>, String> {
@@ -2568,7 +2568,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
     let llvm_i32 = ctx.ctx.i32_type();
     let llvm_usize = generator.get_size_type(ctx.ctx);
 
-    let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) else {
+    let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims_ty) else {
         codegen_unreachable!(ctx)
     };
 
@@ -2601,10 +2601,6 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
         _ => 1,
     };
 
-    let ndarray_ndims_ty = ctx.unifier.get_fresh_literal(
-        ndims.iter().map(|v| SymbolValue::U64(v - subscripted_dims)).collect(),
-        None,
-    );
     let llvm_ndarray_data_t = ctx.get_llvm_type(generator, ty).as_basic_type_enum();
     let sizeof_elem = llvm_ndarray_data_t.size_of().unwrap();
 
@@ -2798,28 +2794,19 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
                 // Accessing an element from a multi-dimensional `ndarray`
                 let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
 
+                let num_dims = extract_ndims(&ctx.unifier, ndims_ty) - 1;
+
                 // Create a new array, remove the top dimension from the dimension-size-list, and copy the
                 // elements over
-                let ndarray = NDArrayType::new(
+                let ndarray =
+                    NDArrayType::new(generator, ctx.ctx, llvm_ndarray_data_t, Some(num_dims))
+                        .construct_uninitialized(
                             generator,
-                    ctx.ctx,
-                    llvm_ndarray_data_t,
-                    Some(extract_ndims(&ctx.unifier, ndarray_ndims_ty)),
-                )
-                .alloca(generator, ctx, None);
-
-                let num_dims = v.load_ndims(ctx);
-                ndarray.store_ndims(
                             ctx,
-                    generator,
+                            llvm_usize.const_int(num_dims, false),
-                    ctx.builder
+                            None,
-                        .build_int_sub(num_dims, llvm_usize.const_int(1, false), "")
-                        .unwrap(),
                         );
 
-                let ndarray_num_dims = ndarray.load_ndims(ctx);
-                ndarray.create_shape(ctx, llvm_usize, ndarray_num_dims);
-
                 let ndarray_num_dims = ctx
                     .builder
                     .build_int_z_extend_or_bit_cast(
@@ -2857,7 +2844,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
                     .builder
                     .build_int_z_extend_or_bit_cast(ndarray_num_elems, sizeof_elem.get_type(), "")
                     .unwrap();
-                ndarray.create_data(generator, ctx, llvm_ndarray_data_t, ndarray_num_elems);
+                unsafe { ndarray.create_data(generator, ctx) };
 
                 let v_data_src_ptr = v.data().ptr_offset(ctx, generator, &index_addr, None);
                 call_memcpy_generic(

nac3core/src/codegen/numpy.rs

@@ -40,25 +40,6 @@ use crate::{
     },
 };
 
-/// Creates an uninitialized `NDArray` instance.
-fn create_ndarray_uninitialized<'ctx, G: CodeGenerator + ?Sized>(
-    generator: &mut G,
-    ctx: &mut CodeGenContext<'ctx, '_>,
-    elem_ty: Type,
-) -> Result<NDArrayValue<'ctx>, String> {
-    let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
-    let llvm_usize = generator.get_size_type(ctx.ctx);
-
-    let llvm_ndarray_t = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None)
-        .as_base_type()
-        .get_element_type()
-        .into_struct_type();
-
-    let ndarray = generator.gen_var_alloc(ctx, llvm_ndarray_t.into(), None)?;
-
-    Ok(NDArrayValue::from_pointer_value(ndarray, llvm_elem_ty, None, llvm_usize, None))
-}
-
 /// Creates an `NDArray` instance from a dynamic shape.
 ///
 /// * `elem_ty` - The element type of the `NDArray`.
@@ -84,6 +65,7 @@ where
     ) -> Result<IntValue<'ctx>, String>,
 {
     let llvm_usize = generator.get_size_type(ctx.ctx);
+    let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
 
     // Assert that all dimensions are non-negative
     let shape_len = shape_len_fn(generator, ctx, shape)?;
@@ -123,13 +105,10 @@ where
         llvm_usize.const_int(1, false),
     )?;
 
-    let ndarray = create_ndarray_uninitialized(generator, ctx, elem_ty)?;
-
     let num_dims = shape_len_fn(generator, ctx, shape)?;
-    ndarray.store_ndims(ctx, generator, num_dims);
 
-    let ndarray_num_dims = ndarray.load_ndims(ctx);
+    let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None)
-    ndarray.create_shape(ctx, llvm_usize, ndarray_num_dims);
+        .construct_uninitialized(generator, ctx, num_dims, None);
 
     // Copy the dimension sizes from shape to ndarray.dims
     let shape_len = shape_len_fn(generator, ctx, shape)?;
@@ -154,7 +133,7 @@ where
         llvm_usize.const_int(1, false),
     )?;
 
-    let ndarray = ndarray_init_data(generator, ctx, elem_ty, ndarray);
+    unsafe { ndarray.create_data(generator, ctx) };
 
     Ok(ndarray)
 }
@@ -194,32 +173,11 @@ pub fn create_ndarray_const_shape<'ctx, G: CodeGenerator + ?Sized>(
 
     let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_dtype, Some(shape.len() as u64))
         .construct_dyn_shape(generator, ctx, shape, None);
-    let ndarray = ndarray_init_data(generator, ctx, elem_ty, ndarray);
+    unsafe { ndarray.create_data(generator, ctx) };
 
     Ok(ndarray)
 }
 
-/// Initializes the `data` field of [`NDArrayValue`] based on the `ndims` and `dim_sz` fields.
-fn ndarray_init_data<'ctx, G: CodeGenerator + ?Sized>(
-    generator: &mut G,
-    ctx: &mut CodeGenContext<'ctx, '_>,
-    elem_ty: Type,
-    ndarray: NDArrayValue<'ctx>,
-) -> NDArrayValue<'ctx> {
-    let llvm_ndarray_data_t = ctx.get_llvm_type(generator, elem_ty).as_basic_type_enum();
-    assert!(llvm_ndarray_data_t.is_sized());
-
-    let ndarray_num_elems = call_ndarray_calc_size(
-        generator,
-        ctx,
-        &ndarray.shape().as_slice_value(ctx, generator),
-        (None, None),
-    );
-    ndarray.create_data(generator, ctx, llvm_ndarray_data_t, ndarray_num_elems);
-
-    ndarray
-}
-
 fn ndarray_zero_value<'ctx, G: CodeGenerator + ?Sized>(
     generator: &mut G,
     ctx: &mut CodeGenContext<'ctx, '_>,
@@ -1263,6 +1221,7 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>(
 ) -> Result<NDArrayValue<'ctx>, String> {
     let llvm_i32 = ctx.ctx.i32_type();
     let llvm_usize = generator.get_size_type(ctx.ctx);
+    let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
 
     let ndarray = if slices.is_empty() {
         create_ndarray_dyn_shape(
@@ -1276,8 +1235,8 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>(
             },
         )?
     } else {
-        let ndarray = create_ndarray_uninitialized(generator, ctx, elem_ty)?;
+        let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None)
-        ndarray.store_ndims(ctx, generator, this.load_ndims(ctx));
+            .construct_uninitialized(generator, ctx, this.load_ndims(ctx), None);
 
         let ndims = this.load_ndims(ctx);
         ndarray.create_shape(ctx, llvm_usize, ndims);
@@ -1340,7 +1299,9 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>(
         )
         .unwrap();
 
-        ndarray_init_data(generator, ctx, elem_ty, ndarray)
+        unsafe { ndarray.create_data(generator, ctx) };
+
+        ndarray
     };
 
     ndarray_sliced_copyto_impl(

nac3core/src/codegen/types/ndarray/mod.rs

@@ -214,7 +214,7 @@ impl<'ctx> NDArrayType<'ctx> {
 
         let itemsize = ctx
             .builder
-            .build_int_z_extend_or_bit_cast(self.dtype.size_of().unwrap(), self.llvm_usize, "")
+            .build_int_truncate_or_bit_cast(self.dtype.size_of().unwrap(), self.llvm_usize, "")
             .unwrap();
         ndarray.store_itemsize(ctx, generator, itemsize);
 

nac3core/src/codegen/values/ndarray/mod.rs

@@ -10,7 +10,7 @@ use super::{
 };
 use crate::codegen::{
     irrt,
-    llvm_intrinsics::call_int_umin,
+    llvm_intrinsics::{call_int_umin, call_memcpy_generic_array},
     stmt::gen_for_callback_incrementing,
     type_aligned_alloca,
     types::{structure::StructField, NDArrayType},
@@ -186,21 +186,23 @@ impl<'ctx> NDArrayValue<'ctx> {
 
     /// Convenience method for creating a new array storing data elements with the given element
     /// type `elem_ty` and `size`.
-    pub fn create_data<G: CodeGenerator + ?Sized>(
+    ///
+    /// The data buffer will be allocated on the stack, and is considered to be owned by this ndarray instance.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `shape` and `itemsize` of this ndarray instance is initialized.
+    pub unsafe fn create_data<G: CodeGenerator + ?Sized>(
         &self,
         generator: &mut G,
         ctx: &mut CodeGenContext<'ctx, '_>,
-        elem_ty: BasicTypeEnum<'ctx>,
-        size: IntValue<'ctx>,
     ) {
-        let itemsize = ctx
+        let nbytes = self.nbytes(generator, ctx);
-            .builder
-            .build_int_z_extend_or_bit_cast(elem_ty.size_of().unwrap(), size.get_type(), "")
-            .unwrap();
-        let nbytes = ctx.builder.build_int_mul(size, itemsize, "").unwrap();
 
-        let data = type_aligned_alloca(generator, ctx, elem_ty, nbytes, None);
+        let data = type_aligned_alloca(generator, ctx, self.dtype, nbytes, None);
         self.store_data(ctx, data);
+
+        self.set_strides_contiguous(generator, ctx);
     }
 
     /// Returns a proxy object to the field storing the data of this `NDArray`.
@@ -208,6 +210,173 @@ impl<'ctx> NDArrayValue<'ctx> {
     pub fn data(&self) -> NDArrayDataProxy<'ctx, '_> {
         NDArrayDataProxy(self)
     }
+
+    /// Copy shape dimensions from an array.
+    pub fn copy_shape_from_array<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+        shape: PointerValue<'ctx>,
+    ) {
+        let num_items = self.load_ndims(ctx);
+
+        call_memcpy_generic_array(
+            ctx,
+            self.shape().base_ptr(ctx, generator),
+            shape,
+            num_items,
+            ctx.ctx.bool_type().const_zero(),
+        );
+    }
+
+    /// Copy shape dimensions from an ndarray.
+    /// Panics if `ndims` mismatches.
+    pub fn copy_shape_from_ndarray<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &mut G,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        src_ndarray: NDArrayValue<'ctx>,
+    ) {
+        if self.ndims.is_some() && src_ndarray.ndims.is_some() {
+            assert_eq!(self.ndims, src_ndarray.ndims);
+        } else {
+            let self_ndims = self.load_ndims(ctx);
+            let src_ndims = src_ndarray.load_ndims(ctx);
+
+            ctx.make_assert(
+                generator,
+                ctx.builder.build_int_compare(
+                    IntPredicate::EQ,
+                    self_ndims,
+                    src_ndims,
+                    ""
+                ).unwrap(),
+                "0:AssertionError",
+                "NDArrayValue::copy_shape_from_ndarray: Expected self.ndims ({0}) == src_ndarray.ndims ({1})",
+                [Some(self_ndims), Some(src_ndims), None],
+                ctx.current_loc
+            );
+        }
+
+        let src_shape = src_ndarray.shape().base_ptr(ctx, generator);
+        self.copy_shape_from_array(generator, ctx, src_shape);
+    }
+
+    /// Copy strides dimensions from an array.
+    pub fn copy_strides_from_array<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+        strides: PointerValue<'ctx>,
+    ) {
+        let num_items = self.load_ndims(ctx);
+
+        call_memcpy_generic_array(
+            ctx,
+            self.strides().base_ptr(ctx, generator),
+            strides,
+            num_items,
+            ctx.ctx.bool_type().const_zero(),
+        );
+    }
+
+    /// Copy strides dimensions from an ndarray.
+    /// Panics if `ndims` mismatches.
+    pub fn copy_strides_from_ndarray<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &mut G,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        src_ndarray: NDArrayValue<'ctx>,
+    ) {
+        if self.ndims.is_some() && src_ndarray.ndims.is_some() {
+            assert_eq!(self.ndims, src_ndarray.ndims);
+        } else {
+            let self_ndims = self.load_ndims(ctx);
+            let src_ndims = src_ndarray.load_ndims(ctx);
+
+            ctx.make_assert(
+                generator,
+                ctx.builder.build_int_compare(
+                    IntPredicate::EQ,
+                    self_ndims,
+                    src_ndims,
+                    ""
+                ).unwrap(),
+                "0:AssertionError",
+                "NDArrayValue::copy_shape_from_ndarray: Expected self.ndims ({0}) == src_ndarray.ndims ({1})",
+                [Some(self_ndims), Some(src_ndims), None],
+                ctx.current_loc
+            );
+        }
+
+        let src_strides = src_ndarray.strides().base_ptr(ctx, generator);
+        self.copy_strides_from_array(generator, ctx, src_strides);
+    }
+
+    /// Get the `np.size()` of this ndarray.
+    pub fn size<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+    ) -> IntValue<'ctx> {
+        irrt::ndarray::call_nac3_ndarray_size(generator, ctx, *self)
+    }
+
+    /// Get the `ndarray.nbytes` of this ndarray.
+    pub fn nbytes<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+    ) -> IntValue<'ctx> {
+        irrt::ndarray::call_nac3_ndarray_nbytes(generator, ctx, *self)
+    }
+
+    /// Get the `len()` of this ndarray.
+    pub fn len<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+    ) -> IntValue<'ctx> {
+        irrt::ndarray::call_nac3_ndarray_len(generator, ctx, *self)
+    }
+
+    /// Check if this ndarray is C-contiguous.
+    ///
+    /// See NumPy's `flags["C_CONTIGUOUS"]`: <https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flags.html#numpy.ndarray.flags>
+    pub fn is_c_contiguous<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+    ) -> IntValue<'ctx> {
+        irrt::ndarray::call_nac3_ndarray_is_c_contiguous(generator, ctx, *self)
+    }
+
+    /// Call [`call_nac3_ndarray_set_strides_by_shape`] on this ndarray to update `strides`.
+    ///
+    /// Update the ndarray's strides to make the ndarray contiguous.
+    pub fn set_strides_contiguous<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+    ) {
+        irrt::ndarray::call_nac3_ndarray_set_strides_by_shape(generator, ctx, *self);
+    }
+
+    /// Copy data from another ndarray.
+    ///
+    /// This ndarray and `src` is that their `np.size()` should be the same. Their shapes
+    /// do not matter. The copying order is determined by how their flattened views look.
+    ///
+    /// Panics if the `dtype`s of ndarrays are different.
+    pub fn copy_data_from<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &G,
+        ctx: &CodeGenContext<'ctx, '_>,
+        src: NDArrayValue<'ctx>,
+    ) {
+        assert_eq!(self.dtype, src.dtype, "self and src dtype should match");
+        irrt::ndarray::call_nac3_ndarray_copy_data(generator, ctx, src, *self);
+    }
 }
 
 impl<'ctx> ProxyValue<'ctx> for NDArrayValue<'ctx> {