diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs index c7bf8cbc..afc24e68 100644 --- a/nac3artiq/src/codegen.rs +++ b/nac3artiq/src/codegen.rs @@ -41,7 +41,6 @@ use std::{ mem, sync::Arc, }; -use inkwell::values::IntValue; /// The parallelism mode within a block. #[derive(Copy, Clone, Eq, PartialEq)] @@ -568,157 +567,183 @@ fn format_rpc_ret<'ctx>( let llvm_ret_ty = ctx.get_llvm_abi_type(generator, ret_ty); let result = match &*ctx.unifier.get_ty_immutable(ret_ty) { - // TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => { - // let llvm_i1 = ctx.ctx.bool_type(); - // let llvm_usize = generator.get_size_type(ctx.ctx); - // - // 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 llvm_ret_ty = - // NDArrayType::new(generator, ctx.ctx, ctx.get_llvm_type(generator, elem_ty)); - // let ndarray = llvm_ret_ty.new_value(generator, ctx, None); - // ctx.builder.build_store(slot, ndarray.as_base_value()).unwrap(); - // - // let ndims = - // if let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) { - // assert_eq!(values.len(), 1); - // - // u64::try_from(values[0].clone()).unwrap() - // } else { - // unreachable!(); - // }; - // ndarray.store_ndims(ctx, generator, llvm_usize.const_int(ndims, false)); - // ndarray.create_dim_sizes(ctx, llvm_usize, ndarray.load_ndims(ctx)); - // - // let llvm_usize_sizeof = ctx - // .builder - // .build_int_truncate_or_bit_cast(llvm_ret_ty.size_type().size_of(), llvm_usize, "") - // .unwrap(); - // let llvm_pdata_sizeof = ctx - // .builder - // .build_int_truncate_or_bit_cast( - // llvm_ret_ty.element_type().size_of().unwrap(), - // llvm_usize, - // "", - // ) - // .unwrap(); - // - // let dims_buf_sz = - // ctx.builder.build_int_mul(ndarray.load_ndims(ctx), llvm_usize_sizeof, "").unwrap(); - // - // let buffer_size = - // ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap(); - // let buffer_size = ctx.builder.build_left_shift( - // ctx.builder.build_int_add( - // ctx.builder.build_right_shift( - // buffer_size, - // llvm_usize.const_int(2, false), - // false, - // "", - // ).unwrap(), - // llvm_usize.const_int(1, false), - // "", - // ).unwrap(), - // llvm_usize.const_int(2, false), - // "", - // ).unwrap(); - // - // let buffer = - // 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 i_addr = ctx.builder.build_alloca(llvm_usize, "i.addr").unwrap(); - // ctx.builder.build_unconditional_branch(head_bb).unwrap(); - // - // ctx.builder.position_at_end(head_bb); - // let phi = ctx.builder.build_phi(llvm_pi8, "rpc.ptr").unwrap(); - // phi.add_incoming(&[(&buffer.base_ptr(ctx, generator), prehead_bb)]); - // let alloc_size = ctx - // .build_call_or_invoke(rpc_recv, &[phi.as_basic_value()], "rpc.size.next") - // .map(BasicValueEnum::into_int_value) - // .unwrap(); - // - // // Parse metadata block(s) for ndarrays - // gen_if_callback( - // generator, - // ctx, - // |_, ctx| { - // let i = ctx - // .builder - // .build_load(i_addr, "") - // .map(BasicValueEnum::into_int_value) - // .unwrap(); - // - // Ok(ctx - // .builder - // .build_int_compare(IntPredicate::EQ, i, llvm_usize.const_zero(), "") - // .unwrap()) - // }, - // |generator, ctx| { - // let ppdata = - // generator.gen_var_alloc(ctx, llvm_ret_ty.element_type(), None).unwrap(); - // ctx.builder - // .build_store(ppdata, ndarray.data().base_ptr(ctx, generator)) - // .unwrap(); - // call_memcpy_generic( - // ctx, - // ppdata, - // buffer.base_ptr(ctx, generator), - // llvm_pdata_sizeof, - // llvm_i1.const_zero(), - // ); - // - // let pbuffer_dims_begin = unsafe { - // buffer.ptr_offset_unchecked(ctx, generator, &llvm_pdata_sizeof, None) - // }; - // call_memcpy_generic( - // ctx, - // ndarray.dim_sizes().base_ptr(ctx, generator), - // pbuffer_dims_begin, - // dims_buf_sz, - // llvm_i1.const_zero(), - // ); - // - // Ok(()) - // }, - // |_, _| Ok(()), - // ) - // .unwrap(); - // - // let is_done = ctx - // .builder - // .build_int_compare(IntPredicate::EQ, llvm_i32.const_zero(), alloc_size, "rpc.done") - // .unwrap(); - // - // ctx.builder.build_conditional_branch(is_done, tail_bb, alloc_bb).unwrap(); - // ctx.builder.position_at_end(alloc_bb); - // - // let alloc_size = ctx.builder.build_left_shift( - // ctx.builder.build_int_add( - // ctx.builder.build_right_shift( - // alloc_size, - // llvm_usize.const_int(2, false), - // false, - // "", - // ).unwrap(), - // llvm_usize.const_int(1, false), - // "", - // ).unwrap(), - // llvm_usize.const_int(2, false), - // "", - // ).unwrap(); - // let alloc_ptr = - // ctx.builder.build_array_alloca(llvm_i8, alloc_size, "rpc.alloc").unwrap(); - // phi.add_incoming(&[(&alloc_ptr, alloc_bb)]); - // let i = - // 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(); - // ctx.builder.build_store(i_addr, i).unwrap(); - // ctx.builder.build_unconditional_branch(head_bb).unwrap(); - // - // ctx.builder.position_at_end(tail_bb); - // ctx.builder.build_load(slot, "rpc.result").unwrap() - // } + TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => { + let llvm_i1 = ctx.ctx.bool_type(); + let llvm_usize = generator.get_size_type(ctx.ctx); + + let (elem_ty, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ret_ty); + let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty); + let llvm_ret_ty = NDArrayType::new(generator, ctx.ctx, llvm_elem_ty); + let ndarray = llvm_ret_ty.new_value(generator, ctx, Some("rpc.result")); + + let ndims = + if let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) { + assert_eq!(values.len(), 1); + + u64::try_from(values[0].clone()).unwrap() + } else { + unreachable!(); + }; + ndarray.store_ndims(ctx, generator, llvm_usize.const_int(ndims, false)); + ndarray.create_dim_sizes(ctx, llvm_usize, ndarray.load_ndims(ctx)); + + let llvm_usize_sizeof = ctx + .builder + .build_int_truncate_or_bit_cast(llvm_ret_ty.size_type().size_of(), llvm_usize, "") + .unwrap(); + let llvm_pdata_sizeof = ctx + .builder + .build_int_truncate_or_bit_cast( + llvm_ret_ty.element_type().size_of().unwrap(), + llvm_usize, + "", + ) + .unwrap(); + + let dims_buf_sz = + ctx.builder.build_int_mul(ndarray.load_ndims(ctx), llvm_usize_sizeof, "").unwrap(); + + let buffer_size = llvm_usize.const_int(8, false); + // let buffer_size = + // ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap(); + // let buffer_size = ctx.builder.build_left_shift( + // ctx.builder.build_int_add( + // ctx.builder.build_right_shift( + // buffer_size, + // llvm_usize.const_int(2, false), + // false, + // "", + // ).unwrap(), + // llvm_usize.const_int(1, false), + // "", + // ).unwrap(), + // llvm_usize.const_int(2, false), + // "", + // ).unwrap(); + + let buffer = + 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 i_addr = ctx.builder.build_alloca(llvm_usize, "i.addr").unwrap(); + ctx.builder.build_unconditional_branch(head_bb).unwrap(); + + ctx.builder.position_at_end(head_bb); + let phi = ctx.builder.build_phi(llvm_pi8, "rpc.ptr").unwrap(); + phi.add_incoming(&[(&buffer.base_ptr(ctx, generator), prehead_bb)]); + let alloc_size = ctx + .build_call_or_invoke(rpc_recv, &[phi.as_basic_value()], "rpc.size.next") + .map(BasicValueEnum::into_int_value) + .unwrap(); + + // Parse metadata block(s) for ndarrays + gen_if_callback( + generator, + ctx, + |_, ctx| { + let i = ctx + .builder + .build_load(i_addr, "") + .map(BasicValueEnum::into_int_value) + .unwrap(); + + Ok(ctx + .builder + .build_int_compare(IntPredicate::EQ, i, llvm_usize.const_zero(), "") + .unwrap()) + }, + |generator, ctx| { + // let ppdata = + // generator.gen_var_alloc(ctx, llvm_ret_ty.element_type(), None).unwrap(); + // ctx.builder + // .build_store(ppdata, ndarray.data().base_ptr(ctx, generator)) + // .unwrap(); + // call_memcpy_generic( + // ctx, + // ppdata, + // buffer.base_ptr(ctx, generator), + // llvm_pdata_sizeof, + // llvm_i1.const_zero(), + // ); + + // let pbuffer_dims_begin = unsafe { + // buffer.ptr_offset_unchecked(ctx, generator, &llvm_pdata_sizeof, None) + // }; + // call_memcpy_generic( + // ctx, + // ndarray.dim_sizes().base_ptr(ctx, generator), + // pbuffer_dims_begin, + // dims_buf_sz, + // llvm_i1.const_zero(), + // ); + + // This works (with uninitialized data) + // unsafe { + // ndarray.dim_sizes().set_typed_unchecked( + // ctx, + // generator, + // &llvm_usize.const_zero(), + // llvm_usize.const_int(3, false), + // ); + // } + + let pbuffer_dims_begin = unsafe { + buffer.ptr_offset_unchecked(ctx, generator, &llvm_usize.const_int(4, false), None) + }; + call_memcpy_generic( + ctx, + ndarray.dim_sizes().base_ptr(ctx, generator), + pbuffer_dims_begin, + dims_buf_sz, + llvm_i1.const_zero(), + ); + + ndarray.create_data( + ctx, + llvm_elem_ty, + call_ndarray_calc_size(generator, ctx, &ndarray.dim_sizes(), (None, None)), + ); + + Ok(()) + }, + |_, _| Ok(()), + ) + .unwrap(); + + let is_done = ctx + .builder + .build_int_compare(IntPredicate::EQ, llvm_i32.const_zero(), alloc_size, "rpc.done") + .unwrap(); + + ctx.builder.build_conditional_branch(is_done, tail_bb, alloc_bb).unwrap(); + ctx.builder.position_at_end(alloc_bb); + + let alloc_size = ctx.builder.build_left_shift( + ctx.builder.build_int_add( + ctx.builder.build_right_shift( + alloc_size, + llvm_usize.const_int(2, false), + false, + "", + ).unwrap(), + llvm_usize.const_int(1, false), + "", + ).unwrap(), + llvm_usize.const_int(2, false), + "", + ).unwrap(); + let alloc_ptr = + ctx.builder.build_array_alloca(llvm_i8, alloc_size, "rpc.alloc").unwrap(); + phi.add_incoming(&[(&alloc_ptr, alloc_bb)]); + let i = + 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(); + ctx.builder.build_store(i_addr, i).unwrap(); + ctx.builder.build_unconditional_branch(head_bb).unwrap(); + + ctx.builder.position_at_end(tail_bb); + ndarray.as_base_value().into() + } _ => { let slot = ctx.builder.build_alloca(llvm_ret_ty, "rpc.ret.slot").unwrap();