From 0c4aae2eea31848c2f3258224d4fa3cc5f4d9bda Mon Sep 17 00:00:00 2001 From: David Mak Date: Tue, 12 Nov 2024 17:00:45 +0800 Subject: [PATCH] WIP --- nac3artiq/src/codegen.rs | 4 +- nac3core/src/codegen/expr.rs | 47 ++---- nac3core/src/codegen/numpy.rs | 195 ++++++++++--------------- nac3core/src/codegen/types/ndarray.rs | 34 +++-- nac3core/src/codegen/values/ndarray.rs | 147 ++++++++++++++++++- 5 files changed, 258 insertions(+), 169 deletions(-) diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs index a89d018c..95a27012 100644 --- a/nac3artiq/src/codegen.rs +++ b/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 = diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs index 1e20ca7a..2dcb4ff3 100644 --- a/nac3core/src/codegen/expr.rs +++ b/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::{ - helper::PrimDef, - numpy::{make_ndarray_ty, unpack_ndarray_var_tys}, - DefinitionId, TopLevelDef, - }, + toplevel::{helper::PrimDef, numpy::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 = - generator.gen_var_alloc(ctx, llvm_ndarray_t.into(), 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 = NDArrayType::new(generator, ctx.ctx, llvm_ndarray_data_t) + .construct_uninitialized(generator, ctx, num_dims, None); 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( diff --git a/nac3core/src/codegen/numpy.rs b/nac3core/src/codegen/numpy.rs index e65bc507..096e25b5 100644 --- a/nac3core/src/codegen/numpy.rs +++ b/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, 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, 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,85 +1223,89 @@ 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() { - create_ndarray_dyn_shape( - generator, - ctx, - elem_ty, - &this, - |_, ctx, shape| Ok(shape.load_ndims(ctx)), - |generator, ctx, shape, idx| unsafe { - Ok(shape.shape().get_typed_unchecked(ctx, generator, &idx, None)) - }, - )? - } else { - let ndarray = create_ndarray_uninitialized(generator, ctx, elem_ty)?; - ndarray.store_ndims(ctx, generator, this.load_ndims(ctx)); + let ndarray = + if slices.is_empty() { + create_ndarray_dyn_shape( + generator, + ctx, + elem_ty, + &this, + |_, ctx, shape| Ok(shape.load_ndims(ctx)), + |generator, ctx, shape, idx| unsafe { + Ok(shape.shape().get_typed_unchecked(ctx, generator, &idx, None)) + }, + )? + } else { + let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty) + .construct_uninitialized(generator, ctx, this.load_ndims(ctx), None); - let ndims = this.load_ndims(ctx); - ndarray.create_shape(ctx, llvm_usize, ndims); + let ndims = this.load_ndims(ctx); + ndarray.create_shape(ctx, llvm_usize, ndims); - // Populate the first slices.len() dimensions by computing the size of each dim slice - for (i, (start, stop, step)) in slices.iter().enumerate() { - // HACK: workaround calculate_len_for_slice_range requiring exclusive stop - let stop = ctx - .builder - .build_select( - ctx.builder - .build_int_compare( - IntPredicate::SLT, - *step, - llvm_i32.const_zero(), - "is_neg", - ) - .unwrap(), - ctx.builder - .build_int_sub(*stop, llvm_i32.const_int(1, true), "e_min_one") - .unwrap(), - ctx.builder - .build_int_add(*stop, llvm_i32.const_int(1, true), "e_add_one") - .unwrap(), - "final_e", - ) - .map(BasicValueEnum::into_int_value) - .unwrap(); + // Populate the first slices.len() dimensions by computing the size of each dim slice + for (i, (start, stop, step)) in slices.iter().enumerate() { + // HACK: workaround calculate_len_for_slice_range requiring exclusive stop + let stop = ctx + .builder + .build_select( + ctx.builder + .build_int_compare( + IntPredicate::SLT, + *step, + llvm_i32.const_zero(), + "is_neg", + ) + .unwrap(), + ctx.builder + .build_int_sub(*stop, llvm_i32.const_int(1, true), "e_min_one") + .unwrap(), + ctx.builder + .build_int_add(*stop, llvm_i32.const_int(1, true), "e_add_one") + .unwrap(), + "final_e", + ) + .map(BasicValueEnum::into_int_value) + .unwrap(); - let slice_len = calculate_len_for_slice_range(generator, ctx, *start, stop, *step); - let slice_len = - ctx.builder.build_int_z_extend_or_bit_cast(slice_len, llvm_usize, "").unwrap(); + let slice_len = calculate_len_for_slice_range(generator, ctx, *start, stop, *step); + let slice_len = + ctx.builder.build_int_z_extend_or_bit_cast(slice_len, llvm_usize, "").unwrap(); - unsafe { - ndarray.shape().set_typed_unchecked( - ctx, - generator, - &llvm_usize.const_int(i as u64, false), - slice_len, - ); - } - } - - // Populate the rest by directly copying the dim size from the source array - gen_for_callback_incrementing( - generator, - ctx, - None, - llvm_usize.const_int(slices.len() as u64, false), - (this.load_ndims(ctx), false), - |generator, ctx, _, idx| { unsafe { - let dim_sz = this.shape().get_typed_unchecked(ctx, generator, &idx, None); - ndarray.shape().set_typed_unchecked(ctx, generator, &idx, dim_sz); + ndarray.shape().set_typed_unchecked( + ctx, + generator, + &llvm_usize.const_int(i as u64, false), + slice_len, + ); } + } - Ok(()) - }, - llvm_usize.const_int(1, false), - ) - .unwrap(); + // Populate the rest by directly copying the dim size from the source array + gen_for_callback_incrementing( + generator, + ctx, + None, + llvm_usize.const_int(slices.len() as u64, false), + (this.load_ndims(ctx), false), + |generator, ctx, _, idx| { + unsafe { + let dim_sz = this.shape().get_typed_unchecked(ctx, generator, &idx, None); + ndarray.shape().set_typed_unchecked(ctx, generator, &idx, dim_sz); + } - ndarray_init_data(generator, ctx, elem_ty, ndarray) - }; + Ok(()) + }, + llvm_usize.const_int(1, false), + ) + .unwrap(); + + unsafe { ndarray.create_data(generator, ctx) }; + + ndarray + }; ndarray_sliced_copyto_impl( generator, diff --git a/nac3core/src/codegen/types/ndarray.rs b/nac3core/src/codegen/types/ndarray.rs index 26f6d066..3c435b7a 100644 --- a/nac3core/src/codegen/types/ndarray.rs +++ b/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>, ) -> >::Value { let ndarray = self.new_value(generator, ctx, name); - let itemsize = ctx - .builder - .build_int_z_extend_or_bit_cast(self.dtype.size_of().unwrap(), self.llvm_usize, "") - .unwrap(); + let itemsize = + ctx.builder.build_int_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_val); + ndarray.store_ndims(ctx, generator, ndims); - ndarray.create_shape(ctx, self.llvm_usize, ndims_val); - ndarray.create_strides(ctx, self.llvm_usize, ndims_val); + ndarray.create_shape(ctx, self.llvm_usize, ndims); + ndarray.create_strides(ctx, self.llvm_usize, ndims); ndarray } @@ -221,7 +219,14 @@ impl<'ctx> NDArrayType<'ctx> { shape: &[u64], name: Option<&'ctx str>, ) -> >::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>, ) -> >::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(); diff --git a/nac3core/src/codegen/values/ndarray.rs b/nac3core/src/codegen/values/ndarray.rs index 9ace252b..1b082843 100644 --- a/nac3core/src/codegen/values/ndarray.rs +++ b/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( + /// + /// 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_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`. @@ -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( + &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( + &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( + &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( + &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( + &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( + &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( + &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"]`: + pub fn is_c_contiguous( + &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( + &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( + &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> {