M-Labs/nac3
M-Labs
/
nac3
8
0
Fork 10
Code Issues 98 Pull Requests 18 Releases Wiki Activity

[core] codegen: Implement NDArray functions from a0a1f35b

This commit is contained in:
David Mak 2024-11-27 14:45:13 +08:00
parent 624e943cd6
commit 57da0f67d1
5 changed files with 205 additions and 88 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nac3artiq/src/codegen.rs
View File

@ -732,7 +732,7 @@ fn format_rpc_ret<'ctx>(
);
}
ndarray.create_data(generator, ctx, llvm_elem_ty, num_elements);
unsafe { ndarray.create_data(generator, ctx) };
let ndarray_data = ndarray.data().base_ptr(ctx, generator);
let ndarray_data_i8 =

39
nac3core/src/codegen/expr.rs
View File

@ -2560,7 +2560,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
ty: Type,
ndims: Type,
ndims_ty: Type,
v: NDArrayValue<'ctx>,
slice: &Expr<Option<Type>>,
) -> Result<Option<ValueEnum<'ctx>>, String> {
@ -2568,7 +2568,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
let llvm_i32 = ctx.ctx.i32_type();
let llvm_usize = generator.get_size_type(ctx.ctx);
let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) else {
let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims_ty) else {
codegen_unreachable!(ctx)
};
@ -2601,10 +2601,6 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
_ => 1,
};
let ndarray_ndims_ty = ctx.unifier.get_fresh_literal(
ndims.iter().map(|v| SymbolValue::U64(v - subscripted_dims)).collect(),
None,
);
let llvm_ndarray_data_t = ctx.get_llvm_type(generator, ty).as_basic_type_enum();
let sizeof_elem = llvm_ndarray_data_t.size_of().unwrap();
@ -2798,27 +2794,18 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
// Accessing an element from a multi-dimensional `ndarray`
let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
let num_dims = extract_ndims(&ctx.unifier, ndims_ty) - 1;
// Create a new array, remove the top dimension from the dimension-size-list, and copy the
// elements over
let ndarray = NDArrayType::new(
generator,
ctx.ctx,
llvm_ndarray_data_t,
Some(extract_ndims(&ctx.unifier, ndarray_ndims_ty)),
)
.alloca(generator, ctx, None);
let num_dims = v.load_ndims(ctx);
ndarray.store_ndims(
ctx,
generator,
ctx.builder
.build_int_sub(num_dims, llvm_usize.const_int(1, false), "")
.unwrap(),
);
let ndarray_num_dims = ndarray.load_ndims(ctx);
ndarray.create_shape(ctx, llvm_usize, ndarray_num_dims);
let ndarray =
NDArrayType::new(generator, ctx.ctx, llvm_ndarray_data_t, Some(num_dims))
.construct_uninitialized(
generator,
ctx,
llvm_usize.const_int(num_dims, false),
None,
);
let ndarray_num_dims = ctx
.builder
@ -2857,7 +2844,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
.builder
.build_int_z_extend_or_bit_cast(ndarray_num_elems, sizeof_elem.get_type(), "")
.unwrap();
ndarray.create_data(generator, ctx, llvm_ndarray_data_t, ndarray_num_elems);
unsafe { ndarray.create_data(generator, ctx) };
let v_data_src_ptr = v.data().ptr_offset(ctx, generator, &index_addr, None);
call_memcpy_generic(

61
nac3core/src/codegen/numpy.rs
View File

@ -40,25 +40,6 @@ use crate::{
},
};
/// Creates an uninitialized `NDArray` instance.
fn create_ndarray_uninitialized<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
elem_ty: Type,
) -> Result<NDArrayValue<'ctx>, String> {
let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_ndarray_t = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None)
.as_base_type()
.get_element_type()
.into_struct_type();
let ndarray = generator.gen_var_alloc(ctx, llvm_ndarray_t.into(), None)?;
Ok(NDArrayValue::from_pointer_value(ndarray, llvm_elem_ty, None, llvm_usize, None))
}
/// Creates an `NDArray` instance from a dynamic shape.
///
/// * `elem_ty` - The element type of the `NDArray`.
@ -84,6 +65,7 @@ where
) -> Result<IntValue<'ctx>, String>,
{
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
// Assert that all dimensions are non-negative
let shape_len = shape_len_fn(generator, ctx, shape)?;
@ -123,13 +105,10 @@ where
llvm_usize.const_int(1, false),
)?;
let ndarray = create_ndarray_uninitialized(generator, ctx, elem_ty)?;
let num_dims = shape_len_fn(generator, ctx, shape)?;
ndarray.store_ndims(ctx, generator, num_dims);
let ndarray_num_dims = ndarray.load_ndims(ctx);
ndarray.create_shape(ctx, llvm_usize, ndarray_num_dims);
let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None)
.construct_uninitialized(generator, ctx, num_dims, None);
// Copy the dimension sizes from shape to ndarray.dims
let shape_len = shape_len_fn(generator, ctx, shape)?;
@ -154,7 +133,7 @@ where
llvm_usize.const_int(1, false),
)?;
let ndarray = ndarray_init_data(generator, ctx, elem_ty, ndarray);
unsafe { ndarray.create_data(generator, ctx) };
Ok(ndarray)
}
@ -194,32 +173,11 @@ pub fn create_ndarray_const_shape<'ctx, G: CodeGenerator + ?Sized>(
let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_dtype, Some(shape.len() as u64))
.construct_dyn_shape(generator, ctx, shape, None);
let ndarray = ndarray_init_data(generator, ctx, elem_ty, ndarray);
unsafe { ndarray.create_data(generator, ctx) };
Ok(ndarray)
}
/// Initializes the `data` field of [`NDArrayValue`] based on the `ndims` and `dim_sz` fields.
fn ndarray_init_data<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
elem_ty: Type,
ndarray: NDArrayValue<'ctx>,
) -> NDArrayValue<'ctx> {
let llvm_ndarray_data_t = ctx.get_llvm_type(generator, elem_ty).as_basic_type_enum();
assert!(llvm_ndarray_data_t.is_sized());
let ndarray_num_elems = call_ndarray_calc_size(
generator,
ctx,
&ndarray.shape().as_slice_value(ctx, generator),
(None, None),
);
ndarray.create_data(generator, ctx, llvm_ndarray_data_t, ndarray_num_elems);
ndarray
}
fn ndarray_zero_value<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
@ -1263,6 +1221,7 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>(
) -> Result<NDArrayValue<'ctx>, String> {
let llvm_i32 = ctx.ctx.i32_type();
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
let ndarray = if slices.is_empty() {
create_ndarray_dyn_shape(
@ -1276,8 +1235,8 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>(
},
)?
} else {
let ndarray = create_ndarray_uninitialized(generator, ctx, elem_ty)?;
ndarray.store_ndims(ctx, generator, this.load_ndims(ctx));
let ndarray = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty, None)
.construct_uninitialized(generator, ctx, this.load_ndims(ctx), None);
let ndims = this.load_ndims(ctx);
ndarray.create_shape(ctx, llvm_usize, ndims);
@ -1340,7 +1299,9 @@ pub fn ndarray_sliced_copy<'ctx, G: CodeGenerator + ?Sized>(
)
.unwrap();
ndarray_init_data(generator, ctx, elem_ty, ndarray)
unsafe { ndarray.create_data(generator, ctx) };
ndarray
};
ndarray_sliced_copyto_impl(

2
nac3core/src/codegen/types/ndarray/mod.rs
View File

@ -214,7 +214,7 @@ impl<'ctx> NDArrayType<'ctx> {
let itemsize = ctx
.builder
.build_int_z_extend_or_bit_cast(self.dtype.size_of().unwrap(), self.llvm_usize, "")
.build_int_truncate_or_bit_cast(self.dtype.size_of().unwrap(), self.llvm_usize, "")
.unwrap();
ndarray.store_itemsize(ctx, generator, itemsize);

189
nac3core/src/codegen/values/ndarray/mod.rs
View File

@ -10,7 +10,7 @@ use super::{
};
use crate::codegen::{
irrt,
llvm_intrinsics::call_int_umin,
llvm_intrinsics::{call_int_umin, call_memcpy_generic_array},
stmt::gen_for_callback_incrementing,
type_aligned_alloca,
types::{structure::StructField, NDArrayType},
@ -186,21 +186,23 @@ impl<'ctx> NDArrayValue<'ctx> {
/// Convenience method for creating a new array storing data elements with the given element
/// type `elem_ty` and `size`.
pub fn create_data<G: CodeGenerator + ?Sized>(
///
/// The data buffer will be allocated on the stack, and is considered to be owned by this ndarray instance.
///
/// # Safety
///
/// The caller must ensure that `shape` and `itemsize` of this ndarray instance is initialized.
pub unsafe fn create_data<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
elem_ty: BasicTypeEnum<'ctx>,
size: IntValue<'ctx>,
) {
let itemsize = ctx
.builder
.build_int_z_extend_or_bit_cast(elem_ty.size_of().unwrap(), size.get_type(), "")
.unwrap();
let nbytes = ctx.builder.build_int_mul(size, itemsize, "").unwrap();
let nbytes = self.nbytes(generator, ctx);
let data = type_aligned_alloca(generator, ctx, elem_ty, nbytes, None);
let data = type_aligned_alloca(generator, ctx, self.dtype, nbytes, None);
self.store_data(ctx, data);
self.set_strides_contiguous(generator, ctx);
}
/// Returns a proxy object to the field storing the data of this `NDArray`.
@ -208,6 +210,173 @@ impl<'ctx> NDArrayValue<'ctx> {
pub fn data(&self) -> NDArrayDataProxy<'ctx, '_> {
NDArrayDataProxy(self)
}
/// Copy shape dimensions from an array.
pub fn copy_shape_from_array<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
shape: PointerValue<'ctx>,
) {
let num_items = self.load_ndims(ctx);
call_memcpy_generic_array(
ctx,
self.shape().base_ptr(ctx, generator),
shape,
num_items,
ctx.ctx.bool_type().const_zero(),
);
}
/// Copy shape dimensions from an ndarray.
/// Panics if `ndims` mismatches.
pub fn copy_shape_from_ndarray<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
src_ndarray: NDArrayValue<'ctx>,
) {
if self.ndims.is_some() && src_ndarray.ndims.is_some() {
assert_eq!(self.ndims, src_ndarray.ndims);
} else {
let self_ndims = self.load_ndims(ctx);
let src_ndims = src_ndarray.load_ndims(ctx);
ctx.make_assert(
generator,
ctx.builder.build_int_compare(
IntPredicate::EQ,
self_ndims,
src_ndims,
""
).unwrap(),
"0:AssertionError",
"NDArrayValue::copy_shape_from_ndarray: Expected self.ndims ({0}) == src_ndarray.ndims ({1})",
[Some(self_ndims), Some(src_ndims), None],
ctx.current_loc
);
}
let src_shape = src_ndarray.shape().base_ptr(ctx, generator);
self.copy_shape_from_array(generator, ctx, src_shape);
}
/// Copy strides dimensions from an array.
pub fn copy_strides_from_array<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
strides: PointerValue<'ctx>,
) {
let num_items = self.load_ndims(ctx);
call_memcpy_generic_array(
ctx,
self.strides().base_ptr(ctx, generator),
strides,
num_items,
ctx.ctx.bool_type().const_zero(),
);
}
/// Copy strides dimensions from an ndarray.
/// Panics if `ndims` mismatches.
pub fn copy_strides_from_ndarray<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
src_ndarray: NDArrayValue<'ctx>,
) {
if self.ndims.is_some() && src_ndarray.ndims.is_some() {
assert_eq!(self.ndims, src_ndarray.ndims);
} else {
let self_ndims = self.load_ndims(ctx);
let src_ndims = src_ndarray.load_ndims(ctx);
ctx.make_assert(
generator,
ctx.builder.build_int_compare(
IntPredicate::EQ,
self_ndims,
src_ndims,
""
).unwrap(),
"0:AssertionError",
"NDArrayValue::copy_shape_from_ndarray: Expected self.ndims ({0}) == src_ndarray.ndims ({1})",
[Some(self_ndims), Some(src_ndims), None],
ctx.current_loc
);
}
let src_strides = src_ndarray.strides().base_ptr(ctx, generator);
self.copy_strides_from_array(generator, ctx, src_strides);
}
/// Get the `np.size()` of this ndarray.
pub fn size<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
) -> IntValue<'ctx> {
irrt::ndarray::call_nac3_ndarray_size(generator, ctx, *self)
}
/// Get the `ndarray.nbytes` of this ndarray.
pub fn nbytes<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
) -> IntValue<'ctx> {
irrt::ndarray::call_nac3_ndarray_nbytes(generator, ctx, *self)
}
/// Get the `len()` of this ndarray.
pub fn len<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
) -> IntValue<'ctx> {
irrt::ndarray::call_nac3_ndarray_len(generator, ctx, *self)
}
/// Check if this ndarray is C-contiguous.
///
/// See NumPy's `flags["C_CONTIGUOUS"]`: <https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flags.html#numpy.ndarray.flags>
pub fn is_c_contiguous<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
) -> IntValue<'ctx> {
irrt::ndarray::call_nac3_ndarray_is_c_contiguous(generator, ctx, *self)
}
/// Call [`call_nac3_ndarray_set_strides_by_shape`] on this ndarray to update `strides`.
///
/// Update the ndarray's strides to make the ndarray contiguous.
pub fn set_strides_contiguous<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
) {
irrt::ndarray::call_nac3_ndarray_set_strides_by_shape(generator, ctx, *self);
}
/// Copy data from another ndarray.
///
/// This ndarray and `src` is that their `np.size()` should be the same. Their shapes
/// do not matter. The copying order is determined by how their flattened views look.
///
/// Panics if the `dtype`s of ndarrays are different.
pub fn copy_data_from<G: CodeGenerator + ?Sized>(
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
src: NDArrayValue<'ctx>,
) {
assert_eq!(self.dtype, src.dtype, "self and src dtype should match");
irrt::ndarray::call_nac3_ndarray_copy_data(generator, ctx, src, *self);
}
}
impl<'ctx> ProxyValue<'ctx> for NDArrayValue<'ctx> {