sklam · Oct 6, 2021 · Sep 13, 2021 · Sep 17, 2021 · Sep 17, 2021 · Sep 17, 2021
diff --git a/numba/parfors/array_analysis.py b/numba/parfors/array_analysis.py
@@ -1325,6 +1325,20 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
                     shape = gvalue
                 elif isinstance(gvalue, int):
                     shape = (gvalue,)
+            elif isinstance(inst.value, ir.Arg):
+                if (
+                    isinstance(typ, types.containers.UniTuple)
+                    and isinstance(typ.dtype, types.Integer)
+                ):
+                    shape = inst.value
+                elif (
+                    isinstance(typ, types.containers.Tuple)
+                    and all([isinstance(x,
+                            (types.Integer, types.IntegerLiteral))
+                        for x in typ.types]
+                    )
+                ):
+                    shape = inst.value
 
             if isinstance(shape, ir.Const):
                 if isinstance(shape.value, tuple):
@@ -1379,6 +1393,16 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
                     shape = self._gen_shape_call(
                         equiv_set, lhs, len(typ), shape, post
                     )
+            elif (
+                isinstance(typ, types.containers.Tuple)
+                and all([isinstance(x,
+                        (types.Integer, types.IntegerLiteral))
+                    for x in typ.types]
+                )
+            ):
+                shape = self._gen_shape_call(
+                    equiv_set, lhs, len(typ), shape, post
+                )
 
             """ See the comment on the define() function.
 
@@ -3128,11 +3152,16 @@ def _gen_shape_call(self, equiv_set, var, ndims, shape, post):
         # attr call: A_sh_attr = getattr(A, shape)
         if isinstance(shape, ir.Var):
             shape = equiv_set.get_shape(shape)
+
         # already a tuple variable that contains size
         if isinstance(shape, ir.Var):
             attr_var = shape
             shape_attr_call = None
             shape = None
+        elif isinstance(shape, ir.Arg):
+            attr_var = var
+            shape_attr_call = None
+            shape = None
         else:
             shape_attr_call = ir.Expr.getattr(var, "shape", var.loc)
             attr_var = ir.Var(

diff --git a/numba/parfors/parfor.py b/numba/parfors/parfor.py
@@ -1873,10 +1873,18 @@ def _type_getitem(self, args):
         return self.pass_states.typingctx.resolve_function_type(fnty, tuple(args), {})
 
 
-def _make_index_var(typemap, scope, index_vars, body_block):
+def _make_index_var(typemap, scope, index_vars, body_block, force_tuple=False):
+    """ When generating a SetItem call to an array in a parfor, the general
+    strategy is to generate a tuple if the array is more than 1 dimension.
+    If it is 1 dimensional then you can use a simple variable.  This routine
+    is also used when converting pndindex to parfor but pndindex requires a
+    tuple even if the iteration space is 1 dimensional.  The pndindex use of
+    this function will use force_tuple to make the output index a tuple even
+    if it is one dimensional.
+    """
     ndims = len(index_vars)
     loc = body_block.loc
-    if ndims > 1:
+    if ndims > 1 or force_tuple:
         tuple_var = ir.Var(scope, mk_unique_var(
             "$parfor_index_tuple_var"), loc)
         typemap[tuple_var.name] = types.containers.UniTuple(
@@ -2446,11 +2454,20 @@ def find_mask_from_size(size_var):
                             in_arr, mask_var, mask_typ, mask_indices = result
                         else:
                             in_arr = args[0]
-                        size_vars = equiv_set.get_shape(in_arr
-                                        if mask_indices is None else mask_var)
-                        index_vars, loops = _mk_parfor_loops(
-                            pass_states.typemap, size_vars, scope, loc,
-                        )
+                        assert(isinstance(in_arr, ir.Var))
+                        in_arr_typ = pass_states.typemap[in_arr.name]
+                        if isinstance(in_arr_typ, types.Integer):
+                            index_var = ir.Var(scope, mk_unique_var("parfor_index"), loc)
+                            pass_states.typemap[index_var.name] = types.uintp
+                            loops = [LoopNest(index_var, 0, in_arr, 1)]
+                            index_vars = [index_var]
+                        else:
+                            size_vars = equiv_set.get_shape(in_arr
+                                          if mask_indices is None else mask_var)
+                            index_vars, loops = _mk_parfor_loops(
+                                pass_states.typemap, size_vars, scope, loc,
+                            )
+                        assert(len(loops) > 0)
                         orig_index = index_vars
                         if mask_indices:
                             # replace mask indices if required;
@@ -2462,6 +2479,7 @@ def find_mask_from_size(size_var):
                         body_block = ir.Block(scope, loc)
                         index_var, index_var_typ = _make_index_var(
                             pass_states.typemap, scope, index_vars, body_block,
+                            force_tuple=True
                         )
                         body = body_block.body + first_body_block.body
                         first_body_block.body = body

diff --git a/numba/tests/test_parfors.py b/numba/tests/test_parfors.py
@@ -1850,6 +1850,88 @@ def test_impl(x):
         x = np.ones((2, 2, 2, 2, 2, 15))
         self.check(test_impl, x)
 
+    def test_tuple_arg(self):
+        def test_impl(x, sz):
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (10, 5)
+        self.check(test_impl, np.empty(sz), sz)
+
+    def test_tuple_arg_not_whole_array(self):
+        def test_impl(x, sz):
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (10, 5)
+        self.check(test_impl, np.zeros(sz), (10, 3))
+
+    def test_tuple_for_pndindex(self):
+        def test_impl(x):
+            sz = (10, 5)
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (10, 5)
+        self.check(test_impl, np.zeros(sz))
+
+    def test_tuple_arg_literal(self):
+        def test_impl(x, first):
+            sz = (first, 5)
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (10, 5)
+        self.check(test_impl, np.zeros(sz), 10)
+
+    def test_tuple_of_literal_nonliteral(self):
+        # This test has to be done manually as the self.check uses
+        # compile_isolated and one function cannot "see" the other
+
+        def test_impl(x, sz):
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+
+        def call(x, fn):
+            return fn(x, (10, 3)) # Only want to iterate to the 3rd
+
+        get_input = lambda: np.zeros((10, 10))
+        expected = call(get_input(), test_impl)
+
+        def check(dec):
+            f1 = dec(test_impl)
+            f2 = njit(call) # no parallel semantics in the caller
+            got = f2(get_input(), f1)
+            self.assertPreciseEqual(expected, got)
+
+        for d in (njit, njit(parallel=True)):
+            check(d)
+
+    def test_tuple_arg_1d(self):
+        def test_impl(x, sz):
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (10,)
+        self.check(test_impl, np.zeros(sz), sz)
+
+    def test_tuple_arg_1d_literal(self):
+        def test_impl(x):
+            sz = (10,)
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (10,)
+        self.check(test_impl, np.zeros(sz))
+
+    def test_int_arg_pndindex(self):
+        def test_impl(x, sz):
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        self.check(test_impl, np.zeros((10, 10)), 3)
+
 
 @skip_parfors_unsupported
 class TestParforsLeaks(MemoryLeakMixin, TestParforsBase):