core/ndstrides: add NDArray with strides definition

This commit is contained in:
lyken 2024-07-28 14:23:31 +08:00
parent 9940d00244
commit b2b46939b5
No known key found for this signature in database
GPG Key ID: 3BD5FC6AC8325DD8
4 changed files with 81 additions and 1 deletions

View File

@ -0,0 +1,44 @@
#pragma once
namespace {
/**
* @brief The NDArray object
*
* The official numpy implementations: https://github.com/numpy/numpy/blob/735a477f0bc2b5b84d0e72d92f224bde78d4e069/doc/source/reference/c-api/types-and-structures.rst
*/
template <typename SizeT>
struct NDArray {
/**
* @brief The underlying data this `ndarray` is pointing to.
*
* Must be set to `nullptr` to indicate that this NDArray's `data` is uninitialized.
*/
uint8_t* data;
/**
* @brief The number of bytes of a single element in `data`.
*/
SizeT itemsize;
/**
* @brief The number of dimensions of this shape.
*/
SizeT ndims;
/**
* @brief The NDArray shape, with length equal to `ndims`.
*
* Note that it may contain 0.
*/
SizeT* shape;
/**
* @brief Array strides, with length equal to `ndims`
*
* The stride values are in units of bytes, not number of elements.
*
* Note that `strides` can have negative values.
*/
SizeT* strides;
};
} // namespace

View File

@ -3,4 +3,5 @@
#include <irrt/core.hpp>
#include <irrt/exception.hpp>
#include <irrt/int_defs.hpp>
#include <irrt/ndarray/def.hpp>
#include <irrt/util.hpp>

View File

@ -1,2 +1,3 @@
pub mod cslice;
pub mod exception;
pub mod ndarray;

View File

@ -0,0 +1,34 @@
use crate::codegen::*;
pub struct NpArrayFields<'ctx, F: FieldTraversal<'ctx>> {
pub data: F::Out<PtrModel<IntModel<Byte>>>,
pub itemsize: F::Out<IntModel<SizeT>>,
pub ndims: F::Out<IntModel<SizeT>>,
pub shape: F::Out<PtrModel<IntModel<SizeT>>>,
pub strides: F::Out<PtrModel<IntModel<SizeT>>>,
}
// TODO: Rename to `NDArray` when the old NDArray is removed.
#[derive(Debug, Clone, Copy, Default)]
pub struct NpArray;
impl<'ctx> StructKind<'ctx> for NpArray {
type Fields<F: FieldTraversal<'ctx>> = NpArrayFields<'ctx, F>;
fn traverse_fields<F: FieldTraversal<'ctx>>(&self, traversal: &mut F) -> Self::Fields<F> {
Self::Fields {
data: traversal.add_auto("data"),
itemsize: traversal.add_auto("itemsize"),
ndims: traversal.add_auto("ndims"),
shape: traversal.add_auto("shape"),
strides: traversal.add_auto("strides"),
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct NDArrayObject<'ctx> {
pub dtype: Type,
pub ndims: Type,
pub value: Ptr<'ctx, StructModel<NpArray>>,
}