Aadityavardhan
/
nac3_sca
forked from M-Labs/nac3
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

core: Use more typed slices in APIs

This commit is contained in:
David Mak 2024-03-22 16:57:06 +08:00
parent d5f4817134
commit 26a01b14d5
3 changed files with 26 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

96
nac3core/src/codegen/classes.rs
View File

@ -1,12 +1,12 @@
use inkwell::{
IntPredicate,
types::{AnyTypeEnum, BasicTypeEnum, IntType, PointerType},
values::{ArrayValue, BasicValueEnum, IntValue, PointerValue},
values::{BasicValueEnum, IntValue, PointerValue},
};
use crate::codegen::{
CodeGenContext,
CodeGenerator,
irrt::{call_ndarray_calc_size, call_ndarray_flatten_index, call_ndarray_flatten_index_const},
irrt::{call_ndarray_calc_size, call_ndarray_flatten_index},
llvm_intrinsics::call_int_umin,
stmt::gen_for_callback_incrementing,
};
@ -1162,98 +1162,6 @@ impl<'ctx> ArrayLikeIndexer<'ctx> for NDArrayDataProxy<'ctx, '_> {
impl<'ctx> UntypedArrayLikeAccessor<'ctx, IntValue<'ctx>> for NDArrayDataProxy<'ctx, '_> {}
impl<'ctx> UntypedArrayLikeMutator<'ctx, IntValue<'ctx>> for NDArrayDataProxy<'ctx, '_> {}
impl<'ctx> ArrayLikeIndexer<'ctx, ArrayValue<'ctx>> for NDArrayDataProxy<'ctx, '_> {
unsafe fn ptr_offset_unchecked<G: CodeGenerator + ?Sized>(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
indices: ArrayValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
let index = call_ndarray_flatten_index_const(
generator,
ctx,
*self.0,
indices,
);
unsafe {
ctx.builder.build_in_bounds_gep(
self.base_ptr(ctx, generator),
&[index],
name.unwrap_or_default(),
)
}.unwrap()
}
fn ptr_offset<G: CodeGenerator + ?Sized>(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
indices: ArrayValue<'ctx>,
name: Option<&str>,
) -> PointerValue<'ctx> {
let llvm_usize = generator.get_size_type(ctx.ctx);
let indices_elem_ty = indices.get_type().get_element_type();
let Ok(indices_elem_ty) = IntType::try_from(indices_elem_ty) else {
panic!("Expected [int32] but got [{indices_elem_ty}]")
};
assert_eq!(indices_elem_ty.get_bit_width(), 32, "Expected [int32] but got [{indices_elem_ty}]");
let nidx_leq_ndims = ctx.builder.build_int_compare(
IntPredicate::SLE,
llvm_usize.const_int(indices.get_type().len() as u64, false),
self.0.load_ndims(ctx),
""
).unwrap();
ctx.make_assert(
generator,
nidx_leq_ndims,
"0:IndexError",
"invalid index to scalar variable",
[None, None, None],
ctx.current_loc,
);
for idx in 0..indices.get_type().len() {
let i = llvm_usize.const_int(idx as u64, false);
let dim_idx = ctx.builder
.build_extract_value(indices, idx, "")
.map(BasicValueEnum::into_int_value)
.map(|v| ctx.builder.build_int_z_extend_or_bit_cast(v, llvm_usize, "").unwrap())
.unwrap();
let dim_sz = unsafe {
self.0.dim_sizes().get_typed_unchecked(ctx, generator, i, None)
};
let dim_lt = ctx.builder.build_int_compare(
IntPredicate::SLT,
dim_idx,
dim_sz,
""
).unwrap();
ctx.make_assert(
generator,
dim_lt,
"0:IndexError",
"index {0} is out of bounds for axis 0 with size {1}",
[Some(dim_idx), Some(dim_sz), None],
ctx.current_loc,
);
}
unsafe {
self.ptr_offset_unchecked(ctx, generator, indices, name)
}
}
}
impl<'ctx> UntypedArrayLikeAccessor<'ctx, ArrayValue<'ctx>> for NDArrayDataProxy<'ctx, '_> {}
impl<'ctx> UntypedArrayLikeMutator<'ctx, ArrayValue<'ctx>> for NDArrayDataProxy<'ctx, '_> {}
impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index> for NDArrayDataProxy<'ctx, '_> {
unsafe fn ptr_offset_unchecked<G: CodeGenerator + ?Sized>(
&self,

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

@ -1,7 +1,15 @@
use crate::typecheck::typedef::Type;
use super::{
classes::{ArrayLikeIndexer, ArrayLikeValue, ListValue, NDArrayValue, UntypedArrayLikeMutator},
classes::{
ArrayLikeIndexer,
ArrayLikeValue,
ArraySliceValue,
ListValue,
NDArrayValue,
TypedArrayLikeAdapter,
UntypedArrayLikeMutator,
},
CodeGenContext,
CodeGenerator,
};
@ -11,7 +19,7 @@ use inkwell::{
memory_buffer::MemoryBuffer,
module::Module,
types::{BasicTypeEnum, IntType},
values::{ArrayValue, BasicValueEnum, CallSiteValue, FloatValue, IntValue, PointerValue},
values::{ArrayValue, BasicValueEnum, CallSiteValue, FloatValue, IntValue},
AddressSpace, IntPredicate,
};
use itertools::Either;
@ -630,7 +638,7 @@ pub fn call_ndarray_calc_nd_indices<'ctx, G: CodeGenerator + ?Sized>(
ctx: &mut CodeGenContext<'ctx, '_>,
index: IntValue<'ctx>,
ndarray: NDArrayValue<'ctx>,
) -> PointerValue<'ctx> {
) -> TypedArrayLikeAdapter<'ctx, IntValue<'ctx>> {
let llvm_void = ctx.ctx.void_type();
let llvm_i32 = ctx.ctx.i32_type();
let llvm_usize = generator.get_size_type(ctx.ctx);
@ -678,7 +686,11 @@ pub fn call_ndarray_calc_nd_indices<'ctx, G: CodeGenerator + ?Sized>(
)
.unwrap();
indices
TypedArrayLikeAdapter::from(
ArraySliceValue::from_ptr_val(indices, ndarray_num_dims, None),
Box::new(|_, v| v.into_int_value()),
Box::new(|_, v| v.into()),
)
}
fn call_ndarray_flatten_index_impl<'ctx, G, Indices>(

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

@ -12,6 +12,7 @@ use crate::{
ListValue,
NDArrayValue,
TypedArrayLikeAccessor,
TypedArrayLikeAdapter,
UntypedArrayLikeAccessor,
},
CodeGenContext,
@ -324,7 +325,7 @@ fn ndarray_fill_indexed<'ctx, G, ValueFn>(
) -> Result<(), String>
where
G: CodeGenerator + ?Sized,
ValueFn: Fn(&mut G, &mut CodeGenContext<'ctx, '_>, PointerValue<'ctx>) -> Result<BasicValueEnum<'ctx>, String>,
ValueFn: Fn(&mut G, &mut CodeGenContext<'ctx, '_>, TypedArrayLikeAdapter<'ctx, IntValue<'ctx>>) -> Result<BasicValueEnum<'ctx>, String>,
{
ndarray_fill_flattened(
generator,
@ -499,16 +500,12 @@ fn call_ndarray_eye_impl<'ctx, G: CodeGenerator + ?Sized>(
ctx,
ndarray,
|generator, ctx, indices| {
let row = ctx.build_gep_and_load(
indices,
&[llvm_usize.const_int(0, false)],
None,
).into_int_value();
let col = ctx.build_gep_and_load(
indices,
&[llvm_usize.const_int(1, false)],
None,
).into_int_value();
let (row, col) = unsafe {
(
indices.get_typed_unchecked(ctx, generator, llvm_usize.const_zero(), None),
indices.get_typed_unchecked(ctx, generator, llvm_usize.const_int(1, false), None),
)
};
let col_with_offset = ctx.builder
.build_int_add(