WIP - [artiq] Implement alignment

nac3artiq/src/codegen.rs

@@ -32,9 +32,11 @@ use pyo3::{
 
 use crate::{symbol_resolver::InnerResolver, timeline::TimeFns};
 
+use inkwell::values::IntValue;
 use itertools::Itertools;
 use nac3core::codegen::classes::{ProxyType, ProxyValue, TypedArrayLikeMutator};
 use nac3core::codegen::llvm_intrinsics::call_memset;
+use nac3core::codegen::stmt::gen_for_callback;
 use std::{
     collections::{hash_map::DefaultHasher, HashMap},
     hash::{Hash, Hasher},
@@ -549,7 +551,6 @@ fn format_rpc_ret<'ctx>(
     let llvm_i8 = ctx.ctx.i8_type();
     let llvm_i32 = ctx.ctx.i32_type();
     let llvm_pi8 = llvm_i8.ptr_type(AddressSpace::default());
-    let llvm_ppi8 = llvm_pi8.ptr_type(AddressSpace::default());
 
     let rpc_recv = ctx.module.get_function("rpc_recv").unwrap_or_else(|| {
         ctx.module.add_function("rpc_recv", llvm_i32.fn_type(&[llvm_pi8.into()], false), None)
@@ -572,6 +573,7 @@ fn format_rpc_ret<'ctx>(
         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 llvm_ppi8 = llvm_pi8.ptr_type(AddressSpace::default());
 
             let (elem_ty, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ret_ty);
             let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
@@ -589,6 +591,95 @@ fn format_rpc_ret<'ctx>(
             ndarray.store_ndims(ctx, generator, llvm_usize.const_int(ndims, false));
             ndarray.create_dim_sizes(ctx, llvm_usize, ndarray.load_ndims(ctx));
 
+            let round_size_to_alignment =
+                |generator: &mut dyn CodeGenerator,
+                 ctx: &mut CodeGenContext<'ctx, '_>,
+                 size: IntValue<'ctx>,
+                 alignment: IntType<'ctx>| {
+                    // https://graphics.stanford.edu/%7Eseander/bithacks.html#RoundUpPowerOf2
+                    let llvm_size_ty = size.get_type();
+
+                    let aligned_size_addr = ctx.builder.build_alloca(llvm_size_ty, "").unwrap();
+                    let size_min_one = ctx
+                        .builder
+                        .build_int_sub(size, llvm_size_ty.const_int(1, false), "")
+                        .unwrap();
+                    ctx.builder.build_store(aligned_size_addr, size_min_one).unwrap();
+
+                    gen_for_callback(
+                        generator,
+                        ctx,
+                        None,
+                        |_, ctx| {
+                            let pow2_addr = ctx.builder.build_alloca(llvm_i32, "").unwrap();
+                            ctx.builder
+                                .build_store(pow2_addr, llvm_i32.const_int(1, false))
+                                .unwrap();
+
+                            Ok(pow2_addr)
+                        },
+                        |_, ctx, pow2_addr| {
+                            let pow2 =
+                                ctx.builder.build_load(pow2_addr, "").unwrap().into_int_value();
+
+                            Ok(ctx
+                                .builder
+                                .build_int_compare(
+                                    IntPredicate::ULT,
+                                    pow2,
+                                    llvm_i32.const_int(alignment.get_bit_width() as u64, false),
+                                    "",
+                                )
+                                .unwrap())
+                        },
+                        |_, ctx, _, pow2_addr| {
+                            let aligned_size = ctx
+                                .builder
+                                .build_load(aligned_size_addr, "")
+                                .unwrap()
+                                .into_int_value();
+                            let pow2 =
+                                ctx.builder.build_load(pow2_addr, "").unwrap().into_int_value();
+
+                            let tmp = ctx
+                                .builder
+                                .build_or(
+                                    aligned_size,
+                                    ctx.builder
+                                        .build_right_shift(aligned_size, pow2, false, "")
+                                        .unwrap(),
+                                    "",
+                                )
+                                .unwrap();
+
+                            ctx.builder.build_store(aligned_size_addr, tmp).unwrap();
+
+                            Ok(())
+                        },
+                        |_, ctx, pow2_addr| {
+                            let pow2 =
+                                ctx.builder.build_load(pow2_addr, "").unwrap().into_int_value();
+                            let pow2 = ctx
+                                .builder
+                                .build_left_shift(pow2, llvm_i32.const_int(1, false), "")
+                                .unwrap();
+                            ctx.builder.build_store(pow2_addr, pow2).unwrap();
+
+                            Ok(())
+                        },
+                    )
+                    .unwrap();
+
+                    let aligned_size =
+                        ctx.builder.build_load(aligned_size_addr, "").unwrap().into_int_value();
+                    let aligned_size = ctx
+                        .builder
+                        .build_int_add(aligned_size, llvm_size_ty.const_int(1, false), "")
+                        .unwrap();
+
+                    aligned_size
+                };
+
             // TODO: Add alignment
             let llvm_usize_sizeof = ctx
                 .builder
@@ -602,31 +693,24 @@ fn format_rpc_ret<'ctx>(
                     "",
                 )
                 .unwrap();
-            let llvm_elem_sizeof = ctx
-                .builder
-                .build_int_truncate_or_bit_cast(llvm_elem_ty.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 = round_size_to_alignment(generator, ctx, buffer_size, llvm_usize);
+            let alloc_size =
+                ctx.builder.build_int_unsigned_div(buffer_size, llvm_usize.size_of(), "").unwrap();
 
             let buffer =
-                ctx.builder.build_array_alloca(llvm_pi8, buffer_size, "rpc.buffer").unwrap();
+                ctx.builder.build_array_alloca(llvm_pi8, alloc_size, "rpc.buffer").unwrap();
             let buffer = ctx
                 .builder
                 .build_bitcast(buffer, llvm_pi8, "")
                 .map(BasicValueEnum::into_pointer_value)
                 .unwrap();
-            let buffer = ArraySliceValue::from_ptr_val(
+            let buffer = ArraySliceValue::from_ptr_val(buffer, buffer_size, Some("rpc.buffer.ptr"));
-                buffer,
-                ctx.builder
-                    .build_left_shift(buffer_size, llvm_usize.const_int(2, false), "")
-                    .unwrap(),
-                Some("rpc.buffer.ptr"),
-            );
 
             let i_addr = ctx.builder.build_alloca(llvm_usize, "i.addr").unwrap();
             ctx.builder.build_store(i_addr, llvm_usize.const_zero()).unwrap();
@@ -654,20 +738,9 @@ fn format_rpc_ret<'ctx>(
                     Ok(ctx
                         .builder
                         .build_int_compare(IntPredicate::EQ, i, llvm_usize.const_zero(), "")
-                        // .build_int_compare(IntPredicate::ULT, i, llvm_usize.const_int(2, false), "")
                         .unwrap())
                 },
                 |generator, ctx| {
-                    // let data_ptr = ctx
-                    //     .builder
-                    //     .build_bitcast(
-                    //         phi.as_basic_value().into_pointer_value(),
-                    //         llvm_ret_ty.element_type(),
-                    //         "",
-                    //     )
-                    //     .unwrap()
-                    //     .into_pointer_value();
-                    // ndarray.store_data(ctx, data_ptr);
                     let data_addr = phi.as_basic_value().into_pointer_value();
                     let data_addr = ctx
                         .builder