WIP - [artiq] Fix RPC of ndarrays from host

nac3artiq/src/codegen.rs

@@ -600,46 +600,16 @@ fn format_rpc_ret<'ctx>(
                 unreachable!();
             };
             llvm_ndarray.store_ndims(ctx, generator, llvm_usize.const_int(ndims, false));
-            llvm_ndarray.create_dim_sizes(ctx, llvm_usize, llvm_ndarray.load_ndims(ctx));
+            llvm_ndarray.create_dim_sizes(ctx, llvm_usize, llvm_usize.const_int(ndims, false));
 
-            let buffer = ctx.builder.build_array_alloca(llvm_i8, llvm_usize_sizeof, "rpc.buffer").unwrap();
-            let buffer = ArraySliceValue::from_ptr_val(buffer, llvm_usize_sizeof, Some("rpc.buffer"));
+            let dims_buf_sz =
+                ctx.builder.build_int_mul(llvm_ndarray.load_ndims(ctx), llvm_usize_sizeof, "").unwrap();
 
-            gen_for_callback_incrementing(
-                generator,
-                ctx,
-                None,
-                llvm_usize.const_zero(),
-                (llvm_ndarray.load_ndims(ctx), false),
-                |generator, ctx, _, i| {
-                    ctx
-                        .build_call_or_invoke(rpc_recv, &[buffer.base_ptr(ctx, generator).into()], "rpc.size.next")
-                        .map(BasicValueEnum::into_int_value)
-                        .unwrap();
+            let buffer_size =
+                ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap();
 
-                    let offset = ctx.builder.build_int_add(
-                        llvm_pdata_sizeof,
-                        ctx.builder.build_int_mul(i, llvm_usize_sizeof, "").unwrap(),
-                        "",
-                    ).unwrap();
-                    let dims_ptr = unsafe { llvm_ndarray.dim_sizes().ptr_offset_unchecked(ctx, generator, &offset, None) };
-                    let pbuffer_dims_ptr =
-                        unsafe { buffer.ptr_offset_unchecked(ctx, generator, &offset, None) };
-                    call_memcpy_generic(
-                        ctx,
-                        dims_ptr,
-                        pbuffer_dims_ptr,
-                        llvm_usize_sizeof,
-                        llvm_i1.const_zero(),
-                    ); 
-
-                    Ok(())
-                },
-                llvm_usize.const_int(1, false),
-            ).unwrap();
-
-            let buffer = ctx.builder.build_array_alloca(llvm_i8, llvm_pdata_sizeof, "rpc.buffer").unwrap();
-            let buffer = ArraySliceValue::from_ptr_val(buffer, llvm_pdata_sizeof, Some("rpc.buffer"));
+            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"));
 
             // recv [*data, dim_sz[..]]
             let alloc_size = ctx
@@ -657,6 +627,16 @@ fn format_rpc_ret<'ctx>(
                 llvm_i1.const_zero(),
             );
 
+            let pbuffer_dims_begin =
+                unsafe { buffer.ptr_offset_unchecked(ctx, generator, &llvm_pdata_sizeof, None) };
+            call_memcpy_generic(
+                ctx,
+                llvm_ndarray.dim_sizes().base_ptr(ctx, generator),
+                pbuffer_dims_begin,
+                dims_buf_sz,
+                llvm_i1.const_zero(),
+            );
+
             let is_done = ctx
                 .builder
                 .build_int_compare(IntPredicate::EQ, llvm_i32.const_zero(), alloc_size, "rpc.done")
@@ -664,13 +644,12 @@ fn format_rpc_ret<'ctx>(
 
             let alloc_ptr = ctx.builder.build_array_alloca(llvm_pi8, alloc_size, "rpc.alloc").unwrap();
             let alloc_ptr = ctx.builder.build_bitcast(alloc_ptr, llvm_pi8, "rpc.alloc.ptr").map(BasicValueEnum::into_pointer_value).unwrap();
-            let prehead_metadata_bb = ctx.builder.get_insert_block().unwrap();
 
             ctx.builder.build_conditional_branch(is_done, tail_bb, head_bb).unwrap();
 
             ctx.builder.position_at_end(head_bb);
             let phi = ctx.builder.build_phi(llvm_pi8, "rpc.ptr").unwrap();
-            phi.add_incoming(&[(&alloc_ptr, prehead_metadata_bb)]);
+            phi.add_incoming(&[(&alloc_ptr, prehead_bb)]);
             let alloc_size = ctx
                 .build_call_or_invoke(rpc_recv, &[phi.as_basic_value()], "rpc.size.next")
                 .unwrap()