diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs index 0db1501457..be68e10428 100644 --- a/nac3artiq/src/codegen.rs +++ b/nac3artiq/src/codegen.rs @@ -722,7 +722,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 = diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs index 402cfbe554..158dfe9607 100644 --- a/nac3core/src/codegen/expr.rs +++ b/nac3core/src/codegen/expr.rs @@ -2521,7 +2521,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>, ) -> Result>, String> { @@ -2529,7 +2529,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) }; @@ -2562,10 +2562,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(); @@ -2759,27 +2755,13 @@ 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( - 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, - 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); + let ndarray = + NDArrayType::new(generator, ctx.ctx, llvm_ndarray_data_t, Some(num_dims)) + .construct_uninitialized(generator, ctx, None); let ndarray_num_dims = ctx .builder @@ -2818,7 +2800,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( diff --git a/nac3core/src/codegen/numpy.rs b/nac3core/src/codegen/numpy.rs index 9d7f79de79..cd113aae0b 100644 --- a/nac3core/src/codegen/numpy.rs +++ b/nac3core/src/codegen/numpy.rs @@ -61,6 +61,7 @@ where ) -> Result, 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)?; @@ -100,15 +101,10 @@ where llvm_usize.const_int(1, false), )?; - let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty); - let ndarray = - NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None).alloca(generator, ctx, None); - let num_dims = shape_len_fn(generator, ctx, shape)?; - ndarray.store_ndims(ctx, generator, num_dims); - let ndarray_num_dims = ndarray.load_ndims(ctx); - ndarray.create_shape(ctx, llvm_usize, ndarray_num_dims); + let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None) + .construct_dyn_ndims(generator, ctx, num_dims, None); // Copy the dimension sizes from shape to ndarray.dims let shape_len = shape_len_fn(generator, ctx, shape)?; @@ -133,7 +129,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) } @@ -173,32 +169,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 `shape` 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, '_>, @@ -1206,6 +1181,7 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>( ) -> Result, 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() { @@ -1220,7 +1196,6 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>( }, )? } else { - let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty); let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None) .construct_dyn_ndims(generator, ctx, this.load_ndims(ctx), None); @@ -1282,7 +1257,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( diff --git a/nac3core/src/codegen/types/ndarray/mod.rs b/nac3core/src/codegen/types/ndarray/mod.rs index d81f673d27..d485eeb47a 100644 --- a/nac3core/src/codegen/types/ndarray/mod.rs +++ b/nac3core/src/codegen/types/ndarray/mod.rs @@ -198,7 +198,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); diff --git a/nac3core/src/codegen/values/ndarray/mod.rs b/nac3core/src/codegen/values/ndarray/mod.rs index 08f6a5b161..61504c2d58 100644 --- a/nac3core/src/codegen/values/ndarray/mod.rs +++ b/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::{ndarray::NDArrayType, structure::StructField}, @@ -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( + /// + /// 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( &self, generator: &mut G, ctx: &mut CodeGenContext<'ctx, '_>, - elem_ty: BasicTypeEnum<'ctx>, - size: IntValue<'ctx>, ) { - let itemsize = 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 nbytes = self.nbytes(generator, ctx); - 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( + &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( + &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( + &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( + &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( + &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( + &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( + &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"]`: + pub fn is_c_contiguous( + &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( + &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( + &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> {