WIP

nac3artiq/src/codegen.rs

@@ -721,7 +721,9 @@ 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

@@ -32,7 +32,7 @@ use super::{
         gen_for_callback_incrementing, gen_if_callback, gen_if_else_expr_callback, gen_raise,
         gen_var,
     },
-    types::{ListType, ProxyType},
+    types::{ListType, NDArrayType, ProxyType},
     values::{
         ArrayLikeIndexer, ArrayLikeValue, ListValue, NDArrayValue, ProxyValue, RangeValue,
         TypedArrayLikeAccessor, UntypedArrayLikeAccessor,
@@ -41,11 +41,7 @@ use super::{
 };
 use crate::{
     symbol_resolver::{SymbolValue, ValueEnum},
-    toplevel::{
+    toplevel::{helper::PrimDef, numpy::unpack_ndarray_var_tys, DefinitionId, TopLevelDef},
-        helper::PrimDef,
-        numpy::{make_ndarray_ty, unpack_ndarray_var_tys},
-        DefinitionId, TopLevelDef,
-    },
     typecheck::{
         magic_methods::{Binop, BinopVariant, HasOpInfo},
         typedef::{FunSignature, FuncArg, Type, TypeEnum, TypeVarId, Unifier, VarMap},
@@ -2595,14 +2591,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 ndarray_ty =
-        make_ndarray_ty(&mut ctx.unifier, &ctx.primitives, Some(ty), Some(ndarray_ndims_ty));
-    let llvm_pndarray_t = ctx.get_llvm_type(generator, ndarray_ty).into_pointer_type();
-    let llvm_ndarray_t = llvm_pndarray_t.get_element_type().into_struct_type();
     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();
 
@@ -2797,26 +2785,16 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
 
                 let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
 
+                let num_dims = v.load_ndims(ctx);
+                let num_dims = ctx
+                    .builder
+                    .build_int_sub(num_dims, llvm_usize.const_int(1, false), "")
+                    .unwrap();
+
                 // Create a new array, remove the top dimension from the dimension-size-list, and copy the
                 // elements over
-                let subscripted_ndarray =
+                let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_ndarray_data_t)
-                    generator.gen_var_alloc(ctx, llvm_ndarray_t.into(), None)?;
+                    .construct_uninitialized(generator, ctx, num_dims, None);
-                let ndarray = NDArrayValue::from_pointer_value(
-                    subscripted_ndarray,
-                    llvm_ndarray_data_t,
-                    None,
-                    llvm_usize,
-                    None,
-                );
-
-                let num_dims = v.load_ndims(ctx);
-                ndarray.store_ndims(
-                    ctx,
-                    generator,
-                    ctx.builder
-                        .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);
@@ -2856,9 +2834,10 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
                 );
                 let ndarray_num_elems = ctx
                     .builder
-                    .build_int_z_extend_or_bit_cast(ndarray_num_elems, sizeof_elem.get_type(), "")
+                    .build_int_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)
-        .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,10 +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 = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty)
+        .construct_uninitialized(generator, ctx, num_dims, None);
 
     let ndarray_num_dims = ndarray.load_ndims(ctx);
     ndarray.create_shape(ctx, llvm_usize, ndarray_num_dims);
@@ -154,7 +136,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 +176,11 @@ pub fn create_ndarray_const_shape<'ctx, G: CodeGenerator + ?Sized>(
 
     let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_dtype)
         .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, '_>,
@@ -1262,8 +1223,10 @@ 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() {
+    let ndarray =
+        if slices.is_empty() {
             create_ndarray_dyn_shape(
                 generator,
                 ctx,
@@ -1275,8 +1238,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)
-        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);
@@ -1339,7 +1302,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.rs

@@ -190,22 +190,20 @@ impl<'ctx> NDArrayType<'ctx> {
         &self,
         generator: &mut G,
         ctx: &mut CodeGenContext<'ctx, '_>,
-        ndims: u64,
+        // ndims: u64,
+        ndims: IntValue<'ctx>,
         name: Option<&'ctx str>,
     ) -> <Self as ProxyType<'ctx>>::Value {
         let ndarray = self.new_value(generator, ctx, name);
 
-        let itemsize = ctx
+        let itemsize =
-            .builder
+            ctx.builder.build_int_cast(self.dtype.size_of().unwrap(), self.llvm_usize, "").unwrap();
-            .build_int_z_extend_or_bit_cast(self.dtype.size_of().unwrap(), self.llvm_usize, "")
-            .unwrap();
         ndarray.store_itemsize(ctx, generator, itemsize);
 
-        let ndims_val = self.llvm_usize.const_int(ndims, false);
+        ndarray.store_ndims(ctx, generator, ndims);
-        ndarray.store_ndims(ctx, generator, ndims_val);
 
-        ndarray.create_shape(ctx, self.llvm_usize, ndims_val);
+        ndarray.create_shape(ctx, self.llvm_usize, ndims);
-        ndarray.create_strides(ctx, self.llvm_usize, ndims_val);
+        ndarray.create_strides(ctx, self.llvm_usize, ndims);
 
         ndarray
     }
@@ -221,7 +219,14 @@ impl<'ctx> NDArrayType<'ctx> {
         shape: &[u64],
         name: Option<&'ctx str>,
     ) -> <Self as ProxyType<'ctx>>::Value {
-        let ndarray = self.construct_uninitialized(generator, ctx, shape.len() as u64, name);
+        let llvm_usize = generator.get_size_type(ctx.ctx);
+
+        let ndarray = self.construct_uninitialized(
+            generator,
+            ctx,
+            llvm_usize.const_int(shape.len() as u64, false),
+            name,
+        );
 
         // Write shape
         let ndarray_shape = ndarray.shape();
@@ -251,7 +256,14 @@ impl<'ctx> NDArrayType<'ctx> {
         shape: &[IntValue<'ctx>],
         name: Option<&'ctx str>,
     ) -> <Self as ProxyType<'ctx>>::Value {
-        let ndarray = self.construct_uninitialized(generator, ctx, shape.len() as u64, name);
+        let llvm_usize = generator.get_size_type(ctx.ctx);
+
+        let ndarray = self.construct_uninitialized(
+            generator,
+            ctx,
+            llvm_usize.const_int(shape.len() as u64, false),
+            name,
+        );
 
         // Write shape
         let ndarray_shape = ndarray.shape();

nac3core/src/codegen/values/ndarray.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},
@@ -179,19 +179,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 =
+        let nbytes = self.nbytes(generator, ctx);
-            ctx.builder.build_int_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`.
@@ -199,6 +203,133 @@ 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: &mut G,
+        ctx: &mut 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>,
+    ) {
+        assert_eq!(self.ndims, src_ndarray.ndims);
+        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: &mut G,
+        ctx: &mut 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>,
+    ) {
+        assert_eq!(self.ndims, src_ndarray.ndims);
+        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: &mut G,
+        ctx: &mut 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: &mut G,
+        ctx: &mut 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: &mut G,
+        ctx: &mut 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: &mut G,
+        ctx: &mut 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: &mut G,
+        ctx: &mut 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: &mut G,
+        ctx: &mut 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> {