From 48fb3ff5add10ea1ef2e3d68449db02a81cc6f90 Mon Sep 17 00:00:00 2001 From: lyken Date: Tue, 20 Aug 2024 12:00:31 +0800 Subject: [PATCH] core/ndstrides: add ndarray iterator (NDIter) --- nac3core/irrt/irrt.cpp | 1 + nac3core/irrt/irrt/ndarray/iter.hpp | 142 +++++++++++++++ nac3core/src/codegen/irrt/mod.rs | 31 +++- nac3core/src/codegen/object/ndarray/mod.rs | 2 + nac3core/src/codegen/object/ndarray/nditer.rs | 168 ++++++++++++++++++ 5 files changed, 343 insertions(+), 1 deletion(-) create mode 100644 nac3core/irrt/irrt/ndarray/iter.hpp create mode 100644 nac3core/src/codegen/object/ndarray/nditer.rs diff --git a/nac3core/irrt/irrt.cpp b/nac3core/irrt/irrt.cpp index 1a4dcf93..c372735f 100644 --- a/nac3core/irrt/irrt.cpp +++ b/nac3core/irrt/irrt.cpp @@ -3,4 +3,5 @@ #include #include #include +#include #include \ No newline at end of file diff --git a/nac3core/irrt/irrt/ndarray/iter.hpp b/nac3core/irrt/irrt/ndarray/iter.hpp new file mode 100644 index 00000000..d27e0454 --- /dev/null +++ b/nac3core/irrt/irrt/ndarray/iter.hpp @@ -0,0 +1,142 @@ +#pragma once + +#include +#include + +namespace +{ +/** + * @brief Helper struct to enumerate through all indices under a shape. + * + * i.e., If `shape` is `[3, 2]`, by repeating `next()`, then you get: + * - `[0, 0]` + * - `[0, 1]` + * - `[1, 0]` + * - `[1, 1]` + * - `[2, 0]` + * - `[2, 1]` + * - end. + * + * Interesting cases: + * - If ndims == 0, there is one enumeration. + * - If shape contains zeroes, there are no enumerations. + */ +template struct NDIter +{ + SizeT ndims; + SizeT *shape; + SizeT *strides; + + /** + * @brief The current indices. + * + * Must be allocated by the caller. + */ + SizeT *indices; + + /** + * @brief The nth (0-based) index of the current indices. + */ + SizeT nth; + + /** + * @brief Pointer to the current element. + */ + uint8_t *element; + + /** + * @brief The product of shape. + */ + SizeT size; + + // TODO:: There is something called backstrides to speedup iteration. + // See https://ajcr.net/stride-guide-part-1/, and https://docs.scipy.org/doc/numpy-1.13.0/reference/c-api.types-and-structures.html#c.PyArrayIterObject.PyArrayIterObject.backstrides. + // Maybe LLVM is clever and knows how to optimize. + + void initialize(SizeT ndims, SizeT *shape, SizeT *strides, uint8_t *element, SizeT *indices) + { + this->ndims = ndims; + this->shape = shape; + this->strides = strides; + + this->indices = indices; + this->element = element; + + // Compute size and backstrides + this->size = 1; + for (SizeT i = 0; i < ndims; i++) + { + this->size *= shape[i]; + } + + for (SizeT axis = 0; axis < ndims; axis++) + indices[axis] = 0; + nth = 0; + } + + void initialize_by_ndarray(NDArray *ndarray, SizeT *indices) + { + this->initialize(ndarray->ndims, ndarray->shape, ndarray->strides, ndarray->data, indices); + } + + bool has_next() + { + return nth < size; + } + + void next() + { + for (SizeT i = 0; i < ndims; i++) + { + SizeT axis = ndims - i - 1; + indices[axis]++; + if (indices[axis] >= shape[axis]) + { + indices[axis] = 0; + + // TODO: Can be optimized with backstrides. + element -= strides[axis] * (shape[axis] - 1); + } + else + { + element += strides[axis]; + break; + } + } + nth++; + } +}; +} // namespace + +extern "C" +{ + void __nac3_nditer_initialize(NDIter *iter, NDArray *ndarray, int32_t *indices) + { + iter->initialize_by_ndarray(ndarray, indices); + } + + void __nac3_nditer_initialize64(NDIter *iter, NDArray *ndarray, int64_t *indices) + { + iter->initialize_by_ndarray(ndarray, indices); + } + + bool __nac3_nditer_has_next(NDIter *iter) + { + return iter->has_next(); + } + + bool __nac3_nditer_has_next64(NDIter *iter) + { + return iter->has_next(); + } + + void __nac3_nditer_next(NDIter *iter) + { + iter->next(); + } + + void __nac3_nditer_next64(NDIter *iter) + { + iter->next(); + } +} \ No newline at end of file diff --git a/nac3core/src/codegen/irrt/mod.rs b/nac3core/src/codegen/irrt/mod.rs index 2ad905d2..dedc5563 100644 --- a/nac3core/src/codegen/irrt/mod.rs +++ b/nac3core/src/codegen/irrt/mod.rs @@ -7,7 +7,7 @@ use super::{ }, llvm_intrinsics, model::*, - object::ndarray::NDArray, + object::ndarray::{nditer::NDIter, NDArray}, CodeGenContext, CodeGenerator, }; use crate::codegen::classes::TypedArrayLikeAccessor; @@ -1090,3 +1090,32 @@ pub fn call_nac3_ndarray_copy_data<'ctx, G: CodeGenerator + ?Sized>( let name = get_sizet_dependent_function_name(generator, ctx, "__nac3_ndarray_copy_data"); CallFunction::begin(generator, ctx, &name).arg(src_ndarray).arg(dst_ndarray).returning_void(); } + +pub fn call_nac3_nditer_initialize<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + iter: Instance<'ctx, Ptr>>, + ndarray: Instance<'ctx, Ptr>>, + indices: Instance<'ctx, Ptr>>, +) { + let name = get_sizet_dependent_function_name(generator, ctx, "__nac3_nditer_initialize"); + CallFunction::begin(generator, ctx, &name).arg(iter).arg(ndarray).arg(indices).returning_void(); +} + +pub fn call_nac3_nditer_has_next<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + iter: Instance<'ctx, Ptr>>, +) -> Instance<'ctx, Int> { + let name = get_sizet_dependent_function_name(generator, ctx, "__nac3_nditer_has_next"); + CallFunction::begin(generator, ctx, &name).arg(iter).returning_auto("has_next") +} + +pub fn call_nac3_nditer_next<'ctx, G: CodeGenerator + ?Sized>( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + iter: Instance<'ctx, Ptr>>, +) { + let name = get_sizet_dependent_function_name(generator, ctx, "__nac3_nditer_next"); + CallFunction::begin(generator, ctx, &name).arg(iter).returning_void(); +} diff --git a/nac3core/src/codegen/object/ndarray/mod.rs b/nac3core/src/codegen/object/ndarray/mod.rs index b2070b02..b270d714 100644 --- a/nac3core/src/codegen/object/ndarray/mod.rs +++ b/nac3core/src/codegen/object/ndarray/mod.rs @@ -1,3 +1,5 @@ +pub mod nditer; + use inkwell::{context::Context, types::BasicType, values::PointerValue, AddressSpace}; use crate::{ diff --git a/nac3core/src/codegen/object/ndarray/nditer.rs b/nac3core/src/codegen/object/ndarray/nditer.rs new file mode 100644 index 00000000..29613949 --- /dev/null +++ b/nac3core/src/codegen/object/ndarray/nditer.rs @@ -0,0 +1,168 @@ +use inkwell::{types::BasicType, values::PointerValue, AddressSpace}; + +use crate::codegen::{ + irrt::{call_nac3_nditer_has_next, call_nac3_nditer_initialize, call_nac3_nditer_next}, + model::*, + object::any::AnyObject, + stmt::{gen_for_callback, BreakContinueHooks}, + CodeGenContext, CodeGenerator, +}; + +use super::NDArrayObject; + +/// Fields of [`NDIter`] +pub struct NDIterFields<'ctx, F: FieldTraversal<'ctx>> { + pub ndims: F::Out>, + pub shape: F::Out>>, + pub strides: F::Out>>, + + pub indices: F::Out>>, + pub nth: F::Out>, + pub element: F::Out>>, + + pub size: F::Out>, +} + +/// An IRRT helper structure used to iterate through an ndarray. +#[derive(Debug, Clone, Copy, Default)] +pub struct NDIter; + +impl<'ctx> StructKind<'ctx> for NDIter { + type Fields> = NDIterFields<'ctx, F>; + + fn traverse_fields>(&self, traversal: &mut F) -> Self::Fields { + Self::Fields { + ndims: traversal.add_auto("ndims"), + shape: traversal.add_auto("shape"), + strides: traversal.add_auto("strides"), + + indices: traversal.add_auto("indices"), + nth: traversal.add_auto("nth"), + element: traversal.add_auto("element"), + + size: traversal.add_auto("size"), + } + } +} + +/// A helper structure containing extra details of an [`NDIter`]. +#[derive(Debug, Clone)] +pub struct NDIterHandle<'ctx> { + instance: Instance<'ctx, Ptr>>, + /// The ndarray this [`NDIter`] to iterating over. + ndarray: NDArrayObject<'ctx>, + /// The current indices of [`NDIter`]. + indices: Instance<'ctx, Ptr>>, +} + +impl<'ctx> NDIterHandle<'ctx> { + /// Allocate an [`NDIter`] that iterates through an ndarray. + pub fn new( + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ndarray: NDArrayObject<'ctx>, + ) -> Self { + let nditer = Struct(NDIter).alloca(generator, ctx); + let ndims = ndarray.ndims_llvm(generator, ctx.ctx); + + // The caller has the responsibility to allocate 'indices' for `NDIter`. + let indices = Int(SizeT).array_alloca(generator, ctx, ndims.value); + call_nac3_nditer_initialize(generator, ctx, nditer, ndarray.instance, indices); + + NDIterHandle { ndarray, instance: nditer, indices } + } + + #[must_use] + pub fn has_next( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ) -> Instance<'ctx, Int> { + call_nac3_nditer_has_next(generator, ctx, self.instance) + } + + pub fn next( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ) { + call_nac3_nditer_next(generator, ctx, self.instance); + } + + /// Get pointer to the current element. + #[must_use] + pub fn get_pointer( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ) -> PointerValue<'ctx> { + let elem_ty = ctx.get_llvm_type(generator, self.ndarray.dtype); + + let p = self.instance.get(generator, ctx, |f| f.element); + ctx.builder + .build_pointer_cast(p.value, elem_ty.ptr_type(AddressSpace::default()), "element") + .unwrap() + } + + /// Get the value of the current element. + #[must_use] + pub fn get_scalar( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ) -> AnyObject<'ctx> { + let p = self.get_pointer(generator, ctx); + let value = ctx.builder.build_load(p, "value").unwrap(); + AnyObject { ty: self.ndarray.dtype, value } + } + + /// Get the index of the current element. + #[must_use] + pub fn get_index( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, '_>, + ) -> Instance<'ctx, Int> { + self.instance.get(generator, ctx, |f| f.nth) + } + + /// Get the indices of the current element. + #[must_use] + pub fn get_indices(&self) -> Instance<'ctx, Ptr>> { + self.indices + } +} + +impl<'ctx> NDArrayObject<'ctx> { + /// Iterate through every element in the ndarray. + /// + /// `body` also access to [`BreakContinueHooks`] to short-circuit. + pub fn foreach<'a, G, F>( + &self, + generator: &mut G, + ctx: &mut CodeGenContext<'ctx, 'a>, + body: F, + ) -> Result<(), String> + where + G: CodeGenerator + ?Sized, + F: FnOnce( + &mut G, + &mut CodeGenContext<'ctx, 'a>, + BreakContinueHooks<'ctx>, + NDIterHandle<'ctx>, + ) -> Result<(), String>, + { + gen_for_callback( + generator, + ctx, + Some("ndarray_foreach"), + |generator, ctx| Ok(NDIterHandle::new(generator, ctx, *self)), + |generator, ctx, nditer| Ok(nditer.has_next(generator, ctx).value), + |generator, ctx, hooks, nditer| body(generator, ctx, hooks, nditer), + |generator, ctx, nditer| { + nditer.next(generator, ctx); + Ok(()) + }, + ) + } +}