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

core/ndstrides: implement ndarray.fill() and .copy()

This commit is contained in:
lyken 2024-08-20 15:18:07 +08:00 committed by David Mak
parent 532fccf1fb
commit f57df4c6a7
2 changed files with 11 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1958,20 +1958,14 @@ pub fn gen_ndarray_copy<'ctx>(
assert!(obj.is_some());
assert!(args.is_empty());
let llvm_usize = generator.get_size_type(context.ctx);
let this_ty = obj.as_ref().unwrap().0;
let (this_elem_ty, _) = unpack_ndarray_var_tys(&mut context.unifier, this_ty);
let this_arg =
obj.as_ref().unwrap().1.clone().to_basic_value_enum(context, generator, this_ty)?;
ndarray_copy_impl(
generator,
context,
this_elem_ty,
NDArrayValue::from_ptr_val(this_arg.into_pointer_value(), llvm_usize, None),
)
.map(NDArrayValue::into)
let this = AnyObject { value: this_arg, ty: this_ty };
let this = NDArrayObject::from_object(generator, context, this);
let ndarray = this.make_copy(generator, context);
Ok(ndarray.instance.value)
}
/// Generates LLVM IR for `ndarray.fill`.
@ -1985,48 +1979,15 @@ pub fn gen_ndarray_fill<'ctx>(
assert!(obj.is_some());
assert_eq!(args.len(), 1);
let llvm_usize = generator.get_size_type(context.ctx);
let this_ty = obj.as_ref().unwrap().0;
let this_arg = obj
.as_ref()
.unwrap()
.1
.clone()
.to_basic_value_enum(context, generator, this_ty)?
.into_pointer_value();
let this_arg =
obj.as_ref().unwrap().1.clone().to_basic_value_enum(context, generator, this_ty)?;
let value_ty = fun.0.args[0].ty;
let value_arg = args[0].1.clone().to_basic_value_enum(context, generator, value_ty)?;
ndarray_fill_flattened(
generator,
context,
NDArrayValue::from_ptr_val(this_arg, llvm_usize, None),
|generator, ctx, _| {
let value = if value_arg.is_pointer_value() {
let llvm_i1 = ctx.ctx.bool_type();
let copy = generator.gen_var_alloc(ctx, value_arg.get_type(), None)?;
call_memcpy_generic(
ctx,
copy,
value_arg.into_pointer_value(),
value_arg.get_type().size_of().map(Into::into).unwrap(),
llvm_i1.const_zero(),
);
copy.into()
} else if value_arg.is_int_value() || value_arg.is_float_value() {
value_arg
} else {
codegen_unreachable!(ctx)
};
Ok(value)
},
)?;
let this = AnyObject { value: this_arg, ty: this_ty };
let this = NDArrayObject::from_object(generator, context, this);
this.fill(generator, context, value_arg);
Ok(())
}

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

@ -417,6 +417,8 @@ impl<'ctx> NDArrayObject<'ctx> {
ctx: &mut CodeGenContext<'ctx, '_>,
value: BasicValueEnum<'ctx>,
) {
// TODO: It is possible to optimize this by exploiting contiguous strides with memset.
// Probably best to implement in IRRT.
self.foreach(generator, ctx, |generator, ctx, _hooks, nditer| {
let p = nditer.get_pointer(generator, ctx);
ctx.builder.build_store(p, value).unwrap();