dist divide ceiling to be more even

Author: Ehsan Totoni

hpat/_distributed.c

@@ -1,13 +1,14 @@
 #include <stdbool.h>
 #include "mpi.h"
+#include <cmath>
+#include <algorithm>
 #include <Python.h>
 
 int hpat_dist_get_rank();
 int hpat_dist_get_size();
-int64_t hpat_dist_get_end(int64_t total, int64_t div_chunk, int num_pes,
-                            int node_id);
-int64_t hpat_dist_get_node_portion(int64_t total, int64_t div_chunk,
-                                    int num_pes, int node_id);
+int64_t hpat_dist_get_start(int64_t total, int num_pes, int node_id);
+int64_t hpat_dist_get_end(int64_t total, int num_pes, int node_id);
+int64_t hpat_dist_get_node_portion(int64_t total, int num_pes, int node_id);
 double hpat_dist_get_time();
 MPI_Datatype get_MPI_typ(int typ_enum);
 int get_elem_size(int type_enum);
@@ -43,6 +44,8 @@ PyMODINIT_FUNC PyInit_hdist(void) {
                             PyLong_FromVoidPtr((void*)(&hpat_dist_get_rank)));
     PyObject_SetAttrString(m, "hpat_dist_get_size",
                             PyLong_FromVoidPtr((void*)(&hpat_dist_get_size)));
+    PyObject_SetAttrString(m, "hpat_dist_get_start",
+                            PyLong_FromVoidPtr((void*)(&hpat_dist_get_start)));
     PyObject_SetAttrString(m, "hpat_dist_get_end",
                             PyLong_FromVoidPtr((void*)(&hpat_dist_get_end)));
     PyObject_SetAttrString(m, "hpat_dist_get_node_portion",
@@ -103,17 +106,29 @@ int hpat_dist_get_size()
     return size;
 }
 
-int64_t hpat_dist_get_end(int64_t total, int64_t div_chunk, int num_pes,
-                            int node_id)
+int64_t hpat_dist_get_start(int64_t total, int num_pes, int node_id)
 {
-    return ((node_id==num_pes-1) ? total : (node_id+1)*div_chunk);
+    int64_t div_chunk = (int64_t)ceil(total/((double)num_pes));
+    int64_t start = std::min(total, node_id*div_chunk);
+    // printf("rank %d start:%lld\n", node_id, start);
+    return start;
 }
 
-int64_t hpat_dist_get_node_portion(int64_t total, int64_t div_chunk,
-                                    int num_pes, int node_id)
+int64_t hpat_dist_get_end(int64_t total, int num_pes, int node_id)
 {
-    int64_t portion = ((node_id==num_pes-1) ? total-node_id*div_chunk : div_chunk);
-    // printf("portion:%lld\n", portion);
+    int64_t div_chunk = (int64_t)ceil(total/((double)num_pes));
+    int64_t end = std::min(total, (node_id+1)*div_chunk);
+    // printf("rank %d end:%lld\n", node_id, end);
+    return end;
+}
+
+int64_t hpat_dist_get_node_portion(int64_t total, int num_pes, int node_id)
+{
+    int64_t div_chunk = (int64_t)ceil(total/((double)num_pes));
+    int64_t start = std::min(total, node_id*div_chunk);
+    int64_t end = std::min(total, (node_id+1)*div_chunk);
+    int64_t portion = end-start;
+    // printf("rank %d portion:%lld\n", node_id, portion);
     return portion;
 }
 

hpat/distributed.py

@@ -852,17 +852,22 @@ def _gen_1D_div(self, size_var, scope, loc, prefix, end_call_name, end_call):
             size_assign = ir.Assign(ir.Const(size_var, loc), new_size_var, loc)
             div_nodes.append(size_assign)
             size_var = new_size_var
-        div_var = ir.Var(scope, mk_unique_var(prefix+"_div_var"), loc)
-        self.typemap[div_var.name] = types.int64
-        div_expr = ir.Expr.binop('//', size_var, self._size_var, loc)
-        self.calltypes[div_expr] = find_op_typ('//', [types.int64, types.int32])
-        div_assign = ir.Assign(div_expr, div_var, loc)
 
+        # attr call: start_attr = getattr(g_dist_var, get_start)
+        start_attr_call = ir.Expr.getattr(self._g_dist_var, "get_start", loc)
+        start_attr_var = ir.Var(scope, mk_unique_var("$get_start_attr"), loc)
+        self.typemap[start_attr_var.name] = get_global_func_typ(distributed_api.get_start)
+        start_attr_assign = ir.Assign(start_attr_call, start_attr_var, loc)
+
+        # start_var = get_start(size, rank, pes)
         start_var = ir.Var(scope, mk_unique_var(prefix+"_start_var"), loc)
         self.typemap[start_var.name] = types.int64
-        start_expr = ir.Expr.binop('*', div_var, self._rank_var, loc)
-        self.calltypes[start_expr] = find_op_typ('*', [types.int64, types.int32])
+        start_expr = ir.Expr.call(start_attr_var, [size_var,
+            self._size_var, self._rank_var], (), loc)
+        self.calltypes[start_expr] = self.typemap[start_attr_var.name].get_call_type(
+            typing.Context(), [types.int64, types.int32, types.int32], {})
         start_assign = ir.Assign(start_expr, start_var, loc)
+
         # attr call: end_attr = getattr(g_dist_var, get_end)
         end_attr_call = ir.Expr.getattr(self._g_dist_var, end_call_name, loc)
         end_attr_var = ir.Var(scope, mk_unique_var("$get_end_attr"), loc)
@@ -871,12 +876,12 @@ def _gen_1D_div(self, size_var, scope, loc, prefix, end_call_name, end_call):
 
         end_var = ir.Var(scope, mk_unique_var(prefix+"_end_var"), loc)
         self.typemap[end_var.name] = types.int64
-        end_expr = ir.Expr.call(end_attr_var, [size_var, div_var,
+        end_expr = ir.Expr.call(end_attr_var, [size_var,
             self._size_var, self._rank_var], (), loc)
         self.calltypes[end_expr] = self.typemap[end_attr_var.name].get_call_type(
-            typing.Context(), [types.int64, types.int64, types.int32, types.int32], {})
+            typing.Context(), [types.int64, types.int32, types.int32], {})
         end_assign = ir.Assign(end_expr, end_var, loc)
-        div_nodes += [div_assign, start_assign, end_attr_assign, end_assign]
+        div_nodes += [start_attr_assign, start_assign, end_attr_assign, end_assign]
         return div_nodes, start_var, end_var
 
     def _get_ind_sub(self, ind_var, start_var):

hpat/distributed_api.py

@@ -11,13 +11,17 @@ def get_size():
     """dummy function for C mpi get_size"""
     return 0
 
-def get_end(total_size, div, pes, rank):
+def get_start(total_size, pes, rank):
     """get end point of range for parfor division"""
-    return total_size if rank==pes-1 else (rank+1)*div
+    return 0
 
-def get_node_portion(total_size, div, pes, rank):
+def get_end(total_size, pes, rank):
+    """get end point of range for parfor division"""
+    return 0
+
+def get_node_portion(total_size, pes, rank):
     """get portion of size for alloc division"""
-    return total_size-div*rank if rank==pes-1 else div
+    return 0
 
 def dist_reduce(value):
     """dummy to implement simple reductions"""
@@ -65,18 +69,25 @@ def generic(self, args, kws):
         assert len(args)==0
         return signature(types.int32, *args)
 
+@infer_global(get_start)
+class DistStart(AbstractTemplate):
+    def generic(self, args, kws):
+        assert not kws
+        assert len(args)==3
+        return signature(types.int64, *args)
+
 @infer_global(get_end)
 class DistEnd(AbstractTemplate):
     def generic(self, args, kws):
         assert not kws
-        assert len(args)==4
+        assert len(args)==3
         return signature(types.int64, *args)
 
 @infer_global(get_node_portion)
 class DistPortion(AbstractTemplate):
     def generic(self, args, kws):
         assert not kws
-        assert len(args)==4
+        assert len(args)==3
         return signature(types.int64, *args)
 
 @infer_global(dist_reduce)

hpat/distributed_lower.py

@@ -11,6 +11,7 @@
 import llvmlite.binding as ll
 ll.add_symbol('hpat_dist_get_rank', hdist.hpat_dist_get_rank)
 ll.add_symbol('hpat_dist_get_size', hdist.hpat_dist_get_size)
+ll.add_symbol('hpat_dist_get_start', hdist.hpat_dist_get_start)
 ll.add_symbol('hpat_dist_get_end', hdist.hpat_dist_get_end)
 ll.add_symbol('hpat_dist_get_node_portion', hdist.hpat_dist_get_node_portion)
 ll.add_symbol('hpat_dist_get_time', hdist.hpat_dist_get_time)
@@ -41,19 +42,26 @@ def dist_get_size(context, builder, sig, args):
     fn = builder.module.get_or_insert_function(fnty, name="hpat_dist_get_size")
     return builder.call(fn, [])
 
-@lower_builtin(distributed_api.get_end, types.int64, types.int64, types.int32, types.int32)
+@lower_builtin(distributed_api.get_start, types.int64, types.int32, types.int32)
+def dist_get_start(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(64), [lir.IntType(64),
+                                            lir.IntType(32), lir.IntType(32)])
+    fn = builder.module.get_or_insert_function(fnty, name="hpat_dist_get_start")
+    return builder.call(fn, [args[0], args[1], args[2]])
+
+@lower_builtin(distributed_api.get_end, types.int64, types.int32, types.int32)
 def dist_get_end(context, builder, sig, args):
-    fnty = lir.FunctionType(lir.IntType(64), [lir.IntType(64), lir.IntType(64),
+    fnty = lir.FunctionType(lir.IntType(64), [lir.IntType(64),
                                             lir.IntType(32), lir.IntType(32)])
     fn = builder.module.get_or_insert_function(fnty, name="hpat_dist_get_end")
-    return builder.call(fn, [args[0], args[1], args[2], args[3]])
+    return builder.call(fn, [args[0], args[1], args[2]])
 
-@lower_builtin(distributed_api.get_node_portion, types.int64, types.int64, types.int32, types.int32)
+@lower_builtin(distributed_api.get_node_portion, types.int64, types.int32, types.int32)
 def dist_get_portion(context, builder, sig, args):
-    fnty = lir.FunctionType(lir.IntType(64), [lir.IntType(64), lir.IntType(64),
+    fnty = lir.FunctionType(lir.IntType(64), [lir.IntType(64),
                                             lir.IntType(32), lir.IntType(32)])
     fn = builder.module.get_or_insert_function(fnty, name="hpat_dist_get_node_portion")
-    return builder.call(fn, [args[0], args[1], args[2], args[3]])
+    return builder.call(fn, [args[0], args[1], args[2]])
 
 @lower_builtin(distributed_api.dist_reduce, types.int64)
 @lower_builtin(distributed_api.dist_reduce, types.int32)