ramtej/nac3
1
0
Fork 0
forked from M-Labs/nac3
Code Pull Requests Activity

artiq: reimplement polymorphic_print for ndarray

Browse Source
This commit is contained in:
lyken 2024-08-22 13:19:39 +08:00 committed by David Mak
parent 2e75d5a730
commit 2a6ee503ba
1 changed files with 34 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

79
nac3artiq/src/codegen.rs
View File

@ -14,9 +14,8 @@ use pyo3::{
use nac3core::{ use nac3core::{
codegen::{ codegen::{
classes::{ListValue, NDArrayValue, RangeValue, UntypedArrayLikeAccessor}, classes::{ListValue, RangeValue, UntypedArrayLikeAccessor},
expr::{destructure_range, gen_call}, expr::{destructure_range, gen_call},
irrt::call_ndarray_calc_size,
llvm_intrinsics::{call_int_smax, call_stackrestore, call_stacksave}, llvm_intrinsics::{call_int_smax, call_stackrestore, call_stacksave},
model::*, model::*,
object::{any::AnyObject, ndarray::NDArrayObject}, object::{any::AnyObject, ndarray::NDArrayObject},
@ -1311,56 +1310,46 @@ fn polymorphic_print<'ctx>(
} }
TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => { TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty);
fmt.push_str("array(["); fmt.push_str("array([");
flush(ctx, generator, &mut fmt, &mut args); flush(ctx, generator, &mut fmt, &mut args);
let val = NDArrayValue::from_ptr_val(value.into_pointer_value(), llvm_usize, None); let ndarray = AnyObject { ty, value };
let len = call_ndarray_calc_size(generator, ctx, &val.dim_sizes(), (None, None)); let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let last =
ctx.builder.build_int_sub(len, llvm_usize.const_int(1, false), "").unwrap();
gen_for_callback_incrementing( let num_0 = Int(SizeT).const_0(generator, ctx.ctx);
generator,
ctx,
None,
llvm_usize.const_zero(),
(len, false),
|generator, ctx, _, i| {
let elem = unsafe { val.data().get_unchecked(ctx, generator, &i, None) };
polymorphic_print( // Print `ndarray` as a flat list delimited by interspersed with ", \0"
ctx, ndarray.foreach(generator, ctx, |generator, ctx, _, hdl| {
generator, let i = hdl.get_index(generator, ctx);
&[(elem_ty, elem.into())], let scalar = hdl.get_scalar(generator, ctx);
"",
None,
true,
as_rtio,
)?;
gen_if_callback( // if (i != 0) { puts(", "); }
generator, gen_if_callback(
ctx, generator,
|_, ctx| { ctx,
Ok(ctx |_, ctx| {
.builder let not_first = i.compare(ctx, IntPredicate::NE, num_0);
.build_int_compare(IntPredicate::ULT, i, last, "") Ok(not_first.value)
.unwrap()) },
}, |generator, ctx| {
|generator, ctx| { printf(ctx, generator, ", \0".into(), Vec::default());
printf(ctx, generator, ", \0".into(), Vec::default()); Ok(())
},
|_, _| Ok(()),
)?;
Ok(()) // Print element
}, polymorphic_print(
|_, _| Ok(()), ctx,
)?; generator,
&[(scalar.ty, scalar.value.into())],
Ok(()) "",
}, None,
llvm_usize.const_int(1, false), true,
)?; as_rtio,
)?;
Ok(())
})?;
fmt.push_str(")]"); fmt.push_str(")]");
flush(ctx, generator, &mut fmt, &mut args); flush(ctx, generator, &mut fmt, &mut args);