From d5880b119ab55df1bf14bcb692830b6aa248a087 Mon Sep 17 00:00:00 2001 From: lyken Date: Mon, 29 Jul 2024 13:34:21 +0800 Subject: [PATCH] core/ndstrides: move functions to numpy_new/util.rs --- nac3core/src/codegen/expr.rs | 7 +- .../src/codegen/irrt/ndarray/allocation.rs | 83 -------- nac3core/src/codegen/irrt/ndarray/mod.rs | 1 - nac3core/src/codegen/numpy_new/control.rs | 49 ----- nac3core/src/codegen/numpy_new/factory.rs | 6 +- nac3core/src/codegen/numpy_new/mod.rs | 2 +- nac3core/src/codegen/numpy_new/util.rs | 184 ++++++++++++++++++ nac3core/src/codegen/numpy_new/view.rs | 2 +- 8 files changed, 191 insertions(+), 143 deletions(-) delete mode 100644 nac3core/src/codegen/irrt/ndarray/allocation.rs delete mode 100644 nac3core/src/codegen/numpy_new/control.rs create mode 100644 nac3core/src/codegen/numpy_new/util.rs diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs index e1fdf3d7..61002848 100644 --- a/nac3core/src/codegen/expr.rs +++ b/nac3core/src/codegen/expr.rs @@ -18,6 +18,7 @@ use crate::{ call_memcpy_generic, }, need_sret, numpy, + numpy_new::util::alloca_ndarray, stmt::{ gen_for_callback_incrementing, gen_if_callback, gen_if_else_expr_callback, gen_raise, gen_var, @@ -42,10 +43,8 @@ use nac3parser::ast::{ self, Boolop, Cmpop, Comprehension, Constant, Expr, ExprKind, Located, Location, Operator, StrRef, Unaryop, }; -use ndarray::{ - allocation::alloca_ndarray, - indexing::{call_nac3_ndarray_index, RustNDIndex}, -}; + +use ndarray::indexing::{call_nac3_ndarray_index, RustNDIndex}; use super::{ model::*, diff --git a/nac3core/src/codegen/irrt/ndarray/allocation.rs b/nac3core/src/codegen/irrt/ndarray/allocation.rs deleted file mode 100644 index ff94afeb..00000000 --- a/nac3core/src/codegen/irrt/ndarray/allocation.rs +++ /dev/null @@ -1,83 +0,0 @@ -use crate::codegen::model::*; -use crate::codegen::util::array_writer::ArrayWriter; -use crate::codegen::{structure::ndarray::NpArray, CodeGenContext, CodeGenerator}; - -use super::basic::{ - call_nac3_ndarray_nbytes, call_nac3_ndarray_set_strides_by_shape, - call_nac3_ndarray_util_assert_shape_no_negative, -}; - -/** - Allocate an ndarray on the stack given its `ndims`. - - `shape` and `strides` will be automatically allocated on the stack. - - The returned ndarray's content will be: - - `data`: `nullptr` - - `itemsize`: **uninitialized** value - - `ndims`: initialized value, set to the input `ndims` - - `shape`: initialized pointer to an allocated stack with **uninitialized** values - - `strides`: initialized pointer to an allocated stack with **uninitialized** values -*/ -pub fn alloca_ndarray<'ctx, G>( - generator: &mut G, - ctx: &mut CodeGenContext<'ctx, '_>, - ndims: Int<'ctx, SizeT>, - name: &str, -) -> Result>, String> -where - G: CodeGenerator + ?Sized, -{ - let tyctx = generator.type_context(ctx.ctx); - - let sizet_model = IntModel(SizeT); - let ndarray_model = StructModel(NpArray); - let ndarray_data_model = PtrModel(IntModel(Byte)); - - // Setup ndarray - let ndarray_ptr = ndarray_model.alloca(tyctx, ctx, name); - let shape = sizet_model.array_alloca(tyctx, ctx, ndims.value, "shape"); - let strides = sizet_model.array_alloca(tyctx, ctx, ndims.value, "strides"); - - ndarray_ptr.gep(ctx, |f| f.data).store(ctx, ndarray_data_model.nullptr(tyctx, ctx.ctx)); - ndarray_ptr.gep(ctx, |f| f.ndims).store(ctx, ndims); - ndarray_ptr.gep(ctx, |f| f.shape).store(ctx, shape); - ndarray_ptr.gep(ctx, |f| f.strides).store(ctx, strides); - - Ok(ndarray_ptr) -} - -/// Initialize an ndarray's `shape` and asserts on. -/// `shape`'s values and prohibit illegal inputs like negative dimensions. -pub fn init_ndarray_shape<'ctx, G: CodeGenerator + ?Sized>( - generator: &mut G, - ctx: &mut CodeGenContext<'ctx, '_>, - pndarray: Ptr<'ctx, StructModel>, - shape_writer: &ArrayWriter<'ctx, G, SizeT, IntModel>, -) -> Result<(), String> { - let tyctx = generator.type_context(ctx.ctx); - let shape = pndarray.gep(ctx, |f| f.shape).load(tyctx, ctx, "shape"); - (shape_writer.write)(generator, ctx, shape)?; - call_nac3_ndarray_util_assert_shape_no_negative(generator, ctx, shape_writer.len, shape); - Ok(()) -} - -/// Initialize an ndarray's `data` by allocating a buffer on the stack. -/// The allocated data buffer is considered to be *owned* by the ndarray. -/// -/// `strides` of the ndarray will also be updated with `set_strides_by_shape`. -/// -/// `shape` and `itemsize` of the ndarray ***must*** be initialized first. -pub fn init_ndarray_data_by_alloca<'ctx, G: CodeGenerator + ?Sized>( - generator: &mut G, - ctx: &mut CodeGenContext<'ctx, '_>, - pndarray: Ptr<'ctx, StructModel>, -) { - let tyctx = generator.type_context(ctx.ctx); - let ndarray_data_model = IntModel(Byte); - - let nbytes = call_nac3_ndarray_nbytes(generator, ctx, pndarray); - let data = ndarray_data_model.array_alloca(tyctx, ctx, nbytes.value, "data"); - pndarray.gep(ctx, |f| f.data).store(ctx, data); - call_nac3_ndarray_set_strides_by_shape(generator, ctx, pndarray); -} diff --git a/nac3core/src/codegen/irrt/ndarray/mod.rs b/nac3core/src/codegen/irrt/ndarray/mod.rs index 6498e1c2..f6c7c8a1 100644 --- a/nac3core/src/codegen/irrt/ndarray/mod.rs +++ b/nac3core/src/codegen/irrt/ndarray/mod.rs @@ -1,4 +1,3 @@ -pub mod allocation; pub mod basic; pub mod indexing; pub mod reshape; diff --git a/nac3core/src/codegen/numpy_new/control.rs b/nac3core/src/codegen/numpy_new/control.rs deleted file mode 100644 index ee99f3ac..00000000 --- a/nac3core/src/codegen/numpy_new/control.rs +++ /dev/null @@ -1,49 +0,0 @@ -use crate::codegen::{ - irrt::ndarray::basic::{call_nac3_ndarray_get_nth_pelement, call_nac3_ndarray_size}, - model::*, - stmt::BreakContinueHooks, - structure::ndarray::NpArray, - util::control::gen_model_for, - CodeGenContext, CodeGenerator, -}; - -/// Iterate through all elements in an ndarray. -/// -/// `body` is given the index of an element and an opaque pointer (as an `uint8_t*`, you might want to cast it) to the element. -/// -/// Short-circuiting is possible with the given [`BreakContinueHooks`]. -pub fn gen_foreach_ndarray_elements<'ctx, G, F>( - generator: &mut G, - ctx: &mut CodeGenContext<'ctx, '_>, - pndarray: Ptr<'ctx, StructModel>, - body: F, -) -> Result<(), String> -where - G: CodeGenerator + ?Sized, - F: Fn( - &mut G, - &mut CodeGenContext<'ctx, '_>, - BreakContinueHooks<'ctx>, - Int<'ctx, SizeT>, - Ptr<'ctx, IntModel>, - ) -> Result<(), String>, -{ - // TODO: Make this more efficient - use a special NDArray iterator? - - let tyctx = generator.type_context(ctx.ctx); - - let sizet_model = IntModel(SizeT); - let size = call_nac3_ndarray_size(generator, ctx, pndarray); - - gen_model_for( - generator, - ctx, - sizet_model.const_0(tyctx, ctx.ctx), - size, - sizet_model.const_1(tyctx, ctx.ctx), - |generator, ctx, hooks, index| { - let pelement = call_nac3_ndarray_get_nth_pelement(generator, ctx, pndarray, index); - body(generator, ctx, hooks, index, pelement) - }, - ) -} diff --git a/nac3core/src/codegen/numpy_new/factory.rs b/nac3core/src/codegen/numpy_new/factory.rs index 817fd5b5..2c52d9a7 100644 --- a/nac3core/src/codegen/numpy_new/factory.rs +++ b/nac3core/src/codegen/numpy_new/factory.rs @@ -7,10 +7,8 @@ use nac3parser::ast::StrRef; use crate::{ codegen::{ - irrt::ndarray::allocation::{ - alloca_ndarray, init_ndarray_data_by_alloca, init_ndarray_shape, - }, model::*, + numpy_new::util::{alloca_ndarray, init_ndarray_data_by_alloca, init_ndarray_shape}, structure::ndarray::NpArray, util::shape::make_shape_writer, CodeGenContext, CodeGenerator, @@ -20,7 +18,7 @@ use crate::{ typecheck::typedef::{FunSignature, Type}, }; -use super::control::gen_foreach_ndarray_elements; +use super::util::gen_foreach_ndarray_elements; /// Helper function to create an ndarray with uninitialized values /// diff --git a/nac3core/src/codegen/numpy_new/mod.rs b/nac3core/src/codegen/numpy_new/mod.rs index 96c68286..95e205a3 100644 --- a/nac3core/src/codegen/numpy_new/mod.rs +++ b/nac3core/src/codegen/numpy_new/mod.rs @@ -1,3 +1,3 @@ -pub mod control; pub mod factory; +pub mod util; pub mod view; diff --git a/nac3core/src/codegen/numpy_new/util.rs b/nac3core/src/codegen/numpy_new/util.rs new file mode 100644 index 00000000..a9cf7c73 --- /dev/null +++ b/nac3core/src/codegen/numpy_new/util.rs @@ -0,0 +1,184 @@ +use inkwell::{types::BasicType, values::BasicValueEnum}; + +use crate::{ + codegen::{ + irrt::ndarray::basic::{ + call_nac3_ndarray_get_nth_pelement, call_nac3_ndarray_nbytes, + call_nac3_ndarray_set_strides_by_shape, call_nac3_ndarray_size, + call_nac3_ndarray_util_assert_shape_no_negative, + }, + model::*, + stmt::BreakContinueHooks, + structure::ndarray::NpArray, + util::{array_writer::ArrayWriter, control::gen_model_for}, + CodeGenContext, CodeGenerator, + }, + toplevel::numpy::unpack_ndarray_var_tys, + typecheck::typedef::{Type, TypeEnum}, +}; + +/** + Allocate an ndarray on the stack given its `ndims`. + + `shape` and `strides` will be automatically allocated on the stack. + + The returned ndarray's content will be: + - `data`: `nullptr` + - `itemsize`: **uninitialized** value + - `ndims`: initialized value, set to the input `ndims` + - `shape`: initialized pointer to an allocated stack with **uninitialized** values + - `strides`: initialized pointer to an allocated stack with **uninitialized** values +*/ +pub fn alloca_ndarray<'ctx, G>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ndims: Int<'ctx, SizeT>, + name: &str, +) -> Result>, String> +where + G: CodeGenerator + ?Sized, +{ + let tyctx = generator.type_context(ctx.ctx); + + let sizet_model = IntModel(SizeT); + let ndarray_model = StructModel(NpArray); + let ndarray_data_model = PtrModel(IntModel(Byte)); + + // Setup ndarray + let ndarray_ptr = ndarray_model.alloca(tyctx, ctx, name); + let shape = sizet_model.array_alloca(tyctx, ctx, ndims.value, "shape"); + let strides = sizet_model.array_alloca(tyctx, ctx, ndims.value, "strides"); + + ndarray_ptr.gep(ctx, |f| f.data).store(ctx, ndarray_data_model.nullptr(tyctx, ctx.ctx)); + ndarray_ptr.gep(ctx, |f| f.ndims).store(ctx, ndims); + ndarray_ptr.gep(ctx, |f| f.shape).store(ctx, shape); + ndarray_ptr.gep(ctx, |f| f.strides).store(ctx, strides); + + Ok(ndarray_ptr) +} + +/// Initialize an ndarray's `shape` and asserts on. +/// `shape`'s values and prohibit illegal inputs like negative dimensions. +pub fn init_ndarray_shape<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + pndarray: Ptr<'ctx, StructModel>, + shape_writer: &ArrayWriter<'ctx, G, SizeT, IntModel>, +) -> Result<(), String> { + let tyctx = generator.type_context(ctx.ctx); + let shape = pndarray.gep(ctx, |f| f.shape).load(tyctx, ctx, "shape"); + (shape_writer.write)(generator, ctx, shape)?; + call_nac3_ndarray_util_assert_shape_no_negative(generator, ctx, shape_writer.len, shape); + Ok(()) +} + +/// Initialize an ndarray's `data` by allocating a buffer on the stack. +/// The allocated data buffer is considered to be *owned* by the ndarray. +/// +/// `strides` of the ndarray will also be updated with `set_strides_by_shape`. +/// +/// `shape` and `itemsize` of the ndarray ***must*** be initialized first. +pub fn init_ndarray_data_by_alloca<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + pndarray: Ptr<'ctx, StructModel>, +) { + let tyctx = generator.type_context(ctx.ctx); + let ndarray_data_model = IntModel(Byte); + + let nbytes = call_nac3_ndarray_nbytes(generator, ctx, pndarray); + let data = ndarray_data_model.array_alloca(tyctx, ctx, nbytes.value, "data"); + pndarray.gep(ctx, |f| f.data).store(ctx, data); + call_nac3_ndarray_set_strides_by_shape(generator, ctx, pndarray); +} + +/// Convert `input` to an ndarray - behaves similarly to `np.asarray`. +/// +/// Returns the ndarray interpretation of `input` and **the element type** of the ndarray. +/// +/// Here are the exact details: +/// - If `input` is an ndarray, the function returns back the **same** ndarray and the `dtype` +/// of the ndarray. +/// - If `input` is not an ndarray, the function creates an ndarray with a single element `input`, +/// and returns the created ndarray and `input_ty`. Note that the created ndarray's `ndims` will +/// be `0` (an *unsized* ndarray). +pub fn as_ndarray<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + input: BasicValueEnum<'ctx>, + input_ty: Type, +) -> Result<(Ptr<'ctx, StructModel>, Type), String> { + let tyctx = generator.type_context(ctx.ctx); + let sizet_model = IntModel(SizeT); + let pbyte_model = PtrModel(IntModel(Byte)); + let pndarray_model = PtrModel(StructModel(NpArray)); + + let input_ty_enum = ctx.unifier.get_ty(input_ty); + match &*input_ty_enum { + TypeEnum::TObj { obj_id, .. } + if *obj_id == ctx.primitives.ndarray.obj_id(&ctx.unifier).unwrap() => + { + let pndarray = pndarray_model.check_value(tyctx, ctx.ctx, input).unwrap(); + let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, input_ty); + Ok((pndarray, elem_ty)) + } + _ => { + let ndims = sizet_model.const_0(tyctx, ctx.ctx); + let pndarray = alloca_ndarray(generator, ctx, ndims, "ndarray")?; + + // We have to put `input` onto the stack to get a data pointer. + let data = ctx.builder.build_alloca(input.get_type(), "as_ndarray_scalar").unwrap(); + ctx.builder.build_store(data, input).unwrap(); + + let data = pbyte_model.transmute(tyctx, ctx, data, "data"); + pndarray.gep(ctx, |f| f.data).store(ctx, data); + + let itemsize = input.get_type().size_of().unwrap(); + let itemsize = sizet_model.check_value(tyctx, ctx.ctx, itemsize).unwrap(); + pndarray.gep(ctx, |f| f.itemsize).store(ctx, itemsize); + + Ok((pndarray, input_ty)) + } + } +} + +/// Iterate through all elements in an ndarray. +/// +/// `body` is given the index of an element and an opaque pointer (as an `uint8_t*`, you might want to cast it) to the element. +/// +/// Short-circuiting is possible with the given [`BreakContinueHooks`]. +pub fn gen_foreach_ndarray_elements<'ctx, G, F>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + pndarray: Ptr<'ctx, StructModel>, + body: F, +) -> Result<(), String> +where + G: CodeGenerator + ?Sized, + F: Fn( + &mut G, + &mut CodeGenContext<'ctx, '_>, + BreakContinueHooks<'ctx>, + Int<'ctx, SizeT>, + Ptr<'ctx, IntModel>, + ) -> Result<(), String>, +{ + // TODO: Make this more efficient - use a special NDArray iterator? + + let tyctx = generator.type_context(ctx.ctx); + + let sizet_model = IntModel(SizeT); + let size = call_nac3_ndarray_size(generator, ctx, pndarray); + + gen_model_for( + generator, + ctx, + sizet_model.const_0(tyctx, ctx.ctx), + size, + sizet_model.const_1(tyctx, ctx.ctx), + |generator, ctx, hooks, index| { + let pelement = call_nac3_ndarray_get_nth_pelement(generator, ctx, pndarray, index); + body(generator, ctx, hooks, index, pelement) + }, + ) +} diff --git a/nac3core/src/codegen/numpy_new/view.rs b/nac3core/src/codegen/numpy_new/view.rs index 805e3568..df16431a 100644 --- a/nac3core/src/codegen/numpy_new/view.rs +++ b/nac3core/src/codegen/numpy_new/view.rs @@ -4,7 +4,6 @@ use nac3parser::ast::StrRef; use crate::{ codegen::{ irrt::ndarray::{ - allocation::{alloca_ndarray, init_ndarray_shape}, basic::{ call_nac3_ndarray_is_c_contiguous, call_nac3_ndarray_nbytes, call_nac3_ndarray_set_strides_by_shape, call_nac3_ndarray_size, @@ -12,6 +11,7 @@ use crate::{ reshape::call_nac3_ndarray_resolve_and_check_new_shape, }, model::*, + numpy_new::util::{alloca_ndarray, init_ndarray_shape}, structure::ndarray::NpArray, util::{array_writer::ArrayWriter, shape::make_shape_writer}, CodeGenContext, CodeGenerator,