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

WIP - [artiq] Revert to unified

This commit is contained in:
David Mak 2024-08-15 04:03:26 +08:00
parent 4522077644
commit b3d86f7187
1 changed files with 151 additions and 151 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

302
nac3artiq/src/codegen.rs
View File

@ -568,157 +568,157 @@ fn format_rpc_ret<'ctx>(
let llvm_ret_ty = ctx.get_llvm_abi_type(generator, ret_ty); let llvm_ret_ty = ctx.get_llvm_abi_type(generator, ret_ty);
let result = match &*ctx.unifier.get_ty_immutable(ret_ty) { let result = match &*ctx.unifier.get_ty_immutable(ret_ty) {
TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => { // TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
let llvm_i1 = ctx.ctx.bool_type(); // let llvm_i1 = ctx.ctx.bool_type();
let llvm_usize = generator.get_size_type(ctx.ctx); // let llvm_usize = generator.get_size_type(ctx.ctx);
//
let slot = ctx.builder.build_alloca(llvm_ret_ty, "rpc.ret.slot").unwrap(); // let slot = ctx.builder.build_alloca(llvm_ret_ty, "rpc.ret.slot").unwrap();
let (elem_ty, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ret_ty); // let (elem_ty, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ret_ty);
let llvm_ret_ty = // let llvm_ret_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));
let ndarray = llvm_ret_ty.new_value(generator, ctx, None); // let ndarray = llvm_ret_ty.new_value(generator, ctx, None);
ctx.builder.build_store(slot, ndarray.as_base_value()).unwrap(); // ctx.builder.build_store(slot, ndarray.as_base_value()).unwrap();
//
let ndims = // let ndims =
if let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) { // if let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) {
assert_eq!(values.len(), 1); // assert_eq!(values.len(), 1);
//
u64::try_from(values[0].clone()).unwrap() // u64::try_from(values[0].clone()).unwrap()
} else { // } else {
unreachable!(); // unreachable!();
}; // };
ndarray.store_ndims(ctx, generator, llvm_usize.const_int(ndims, false)); // ndarray.store_ndims(ctx, generator, llvm_usize.const_int(ndims, false));
ndarray.create_dim_sizes(ctx, llvm_usize, ndarray.load_ndims(ctx)); // ndarray.create_dim_sizes(ctx, llvm_usize, ndarray.load_ndims(ctx));
//
let llvm_usize_sizeof = ctx // let llvm_usize_sizeof = ctx
.builder // .builder
.build_int_truncate_or_bit_cast(llvm_ret_ty.size_type().size_of(), llvm_usize, "") // .build_int_truncate_or_bit_cast(llvm_ret_ty.size_type().size_of(), llvm_usize, "")
.unwrap(); // .unwrap();
let llvm_pdata_sizeof = ctx // let llvm_pdata_sizeof = ctx
.builder // .builder
.build_int_truncate_or_bit_cast( // .build_int_truncate_or_bit_cast(
llvm_ret_ty.element_type().size_of().unwrap(), // llvm_ret_ty.element_type().size_of().unwrap(),
llvm_usize, // llvm_usize,
"", // "",
) // )
.unwrap(); // .unwrap();
//
let dims_buf_sz = // let dims_buf_sz =
ctx.builder.build_int_mul(ndarray.load_ndims(ctx), llvm_usize_sizeof, "").unwrap(); // ctx.builder.build_int_mul(ndarray.load_ndims(ctx), llvm_usize_sizeof, "").unwrap();
//
let buffer_size = // let buffer_size =
ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap(); // ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap();
let buffer_size = ctx.builder.build_left_shift( // let buffer_size = ctx.builder.build_left_shift(
ctx.builder.build_int_add( // ctx.builder.build_int_add(
ctx.builder.build_right_shift( // ctx.builder.build_right_shift(
buffer_size, // buffer_size,
llvm_usize.const_int(2, false), // llvm_usize.const_int(2, false),
false, // false,
"", // "",
).unwrap(), // ).unwrap(),
llvm_usize.const_int(1, false), // llvm_usize.const_int(1, false),
"", // "",
).unwrap(), // ).unwrap(),
llvm_usize.const_int(2, false), // llvm_usize.const_int(2, false),
"", // "",
).unwrap(); // ).unwrap();
//
let buffer = // let buffer =
ctx.builder.build_array_alloca(llvm_i8, buffer_size, "rpc.buffer").unwrap(); // ctx.builder.build_array_alloca(llvm_i8, buffer_size, "rpc.buffer").unwrap();
let buffer = ArraySliceValue::from_ptr_val(buffer, buffer_size, Some("rpc.buffer")); // let buffer = ArraySliceValue::from_ptr_val(buffer, buffer_size, Some("rpc.buffer"));
//
let i_addr = ctx.builder.build_alloca(llvm_usize, "i.addr").unwrap(); // let i_addr = ctx.builder.build_alloca(llvm_usize, "i.addr").unwrap();
ctx.builder.build_unconditional_branch(head_bb).unwrap(); // ctx.builder.build_unconditional_branch(head_bb).unwrap();
//
ctx.builder.position_at_end(head_bb); // ctx.builder.position_at_end(head_bb);
let phi = ctx.builder.build_phi(llvm_pi8, "rpc.ptr").unwrap(); // let phi = ctx.builder.build_phi(llvm_pi8, "rpc.ptr").unwrap();
phi.add_incoming(&[(&buffer.base_ptr(ctx, generator), prehead_bb)]); // phi.add_incoming(&[(&buffer.base_ptr(ctx, generator), prehead_bb)]);
let alloc_size = ctx // let alloc_size = ctx
.build_call_or_invoke(rpc_recv, &[phi.as_basic_value()], "rpc.size.next") // .build_call_or_invoke(rpc_recv, &[phi.as_basic_value()], "rpc.size.next")
.map(BasicValueEnum::into_int_value) // .map(BasicValueEnum::into_int_value)
.unwrap(); // .unwrap();
//
// Parse metadata block(s) for ndarrays // // Parse metadata block(s) for ndarrays
gen_if_callback( // gen_if_callback(
generator, // generator,
ctx, // ctx,
|_, ctx| { // |_, ctx| {
let i = ctx // let i = ctx
.builder // .builder
.build_load(i_addr, "") // .build_load(i_addr, "")
.map(BasicValueEnum::into_int_value) // .map(BasicValueEnum::into_int_value)
.unwrap(); // .unwrap();
//
Ok(ctx // Ok(ctx
.builder // .builder
.build_int_compare(IntPredicate::EQ, i, llvm_usize.const_zero(), "") // .build_int_compare(IntPredicate::EQ, i, llvm_usize.const_zero(), "")
.unwrap()) // .unwrap())
}, // },
|generator, ctx| { // |generator, ctx| {
let ppdata = // let ppdata =
generator.gen_var_alloc(ctx, llvm_ret_ty.element_type(), None).unwrap(); // generator.gen_var_alloc(ctx, llvm_ret_ty.element_type(), None).unwrap();
ctx.builder // ctx.builder
.build_store(ppdata, ndarray.data().base_ptr(ctx, generator)) // .build_store(ppdata, ndarray.data().base_ptr(ctx, generator))
.unwrap(); // .unwrap();
call_memcpy_generic( // call_memcpy_generic(
ctx, // ctx,
ppdata, // ppdata,
buffer.base_ptr(ctx, generator), // buffer.base_ptr(ctx, generator),
llvm_pdata_sizeof, // llvm_pdata_sizeof,
llvm_i1.const_zero(), // llvm_i1.const_zero(),
); // );
//
let pbuffer_dims_begin = unsafe { // let pbuffer_dims_begin = unsafe {
buffer.ptr_offset_unchecked(ctx, generator, &llvm_pdata_sizeof, None) // buffer.ptr_offset_unchecked(ctx, generator, &llvm_pdata_sizeof, None)
}; // };
call_memcpy_generic( // call_memcpy_generic(
ctx, // ctx,
ndarray.dim_sizes().base_ptr(ctx, generator), // ndarray.dim_sizes().base_ptr(ctx, generator),
pbuffer_dims_begin, // pbuffer_dims_begin,
dims_buf_sz, // dims_buf_sz,
llvm_i1.const_zero(), // llvm_i1.const_zero(),
); // );
//
Ok(()) // Ok(())
}, // },
|_, _| Ok(()), // |_, _| Ok(()),
) // )
.unwrap(); // .unwrap();
//
let is_done = ctx // let is_done = ctx
.builder // .builder
.build_int_compare(IntPredicate::EQ, llvm_i32.const_zero(), alloc_size, "rpc.done") // .build_int_compare(IntPredicate::EQ, llvm_i32.const_zero(), alloc_size, "rpc.done")
.unwrap(); // .unwrap();
//
ctx.builder.build_conditional_branch(is_done, tail_bb, alloc_bb).unwrap(); // ctx.builder.build_conditional_branch(is_done, tail_bb, alloc_bb).unwrap();
ctx.builder.position_at_end(alloc_bb); // ctx.builder.position_at_end(alloc_bb);
//
let alloc_size = ctx.builder.build_left_shift( // let alloc_size = ctx.builder.build_left_shift(
ctx.builder.build_int_add( // ctx.builder.build_int_add(
ctx.builder.build_right_shift( // ctx.builder.build_right_shift(
alloc_size, // alloc_size,
llvm_usize.const_int(2, false), // llvm_usize.const_int(2, false),
false, // false,
"", // "",
).unwrap(), // ).unwrap(),
llvm_usize.const_int(1, false), // llvm_usize.const_int(1, false),
"", // "",
).unwrap(), // ).unwrap(),
llvm_usize.const_int(2, false), // llvm_usize.const_int(2, false),
"", // "",
).unwrap(); // ).unwrap();
let alloc_ptr = // let alloc_ptr =
ctx.builder.build_array_alloca(llvm_i8, alloc_size, "rpc.alloc").unwrap(); // ctx.builder.build_array_alloca(llvm_i8, alloc_size, "rpc.alloc").unwrap();
phi.add_incoming(&[(&alloc_ptr, alloc_bb)]); // phi.add_incoming(&[(&alloc_ptr, alloc_bb)]);
let i = // let i =
ctx.builder.build_load(i_addr, "i").map(BasicValueEnum::into_int_value).unwrap(); // ctx.builder.build_load(i_addr, "i").map(BasicValueEnum::into_int_value).unwrap();
let i = ctx.builder.build_int_add(i, llvm_usize.const_int(1, false), "").unwrap(); // let i = ctx.builder.build_int_add(i, llvm_usize.const_int(1, false), "").unwrap();
ctx.builder.build_store(i_addr, i).unwrap(); // ctx.builder.build_store(i_addr, i).unwrap();
ctx.builder.build_unconditional_branch(head_bb).unwrap(); // ctx.builder.build_unconditional_branch(head_bb).unwrap();
//
ctx.builder.position_at_end(tail_bb); // ctx.builder.position_at_end(tail_bb);
ctx.builder.build_load(slot, "rpc.result").unwrap() // ctx.builder.build_load(slot, "rpc.result").unwrap()
} // }
_ => { _ => {
let slot = ctx.builder.build_alloca(llvm_ret_ty, "rpc.ret.slot").unwrap(); let slot = ctx.builder.build_alloca(llvm_ret_ty, "rpc.ret.slot").unwrap();