From 3344a2bcd3498fb6052a2a3181af3b0679a1ad25 Mon Sep 17 00:00:00 2001 From: lyken Date: Sun, 14 Jul 2024 23:16:28 +0800 Subject: [PATCH] core: new np_{zeros,ones,fill} + some irrt model additions --- nac3core/irrt/irrt/numpy/ndarray.hpp | 155 +++++++++++++ nac3core/irrt/irrt/numpy/ndarray_util.hpp | 107 +++++++++ nac3core/irrt/irrt_everything.hpp | 6 +- nac3core/irrt/irrt_test.cpp | 2 + nac3core/irrt/test/ndarray.hpp | 44 ++++ nac3core/irrt/test/test_core.hpp | 2 +- nac3core/src/codegen/irrt/mod.rs | 1 + nac3core/src/codegen/irrt/numpy/mod.rs | 5 + nac3core/src/codegen/irrt/numpy/ndarray.rs | 243 +++++++++++++++++++++ nac3core/src/codegen/irrt/numpy/shape.rs | 162 ++++++++++++++ nac3core/src/codegen/mod.rs | 1 + nac3core/src/codegen/model/int.rs | 9 + nac3core/src/codegen/model/pointer.rs | 6 + nac3core/src/codegen/model/slice.rs | 16 +- nac3core/src/codegen/numpy_new.rs | 188 ++++++++++++++++ nac3core/src/toplevel/builtins.rs | 7 +- 16 files changed, 939 insertions(+), 15 deletions(-) create mode 100644 nac3core/irrt/irrt/numpy/ndarray.hpp create mode 100644 nac3core/irrt/irrt/numpy/ndarray_util.hpp create mode 100644 nac3core/irrt/test/ndarray.hpp create mode 100644 nac3core/src/codegen/irrt/numpy/mod.rs create mode 100644 nac3core/src/codegen/irrt/numpy/ndarray.rs create mode 100644 nac3core/src/codegen/irrt/numpy/shape.rs create mode 100644 nac3core/src/codegen/numpy_new.rs diff --git a/nac3core/irrt/irrt/numpy/ndarray.hpp b/nac3core/irrt/irrt/numpy/ndarray.hpp new file mode 100644 index 00000000..5e91877f --- /dev/null +++ b/nac3core/irrt/irrt/numpy/ndarray.hpp @@ -0,0 +1,155 @@ +#pragma once + +#include +#include + +namespace { + +// The NDArray object. `SizeT` is the *signed* size type of this ndarray. +// +// NOTE: The order of fields is IMPORTANT. DON'T TOUCH IT +// +// Some resources you might find helpful: +// - The official numpy implementations: +// - https://github.com/numpy/numpy/blob/735a477f0bc2b5b84d0e72d92f224bde78d4e069/doc/source/reference/c-api/types-and-structures.rst +// - On strides (about reshaping, slicing, C-contagiousness, etc) +// - https://ajcr.net/stride-guide-part-1/. +// - https://ajcr.net/stride-guide-part-2/. +// - https://ajcr.net/stride-guide-part-3/. +template +struct NDArray { + // The underlying data this `ndarray` is pointing to. + // + // NOTE: Formally this should be of type `void *`, but clang + // translates `void *` to `i8 *` when run with `-S -emit-llvm`, + // so we will put `uint8_t *` here for clarity. + // + // This pointer should point to the first element of the ndarray directly + uint8_t *data; + + // The number of bytes of a single element in `data`. + // + // The `SizeT` is treated as `unsigned`. + SizeT itemsize; + + // The number of dimensions of this shape. + // + // The `SizeT` is treated as `unsigned`. + SizeT ndims; + + // Array shape, with length equal to `ndims`. + // + // The `SizeT` is treated as `unsigned`. + // + // NOTE: `shape` can contain 0. + // (those appear when the user makes an out of bounds slice into an ndarray, e.g., `np.zeros((3, 3))[400:].shape == (0, 3)`) + SizeT *shape; + + // Array strides (stride value is in number of bytes, NOT number of elements), with length equal to `ndims`. + // + // The `SizeT` is treated as `signed`. + // + // NOTE: `strides` can have negative numbers. + // (those appear when there is a slice with a negative step, e.g., `my_array[::-1]`) + SizeT *strides; + + // Calculate the size/# of elements of an `ndarray`. + // This function corresponds to `np.size()` or `ndarray.size` + SizeT size() { + return ndarray_util::calc_size_from_shape(ndims, shape); + } + + // Calculate the number of bytes of its content of an `ndarray` *in its view*. + // This function corresponds to `ndarray.nbytes` + SizeT nbytes() { + return this->size() * itemsize; + } + + // Set the strides of the ndarray with `ndarray_util::set_strides_by_shape` + void set_strides_by_shape() { + ndarray_util::set_strides_by_shape(itemsize, ndims, strides, shape); + } + + uint8_t* get_pelement_by_indices(const SizeT *indices) { + uint8_t* element = data; + for (SizeT dim_i = 0; dim_i < ndims; dim_i++) + element += indices[dim_i] * strides[dim_i]; + return element; + } + + uint8_t* get_nth_pelement(SizeT nth) { + SizeT* indices = (SizeT*) __builtin_alloca(sizeof(SizeT) * this->ndims); + ndarray_util::set_indices_by_nth(this->ndims, this->shape, indices, nth); + return get_pelement_by_indices(indices); + } + + // Get the pointer to the nth element of the ndarray as if it were flattened. + uint8_t* checked_get_nth_pelement(ErrorContext* errctx, SizeT nth) { + SizeT arr_size = this->size(); + if (!(0 <= nth && nth < arr_size)) { + errctx->set_error( + errctx->error_ids->index_error, + "index {0} is out of bounds, valid range is {1} <= index < {2}", + nth, 0, arr_size + ); + return 0; + } + return get_nth_pelement(nth); + } + + void set_pelement_value(uint8_t* pelement, const uint8_t* pvalue) { + __builtin_memcpy(pelement, pvalue, itemsize); + } + + // Fill the ndarray with a value + void fill_generic(const uint8_t* pvalue) { + const SizeT size = this->size(); + for (SizeT i = 0; i < size; i++) { + uint8_t* pelement = get_nth_pelement(i); // No need for checked_get_nth_pelement + set_pelement_value(pelement, pvalue); + } + } +}; +} + +extern "C" { +uint32_t __nac3_ndarray_size(NDArray* ndarray) { + return ndarray->size(); +} + +uint64_t __nac3_ndarray_size64(NDArray* ndarray) { + return ndarray->size(); +} + +uint32_t __nac3_ndarray_nbytes(NDArray* ndarray) { + return ndarray->nbytes(); +} + +uint64_t __nac3_ndarray_nbytes64(NDArray* ndarray) { + return ndarray->nbytes(); +} + +void __nac3_ndarray_util_assert_shape_no_negative(ErrorContext* errctx, int32_t ndims, int32_t* shape) { + ndarray_util::assert_shape_no_negative(errctx, ndims, shape); +} + +void __nac3_ndarray_util_assert_shape_no_negative64(ErrorContext* errctx, int64_t ndims, int64_t* shape) { + ndarray_util::assert_shape_no_negative(errctx, ndims, shape); +} + +void __nac3_ndarray_set_strides_by_shape(NDArray* ndarray) { + ndarray->set_strides_by_shape(); +} + +void __nac3_ndarray_set_strides_by_shape64(NDArray* ndarray) { + ndarray->set_strides_by_shape(); +} + +void __nac3_ndarray_fill_generic(NDArray* ndarray, uint8_t* pvalue) { + ndarray->fill_generic(pvalue); +} + +void __nac3_ndarray_fill_generic64(NDArray* ndarray, uint8_t* pvalue) { + ndarray->fill_generic(pvalue); +} +} \ No newline at end of file diff --git a/nac3core/irrt/irrt/numpy/ndarray_util.hpp b/nac3core/irrt/irrt/numpy/ndarray_util.hpp new file mode 100644 index 00000000..e99804c8 --- /dev/null +++ b/nac3core/irrt/irrt/numpy/ndarray_util.hpp @@ -0,0 +1,107 @@ +#pragma once + +#include + +namespace { +namespace ndarray_util { + +// Throw an error if there is an axis with negative dimension +template +void assert_shape_no_negative(ErrorContext* errctx, SizeT ndims, const SizeT* shape) { + for (SizeT axis = 0; axis < ndims; axis++) { + if (shape[axis] < 0) { + errctx->set_error( + errctx->error_ids->value_error, + "negative dimensions are not allowed; axis {0} has dimension {1}", + axis, shape[axis] + ); + return; + } + } +} + +// Compute the size/# of elements of an ndarray given its shape +template +SizeT calc_size_from_shape(SizeT ndims, const SizeT* shape) { + SizeT size = 1; + for (SizeT axis = 0; axis < ndims; axis++) size *= shape[axis]; + return size; +} + +// Compute the strides of an ndarray given an ndarray `shape` +// and assuming that the ndarray is *fully C-contagious*. +// +// You might want to read up on https://ajcr.net/stride-guide-part-1/. +template +void set_strides_by_shape(SizeT itemsize, SizeT ndims, SizeT* dst_strides, const SizeT* shape) { + SizeT stride_product = 1; + for (SizeT i = 0; i < ndims; i++) { + int axis = ndims - i - 1; + dst_strides[axis] = stride_product * itemsize; + stride_product *= shape[axis]; + } +} + +template +void set_indices_by_nth(SizeT ndims, const SizeT* shape, SizeT* indices, SizeT nth) { + for (int32_t i = 0; i < ndims; i++) { + int32_t axis = ndims - i - 1; + int32_t dim = shape[axis]; + + indices[axis] = nth % dim; + nth /= dim; + } +} + +template +bool can_broadcast_shape_to( + const SizeT target_ndims, + const SizeT *target_shape, + const SizeT src_ndims, + const SizeT *src_shape +) { + /* + // See https://numpy.org/doc/stable/user/basics.broadcasting.html + + This function handles this example: + ``` + Image (3d array): 256 x 256 x 3 + Scale (1d array): 3 + Result (3d array): 256 x 256 x 3 + ``` + + Other interesting examples to consider: + - `can_broadcast_shape_to([3], [1, 1, 1, 1, 3]) == true` + - `can_broadcast_shape_to([3], [3, 1]) == false` + - `can_broadcast_shape_to([256, 256, 3], [256, 1, 3]) == true` + + In cases when the shapes contain zero(es): + - `can_broadcast_shape_to([0], [1]) == true` + - `can_broadcast_shape_to([0], [2]) == false` + - `can_broadcast_shape_to([0, 4, 0, 0], [1]) == true` + - `can_broadcast_shape_to([0, 4, 0, 0], [1, 1, 1, 1]) == true` + - `can_broadcast_shape_to([0, 4, 0, 0], [1, 4, 1, 1]) == true` + - `can_broadcast_shape_to([4, 3], [0, 3]) == false` + - `can_broadcast_shape_to([4, 3], [0, 0]) == false` + */ + + // This is essentially doing the following in Python: + // `for target_dim, src_dim in itertools.zip_longest(target_shape[::-1], src_shape[::-1], fillvalue=1)` + for (SizeT i = 0; i < max(target_ndims, src_ndims); i++) { + SizeT target_axis = target_ndims - i - 1; + SizeT src_axis = src_ndims - i - 1; + + bool target_dim_exists = target_axis >= 0; + bool src_dim_exists = src_axis >= 0; + + SizeT target_dim = target_dim_exists ? target_shape[target_axis] : 1; + SizeT src_dim = src_dim_exists ? src_shape[src_axis] : 1; + + bool ok = src_dim == 1 || target_dim == src_dim; + if (!ok) return false; + } + + return true; +} +} +} \ No newline at end of file diff --git a/nac3core/irrt/irrt_everything.hpp b/nac3core/irrt/irrt_everything.hpp index d0480a5b..195649b4 100644 --- a/nac3core/irrt/irrt_everything.hpp +++ b/nac3core/irrt/irrt_everything.hpp @@ -1,6 +1,8 @@ #pragma once #include +#include #include -#include -#include \ No newline at end of file +#include +#include +#include \ No newline at end of file diff --git a/nac3core/irrt/irrt_test.cpp b/nac3core/irrt/irrt_test.cpp index fefcd2a0..0332ef66 100644 --- a/nac3core/irrt/irrt_test.cpp +++ b/nac3core/irrt/irrt_test.cpp @@ -8,9 +8,11 @@ #include #include +#include #include int main() { test_int_exp(); + run_all_tests_ndarray(); return 0; } \ No newline at end of file diff --git a/nac3core/irrt/test/ndarray.hpp b/nac3core/irrt/test/ndarray.hpp new file mode 100644 index 00000000..3554f613 --- /dev/null +++ b/nac3core/irrt/test/ndarray.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include + +void test_calc_size_from_shape_normal() { + // Test shapes with normal values + BEGIN_TEST(); + + int32_t shape[4] = { 2, 3, 5, 7 }; + assert_values_match(210, ndarray_util::calc_size_from_shape(4, shape)); +} + +void test_calc_size_from_shape_has_zero() { + // Test shapes with 0 in them + BEGIN_TEST(); + + int32_t shape[4] = { 2, 0, 5, 7 }; + assert_values_match(0, ndarray_util::calc_size_from_shape(4, shape)); +} + +void test_set_strides_by_shape() { + // Test `set_strides_by_shape()` + BEGIN_TEST(); + + int32_t shape[4] = { 99, 3, 5, 7 }; + int32_t strides[4] = { 0 }; + ndarray_util::set_strides_by_shape((int32_t) sizeof(int32_t), 4, strides, shape); + + int32_t expected_strides[4] = { + 105 * sizeof(int32_t), + 35 * sizeof(int32_t), + 7 * sizeof(int32_t), + 1 * sizeof(int32_t) + }; + assert_arrays_match(4, expected_strides, strides); +} + +void run_all_tests_ndarray() { + test_calc_size_from_shape_normal(); + test_calc_size_from_shape_has_zero(); + test_set_strides_by_shape(); +} \ No newline at end of file diff --git a/nac3core/irrt/test/test_core.hpp b/nac3core/irrt/test/test_core.hpp index 09b0c08a..661ca2b7 100644 --- a/nac3core/irrt/test/test_core.hpp +++ b/nac3core/irrt/test/test_core.hpp @@ -8,4 +8,4 @@ void test_int_exp() { assert_values_match(125, __nac3_int_exp_impl(5, 3)); assert_values_match(3125, __nac3_int_exp_impl(5, 5)); -} \ No newline at end of file +} diff --git a/nac3core/src/codegen/irrt/mod.rs b/nac3core/src/codegen/irrt/mod.rs index 280c7ff1..ebd66cd8 100644 --- a/nac3core/src/codegen/irrt/mod.rs +++ b/nac3core/src/codegen/irrt/mod.rs @@ -1,6 +1,7 @@ use crate::typecheck::typedef::Type; pub mod error_context; +pub mod numpy; mod test; mod util; diff --git a/nac3core/src/codegen/irrt/numpy/mod.rs b/nac3core/src/codegen/irrt/numpy/mod.rs new file mode 100644 index 00000000..817dc114 --- /dev/null +++ b/nac3core/src/codegen/irrt/numpy/mod.rs @@ -0,0 +1,5 @@ +pub mod ndarray; +pub mod shape; + +pub use ndarray::*; +pub use shape::*; diff --git a/nac3core/src/codegen/irrt/numpy/ndarray.rs b/nac3core/src/codegen/irrt/numpy/ndarray.rs new file mode 100644 index 00000000..2e510dce --- /dev/null +++ b/nac3core/src/codegen/irrt/numpy/ndarray.rs @@ -0,0 +1,243 @@ +use inkwell::types::{BasicType, BasicTypeEnum}; + +use crate::codegen::{ + irrt::{ + error_context::{check_error_context, prepare_error_context, ErrorContext}, + util::{get_sized_dependent_function_name, FunctionBuilder}, + }, + model::*, + CodeGenContext, CodeGenerator, +}; + +use super::Producer; + +pub struct NpArrayFields<'ctx> { + pub data: Field, + pub itemsize: Field>, + pub ndims: Field>, + pub shape: Field>>, + pub strides: Field>>, +} + +#[derive(Debug, Clone, Copy)] +pub struct NpArray<'ctx> { + pub sizet: IntModel<'ctx>, +} + +impl<'ctx> IsStruct<'ctx> for NpArray<'ctx> { + type Fields = NpArrayFields<'ctx>; + + fn struct_name(&self) -> &'static str { + "NDArray" + } + + fn build_fields(&self, builder: &mut FieldBuilder<'ctx>) -> Self::Fields { + NpArrayFields { + data: builder.add_field_auto("data"), + itemsize: builder.add_field("itemsize", self.sizet), + ndims: builder.add_field("ndims", self.sizet), + shape: builder.add_field("shape", PointerModel(self.sizet)), + strides: builder.add_field("strides", PointerModel(self.sizet)), + } + } +} + +impl<'ctx> Pointer<'ctx, StructModel>> { + pub fn shape_slice(&self, ctx: &CodeGenContext<'ctx, '_>) -> ArraySlice<'ctx, IntModel<'ctx>> { + let ndims = self.gep(ctx, |f| f.ndims).load(ctx, "ndims"); + let shape_base_ptr = self.gep(ctx, |f| f.shape).load(ctx, "shape"); + ArraySlice { num_elements: ndims, pointer: shape_base_ptr } + } + + pub fn strides_slice( + &self, + ctx: &CodeGenContext<'ctx, '_>, + ) -> ArraySlice<'ctx, IntModel<'ctx>> { + let ndims = self.gep(ctx, |f| f.ndims).load(ctx, "ndims"); + let strides_base_ptr = self.gep(ctx, |f| f.strides).load(ctx, "strides"); + ArraySlice { num_elements: ndims, pointer: strides_base_ptr } + } +} + +pub fn alloca_ndarray<'ctx, G>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + elem_type: BasicTypeEnum<'ctx>, + ndims: &Int<'ctx>, + name: &str, +) -> Result>>, String> +where + G: CodeGenerator + ?Sized, +{ + let sizet = IntModel(generator.get_size_type(ctx.ctx)); + + // Allocate ndarray + let ndarray_ptr = StructModel(NpArray { sizet }).alloca(ctx, name); + + // Set ndims + ndarray_ptr.gep(ctx, |f| f.ndims).store(ctx, ndims); + + // Set itemsize + let itemsize = elem_type.size_of().unwrap(); + let itemsize = + ctx.builder.build_int_s_extend_or_bit_cast(itemsize, sizet.0, "itemsize").unwrap(); + ndarray_ptr.gep(ctx, |f| f.itemsize).store(ctx, &Int(itemsize)); + + // Allocate and set shape + let shape_ptr = ctx.builder.build_array_alloca(sizet.0, ndims.0, "shape").unwrap(); + ndarray_ptr.gep(ctx, |f| f.shape).store(ctx, &Pointer { element: sizet, value: shape_ptr }); + // .store(ctx, &Pointer { addressee_optic: IntLens(sizet), address: shape_ptr }); + + // Allocate and set strides + let strides_ptr = ctx.builder.build_array_alloca(sizet.0, ndims.0, "strides").unwrap(); + ndarray_ptr.gep(ctx, |f| f.strides).store(ctx, &Pointer { element: sizet, value: strides_ptr }); + + Ok(ndarray_ptr) +} + +pub enum NDArrayInitMode<'ctx, G: CodeGenerator + ?Sized> { + NDims { ndims: Int<'ctx> }, + Shape { shape: Producer<'ctx, G, IntModel<'ctx>> }, + ShapeAndAllocaData { shape: Producer<'ctx, G, IntModel<'ctx>> }, +} + +/// TODO: DOCUMENT ME +pub fn alloca_ndarray_and_init<'ctx, G>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + elem_type: BasicTypeEnum<'ctx>, + init_mode: NDArrayInitMode<'ctx, G>, + name: &str, +) -> Result>>, String> +where + G: CodeGenerator + ?Sized, +{ + // It is implemented verbosely in order to make the initialization modes super clear in their intent. + match init_mode { + NDArrayInitMode::NDims { ndims } => { + let ndarray_ptr = alloca_ndarray(generator, ctx, elem_type, &ndims, name)?; + Ok(ndarray_ptr) + } + NDArrayInitMode::Shape { shape } => { + let ndims = shape.count; + let ndarray_ptr = alloca_ndarray(generator, ctx, elem_type, &ndims, name)?; + + // Fill `ndarray.shape` + (shape.write_to_array)(generator, ctx, &ndarray_ptr.shape_slice(ctx))?; + + // Check if `shape` has bad inputs + call_nac3_ndarray_util_assert_shape_no_negative( + generator, + ctx, + &ndims, + &ndarray_ptr.gep(ctx, |f| f.shape).load(ctx, "shape"), + ); + + // NOTE: DO NOT DO `set_strides_by_shape` HERE. + // Simply this is because we specified that `SetShape` wouldn't do `set_strides_by_shape` + + Ok(ndarray_ptr) + } + NDArrayInitMode::ShapeAndAllocaData { shape } => { + let ndims = shape.count; + let ndarray_ptr = alloca_ndarray(generator, ctx, elem_type, &ndims, name)?; + + // Fill `ndarray.shape` + (shape.write_to_array)(generator, ctx, &ndarray_ptr.shape_slice(ctx))?; + + // Check if `shape` has bad inputs + call_nac3_ndarray_util_assert_shape_no_negative( + generator, + ctx, + &ndims, + &ndarray_ptr.gep(ctx, |f| f.shape).load(ctx, "shape"), + ); + + // Now we populate `ndarray.data` by alloca-ing. + // But first, we need to know the size of the ndarray to know how many elements to alloca, + // since calculating nbytes of an ndarray requires `ndarray.shape` to be set. + let ndarray_nbytes = call_nac3_ndarray_nbytes(generator, ctx, &ndarray_ptr); + + // Alloca `data` and assign it to `ndarray.data` + let data_ptr = OpaquePointer( + ctx.builder + .build_array_alloca(ctx.ctx.i8_type(), ndarray_nbytes.0, "data") + .unwrap(), + ); + ndarray_ptr.gep(ctx, |f| f.data).store(ctx, &data_ptr); + + // Finally, do `set_strides_by_shape` + // Check out https://ajcr.net/stride-guide-part-1/ to see what numpy "strides" are. + call_nac3_ndarray_set_strides_by_shape(generator, ctx, &ndarray_ptr); + + Ok(ndarray_ptr) + } + } +} + +fn call_nac3_ndarray_util_assert_shape_no_negative<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ndims: &Int<'ctx>, + shape_ptr: &Pointer<'ctx, IntModel<'ctx>>, +) { + let sizet = IntModel(generator.get_size_type(ctx.ctx)); + + let errctx = prepare_error_context(ctx); + FunctionBuilder::begin( + ctx, + &get_sized_dependent_function_name(sizet, "__nac3_ndarray_util_assert_shape_no_negative"), + ) + .arg("errctx", &PointerModel(StructModel(ErrorContext)), &errctx) + .arg("ndims", &sizet, ndims) + .arg("shape", &PointerModel(sizet), shape_ptr) + .returning_void(); + check_error_context(generator, ctx, &errctx); +} + +fn call_nac3_ndarray_set_strides_by_shape<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ndarray_ptr: &Pointer<'ctx, StructModel>>, +) { + let sizet = IntModel(generator.get_size_type(ctx.ctx)); + + FunctionBuilder::begin( + ctx, + &get_sized_dependent_function_name(sizet, "__nac3_ndarray_util_assert_shape_no_negative"), + ) + .arg("ndarray", &PointerModel(StructModel(NpArray { sizet })), ndarray_ptr) + .returning_void(); +} + +fn call_nac3_ndarray_nbytes<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ndarray_ptr: &Pointer<'ctx, StructModel>>, +) -> Int<'ctx> { + let sizet = IntModel(generator.get_size_type(ctx.ctx)); + + FunctionBuilder::begin( + ctx, + &get_sized_dependent_function_name(sizet, "__nac3_ndarray_util_assert_shape_no_negative"), + ) + .arg("ndarray", &PointerModel(StructModel(NpArray { sizet })), ndarray_ptr) + .returning("nbytes", &sizet) +} + +pub fn call_nac3_ndarray_fill_generic<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ndarray_ptr: &Pointer<'ctx, StructModel>>, + fill_value_ptr: &OpaquePointer<'ctx>, +) { + let sizet = IntModel(generator.get_size_type(ctx.ctx)); + + FunctionBuilder::begin( + ctx, + &get_sized_dependent_function_name(sizet, "__nac3_ndarray_fill_generic"), + ) + .arg("ndarray", &PointerModel(StructModel(NpArray { sizet })), ndarray_ptr) + .arg("pvalue", &OpaquePointerModel, fill_value_ptr) + .returning_void(); +} diff --git a/nac3core/src/codegen/irrt/numpy/shape.rs b/nac3core/src/codegen/irrt/numpy/shape.rs new file mode 100644 index 00000000..d3b01a0b --- /dev/null +++ b/nac3core/src/codegen/irrt/numpy/shape.rs @@ -0,0 +1,162 @@ +use inkwell::values::BasicValueEnum; + +use crate::{ + codegen::{ + classes::{ListValue, UntypedArrayLikeAccessor}, + model::*, + stmt::gen_for_callback_incrementing, + CodeGenContext, CodeGenerator, + }, + typecheck::typedef::{Type, TypeEnum}, +}; + +pub type ProducerWriteToArray<'ctx, G, E> = Box< + dyn Fn(&mut G, &mut CodeGenContext<'ctx, '_>, &ArraySlice<'ctx, E>) -> Result<(), String> + + 'ctx, +>; + +pub struct Producer<'ctx, G: CodeGenerator + ?Sized, E: Model<'ctx>> { + pub count: Int<'ctx>, + pub write_to_array: ProducerWriteToArray<'ctx, G, E>, +} + +/// TODO: UPDATE DOCUMENTATION +/// LLVM-typed implementation for generating a [`Producer`] that sets a list of ints. +/// +/// * `elem_ty` - The element type of the `NDArray`. +/// * `shape` - The `shape` parameter used to construct the `NDArray`. +/// +/// ### Notes on `shape` +/// +/// Just like numpy, the `shape` argument can be: +/// 1. A list of `int32`; e.g., `np.empty([600, 800, 3])` +/// 2. A tuple of `int32`; e.g., `np.empty((600, 800, 3))` +/// 3. A scalar `int32`; e.g., `np.empty(3)`, this is functionally equivalent to `np.empty([3])` +/// +/// See also [`typecheck::type_inferencer::fold_numpy_function_call_shape_argument`] to +/// learn how `shape` gets from being a Python user expression to here. +pub fn parse_input_shape_arg<'ctx, G>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + shape: BasicValueEnum<'ctx>, + shape_ty: Type, +) -> Producer<'ctx, G, IntModel<'ctx>> +where + G: CodeGenerator + ?Sized, +{ + let sizet = IntModel(generator.get_size_type(ctx.ctx)); + + match &*ctx.unifier.get_ty(shape_ty) { + TypeEnum::TObj { obj_id, .. } + if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() => + { + // 1. A list of ints; e.g., `np.empty([600, 800, 3])` + + // A list has to be a PointerValue + let shape_list = ListValue::from_ptr_val(shape.into_pointer_value(), sizet.0, None); + + // Create `Producer` + let ndims = Int(shape_list.load_size(ctx, Some("count"))); + Producer { + count: ndims, + write_to_array: Box::new(move |ctx, generator, dst_array| { + // Basically iterate through the list and write to `dst_slice` accordingly + let init_val = sizet.constant(0).0; + let max_val = (ndims.0, false); + let incr_val = sizet.constant(1).0; + gen_for_callback_incrementing( + ctx, + generator, + init_val, + max_val, + |generator, ctx, _hooks, axis| { + let axis = Int(axis); + + // Get the dimension at `axis` + let dim = shape_list + .data() + .get(ctx, generator, &axis.0, None) + .into_int_value(); + + // Cast `dim` to SizeT + let dim = ctx + .builder + .build_int_s_extend_or_bit_cast(dim, sizet.0, "dim_casted") + .unwrap(); + + // Write + dst_array.ix(generator, ctx, axis, "dim").store(ctx, &Int(dim)); + Ok(()) + }, + incr_val, + ) + }), + } + } + TypeEnum::TTuple { ty: tuple_types } => { + // 2. A tuple of ints; e.g., `np.empty((600, 800, 3))` + + // Get the length/size of the tuple, which also happens to be the value of `ndims`. + let ndims = tuple_types.len(); + + // A tuple has to be a StructValue + // Read [`codegen::expr::gen_expr`] to see how `nac3core` translates a Python tuple into LLVM. + let shape_tuple = shape.into_struct_value(); + + Producer { + count: sizet.constant(ndims as u64), + write_to_array: Box::new(move |generator, ctx, dst_array| { + for axis in 0..ndims { + // Get the dimension at `axis` + let dim = ctx + .builder + .build_extract_value( + shape_tuple, + axis as u32, + format!("dim{axis}").as_str(), + ) + .unwrap() + .into_int_value(); + + // Cast `dim` to SizeT + let dim = ctx + .builder + .build_int_s_extend_or_bit_cast(dim, sizet.0, "dim_casted") + .unwrap(); + + // Write + dst_array + .ix(generator, ctx, sizet.constant(axis as u64), "dim") + .store(ctx, &Int(dim)); + } + Ok(()) + }), + } + } + TypeEnum::TObj { obj_id, .. } + if *obj_id == ctx.primitives.int32.obj_id(&ctx.unifier).unwrap() => + { + // 3. A scalar int; e.g., `np.empty(3)`, this is functionally equivalent to `np.empty([3])` + + // The value has to be an integer + let shape_int = shape.into_int_value(); + + Producer { + count: sizet.constant(1), + write_to_array: Box::new(move |generator, ctx, dst_array| { + // Cast `shape_int` to SizeT + let dim = ctx + .builder + .build_int_s_extend_or_bit_cast(shape_int, sizet.0, "dim_casted") + .unwrap(); + + // Set shape[0] = shape_int + dst_array.ix(generator, ctx, sizet.constant(0), "dim").store(ctx, &Int(dim)); + + Ok(()) + }), + } + } + _ => panic!("parse_input_shape_arg encountered unknown type"), + } +} diff --git a/nac3core/src/codegen/mod.rs b/nac3core/src/codegen/mod.rs index 36741905..f9e42290 100644 --- a/nac3core/src/codegen/mod.rs +++ b/nac3core/src/codegen/mod.rs @@ -45,6 +45,7 @@ pub mod irrt; pub mod llvm_intrinsics; pub mod model; pub mod numpy; +pub mod numpy_new; pub mod stmt; #[cfg(test)] diff --git a/nac3core/src/codegen/model/int.rs b/nac3core/src/codegen/model/int.rs index ebf06408..228ed705 100644 --- a/nac3core/src/codegen/model/int.rs +++ b/nac3core/src/codegen/model/int.rs @@ -8,6 +8,8 @@ use super::core::*; #[derive(Debug, Clone, Copy)] pub struct IntModel<'ctx>(pub IntType<'ctx>); + +#[derive(Debug, Clone, Copy)] pub struct Int<'ctx>(pub IntValue<'ctx>); impl<'ctx> ModelValue<'ctx> for Int<'ctx> { @@ -30,6 +32,13 @@ impl<'ctx> Model<'ctx> for IntModel<'ctx> { } } +impl<'ctx> IntModel<'ctx> { + #[must_use] + pub fn constant(&self, value: u64) -> Int<'ctx> { + Int(self.0.const_int(value, false)) + } +} + #[derive(Debug, Clone, Default)] pub struct FixedIntModel(pub T); pub struct FixedInt<'ctx, T: IsFixedInt> { diff --git a/nac3core/src/codegen/model/pointer.rs b/nac3core/src/codegen/model/pointer.rs index 45c6ac6f..bbe0bb93 100644 --- a/nac3core/src/codegen/model/pointer.rs +++ b/nac3core/src/codegen/model/pointer.rs @@ -72,3 +72,9 @@ impl<'ctx> Model<'ctx> for OpaquePointerModel { OpaquePointer(ptr) } } + +impl<'ctx> OpaquePointer<'ctx> { + pub fn store(&self, ctx: &CodeGenContext<'ctx, '_>, value: BasicValueEnum<'ctx>) { + ctx.builder.build_store(self.0, value).unwrap(); + } +} diff --git a/nac3core/src/codegen/model/slice.rs b/nac3core/src/codegen/model/slice.rs index 8e3cefa0..db6c6406 100644 --- a/nac3core/src/codegen/model/slice.rs +++ b/nac3core/src/codegen/model/slice.rs @@ -1,5 +1,3 @@ -use inkwell::values::IntValue; - use crate::codegen::{CodeGenContext, CodeGenerator}; use super::{Int, Model, Pointer}; @@ -13,11 +11,11 @@ impl<'ctx, E: Model<'ctx>> ArraySlice<'ctx, E> { pub fn ix_unchecked( &self, ctx: &CodeGenContext<'ctx, '_>, - idx: IntValue<'ctx>, + idx: Int<'ctx>, name: &str, ) -> Pointer<'ctx, E> { let element_addr = - unsafe { ctx.builder.build_in_bounds_gep(self.pointer.value, &[idx], name).unwrap() }; + unsafe { ctx.builder.build_in_bounds_gep(self.pointer.value, &[idx.0], name).unwrap() }; Pointer { value: element_addr, element: self.pointer.element.clone() } } @@ -25,12 +23,12 @@ impl<'ctx, E: Model<'ctx>> ArraySlice<'ctx, E> { &self, generator: &mut G, ctx: &mut CodeGenContext<'ctx, '_>, - idx: IntValue<'ctx>, + idx: Int<'ctx>, name: &str, ) -> Pointer<'ctx, E> { let int_type = self.num_elements.0.get_type(); // NOTE: Weird get_type(), see comment under `trait Ixed` - assert_eq!(int_type.get_bit_width(), idx.get_type().get_bit_width()); // Might as well check bit width to catch bugs + assert_eq!(int_type.get_bit_width(), idx.0.get_type().get_bit_width()); // Might as well check bit width to catch bugs // TODO: SGE or UGE? or make it defined by the implementee? @@ -40,7 +38,7 @@ impl<'ctx, E: Model<'ctx>> ArraySlice<'ctx, E> { .build_int_compare( inkwell::IntPredicate::SLE, int_type.const_zero(), - idx, + idx.0, "lower_bounded", ) .unwrap(); @@ -50,7 +48,7 @@ impl<'ctx, E: Model<'ctx>> ArraySlice<'ctx, E> { .builder .build_int_compare( inkwell::IntPredicate::SLT, - idx, + idx.0, self.num_elements.0, "upper_bounded", ) @@ -65,7 +63,7 @@ impl<'ctx, E: Model<'ctx>> ArraySlice<'ctx, E> { bounded, "0:IndexError", "nac3core LLVM codegen attempting to access out of bounds array index {0}. Must satisfy 0 <= index < {2}", - [ Some(idx), Some(self.num_elements.0), None], + [ Some(idx.0), Some(self.num_elements.0), None], ctx.current_loc ); diff --git a/nac3core/src/codegen/numpy_new.rs b/nac3core/src/codegen/numpy_new.rs new file mode 100644 index 00000000..f726d5f8 --- /dev/null +++ b/nac3core/src/codegen/numpy_new.rs @@ -0,0 +1,188 @@ +use inkwell::values::{BasicValue, BasicValueEnum, PointerValue}; +use nac3parser::ast::StrRef; + +use crate::{ + symbol_resolver::ValueEnum, + toplevel::DefinitionId, + typecheck::typedef::{FunSignature, Type}, +}; + +use super::{ + irrt::numpy::{ + alloca_ndarray_and_init, call_nac3_ndarray_fill_generic, parse_input_shape_arg, + NDArrayInitMode, NpArray, + }, + model::*, + CodeGenContext, CodeGenerator, +}; + +/// LLVM-typed implementation for generating the implementation for constructing an empty `NDArray`. +fn call_ndarray_empty_impl<'ctx, G>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + elem_ty: Type, + shape: BasicValueEnum<'ctx>, + shape_ty: Type, + name: &str, +) -> Result>>, String> +where + G: CodeGenerator + ?Sized, +{ + let elem_type = ctx.get_llvm_type(generator, elem_ty); + let shape = parse_input_shape_arg(generator, ctx, shape, shape_ty); + let ndarray_ptr = alloca_ndarray_and_init( + generator, + ctx, + elem_type, + NDArrayInitMode::ShapeAndAllocaData { shape }, + name, + )?; + Ok(ndarray_ptr) +} + +fn call_ndarray_full_impl<'ctx, G>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + elem_ty: Type, + shape: BasicValueEnum<'ctx>, + shape_ty: Type, + fill_value: BasicValueEnum<'ctx>, + name: &str, +) -> Result>>, String> +where + G: CodeGenerator + ?Sized, +{ + let ndarray_ptr = call_ndarray_empty_impl(generator, ctx, elem_ty, shape, shape_ty, name)?; + + // NOTE: fill_value's type is not checked!! so be careful with logics + let fill_value_ptr = + OpaquePointer(ctx.builder.build_alloca(fill_value.get_type(), "fill_value_ptr").unwrap()); + fill_value_ptr.store(ctx, fill_value); + call_nac3_ndarray_fill_generic(generator, ctx, &ndarray_ptr, &fill_value_ptr); + + Ok(ndarray_ptr) +} + +/// Generates LLVM IR for `np.empty`. +pub fn gen_ndarray_empty<'ctx>( + context: &mut CodeGenContext<'ctx, '_>, + obj: &Option<(Type, ValueEnum<'ctx>)>, + fun: (&FunSignature, DefinitionId), + args: &[(Option, ValueEnum<'ctx>)], + generator: &mut dyn CodeGenerator, +) -> Result, String> { + assert!(obj.is_none()); + assert_eq!(args.len(), 1); + + // Parse arguments + let shape_ty = fun.0.args[0].ty; + let shape = args[0].1.clone().to_basic_value_enum(context, generator, shape_ty)?; + + // Implementation + let ndarray_ptr = call_ndarray_empty_impl( + generator, + context, + context.primitives.float, + shape, + shape_ty, + "ndarray", + )?; + Ok(ndarray_ptr.value) +} + +/// Generates LLVM IR for `np.zeros`. +pub fn gen_ndarray_zeros<'ctx>( + context: &mut CodeGenContext<'ctx, '_>, + obj: &Option<(Type, ValueEnum<'ctx>)>, + fun: (&FunSignature, DefinitionId), + args: &[(Option, ValueEnum<'ctx>)], + generator: &mut dyn CodeGenerator, +) -> Result, String> { + assert!(obj.is_none()); + assert_eq!(args.len(), 1); + + // Parse arguments + let shape_ty = fun.0.args[0].ty; + let shape = args[0].1.clone().to_basic_value_enum(context, generator, shape_ty)?; + + // Implementation + // NOTE: Currently nac3's `np.zeros` is always `float64`. + let float64_ty = context.primitives.float; + let float64_llvm_type = context.get_llvm_type(generator, float64_ty).into_float_type(); + + let ndarray_ptr = call_ndarray_full_impl( + generator, + context, + float64_ty, // `elem_ty` is always `float64` + shape, + shape_ty, + float64_llvm_type.const_zero().as_basic_value_enum(), + "ndarray", + )?; + Ok(ndarray_ptr.value) +} + +/// Generates LLVM IR for `np.ones`. +pub fn gen_ndarray_ones<'ctx>( + context: &mut CodeGenContext<'ctx, '_>, + obj: &Option<(Type, ValueEnum<'ctx>)>, + fun: (&FunSignature, DefinitionId), + args: &[(Option, ValueEnum<'ctx>)], + generator: &mut dyn CodeGenerator, +) -> Result, String> { + assert!(obj.is_none()); + assert_eq!(args.len(), 1); + + // Parse arguments + let shape_ty = fun.0.args[0].ty; + let shape = args[0].1.clone().to_basic_value_enum(context, generator, shape_ty)?; + + // Implementation + // NOTE: Currently nac3's `np.ones` is always `float64`. + let float64_ty = context.primitives.float; + let float64_llvm_type = context.get_llvm_type(generator, float64_ty).into_float_type(); + + let ndarray_ptr = call_ndarray_full_impl( + generator, + context, + float64_ty, // `elem_ty` is always `float64` + shape, + shape_ty, + float64_llvm_type.const_float(1.0).as_basic_value_enum(), + "ndarray", + )?; + Ok(ndarray_ptr.value) +} + +/// Generates LLVM IR for `ndarray.full`. +pub fn gen_ndarray_full<'ctx>( + context: &mut CodeGenContext<'ctx, '_>, + obj: &Option<(Type, ValueEnum<'ctx>)>, + fun: (&FunSignature, DefinitionId), + args: &[(Option, ValueEnum<'ctx>)], + generator: &mut dyn CodeGenerator, +) -> Result, String> { + assert!(obj.is_none()); + assert_eq!(args.len(), 2); + + // Parse argument #1 shape + let shape_ty = fun.0.args[0].ty; + let shape_arg = args[0].1.clone().to_basic_value_enum(context, generator, shape_ty)?; + + // Parse argument #2 fill_value + let fill_value_ty = fun.0.args[1].ty; + let fill_value_arg = + args[1].1.clone().to_basic_value_enum(context, generator, fill_value_ty)?; + + // Implementation + let ndarray_ptr = call_ndarray_full_impl( + generator, + context, + fill_value_ty, + shape_arg, + shape_ty, + fill_value_arg, + "ndarray", + )?; + Ok(ndarray_ptr.value) +} diff --git a/nac3core/src/toplevel/builtins.rs b/nac3core/src/toplevel/builtins.rs index 0f6e8f24..84bcbf76 100644 --- a/nac3core/src/toplevel/builtins.rs +++ b/nac3core/src/toplevel/builtins.rs @@ -18,6 +18,7 @@ use crate::{ expr::destructure_range, irrt::*, numpy::*, + numpy_new, stmt::exn_constructor, }, symbol_resolver::SymbolValue, @@ -1194,9 +1195,9 @@ impl<'a> BuiltinBuilder<'a> { &[(self.ndarray_factory_fn_shape_arg_tvar.ty, "shape")], Box::new(move |ctx, obj, fun, args, generator| { let func = match prim { - PrimDef::FunNpNDArray | PrimDef::FunNpEmpty => gen_ndarray_empty, - PrimDef::FunNpZeros => gen_ndarray_zeros, - PrimDef::FunNpOnes => gen_ndarray_ones, + PrimDef::FunNpNDArray | PrimDef::FunNpEmpty => numpy_new::gen_ndarray_empty, + PrimDef::FunNpZeros => numpy_new::gen_ndarray_zeros, + PrimDef::FunNpOnes => numpy_new::gen_ndarray_ones, _ => unreachable!(), }; func(ctx, &obj, fun, &args, generator).map(|val| Some(val.as_basic_value_enum()))