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

core: irrt split ndarray.hpp

Browse Source
This commit is contained in:
lyken 2024-07-14 23:35:50 +08:00
parent 3344a2bcd3
commit d92cccb85e
7 changed files with 293 additions and 264 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,155 +0,0 @@
#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);
}
}

151
nac3core/irrt/irrt/numpy/ndarray_basic.hpp Normal file
View File

@ -0,0 +1,151 @@
#pragma once
#include <irrt/int_defs.hpp>
#include <irrt/error_context.hpp>
#include <irrt/numpy/ndarray_def.hpp>
namespace {
namespace ndarray {
namespace 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/.
//
// This function might be used in isolation without an ndarray. That's
// why it separated out into its own util function.
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;
}
}
}
// Calculate the size/# of elements of an `ndarray`.
// This function corresponds to `np.size(<ndarray>)` or `ndarray.size`
template <typename SizeT>
SizeT size(NDArray<SizeT>* ndarray) {
return ndarray::util::calc_size_from_shape(ndarray->ndims, ndarray->shape);
}
// Calculate the number of bytes of its content of an `ndarray` *in its view*.
// This function corresponds to `ndarray.nbytes`
template <typename SizeT>
SizeT nbytes(NDArray<SizeT>* ndarray) {
return ndarray::size(ndarray) * ndarray->itemsize;
}
// Set the strides of the ndarray with `ndarray_util::set_strides_by_shape`
template <typename SizeT>
void set_strides_by_shape(NDArray<SizeT>* ndarray) {
ndarray::util::set_strides_by_shape(ndarray->itemsize, ndarray->ndims, ndarray->strides, ndarray->shape);
}
template <typename SizeT>
uint8_t* get_pelement_by_indices(NDArray<SizeT>* ndarray, const SizeT *indices) {
uint8_t* element = ndarray->data;
for (SizeT dim_i = 0; dim_i < ndarray->ndims; dim_i++)
element += indices[dim_i] * ndarray->strides[dim_i];
return element;
}
template <typename SizeT>
uint8_t* get_nth_pelement(NDArray<SizeT>* ndarray, SizeT nth) {
SizeT* indices = (SizeT*) __builtin_alloca(sizeof(SizeT) * ndarray->ndims);
ndarray::util::set_indices_by_nth(ndarray->ndims, ndarray->shape, indices, nth);
return ndarray::get_pelement_by_indices(ndarray, indices);
}
// Get the pointer to the nth element of the ndarray as if it were flattened.
template <typename SizeT>
uint8_t* checked_get_nth_pelement(NDArray<SizeT>* ndarray, ErrorContext* errctx, SizeT nth) {
SizeT arr_size = ndarray->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 ndarray::get_nth_pelement(ndarray, nth);
}
template <typename SizeT>
void set_pelement_value(NDArray<SizeT>* ndarray, uint8_t* pelement, const uint8_t* pvalue) {
__builtin_memcpy(pelement, pvalue, ndarray->itemsize);
}
};
}
extern "C" {
uint32_t __nac3_ndarray_size(NDArray<int32_t>* ndarray) {
return ndarray::size(ndarray);
}
uint64_t __nac3_ndarray_size64(NDArray<int64_t>* ndarray) {
return ndarray::size(ndarray);
}
uint32_t __nac3_ndarray_nbytes(NDArray<int32_t>* ndarray) {
return ndarray::nbytes(ndarray);
}
uint64_t __nac3_ndarray_nbytes64(NDArray<int64_t>* ndarray) {
return ndarray::nbytes(ndarray);
}
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(ndarray);
}
void __nac3_ndarray_set_strides_by_shape64(NDArray<int64_t>* ndarray) {
ndarray::set_strides_by_shape(ndarray);
}
}

58
nac3core/irrt/irrt/numpy/ndarray_broadcast.hpp Normal file
View File

@ -0,0 +1,58 @@
#include <irrt/numpy/ndarray_def.hpp>
namespace {
namespace ndarray {
namespace util {
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;
}
}
}
}

52
nac3core/irrt/irrt/numpy/ndarray_def.hpp Normal file
View File

@ -0,0 +1,52 @@
#pragma once
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;
};
}

28
nac3core/irrt/irrt/numpy/ndarray_fill.hpp Normal file
View File

@ -0,0 +1,28 @@
#pragma once
#include <irrt/numpy/ndarray_def.hpp>
#include <irrt/numpy/ndarray_basic.hpp>
namespace {
namespace ndarray {
// Fill the ndarray with a value
template <typename SizeT>
void fill_generic(NDArray<SizeT>* ndarray, const uint8_t* pvalue) {
const SizeT size = ndarray::size(ndarray);
for (SizeT i = 0; i < size; i++) {
uint8_t* pelement = ndarray::get_nth_pelement(ndarray, i); // No need for checked_get_nth_pelement
ndarray::set_pelement_value(ndarray, pelement, pvalue);
}
}
}
}
extern "C" {
void __nac3_ndarray_fill_generic(NDArray<int32_t>* ndarray, uint8_t* pvalue) {
ndarray::fill_generic(ndarray, pvalue);
}
void __nac3_ndarray_fill_generic64(NDArray<int64_t>* ndarray, uint8_t* pvalue) {
ndarray::fill_generic(ndarray, pvalue);
}
}

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

@ -1,107 +0,0 @@
#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

@ -3,6 +3,8 @@
#include <irrt/core.hpp>
#include <irrt/error_context.hpp>
#include <irrt/int_defs.hpp>
#include <irrt/numpy/ndarray.hpp>
#include <irrt/numpy/ndarray_util.hpp>
#include <irrt/numpy/ndarray_def.hpp>
#include <irrt/numpy/ndarray_basic.hpp>
#include <irrt/numpy/ndarray_broadcast.hpp>
#include <irrt/numpy/ndarray_fill.hpp>
#include <irrt/utils.hpp>