support S.head(), dist broadcast of const slice

Author: Ehsan Totoni

hpat/distributed.py

@@ -36,7 +36,8 @@
 import hpat.utils
 from hpat.utils import (is_alloc_callname, is_whole_slice, is_array_container,
                         get_slice_step, is_array, is_np_array, find_build_tuple,
-                        debug_prints, ReplaceFunc, gen_getitem, is_call)
+                        debug_prints, ReplaceFunc, gen_getitem, is_call,
+                        is_const_slice)
 from hpat.distributed_api import Reduce_Type
 from hpat.hiframes.pd_dataframe_ext import DataFrameType
 
@@ -739,6 +740,7 @@ def f(arr, bag, start, count):  # pragma: no cover
             return [assign]
 
         if (fdef == ('get_series_data', 'hpat.hiframes.api')
+                or fdef == ('get_series_index', 'hpat.hiframes.api')
                 or fdef == ('get_dataframe_data', 'hpat.hiframes.pd_dataframe_ext')):
             out = [assign]
             arr = assign.target
@@ -1583,6 +1585,17 @@ def f(A, start, step):
                 out += self._run_call_rebalance_array(lhs.name, full_node, [imb_arr])
                 out[-1].target = lhs
 
+            elif self._is_REP(lhs.name) and guard(
+                    is_const_slice, self.typemap, self.func_ir, index_var):
+                # cases like S.head()
+                # bcast if all in rank 0, otherwise gatherv
+                in_arr = full_node.value.value
+                start = self._array_starts[in_arr.name][0]
+                count = self._array_counts[in_arr.name][0]
+                return self._replace_func(
+                    lambda arr, slice_index, start, count: hpat.distributed_api.const_slice_getitem(
+                        arr, slice_index, start, count), [in_arr, index_var, start, count])
+
         return out
 
     def _run_parfor(self, parfor, namevar_table):

hpat/distributed_analysis.py

@@ -19,7 +19,8 @@
 from hpat.hiframes.pd_series_ext import SeriesType
 from hpat.utils import (get_constant, is_alloc_callname,
                         is_whole_slice, is_array, is_array_container,
-                        is_np_array, find_build_tuple, debug_prints)
+                        is_np_array, find_build_tuple, debug_prints,
+                        is_const_slice)
 from hpat.hiframes.pd_dataframe_ext import DataFrameType
 from enum import Enum
 
@@ -358,8 +359,12 @@ def _analyze_call(self, lhs, rhs, func_var, args, array_dists):
         if fdef == ('isna', 'hpat.hiframes.api'):
             return
 
+        if fdef == ('get_series_name', 'hpat.hiframes.api'):
+            return
+
         # dummy hiframes functions
         if func_mod == 'hpat.hiframes.api' and func_name in ('get_series_data',
+                'get_series_index',
                 'to_arr_from_series', 'ts_series_to_arr_typ',
                 'to_date_series_type', 'dummy_unbox_series',
                 'parallel_fix_df_array'):
@@ -772,6 +777,12 @@ def _analyze_getitem(self, inst, lhs, rhs, array_dists):
             self._meet_array_dists(lhs, rhs.value.name, array_dists)
             return
 
+        # output of operations like S.head() is REP since it's a "small" slice
+        # input can remain 1D
+        if guard(is_const_slice, self.typemap, self.func_ir, index_var):
+            array_dists[lhs] = Distribution.REP
+            return
+
         self._set_REP(inst.list_vars(), array_dists)
         return
 

hpat/distributed_api.py

@@ -276,6 +276,64 @@ def prealloc_impl(arr):
 
     return lambda arr: arr
 
+
+# assuming start and step are None
+def const_slice_getitem(arr, slice_index, start, count):
+    return arr[slice_index]
+
+
+@overload(const_slice_getitem)
+def const_slice_getitem_overload(arr, slice_index, start, count):
+    if arr == string_array_type:
+        reduce_op = Reduce_Type.Sum.value
+        def getitem_str_impl(arr, slice_index, start, count):
+            rank = hpat.distributed_api.get_rank()
+            k = slice_index.stop
+            # get total characters for allocation
+            n_chars = np.uint64(0)
+            if k > count:
+                my_end = min(count, max(k-start, 0))
+                my_arr = arr[:my_end]
+                my_arr = hpat.distributed_api.gatherv(my_arr)
+                n_chars = hpat.distributed_api.dist_reduce(
+                    num_total_chars(my_arr), np.int32(reduce_op))
+                if rank == 0:
+                    out_arr = my_arr
+            else:
+                if rank == 0:
+                    my_arr = arr[:k]
+                    n_chars = num_total_chars(my_arr)
+                    out_arr = my_arr
+                n_chars = bcast_scalar(n_chars)
+            if rank != 0:
+                out_arr = pre_alloc_string_array(k, n_chars)
+
+            # actual communication
+            hpat.distributed_api.bcast(out_arr)
+            return out_arr
+
+        return getitem_str_impl
+
+    def getitem_impl(arr, slice_index, start, count):
+        rank = hpat.distributed_api.get_rank()
+        k = slice_index.stop
+        out_arr = np.empty(k, arr.dtype)
+        if k > count:
+            my_end = min(count, max(k-start, 0))
+            my_arr = arr[:my_end]
+            my_arr = hpat.distributed_api.gatherv(my_arr)
+            if rank == 0:
+                print(my_arr)
+                out_arr = my_arr
+        else:
+            if rank == 0:
+                out_arr = arr[:k]
+        hpat.distributed_api.bcast(out_arr)
+        return out_arr
+
+    return getitem_impl
+
+
 # send_data, recv_data, send_counts, recv_counts, send_disp, recv_disp, typ_enum
 c_alltoallv = types.ExternalFunction("c_alltoallv", types.void(types.voidptr,
     types.voidptr, types.voidptr, types.voidptr, types.voidptr, types.voidptr, types.int32))

hpat/tests/test_series.py

@@ -1168,6 +1168,16 @@ def test_impl():
         hpat_func = hpat.jit(test_impl)
         pd.testing.assert_series_equal(hpat_func(), test_impl())
 
+    def test_series_head_index_parallel1(self):
+        def test_impl(S):
+            return S.head(3)
+
+        S = pd.Series([6,9,2,3,6,4,5], ['a','ab','abc','c','f','hh',''])
+        hpat_func = hpat.jit(distributed={'S'})(test_impl)
+        start, end = get_start_end(len(S))
+        pd.testing.assert_series_equal(hpat_func(S[start:end]), test_impl(S))
+        self.assertTrue(count_array_OneDs()>0)
+
     def test_series_median1(self):
         def test_impl(S):
             return S.median()

hpat/utils.py

@@ -3,7 +3,7 @@
 import numba
 from numba import ir_utils, ir, types, cgutils
 from numba.ir_utils import (guard, get_definition, find_callname, require,
-    add_offset_to_labels, find_topo_order)
+    add_offset_to_labels, find_topo_order, find_const)
 from numba.parfor import wrap_parfor_blocks, unwrap_parfor_blocks
 from numba.typing import signature
 from numba.typing.templates import infer_global, AbstractTemplate
@@ -279,6 +279,23 @@ def is_whole_slice(typemap, func_ir, var, accept_stride=False):
     require(isinstance(arg1_def, ir.Const) and arg1_def.value == None)
     return True
 
+
+def is_const_slice(typemap, func_ir, var, accept_stride=False):
+    """ return True if var can be determined to be a constant size slice """
+    require(typemap[var.name] == types.slice2_type
+            or (accept_stride and typemap[var.name] == types.slice3_type))
+    call_expr = get_definition(func_ir, var)
+    require(isinstance(call_expr, ir.Expr) and call_expr.op == 'call')
+    assert (len(call_expr.args) == 2
+            or (accept_stride and len(call_expr.args) == 3))
+    assert find_callname(func_ir, call_expr) == ('slice', 'builtins')
+    arg0_def = get_definition(func_ir, call_expr.args[0])
+    require(isinstance(arg0_def, ir.Const) and arg0_def.value == None)
+    size_const = find_const(func_ir, call_expr.args[1])
+    require(isinstance(size_const, int))
+    return True
+
+
 def get_slice_step(typemap, func_ir, var):
     require(typemap[var.name] == types.slice3_type)
     call_expr = get_definition(func_ir, var)