Support array shuffling with an index array generated with np.random.permutation #17
Conversation
hpat/distributed_analysis.py
Outdated
| arr_def = guard(get_definition, self.func_ir, index_var) | ||
| if (arr_def.op == 'call' and | ||
| self._call_table[arr_def.func.name] == ['permutation', 'random', np]): | ||
| self._meet_array_dists(lhs, rhs.value.name, array_dists) |
There was a problem hiding this comment.
arr_def could be None or any node (not necessarily Expr). This if needs more checks otherwise it fails.
There was a problem hiding this comment.
We need to check rhs.value to be multi-dimensional.
There was a problem hiding this comment.
call_table is gone now. need to use find_callname.
hpat/distributed.py
Outdated
| is_multi_dim = True | ||
|
|
||
| arr_def = guard(get_definition, self.func_ir, index_var) | ||
| if (arr_def.op == 'call' and |
There was a problem hiding this comment.
same issue with arr_def check before.
hpat/distributed_analysis.py
Outdated
| if fdef == ('permutation', 'np.random'): | ||
| if fdef == ('permutation', 'numpy.random'): | ||
| assert len(args) == 1 | ||
| assert self.typemap[args[0].name] == types.int64 |
hpat/distributed.py
Outdated
| def _run_permutation_array_index(self, lhs, rhs, idx): | ||
| scope, loc = lhs.scope, lhs.loc | ||
| alloc = mk_alloc(self.typemap, self.calltypes, lhs, | ||
| tuple(self._array_sizes[lhs.name]), |
There was a problem hiding this comment.
chunk size of first dimention: self._array_counts[lhs.name][0].
|
Fixed the chunk size. You were wrong about the code path not being exercised. It was being exercised, except when you execute your shuffle algorithm with one node, you end up sending all of the elements to yourself in sorted order. The missing thing was shuffle after MPI_Alltoallv. I added a hack to receive make up a PRNG and shuffle with that, but the proper way is to obtain the seed from numba and use that. |
|
Dead code eliminator removes np.random.seed() call, which makes it hard to write test for this. Fixed it in numba/numba#2873. |
hpat/_distributed.cpp
Outdated
| } | ||
|
|
||
| std::vector<int> find_dest_ranks(int rank, int num_ranks, | ||
| int64_t *idx, int idx_len) |
hpat/_distributed.cpp
Outdated
| MPI_INT64_T, MPI_COMM_WORLD); | ||
|
|
||
| std::mt19937 rng(0); | ||
| std::shuffle(ilhs, ilhs + dest_ranks.size(), rng); |
There was a problem hiding this comment.
shuffle? idx should determine the order.
| B = np.random.permutation(n) | ||
| A = A[B] | ||
| return A.sum() | ||
|
|
There was a problem hiding this comment.
Needs a test where two arrays are permuted with the same index array, to make sure they are permuted in the same way.
hpat/_distributed.cpp
Outdated
|
|
||
| MPI_Alltoallv(send_buf.data(), send_counts.data(), send_disps.data(), | ||
| MPI_INT64_T, ilhs, recv_counts.data(), recv_disps.data(), | ||
| MPI_INT64_T, MPI_COMM_WORLD); |
There was a problem hiding this comment.
Needs support for non-int types.
|
The last patch addresses all the issues, except the unit tests. You can review this, while I work on the unit tests. |
hpat/_distributed.cpp
Outdated
| // |num_ranks|, finds the destination ranks of indices of the |rank|. For | ||
| // example, if |rank| is 1, |num_ranks| is 3, |p_len| is 12, and |p| is the | ||
| // following array [ 9, 8, 6, 4, 11, 7, 2, 3, 5, 0, 1, 10], the function returns | ||
| // [2, 0, 1, 0]. |
There was a problem hiding this comment.
Shouldn't this be [0,2,0,1]?
hpat/_distributed.cpp
Outdated
| void apply_permutation(unsigned char *v, int64_t elem_size, | ||
| std::vector<size_t>& p) | ||
| { | ||
| std::vector<unsigned char> tmp(elem_size); |
| std::vector<unsigned char> tmp(elem_size); | ||
| for (size_t i = 0; i < p.size(); ++i) { | ||
| auto current = i; | ||
| while (i != p[current]) { |
There was a problem hiding this comment.
This loop needs doc/reference.
hpat/_distributed.cpp
Outdated
| } | ||
|
|
||
| MPI_Alltoallv(send_buf.data(), send_counts.data(), send_disps.data(), | ||
| MPI_INT64_T, lhs, recv_counts.data(), recv_disps.data(), |
There was a problem hiding this comment.
Needs MPI data type of size elem_size.
|
|
||
| auto begin = p + hpat_dist_get_start(p_len, num_ranks, rank); | ||
| auto p1 = arg_sort(begin, dest_ranks.size()); | ||
| auto p2 = arg_sort(p1.data(), dest_ranks.size()); |
… on permutation array.
Change the rhs to be a contiguous array before passing it to C++ function. When adding alloc IR node, take into account the size of the lower dimensions.
|
Looks good. Merging. |
3 participants
Currently, it assumes the array that is being indexed is a one-dimensional array of ints and appears to work. Please see if this makes sense. If it does, then I will extend it to support multi-dimensional arrays.