1
0
forked from M-Labs/nac3

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

View File

@ -1,12 +1,12 @@
use inkwell::{ use inkwell::{
IntPredicate, IntPredicate,
types::{AnyTypeEnum, BasicTypeEnum, IntType, PointerType}, types::{AnyTypeEnum, BasicTypeEnum, IntType, PointerType},
values::{ArrayValue, BasicValueEnum, IntValue, PointerValue}, values::{BasicValueEnum, IntValue, PointerValue},
}; };
use crate::codegen::{ use crate::codegen::{
CodeGenContext, CodeGenContext,
CodeGenerator, 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, llvm_intrinsics::call_int_umin,
stmt::gen_for_callback_incrementing, 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> UntypedArrayLikeAccessor<'ctx, IntValue<'ctx>> for NDArrayDataProxy<'ctx, '_> {}
impl<'ctx> UntypedArrayLikeMutator<'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, '_> { impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index> for NDArrayDataProxy<'ctx, '_> {
unsafe fn ptr_offset_unchecked<G: CodeGenerator + ?Sized>( unsafe fn ptr_offset_unchecked<G: CodeGenerator + ?Sized>(
&self, &self,

View File

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

View File

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