WIP

nac3artiq/src/codegen.rs

@@ -15,9 +15,10 @@ use pyo3::{
 use nac3core::{
     codegen::{
         expr::{destructure_range, gen_call},
-        irrt::call_ndarray_calc_size,
+        irrt::ndarray::call_ndarray_calc_size,
         llvm_intrinsics::{call_int_smax, call_memcpy_generic, call_stackrestore, call_stacksave},
         stmt::{gen_block, gen_for_callback_incrementing, gen_if_callback, gen_with},
+        type_aligned_alloca,
         types::{NDArrayType, ProxyType},
         values::{
             ArrayLikeIndexer, ArrayLikeValue, ArraySliceValue, ListValue, NDArrayValue, ProxyValue,
@@ -642,27 +643,12 @@ fn format_rpc_ret<'ctx>(
             // (4 + 4 * ndims) bytes with 8-byte alignment
             let sizeof_dims =
                 ctx.builder.build_int_mul(ndarray.load_ndims(ctx), llvm_usize_sizeof, "").unwrap();
-            let unaligned_buffer_size =
+            let buffer_size =
                 ctx.builder.build_int_add(sizeof_dims, llvm_pdata_sizeof, "").unwrap();
-            let buffer_size = round_up(ctx, unaligned_buffer_size, llvm_usize.const_int(8, false));
 
             let stackptr = call_stacksave(ctx, None);
-            // Just to be absolutely sure, alloca in [i8 x 8] slices to force 8-byte alignment
+            let buffer =
-            let buffer = ctx
+                type_aligned_alloca(generator, ctx, llvm_i8_8, buffer_size, Some("rpc.buffer"));
-                .builder
-                .build_array_alloca(
-                    llvm_i8_8,
-                    ctx.builder
-                        .build_int_unsigned_div(buffer_size, llvm_usize.const_int(8, false), "")
-                        .unwrap(),
-                    "rpc.buffer",
-                )
-                .unwrap();
-            let buffer = ctx
-                .builder
-                .build_bit_cast(buffer, llvm_pi8, "")
-                .map(BasicValueEnum::into_pointer_value)
-                .unwrap();
             let buffer = ArraySliceValue::from_ptr_val(buffer, buffer_size, None);
 
             // The first call to `rpc_recv` reads the top-level ndarray object: [pdata, shape]
@@ -736,7 +722,7 @@ fn format_rpc_ret<'ctx>(
             }
 
             unsafe {
-                ndarray.create_data(generator, ctx, llvm_elem_ty, num_elements);
+                ndarray.create_data(generator, ctx, num_elements);
             }
 
             let ndarray_data = ndarray.data().base_ptr(ctx, generator);

nac3core/src/codegen/builtin_fns.rs

@@ -923,7 +923,8 @@ pub fn call_numpy_max_min<'ctx, G: CodeGenerator + ?Sized>(
             let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
 
             let n = NDArrayValue::from_pointer_value(n, llvm_elem_ty, None, llvm_usize, None);
-            let n_sz = irrt::call_ndarray_calc_size(generator, ctx, &n.shape(), (None, None));
+            let n_sz =
+                irrt::ndarray::call_ndarray_calc_size(generator, ctx, &n.shape(), (None, None));
             if ctx.registry.llvm_options.opt_level == OptimizationLevel::None {
                 let n_sz_eqz = ctx
                     .builder

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,
@@ -43,7 +43,7 @@ use crate::{
     symbol_resolver::{SymbolValue, ValueEnum},
     toplevel::{
         helper::PrimDef,
-        numpy::{make_ndarray_ty, unpack_ndarray_var_tys},
+        numpy::unpack_ndarray_var_tys,
         DefinitionId, TopLevelDef,
     },
     typecheck::{
@@ -2595,14 +2595,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 +2789,15 @@ 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);
@@ -2848,7 +2829,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
                     llvm_i1.const_zero(),
                 );
 
-                let ndarray_num_elems = call_ndarray_calc_size(
+                let ndarray_num_elems = ndarray::call_ndarray_calc_size(
                     generator,
                     ctx,
                     &ndarray.shape().as_slice_value(ctx, generator),
@@ -2859,7 +2840,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
                     .build_int_z_extend_or_bit_cast(ndarray_num_elems, sizeof_elem.get_type(), "")
                     .unwrap();
                 unsafe {
-                    ndarray.create_data(generator, ctx, llvm_ndarray_data_t, ndarray_num_elems);
+                    ndarray.create_data(generator, ctx, ndarray_num_elems);
                 }
 
                 let v_data_src_ptr = v.data().ptr_offset(ctx, generator, &index_addr, None);
@@ -3618,10 +3599,7 @@ pub fn create_fn_and_call<'ctx>(
     configure: Option<&dyn Fn(&FunctionValue<'ctx>)>,
 ) -> Option<BasicValueEnum<'ctx>> {
     let intrinsic_fn = ctx.module.get_function(fn_name).unwrap_or_else(|| {
-        let params = params.iter()
+        let params = params.iter().copied().map(BasicTypeEnum::into).collect_vec();
-            .copied()
-            .map(BasicTypeEnum::into)
-            .collect_vec();
         let fn_type = if let Some(ret_type) = ret_type {
             ret_type.fn_type(params.as_slice(), is_var_args)
         } else {
@@ -3635,10 +3613,7 @@ pub fn create_fn_and_call<'ctx>(
         configure(&intrinsic_fn);
     }
 
-    let args = args.iter()
+    let args = args.iter().copied().map(BasicValueEnum::into).collect_vec();
-        .copied()
-        .map(BasicValueEnum::into)
-        .collect_vec();
     ctx.builder
         .build_call(intrinsic_fn, args.as_slice(), call_value_name.unwrap_or_default())
         .map(CallSiteValue::try_as_basic_value)
@@ -3660,7 +3635,8 @@ pub fn create_and_call_function<'ctx>(
     configure: Option<&dyn Fn(&FunctionValue<'ctx>)>,
 ) -> Option<BasicValueEnum<'ctx>> {
     let param_tys = params.iter().map(|(ty, _)| ty).copied().map(BasicTypeEnum::into).collect_vec();
-    let arg_values = params.iter().map(|(_, value)| value).copied().map(BasicValueEnum::into).collect_vec();
+    let arg_values =
+        params.iter().map(|(_, value)| value).copied().map(BasicValueEnum::into).collect_vec();
 
     create_fn_and_call(
         ctx,
@@ -3686,9 +3662,7 @@ pub fn infer_and_call_function<'ctx>(
     value_name: Option<&str>,
     configure: Option<&dyn Fn(&FunctionValue<'ctx>)>,
 ) -> Option<BasicValueEnum<'ctx>> {
-    let param_tys = args.iter()
+    let param_tys = args.iter().map(BasicValueEnum::get_type).collect_vec();
-        .map(BasicValueEnum::get_type)
-        .collect_vec();
 
     create_fn_and_call(
         ctx,

nac3core/src/codegen/irrt/mod.rs

@@ -13,12 +13,11 @@ use super::{CodeGenContext, CodeGenerator};
 use crate::{symbol_resolver::SymbolResolver, typecheck::typedef::Type};
 pub use list::*;
 pub use math::*;
-pub use ndarray::*;
 pub use slice::*;
 
 mod list;
 mod math;
-mod ndarray;
+pub mod ndarray;
 mod slice;
 
 #[must_use]

nac3core/src/codegen/llvm_intrinsics.rs

@@ -201,6 +201,52 @@ pub fn call_memcpy_generic<'ctx>(
     call_memcpy(ctx, dest, src, len, is_volatile);
 }
 
+/// Invokes the `llvm.memcpy` intrinsic.
+///
+/// Unlike [`call_memcpy`], this function accepts any type of pointer value. If `dest` or `src` is
+/// not a pointer to an integer, the pointer(s) will be cast to `i8*` before invoking `memcpy`.
+/// Moreover, `len` now refers to the number of elements (rather than bytes) to copy.
+pub fn call_memcpy_generic_array<'ctx>(
+    ctx: &CodeGenContext<'ctx, '_>,
+    dest: PointerValue<'ctx>,
+    src: PointerValue<'ctx>,
+    len: IntValue<'ctx>,
+    is_volatile: IntValue<'ctx>,
+) {
+    let llvm_i8 = ctx.ctx.i8_type();
+    let llvm_p0i8 = llvm_i8.ptr_type(AddressSpace::default());
+    let llvm_sizeof_expr_t = llvm_i8.size_of().get_type();
+
+    let dest_elem_t = dest.get_type().get_element_type();
+    let src_elem_t = src.get_type().get_element_type();
+
+    let dest = if matches!(dest_elem_t, IntType(t) if t.get_bit_width() == 8) {
+        dest
+    } else {
+        ctx.builder
+            .build_bit_cast(dest, llvm_p0i8, "")
+            .map(BasicValueEnum::into_pointer_value)
+            .unwrap()
+    };
+    let src = if matches!(src_elem_t, IntType(t) if t.get_bit_width() == 8) {
+        src
+    } else {
+        ctx.builder
+            .build_bit_cast(src, llvm_p0i8, "")
+            .map(BasicValueEnum::into_pointer_value)
+            .unwrap()
+    };
+
+    let len = ctx.builder.build_int_cast(len, llvm_sizeof_expr_t, "").unwrap();
+    let len = ctx.builder.build_int_mul(
+        len,
+        src_elem_t.size_of().unwrap(),
+        ""
+    ).unwrap();
+
+    call_memcpy(ctx, dest, src, len, is_volatile);
+}
+
 /// Macro to find and generate build call for llvm intrinsic (body of llvm intrinsic function)
 ///
 /// Arguments:
@@ -343,3 +389,25 @@ pub fn call_float_powi<'ctx>(
         .map(Either::unwrap_left)
         .unwrap()
 }
+
+/// Invokes the [`llvm.ctpop`](https://llvm.org/docs/LangRef.html#llvm-ctpop-intrinsic) intrinsic.
+pub fn call_int_ctpop<'ctx>(
+    ctx: &CodeGenContext<'ctx, '_>,
+    src: IntValue<'ctx>,
+    name: Option<&str>,
+) -> IntValue<'ctx> {
+    const FN_NAME: &str = "llvm.ctpop";
+
+    let llvm_src_t = src.get_type();
+
+    let intrinsic_fn = Intrinsic::find(FN_NAME)
+        .and_then(|intrinsic| intrinsic.get_declaration(&ctx.module, &[llvm_src_t.into()]))
+        .unwrap();
+
+    ctx.builder
+        .build_call(intrinsic_fn, &[src.into()], name.unwrap_or_default())
+        .map(CallSiteValue::try_as_basic_value)
+        .map(|v| v.map_left(BasicValueEnum::into_int_value))
+        .map(Either::unwrap_left)
+        .unwrap()
+}

nac3core/src/codegen/mod.rs

@@ -1119,3 +1119,106 @@ fn gen_in_range_check<'ctx>(
 fn get_va_count_arg_name(arg_name: StrRef) -> StrRef {
     format!("__{}_va_count", &arg_name).into()
 }
+
+/// Returns the alignment of the type.
+///
+/// This is necessary as `get_alignment` is not implemented as part of [`BasicType`].
+pub fn get_type_alignment<'ctx>(ty: impl Into<BasicTypeEnum<'ctx>>) -> IntValue<'ctx> {
+    match ty.into() {
+        BasicTypeEnum::ArrayType(ty) => ty.get_alignment(),
+        BasicTypeEnum::FloatType(ty) => ty.get_alignment(),
+        BasicTypeEnum::IntType(ty) => ty.get_alignment(),
+        BasicTypeEnum::PointerType(ty) => ty.get_alignment(),
+        BasicTypeEnum::StructType(ty) => ty.get_alignment(),
+        BasicTypeEnum::VectorType(ty) => ty.get_alignment(),
+    }
+}
+
+/// Inserts an `alloca` instruction with allocation `size` given in bytes and the alignment of the
+/// given type.
+///
+/// The returned [`PointerValue`] will have a type of `i8*`, a size of at least `size`, and will be
+/// aligned with the alignment of `align_ty`.
+pub fn type_aligned_alloca<'ctx, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    align_ty: impl Into<BasicTypeEnum<'ctx>>,
+    size: IntValue<'ctx>,
+    name: Option<&str>,
+) -> PointerValue<'ctx> {
+    /// Round `val` up to its modulo `power_of_two`.
+    fn round_up<'ctx>(
+        ctx: &CodeGenContext<'ctx, '_>,
+        val: IntValue<'ctx>,
+        power_of_two: IntValue<'ctx>,
+    ) -> IntValue<'ctx> {
+        debug_assert_eq!(
+            val.get_type().get_bit_width(),
+            power_of_two.get_type().get_bit_width(),
+            "`val` ({}) and `power_of_two` ({}) must be the same type",
+            val.get_type(),
+            power_of_two.get_type(),
+        );
+
+        let llvm_val_t = val.get_type();
+
+        let max_rem =
+            ctx.builder.build_int_sub(power_of_two, llvm_val_t.const_int(1, false), "").unwrap();
+        ctx.builder
+            .build_and(
+                ctx.builder.build_int_add(val, max_rem, "").unwrap(),
+                ctx.builder.build_not(max_rem, "").unwrap(),
+                "",
+            )
+            .unwrap()
+    }
+
+    let llvm_i8 = ctx.ctx.i8_type();
+    let llvm_pi8 = llvm_i8.ptr_type(AddressSpace::default());
+    let llvm_usize = generator.get_size_type(ctx.ctx);
+    let align_ty = align_ty.into();
+
+    let size = ctx.builder.build_int_cast(size, llvm_usize, "").unwrap();
+
+    debug_assert_eq!(
+        size.get_type().get_bit_width(),
+        llvm_usize.get_bit_width(),
+        "Expected size_t ({}) for parameter `size` of `aligned_alloca`, got {}",
+        llvm_usize,
+        size.get_type(),
+    );
+
+    let alignment = get_type_alignment(align_ty);
+    let alignment = ctx.builder.build_int_cast(alignment, llvm_usize, "").unwrap();
+
+    if ctx.registry.llvm_options.opt_level == OptimizationLevel::None {
+        let alignment_bitcount = llvm_intrinsics::call_int_ctpop(ctx, alignment, None);
+
+        ctx.make_assert(
+            generator,
+            ctx.builder
+                .build_int_compare(
+                    IntPredicate::EQ,
+                    alignment_bitcount,
+                    alignment_bitcount.get_type().const_int(1, false),
+                    "",
+                )
+                .unwrap(),
+            "0:AssertionError",
+            "Expected power-of-two alignment for aligned_alloca, got {0}",
+            [Some(alignment), None, None],
+            ctx.current_loc,
+        );
+    }
+
+    let buffer_size = round_up(ctx, size, alignment);
+    let aligned_slices = ctx.builder.build_int_unsigned_div(buffer_size, alignment, "").unwrap();
+
+    // Just to be absolutely sure, alloca in [i8 x alignment] slices
+    let buffer = ctx.builder.build_array_alloca(align_ty, aligned_slices, "").unwrap();
+
+    ctx.builder
+        .build_bit_cast(buffer, llvm_pi8, name.unwrap_or_default())
+        .map(BasicValueEnum::into_pointer_value)
+        .unwrap()
+}

nac3core/src/codegen/numpy.rs

@@ -10,8 +10,11 @@ use nac3parser::ast::{Operator, StrRef};
 use super::{
     expr::gen_binop_expr_with_values,
     irrt::{
-        calculate_len_for_slice_range, call_ndarray_calc_broadcast,
+        calculate_len_for_slice_range,
-        call_ndarray_calc_broadcast_index, call_ndarray_calc_nd_indices, call_ndarray_calc_size,
+        ndarray::{
+            call_ndarray_calc_broadcast, call_ndarray_calc_broadcast_index,
+            call_ndarray_calc_nd_indices, call_ndarray_calc_size,
+        },
     },
     llvm_intrinsics::{self, call_memcpy_generic},
     macros::codegen_unreachable,
@@ -38,6 +41,7 @@ use crate::{
 };
 
 /// Creates an uninitialized `NDArray` instance.
+#[deprecated = "Use NDArrayType::construct_uninitialized instead."]
 fn create_ndarray_uninitialized<'ctx, G: CodeGenerator + ?Sized>(
     generator: &mut G,
     ctx: &mut CodeGenContext<'ctx, '_>,
@@ -81,6 +85,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)?;
@@ -120,10 +125,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);
@@ -213,7 +218,7 @@ fn ndarray_init_data<'ctx, G: CodeGenerator + ?Sized>(
         (None, None),
     );
     unsafe {
-        ndarray.create_data(generator, ctx, llvm_ndarray_data_t, ndarray_num_elems);
+        ndarray.create_data(generator, ctx, ndarray_num_elems);
     }
 
     ndarray
@@ -1261,6 +1266,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(
@@ -1274,8 +1280,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);

nac3core/src/codegen/types/ndarray.rs

@@ -82,7 +82,7 @@ impl<'ctx> NDArrayType<'ctx> {
         Ok(())
     }
 
-    // TODO: Move this into e.g. StructProxyType
+    // TODO: Move this as a member of this Struct
     #[must_use]
     fn fields(
         ctx: impl AsContextRef<'ctx>,
@@ -189,22 +189,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
     }
@@ -220,7 +218,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();
@@ -250,7 +255,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/types/structure.rs

@@ -145,7 +145,7 @@ where
     }
 
     /// Sets the value of this field for a given `obj`.
-    pub fn set_from_value(&self, obj: StructValue<'ctx>, value: Value) {
+    pub fn set_for_value(&self, obj: StructValue<'ctx>, value: Value) {
         obj.set_field_at_index(self.index, value);
     }
 

nac3core/src/codegen/values/ndarray.rs

@@ -9,10 +9,11 @@ use super::{
     UntypedArrayLikeAccessor, UntypedArrayLikeMutator,
 };
 use crate::codegen::{
-    irrt::{call_ndarray_calc_size, call_ndarray_flatten_index},
+    irrt,
-    llvm_intrinsics::call_int_umin,
+    llvm_intrinsics::{call_int_umin, call_memcpy_generic_array},
     stmt::gen_for_callback_incrementing,
-    types::NDArrayType,
+    type_aligned_alloca,
+    types::{structure::StructField, NDArrayType},
     CodeGenContext, CodeGenerator,
 };
 
@@ -77,12 +78,10 @@ impl<'ctx> NDArrayValue<'ctx> {
         ctx.builder.build_load(pndims, "").map(BasicValueEnum::into_int_value).unwrap()
     }
 
-    /// Returns the pointer to the field storing the size of each element of this `NDArray`.
+    fn itemsize(&self, ctx: &CodeGenContext<'ctx, '_>) -> StructField<'ctx, IntValue<'ctx>> {
-    fn ptr_to_itemsize(&self, ctx: &CodeGenContext<'ctx, '_>) -> PointerValue<'ctx> {
         self.get_type()
             .get_fields(ctx.ctx, self.llvm_usize)
             .itemsize
-            .ptr_by_gep(ctx, self.value, self.name)
     }
 
     /// Stores the size of each element `itemsize` into this instance.
@@ -90,18 +89,16 @@ impl<'ctx> NDArrayValue<'ctx> {
         &self,
         ctx: &CodeGenContext<'ctx, '_>,
         generator: &G,
-        ndims: IntValue<'ctx>,
+        itemsize: IntValue<'ctx>,
     ) {
-        debug_assert_eq!(ndims.get_type(), generator.get_size_type(ctx.ctx));
+        debug_assert_eq!(itemsize.get_type(), generator.get_size_type(ctx.ctx));
 
-        let pndims = self.ptr_to_ndims(ctx);
+        self.itemsize(ctx).set(ctx, self.value, itemsize, self.name);
-        ctx.builder.build_store(pndims, ndims).unwrap();
     }
 
     /// Returns the size of each element of this `NDArray` as a value.
     pub fn load_itemsize(&self, ctx: &CodeGenContext<'ctx, '_>) -> IntValue<'ctx> {
-        let pndims = self.ptr_to_ndims(ctx);
+        self.itemsize(ctx).get(ctx, self.value, self.name)
-        ctx.builder.build_load(pndims, "").map(BasicValueEnum::into_int_value).unwrap()
     }
 
     /// Returns the double-indirection pointer to the `shape` array, as if by calling
@@ -145,7 +142,7 @@ impl<'ctx> NDArrayValue<'ctx> {
 
     /// Stores the array of dimension sizes `dims` into this instance.
     fn store_strides(&self, ctx: &CodeGenContext<'ctx, '_>, dims: PointerValue<'ctx>) {
-        ctx.builder.build_store(self.ptr_to_shape(ctx), dims).unwrap();
+        ctx.builder.build_store(self.ptr_to_strides(ctx), dims).unwrap();
     }
 
     /// Convenience method for creating a new array storing the stride with the given `size`.
@@ -155,7 +152,7 @@ impl<'ctx> NDArrayValue<'ctx> {
         llvm_usize: IntType<'ctx>,
         size: IntValue<'ctx>,
     ) {
-        self.store_shape(ctx, ctx.builder.build_array_alloca(llvm_usize, size, "").unwrap());
+        self.store_strides(ctx, ctx.builder.build_array_alloca(llvm_usize, size, "").unwrap());
     }
 
     /// Returns a proxy object to the field storing the stride of each dimension of this `NDArray`.
@@ -193,16 +190,17 @@ impl<'ctx> NDArrayValue<'ctx> {
     pub unsafe fn create_data<G: CodeGenerator + ?Sized>(
         &self,
         generator: &mut G,
-        ctx: &CodeGenContext<'ctx, '_>,
+        ctx: &mut CodeGenContext<'ctx, '_>,
-        elem_ty: BasicTypeEnum<'ctx>,
         size: IntValue<'ctx>,
     ) {
+        // let itemsize =
+        //     ctx.builder.build_int_cast(self.load_itemsize(ctx), size.get_type(), "").unwrap();
         let itemsize =
-            ctx.builder.build_int_cast(elem_ty.size_of().unwrap(), size.get_type(), "").unwrap();
+            ctx.builder.build_int_cast(self.dtype.size_of().unwrap(), size.get_type(), "").unwrap();
         let nbytes = ctx.builder.build_int_mul(size, itemsize, "").unwrap();
+        // let nbytes = self.nbytes(generator, ctx);
 
-        // TODO: What about alignment?
+        let data = type_aligned_alloca(generator, ctx, self.dtype, nbytes, None);
-        let data = ctx.builder.build_array_alloca(ctx.ctx.i8_type(), nbytes, "").unwrap();
         self.store_data(ctx, data);
 
         // self.set_strides_contiguous(generator, ctx);
@@ -221,7 +219,15 @@ impl<'ctx> NDArrayValue<'ctx> {
         ctx: &mut CodeGenContext<'ctx, '_>,
         shape: PointerValue<'ctx>,
     ) {
-        todo!()
+        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.
@@ -232,7 +238,9 @@ impl<'ctx> NDArrayValue<'ctx> {
         ctx: &mut CodeGenContext<'ctx, '_>,
         src_ndarray: NDArrayValue<'ctx>,
     ) {
-        todo!()
+        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.
@@ -242,7 +250,15 @@ impl<'ctx> NDArrayValue<'ctx> {
         ctx: &mut CodeGenContext<'ctx, '_>,
         strides: PointerValue<'ctx>,
     ) {
-        todo!()
+        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.
@@ -253,7 +269,9 @@ impl<'ctx> NDArrayValue<'ctx> {
         ctx: &mut CodeGenContext<'ctx, '_>,
         src_ndarray: NDArrayValue<'ctx>,
     ) {
-        todo!()
+        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.
@@ -262,7 +280,7 @@ impl<'ctx> NDArrayValue<'ctx> {
         generator: &mut G,
         ctx: &mut CodeGenContext<'ctx, '_>,
     ) -> IntValue<'ctx> {
-        todo!()
+        irrt::ndarray::call_nac3_ndarray_size(generator, ctx, *self)
     }
 
     /// Get the `ndarray.nbytes` of this ndarray.
@@ -271,7 +289,7 @@ impl<'ctx> NDArrayValue<'ctx> {
         generator: &mut G,
         ctx: &mut CodeGenContext<'ctx, '_>,
     ) -> IntValue<'ctx> {
-        todo!()
+        irrt::ndarray::call_nac3_ndarray_nbytes(generator, ctx, *self)
     }
 
     /// Get the `len()` of this ndarray.
@@ -280,7 +298,7 @@ impl<'ctx> NDArrayValue<'ctx> {
         generator: &mut G,
         ctx: &mut CodeGenContext<'ctx, '_>,
     ) -> IntValue<'ctx> {
-        todo!()
+        irrt::ndarray::call_nac3_ndarray_len(generator, ctx, *self)
     }
 
     /// Check if this ndarray is C-contiguous.
@@ -291,18 +309,18 @@ impl<'ctx> NDArrayValue<'ctx> {
         generator: &mut G,
         ctx: &mut CodeGenContext<'ctx, '_>,
     ) -> IntValue<'ctx> {
-        todo!()
+        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,
+        &self,
         generator: &mut G,
-        ctx: &CodeGenContext<'ctx, '_>,
+        ctx: &mut CodeGenContext<'ctx, '_>,
     ) {
-        todo!()
+        irrt::ndarray::call_nac3_ndarray_set_strides_by_shape(generator, ctx, *self);
     }
 
     /// Copy data from another ndarray.
@@ -317,7 +335,8 @@ impl<'ctx> NDArrayValue<'ctx> {
         ctx: &mut CodeGenContext<'ctx, '_>,
         src: NDArrayValue<'ctx>,
     ) {
-        todo!()
+        assert_eq!(self.dtype, src.dtype, "self and src dtype should match");
+        irrt::ndarray::call_nac3_ndarray_copy_data(generator, ctx, src, *self);
     }
 }
 
@@ -371,7 +390,12 @@ impl<'ctx> ArrayLikeValue<'ctx> for NDArrayDataProxy<'ctx, '_> {
         ctx: &CodeGenContext<'ctx, '_>,
         generator: &G,
     ) -> IntValue<'ctx> {
-        call_ndarray_calc_size(generator, ctx, &self.as_slice_value(ctx, generator), (None, None))
+        irrt::ndarray::call_ndarray_calc_size(
+            generator,
+            ctx,
+            &self.as_slice_value(ctx, generator),
+            (None, None),
+        )
     }
 }
 
@@ -480,7 +504,7 @@ impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index>
             indices_elem_ty.get_bit_width()
         );
 
-        let index = call_ndarray_flatten_index(generator, ctx, *self.0, indices);
+        let index = irrt::ndarray::call_ndarray_flatten_index(generator, ctx, *self.0, indices);
         let sizeof_elem = ctx
             .builder
             .build_int_truncate_or_bit_cast(
@@ -704,7 +728,7 @@ impl<'ctx> ArrayLikeValue<'ctx> for NDArrayStridesProxy<'ctx, '_> {
         ctx: &CodeGenContext<'ctx, '_>,
         generator: &G,
     ) -> AnyTypeEnum<'ctx> {
-        self.0.shape().base_ptr(ctx, generator).get_type().get_element_type()
+        self.0.strides().base_ptr(ctx, generator).get_type().get_element_type()
     }
 
     fn base_ptr<G: CodeGenerator + ?Sized>(
@@ -712,10 +736,10 @@ impl<'ctx> ArrayLikeValue<'ctx> for NDArrayStridesProxy<'ctx, '_> {
         ctx: &CodeGenContext<'ctx, '_>,
         _: &G,
     ) -> PointerValue<'ctx> {
-        let var_name = self.0.name.map(|v| format!("{v}.data")).unwrap_or_default();
+        let var_name = self.0.name.map(|v| format!("{v}.strides")).unwrap_or_default();
 
         ctx.builder
-            .build_load(self.0.ptr_to_shape(ctx), var_name.as_str())
+            .build_load(self.0.ptr_to_strides(ctx), var_name.as_str())
             .map(BasicValueEnum::into_pointer_value)
             .unwrap()
     }