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

WIP

This commit is contained in:
David Mak 2024-08-13 17:37:48 +08:00
parent 7a8f201061
commit 0c1153e40e
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
nac3artiq/src/codegen.rs
View File

@ -549,6 +549,8 @@ fn rpc_codegen_callback_fn<'ctx>(
ctx.builder.build_store(arg_slot, *arg).unwrap(); ctx.builder.build_store(arg_slot, *arg).unwrap();
let arg_slot = if matches!(&*ctx.unifier.get_ty_immutable(*arg_ty), TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id()) let arg_slot = if matches!(&*ctx.unifier.get_ty_immutable(*arg_ty), TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id())
{ {
println!("Codegen for ndarray RPC");
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, *arg_ty); let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, *arg_ty);
let llvm_arg_ty = let llvm_arg_ty =
NDArrayType::new(generator, ctx.ctx, ctx.get_llvm_type(generator, elem_ty)); NDArrayType::new(generator, ctx.ctx, ctx.get_llvm_type(generator, elem_ty));
@ -557,6 +559,8 @@ fn rpc_codegen_callback_fn<'ctx>(
let llvm_usize_sizeof = llvm_arg_ty.size_type().size_of(); let llvm_usize_sizeof = llvm_arg_ty.size_type().size_of();
let llvm_elem_sizeof = llvm_arg_ty.element_type().size_of().unwrap(); let llvm_elem_sizeof = llvm_arg_ty.element_type().size_of().unwrap();
println!("Calc required buffer size");
let dims_buf_sz = let dims_buf_sz =
ctx.builder.build_int_mul(llvm_arg.load_ndims(ctx), llvm_usize_sizeof, "").unwrap(); ctx.builder.build_int_mul(llvm_arg.load_ndims(ctx), llvm_usize_sizeof, "").unwrap();
let data_buf_sz = ctx let data_buf_sz = ctx
@ -570,9 +574,13 @@ fn rpc_codegen_callback_fn<'ctx>(
let buffer_size = ctx.builder.build_int_add(dims_buf_sz, data_buf_sz, "").unwrap(); let buffer_size = ctx.builder.build_int_add(dims_buf_sz, data_buf_sz, "").unwrap();
println!("Alloc buffer");
let buffer = let buffer =
generator.gen_array_var_alloc(ctx, int8.into(), buffer_size, None).unwrap(); generator.gen_array_var_alloc(ctx, int8.into(), buffer_size, None).unwrap();
println!("memcpy dim_sizes");
llvm_intrinsics::call_memcpy_generic( llvm_intrinsics::call_memcpy_generic(
ctx, ctx,
buffer.base_ptr(ctx, generator), buffer.base_ptr(ctx, generator),
@ -581,6 +589,8 @@ fn rpc_codegen_callback_fn<'ctx>(
int1.const_zero(), int1.const_zero(),
); );
println!("memcpy data");
let pbuffer_data_begin = let pbuffer_data_begin =
unsafe { buffer.ptr_offset_unchecked(ctx, generator, &dims_buf_sz, None) }; unsafe { buffer.ptr_offset_unchecked(ctx, generator, &dims_buf_sz, None) };
llvm_intrinsics::call_memcpy_generic( llvm_intrinsics::call_memcpy_generic(
@ -591,6 +601,8 @@ fn rpc_codegen_callback_fn<'ctx>(
int1.const_zero(), int1.const_zero(),
); );
println!("done");
buffer.base_ptr(ctx, generator) buffer.base_ptr(ctx, generator)
} else { } else {
ctx.builder ctx.builder