def bcast_str_impl(data):
rank = hpat.distributed_api.get_rank()
n_loc = len(data)
n_all_chars = num_total_chars(data)
assert n_loc < INT_MAX
assert n_all_chars < INT_MAX
offset_ptr = get_offset_ptr(data)
data_ptr = get_data_ptr(data)
if rank == MPI_ROOT:
send_arr_lens = np.empty(
n_loc, np.uint32) # XXX offset type is uint32
for i in range(n_loc):
_str = data[i]
send_arr_lens[i] = len(_str)
del_str(_str)
c_bcast(send_arr_lens.ctypes, np.int32(n_loc), int32_typ_enum)
else:
c_bcast(offset_ptr, np.int32(n_loc), int32_typ_enum)
c_bcast(data_ptr, np.int32(n_all_chars), char_typ_enum)
if rank != MPI_ROOT:
convert_len_arr_to_offset(offset_ptr, n_loc)
def update_str_impl(shuffle_meta, node_id, ind, val, is_contig=True):
n_chars = len(val)
shuffle_meta.send_counts[node_id] += 1
shuffle_meta.send_counts_char[node_id] += n_chars
if is_contig:
shuffle_meta.send_arr_lens[ind] = n_chars
del_str(val)
def gatherv_str_arr_impl(data):
rank = hpat.distributed_api.get_rank()
n_loc = len(data)
n_all_chars = num_total_chars(data)
# allocate send lens arrays
send_arr_lens = np.empty(n_loc,
np.uint32) # XXX offset type is uint32
send_data_ptr = get_data_ptr(data)
for i in range(n_loc):
_str = data[i]
send_arr_lens[i] = len(_str)
del_str(_str)
recv_counts = gather_scalar(np.int32(n_loc))
recv_counts_char = gather_scalar(np.int32(n_all_chars))
n_total = recv_counts.sum()
n_total_char = recv_counts_char.sum()
# displacements
all_data = StringArray(['']) # dummy arrays on non-root PEs
displs = np.empty(1, np.int32)
displs_char = np.empty(1, np.int32)
if rank == MPI_ROOT:
all_data = pre_alloc_string_array(n_total, n_total_char)
displs = hpat.hiframes_join.calc_disp(recv_counts)
displs_char = hpat.hiframes_join.calc_disp(recv_counts_char)
# print(rank, n_loc, n_total, recv_counts, displs)
offset_ptr = get_offset_ptr(all_data)
data_ptr = get_data_ptr(all_data)
c_gatherv(send_arr_lens.ctypes, np.int32(n_loc), offset_ptr,
recv_counts.ctypes, displs.ctypes, int32_typ_enum)
c_gatherv(send_data_ptr, np.int32(n_all_chars), data_ptr,
recv_counts_char.ctypes, displs_char.ctypes,
char_typ_enum)
convert_len_arr_to_offset(offset_ptr, n_total)
return all_data
