From e8f94b43933fc4f874093f03303613fe6071d766 Mon Sep 17 00:00:00 2001
From: David Mak <david.18.19.21@gmail.com>
Date: Thu, 15 Aug 2024 15:11:47 +0800
Subject: [PATCH] [artiq] WIP

---
 nac3artiq/src/codegen.rs | 39 +++++++++++----------------------------
 1 file changed, 11 insertions(+), 28 deletions(-)

diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs
index dc07c836..0f2b34bc 100644
--- a/nac3artiq/src/codegen.rs
+++ b/nac3artiq/src/codegen.rs
@@ -544,6 +544,7 @@ fn format_rpc_ret<'ctx>(
 
     let llvm_i8 = ctx.ctx.i8_type();
     let llvm_i32 = ctx.ctx.i32_type();
+    let llvm_i8_8 = ctx.ctx.struct_type(&[llvm_i8.array_type(8).into()], false);
     let llvm_pi8 = llvm_i8.ptr_type(AddressSpace::default());
 
     let rpc_recv = ctx.module.get_function("rpc_recv").unwrap_or_else(|| {
@@ -570,13 +571,13 @@ fn format_rpc_ret<'ctx>(
 
             let round_up = |ctx: &mut CodeGenContext<'ctx, '_>,
                             val: IntValue<'ctx>,
-                            power_of_two: IntType<'ctx>| {
-                debug_assert!((power_of_two.get_bit_width() / 8).is_power_of_two());
-
+                            power_of_two: IntValue<'ctx>| {
                 let llvm_val_t = val.get_type();
 
-                let max_rem = (power_of_two.get_bit_width() / 8) - 1;
-                let max_rem = llvm_val_t.const_int(max_rem as u64, false);
+                let max_rem = ctx
+                    .builder
+                    .build_int_sub(power_of_two, llvm_val_t.const_int(1, false), "")
+                    .unwrap();
                 ctx.builder
                     .build_and(
                         ctx.builder.build_int_add(val, max_rem, "").unwrap(),
@@ -620,10 +621,10 @@ fn format_rpc_ret<'ctx>(
 
             let buffer_size =
                 ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap();
-            let buffer_size = round_up(ctx, buffer_size, ctx.ctx.i64_type());
+            let buffer_size = round_up(ctx, buffer_size, llvm_usize.const_int(8, false));
 
             let buffer =
-                ctx.builder.build_array_alloca(llvm_pi8, buffer_size, "rpc.buffer").unwrap();
+                ctx.builder.build_array_alloca(llvm_i8_8, buffer_size, "rpc.buffer").unwrap();
             let buffer = ctx
                 .builder
                 .build_bitcast(buffer, llvm_pi8, "")
@@ -661,14 +662,10 @@ fn format_rpc_ret<'ctx>(
                 },
                 |generator, ctx| {
                     let phi = phi.as_basic_value().into_pointer_value();
-                    let pbuffer_data_begin = unsafe {
-                        ctx.builder.build_in_bounds_gep(phi, &[llvm_usize.const_int(8, false)], "")
-                    }
-                    .unwrap();
                     call_memcpy_generic(
                         ctx,
                         ndarray.ptr_to_data(ctx),
-                        pbuffer_data_begin,
+                        phi,
                         llvm_pdata_sizeof,
                         llvm_i1.const_zero(),
                     );
@@ -684,21 +681,6 @@ fn format_rpc_ret<'ctx>(
                         llvm_i1.const_zero(),
                     );
 
-                    // // TODO: Testing for buffer
-                    // ndarray.create_data(
-                    //     ctx,
-                    //     llvm_elem_ty,
-                    //     call_ndarray_calc_size(generator, ctx, &ndarray.dim_sizes(), (None, None)),
-                    // );
-                    //
-                    // call_memcpy_generic(
-                    //     ctx,
-                    //     ndarray.data().base_ptr(ctx, generator),
-                    //     buffer.base_ptr(ctx, generator),
-                    //     llvm_usize.const_int(8, false),
-                    //     llvm_i1.const_zero(),
-                    // );
-
                     Ok(())
                 },
                 |_, _| Ok(()),
@@ -713,8 +695,9 @@ fn format_rpc_ret<'ctx>(
             ctx.builder.build_conditional_branch(is_done, tail_bb, alloc_bb).unwrap();
             ctx.builder.position_at_end(alloc_bb);
 
+            let alloc_size = round_up(ctx, alloc_size, llvm_elem_ty.size_of().unwrap());
             let alloc_ptr =
-                ctx.builder.build_array_alloca(llvm_pi8, alloc_size, "rpc.alloc").unwrap();
+                ctx.builder.build_array_alloca(llvm_elem_ty, alloc_size, "rpc.alloc").unwrap();
             let alloc_ptr = ctx
                 .builder
                 .build_bitcast(alloc_ptr, llvm_pi8, "")