@@ -58,9 +58,9 @@ def do_iteration(self, message_size, num_iters, done_future):
5858
5959 if self .gpu_direct and USE_ADDRESS_OPTIMIZATION :
6060 d_data_recv_addr = array .array ('L' , [0 ])
61- d_data_recv_size = array .array ('L ' , [0 ])
61+ d_data_recv_size = array .array ('i ' , [0 ])
6262 d_data_send_addr = array .array ('L' , [0 ])
63- d_data_send_size = array .array ('L ' , [0 ])
63+ d_data_send_size = array .array ('i ' , [0 ])
6464
6565 d_data_recv_addr [0 ] = d_data_recv .__cuda_array_interface__ ['data' ][0 ]
6666 d_data_recv_size [0 ] = d_data_recv .nbytes
@@ -85,9 +85,9 @@ def do_iteration(self, message_size, num_iters, done_future):
8585 charm .lib .CudaStreamSynchronize (stream_address )
8686 else :
8787 if USE_ADDRESS_OPTIMIZATION :
88- partner_channel .send (gpu_src_ptrs = d_data_send_addr , gpu_src_sizes = d_data_send_size )
89- partner_channel .recv (post_buf_addresses = d_data_recv_addr ,
90- post_buf_sizes = d_data_recv_size
88+ partner_channel .send (src_ptrs = d_data_send_addr , src_sizes = d_data_send_size )
89+ partner_channel .recv (post_addresses = d_data_recv_addr ,
90+ post_sizes = d_data_recv_size
9191 )
9292 else :
9393 partner_channel .send (d_data_send )
@@ -103,10 +103,10 @@ def do_iteration(self, message_size, num_iters, done_future):
103103 partner_channel .send (h_data_send )
104104 else :
105105 if USE_ADDRESS_OPTIMIZATION :
106- partner_channel .recv (post_buf_addresses = d_data_recv_addr ,
107- post_buf_sizes = d_data_recv_size
106+ partner_channel .recv (post_addresses = d_data_recv_addr ,
107+ post_sizes = d_data_recv_size
108108 )
109- partner_channel .send (gpu_src_ptrs = d_data_send_addr , gpu_src_sizes = d_data_send_size )
109+ partner_channel .send (src_ptrs = d_data_send_addr , src_sizes = d_data_send_size )
110110 else :
111111 partner_channel .recv (d_data_recv )
112112 partner_channel .send (d_data_send )
0 commit comments