abdul124
/
nac3
forked from M-Labs/nac3
1
0
Fork 0
Code Pull Requests Activity

core: new np_{zeros,ones,fill} + some irrt model additions

This commit is contained in:
lyken 2024-07-14 23:16:28 +08:00
parent 51a099b602
commit 3344a2bcd3
16 changed files with 939 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

155
nac3core/irrt/irrt/numpy/ndarray.hpp Normal file
View File

@ -0,0 +1,155 @@
#pragma once
#include <irrt/int_defs.hpp>
#include <irrt/numpy/ndarray_util.hpp>
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 <typename SizeT>
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(<ndarray>)` 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<int32_t>* ndarray) {
return ndarray->size();
}
uint64_t __nac3_ndarray_size64(NDArray<int64_t>* ndarray) {
return ndarray->size();
}
uint32_t __nac3_ndarray_nbytes(NDArray<int32_t>* ndarray) {
return ndarray->nbytes();
}
uint64_t __nac3_ndarray_nbytes64(NDArray<int64_t>* 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<int32_t>* ndarray) {
ndarray->set_strides_by_shape();
}
void __nac3_ndarray_set_strides_by_shape64(NDArray<int64_t>* ndarray) {
ndarray->set_strides_by_shape();
}
void __nac3_ndarray_fill_generic(NDArray<int32_t>* ndarray, uint8_t* pvalue) {
ndarray->fill_generic(pvalue);
}
void __nac3_ndarray_fill_generic64(NDArray<int64_t>* ndarray, uint8_t* pvalue) {
ndarray->fill_generic(pvalue);
}
}

107
nac3core/irrt/irrt/numpy/ndarray_util.hpp Normal file
View File

@ -0,0 +1,107 @@
#pragma once
#include <irrt/int_defs.hpp>
namespace {
namespace ndarray_util {
// Throw an error if there is an axis with negative dimension
template <typename SizeT>
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 <typename SizeT>
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 <typename SizeT>
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 <typename SizeT>
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 <typename SizeT>
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;
}
}
}

6
nac3core/irrt/irrt_everything.hpp
View File

@ -1,6 +1,8 @@
#pragma once
#include <irrt/core.hpp>
#include <irrt/int_defs.hpp>
#include <irrt/utils.hpp>
#include <irrt/error_context.hpp>
#include <irrt/int_defs.hpp>
#include <irrt/numpy/ndarray.hpp>
#include <irrt/numpy/ndarray_util.hpp>
#include <irrt/utils.hpp>

2
nac3core/irrt/irrt_test.cpp
View File

@ -8,9 +8,11 @@
#include <irrt_everything.hpp>
#include <test/core.hpp>
#include <test/ndarray.hpp>
#include <test/test_core.hpp>
int main() {
test_int_exp();
run_all_tests_ndarray();
return 0;
}

44
nac3core/irrt/test/ndarray.hpp Normal file
View File

@ -0,0 +1,44 @@
#pragma once
#include <test/core.hpp>
#include <irrt/numpy/ndarray.hpp>
#include <irrt/numpy/ndarray_util.hpp>
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<int32_t>(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<int32_t>(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();
}

1
nac3core/src/codegen/irrt/mod.rs
View File

@ -1,6 +1,7 @@
use crate::typecheck::typedef::Type;
pub mod error_context;
pub mod numpy;
mod test;
mod util;

5
nac3core/src/codegen/irrt/numpy/mod.rs Normal file
View File

@ -0,0 +1,5 @@
pub mod ndarray;
pub mod shape;
pub use ndarray::*;
pub use shape::*;

243
nac3core/src/codegen/irrt/numpy/ndarray.rs Normal file
View File

@ -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<OpaquePointerModel>,
pub itemsize: Field<IntModel<'ctx>>,
pub ndims: Field<IntModel<'ctx>>,
pub shape: Field<PointerModel<IntModel<'ctx>>>,
pub strides: Field<PointerModel<IntModel<'ctx>>>,
}
#[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<NpArray<'ctx>>> {
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<Pointer<'ctx, StructModel<NpArray<'ctx>>>, 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<Pointer<'ctx, StructModel<NpArray<'ctx>>>, 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<NpArray<'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_void();
}
fn call_nac3_ndarray_nbytes<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
ndarray_ptr: &Pointer<'ctx, StructModel<NpArray<'ctx>>>,
) -> 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<NpArray<'ctx>>>,
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();
}

162
nac3core/src/codegen/irrt/numpy/shape.rs Normal file
View File

@ -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"),
}
}

1
nac3core/src/codegen/mod.rs
View File

@ -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)]

9
nac3core/src/codegen/model/int.rs
View File

@ -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<T>(pub T);
pub struct FixedInt<'ctx, T: IsFixedInt> {

6
nac3core/src/codegen/model/pointer.rs
View File

@ -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();
}
}

16
nac3core/src/codegen/model/slice.rs
View File

@ -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
);

188
nac3core/src/codegen/numpy_new.rs Normal file
View File

@ -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<Pointer<'ctx, StructModel<NpArray<'ctx>>>, 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<Pointer<'ctx, StructModel<NpArray<'ctx>>>, 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<StrRef>, ValueEnum<'ctx>)],
generator: &mut dyn CodeGenerator,
) -> Result<PointerValue<'ctx>, 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<StrRef>, ValueEnum<'ctx>)],
generator: &mut dyn CodeGenerator,
) -> Result<PointerValue<'ctx>, 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<StrRef>, ValueEnum<'ctx>)],
generator: &mut dyn CodeGenerator,
) -> Result<PointerValue<'ctx>, 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<StrRef>, ValueEnum<'ctx>)],
generator: &mut dyn CodeGenerator,
) -> Result<PointerValue<'ctx>, 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)
}

7
nac3core/src/toplevel/builtins.rs
View File

@ -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()))