Testing... 5342

numba/core/config.py

@@ -173,6 +173,11 @@ def optional_str(x):
         # of external references
         FUNCTION_CACHE_SIZE = _readenv("NUMBA_FUNCTION_CACHE_SIZE", int, 128)
 
+        # Maximum tuple size that parfors will unpack and pass to
+        # internal gufunc.
+        PARFOR_MAX_TUPLE_SIZE = _readenv("NUMBA_PARFOR_MAX_TUPLE_SIZE",
+                                         int, 100)
+
         # Enable logging of cache operation
         DEBUG_CACHE = _readenv("NUMBA_DEBUG_CACHE", int, DEBUG)
 

numba/parfors/parfor.py

@@ -614,6 +614,35 @@ def dump(self, file=None):
             block.dump(file)
         print(("end parfor {}".format(self.id)).center(20, '-'), file=file)
 
+    def validate_params(self, typemap):
+        """
+        Check that Parfors params are of valid types.
+        """
+        if self.params is None:
+            msg = ("Cannot run parameter validation on a Parfor with params "
+                   "not set")
+            raise ValueError(msg)
+        for p in self.params:
+            ty = typemap.get(p)
+            if ty is None:
+                msg = ("Cannot validate parameter %s, there is no type "
+                       "information available")
+                raise ValueError(msg)
+            if isinstance(ty, types.BaseTuple):
+                if ty.count > config.PARFOR_MAX_TUPLE_SIZE:
+                    msg = ("Use of a tuple (%s) of length %d in a parallel region "
+                           "exceeds the maximum supported tuple size.  Since "
+                           "Generalized Universal Functions back parallel regions "
+                           "and those do not support tuples, tuples passed to "
+                           "parallel regions are unpacked if their size is below "
+                           "a certain threshold, currently configured to be %d. "
+                           "This threshold can be modified using the Numba "
+                           "environment variable NUMBA_PARFOR_MAX_TUPLE_SIZE.")
+                    raise errors.UnsupportedParforsError(msg %
+                              (p, ty.count, config.PARFOR_MAX_TUPLE_SIZE),
+                              self.loc)
+
+
 def _analyze_parfor(parfor, equiv_set, typemap, array_analysis):
     """Recursive array analysis for parfor nodes.
     """
@@ -2647,8 +2676,15 @@ def run(self):
             parfor_ids, parfors = get_parfor_params(self.func_ir.blocks,
                                                     self.options.fusion,
                                                     self.nested_fusion_info)
+
+            # Validate reduction in parfors.
             for p in parfors:
                 get_parfor_reductions(self.func_ir, p, p.params, self.calltypes)
+
+            # Validate parameters:
+            for p in parfors:
+                p.validate_params(self.typemap)
+
             if config.DEBUG_ARRAY_OPT_STATS:
                 name = self.func_ir.func_id.func_qualname
                 n_parfors = len(parfor_ids)

numba/parfors/parfor_lowering.py

@@ -48,6 +48,8 @@ def _lower_parfor_parallel(lowerer, parfor):
     orig_typemap = lowerer.fndesc.typemap
     # replace original typemap with copy and restore the original at the end.
     lowerer.fndesc.typemap = copy.copy(orig_typemap)
+    if config.DEBUG_ARRAY_OPT:
+        print("lowerer.fndesc", lowerer.fndesc, type(lowerer.fndesc))
     typemap = lowerer.fndesc.typemap
     varmap = lowerer.varmap
 
@@ -228,7 +230,12 @@ def _lower_parfor_parallel(lowerer, parfor):
         assert typemap[l.index_variable.name] == index_var_typ
     numba.parfors.parfor.sequential_parfor_lowering = True
     try:
-        func, func_args, func_sig, redargstartdim, func_arg_types = _create_gufunc_for_parfor_body(
+        (func,
+         func_args,
+         func_sig,
+         redargstartdim,
+         func_arg_types,
+         exp_name_to_tuple_var) = _create_gufunc_for_parfor_body(
             lowerer, parfor, typemap, typingctx, targetctx, flags, {},
             bool(alias_map), index_var_typ, parfor.races)
     finally:
@@ -274,7 +281,8 @@ def _lower_parfor_parallel(lowerer, parfor):
         redarrs,
         parfor.init_block,
         index_var_typ,
-        parfor.races)
+        parfor.races,
+        exp_name_to_tuple_var)
     if config.DEBUG_ARRAY_OPT:
         sys.stdout.flush()
 
@@ -829,6 +837,66 @@ def _create_gufunc_for_parfor_body(
         print("parfor_inputs = ", parfor_inputs, " ", type(parfor_inputs))
         print("parfor_redvars = ", parfor_redvars, " ", type(parfor_redvars))
 
+    # -------------------------------------------------------------------------
+    # Convert tuples to individual parameters.
+    tuple_expanded_parfor_inputs = []
+    tuple_var_to_expanded_names = {}
+    expanded_name_to_tuple_var = {}
+    next_expanded_tuple_var = 0
+    parfor_tuple_params = []
+    # For each input to the parfor.
+    for pi in parfor_inputs:
+        # Get the type of the input.
+        pi_type = typemap[pi]
+        # If it is a UniTuple or Tuple we will do the conversion.
+        if isinstance(pi_type, types.UniTuple) or isinstance(pi_type, types.NamedUniTuple):
+            # Get the size and dtype of the tuple.
+            tuple_count = pi_type.count
+            tuple_dtype = pi_type.dtype
+            # Only do tuples up to config.PARFOR_MAX_TUPLE_SIZE length.
+            assert(tuple_count <= config.PARFOR_MAX_TUPLE_SIZE)
+            this_var_expansion = []
+            for i in range(tuple_count):
+                # Generate a new name for the individual part of the tuple var.
+                expanded_name = "expanded_tuple_var_" + str(next_expanded_tuple_var)
+                # Add that name to the new list of inputs to the gufunc.
+                tuple_expanded_parfor_inputs.append(expanded_name)
+                this_var_expansion.append(expanded_name)
+                # Remember a mapping from new param name to original tuple
+                # var and the index within the tuple.
+                expanded_name_to_tuple_var[expanded_name] = (pi, i)
+                next_expanded_tuple_var += 1
+                # Set the type of the new parameter.
+                typemap[expanded_name] = tuple_dtype
+            # Remember a mapping from the original tuple var to the
+            # individual parts.
+            tuple_var_to_expanded_names[pi] = this_var_expansion
+            parfor_tuple_params.append(pi)
+        elif isinstance(pi_type, types.Tuple) or isinstance(pi_type, types.NamedTuple):
+            # This is the same as above for UniTuple except that each part of
+            # the tuple can have a different type and we fetch that type with
+            # pi_type.types[offset].
+            tuple_count = pi_type.count
+            tuple_types = pi_type.types
+            # Only do tuples up to config.PARFOR_MAX_TUPLE_SIZE length.
+            assert(tuple_count <= config.PARFOR_MAX_TUPLE_SIZE)
+            this_var_expansion = []
+            for i in range(tuple_count):
+                expanded_name = "expanded_tuple_var_" + str(next_expanded_tuple_var)
+                tuple_expanded_parfor_inputs.append(expanded_name)
+                this_var_expansion.append(expanded_name)
+                expanded_name_to_tuple_var[expanded_name] = (pi, i)
+                next_expanded_tuple_var += 1
+                typemap[expanded_name] = tuple_types[i]
+            tuple_var_to_expanded_names[pi] = this_var_expansion
+            parfor_tuple_params.append(pi)
+        else:
+            tuple_expanded_parfor_inputs.append(pi)
+    parfor_inputs = tuple_expanded_parfor_inputs
+    if config.DEBUG_ARRAY_OPT >= 1:
+        print("parfor_inputs post tuple handling = ", parfor_inputs, " ", type(parfor_inputs))
+    # -------------------------------------------------------------------------
+
     races = races.difference(set(parfor_redvars))
     for race in races:
         msg = ("Variable %s used in parallel loop may be written "
@@ -863,7 +931,7 @@ def _create_gufunc_for_parfor_body(
 
     # Some Var are not legal parameter names so create a dict of potentially illegal
     # param name to guaranteed legal name.
-    param_dict = legalize_names_with_typemap(parfor_params + parfor_redvars, typemap)
+    param_dict = legalize_names_with_typemap(parfor_params + parfor_redvars + parfor_tuple_params, typemap)
     if config.DEBUG_ARRAY_OPT >= 1:
         print(
             "param_dict = ",
@@ -937,6 +1005,68 @@ def _create_gufunc_for_parfor_body(
     gufunc_txt += "def " + gufunc_name + \
         "(sched, " + (", ".join(parfor_params)) + "):\n"
 
+    globls = {"np": np}
+
+    # First thing in the gufunc, we reconstruct tuples from their
+    # individual parts, e.g., orig_tup_name = (part1, part2,).
+    # The rest of the code of the function will use the original tuple name.
+    for tup_var, exp_names in tuple_var_to_expanded_names.items():
+        tup_type = typemap[tup_var]
+        gufunc_txt += "    " + param_dict[tup_var]
+        # Determine if the tuple is a named tuple.
+        if (isinstance(tup_type, types.NamedTuple) or
+            isinstance(tup_type, types.NamedUniTuple)):
+            named_tup = True
+        else:
+            named_tup = False
+
+        if named_tup:
+            # It is a named tuple so try to find the global that defines the
+            # named tuple.
+            func_def = guard(get_definition, lowerer.func_ir, tup_var)
+            named_tuple_def = None
+            if config.DEBUG_ARRAY_OPT:
+                print("func_def:", func_def, type(func_def))
+            if func_def is not None:
+                if (isinstance(func_def, ir.Expr) and
+                    func_def.op == 'call'):
+                    named_tuple_def = guard(get_definition, lowerer.func_ir, func_def.func)
+                    if config.DEBUG_ARRAY_OPT:
+                        print("named_tuple_def:", named_tuple_def, type(named_tuple_def))
+                elif isinstance(func_def, ir.Arg):
+                    named_tuple_def = typemap[func_def.name]
+                    if config.DEBUG_ARRAY_OPT:
+                        print("named_tuple_def:", named_tuple_def,
+                              type(named_tuple_def), named_tuple_def.name)
+            if named_tuple_def is not None:
+                if (isinstance(named_tuple_def, ir.Global) or
+                    isinstance(named_tuple_def, ir.FreeVar)):
+                    gval = named_tuple_def.value
+                    if config.DEBUG_ARRAY_OPT:
+                        print("gval:", gval, type(gval))
+                    globls[named_tuple_def.name] = gval
+                elif isinstance(named_tuple_def, types.containers.BaseNamedTuple):
+                    named_tuple_name = named_tuple_def.name.split('(')[0]
+                    if config.DEBUG_ARRAY_OPT:
+                        print("name:", named_tuple_name,
+                              named_tuple_def.instance_class,
+                              type(named_tuple_def.instance_class))
+                    globls[named_tuple_name] = named_tuple_def.instance_class
+            elif config.DEBUG_ARRAY_OPT:
+                print("Didn't find definition of namedtuple for globls.")
+            gufunc_txt += " = " + tup_type.instance_class.__name__ + "("
+            for name, field_name in zip(exp_names, tup_type.fields):
+                gufunc_txt += field_name + "=" + param_dict[name] + ","
+        else:
+            # Just a regular tuple so use (part0, part1, ...)
+            gufunc_txt += " = (" + ", ".join([param_dict[x] for x in exp_names])
+            if len(exp_names) == 1:
+                # Add comma for tuples with singular values.  We can't unilaterally
+                # add a comma alway because (,) isn't valid.
+                gufunc_txt += ","
+
+        gufunc_txt += ")\n"
+
     for pindex in range(len(parfor_inputs)):
         if ascontig and isinstance(param_types[pindex], types.npytypes.Array):
             gufunc_txt += ("    " + parfor_params_orig[pindex]
@@ -1002,7 +1132,6 @@ def _create_gufunc_for_parfor_body(
     if config.DEBUG_ARRAY_OPT:
         print("gufunc_txt = ", type(gufunc_txt), "\n", gufunc_txt)
     # Force gufunc outline into existence.
-    globls = {"np": np}
     locls = {}
     exec(gufunc_txt, globls, locls)
     gufunc_func = locls[gufunc_name]
@@ -1172,7 +1301,7 @@ def _create_gufunc_for_parfor_body(
     if config.DEBUG_ARRAY_OPT:
         print("finished create_gufunc_for_parfor_body. kernel_sig = ", kernel_sig)
 
-    return kernel_func, parfor_args, kernel_sig, redargstartdim, func_arg_types
+    return kernel_func, parfor_args, kernel_sig, redargstartdim, func_arg_types, expanded_name_to_tuple_var
 
 def replace_var_with_array_in_block(vars, block, typemap, calltypes):
     new_block = []
@@ -1208,7 +1337,8 @@ def replace_var_with_array(vars, loop_body, typemap, calltypes):
         typemap[v] = types.npytypes.Array(el_typ, 1, "C")
 
 def call_parallel_gufunc(lowerer, cres, gu_signature, outer_sig, expr_args, expr_arg_types,
-                         loop_ranges, redvars, reddict, redarrdict, init_block, index_var_typ, races):
+                         loop_ranges, redvars, reddict, redarrdict, init_block, index_var_typ, races,
+                         exp_name_to_tuple_var):
     '''
     Adds the call to the gufunc function from the main function.
     '''
@@ -1221,7 +1351,6 @@ def call_parallel_gufunc(lowerer, cres, gu_signature, outer_sig, expr_args, expr
 
     if config.DEBUG_ARRAY_OPT:
         print("make_parallel_loop")
-        print("args = ", expr_args)
         print("outer_sig = ", outer_sig.args, outer_sig.return_type,
               outer_sig.recvr, outer_sig.pysig)
         print("loop_ranges = ", loop_ranges)
@@ -1371,9 +1500,27 @@ def load_range(v):
                                     types.intp, i * num_dim * 2 + j)])))
             cgutils.printf(builder, "\n")
 
+    def load_potential_tuple_var(x):
+        """Given a variable name, if that variable is not a new name
+           introduced as the extracted part of a tuple then just return
+           the variable loaded from its name.  However, if the variable
+           does represent part of a tuple, as recognized by the name of
+           the variable being present in the exp_name_to_tuple_var dict,
+           then we load the original tuple var instead that we get from
+           the dict and then extract the corresponding element of the
+           tuple, also stored and returned to use in the dict (i.e., offset).
+        """
+        if x in exp_name_to_tuple_var:
+            orig_tup, offset = exp_name_to_tuple_var[x]
+            tup_var = lowerer.loadvar(orig_tup)
+            res = builder.extract_value(tup_var, offset)
+            return res
+        else:
+            return lowerer.loadvar(x)
+
     # ----------------------------------------------------------------------------
     # Prepare arguments: args, shapes, steps, data
-    all_args = [lowerer.loadvar(x) for x in expr_args[:ninouts]] + redarrs
+    all_args = [load_potential_tuple_var(x) for x in expr_args[:ninouts]] + redarrs
     num_args = len(all_args)
     num_inps = len(sin) + 1
     args = cgutils.alloca_once(