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,10 @@ 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):
     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(
@@ -2462,6 +2462,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
@@ -24,7 +24,7 @@
 import subprocess as subp
 
 import numba.parfors.parfor
-from numba import njit, prange, set_num_threads, get_num_threads, typeof
+from numba import njit, prange, set_num_threads, get_num_threads, typeof, literally
 from numba.core import (types, utils, typing, errors, ir, rewrites,
                         typed_passes, inline_closurecall, config, compiler, cpu)
 from numba.extending import (overload_method, register_model,
@@ -1850,6 +1850,37 @@ 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,10)
+        self.check(test_impl, np.empty(sz), sz)
+
+    def test_tuple_arg_literal(self):
+        def test_impl(x, sz):
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (10,numba.literally(10))
 def literally(obj): 
     """Forces Numba to interpret *obj* as an Literal value. 
     *obj* must be either a literal or an argument of the caller function, where 
     the argument must be bound to a literal. The literal requirement 
     propagates up the call stack. 
     This function is intercepted by the compiler to alter the compilation 
     behavior to wrap the corresponding function parameters as ``Literal``. 
     It has **no effect** outside of nopython-mode (interpreter, and objectmode). 
     The current implementation detects literal arguments in two ways: 
     1. Scans for uses of ``literally`` via a compiler pass. 
     2. ``literally`` is overloaded to raise ``numba.errors.ForceLiteralArg`` 
        to signal the dispatcher to treat the corresponding parameter 
        differently. This mode is to support indirect use (via a function call). 
     The execution semantic of this function is equivalent to an identity 
     function. 
     See :ghfile:`numba/tests/test_literal_dispatch.py` for examples. 
     """ 
     return obj 
 def literally(obj): 
     """Forces Numba to interpret *obj* as an Literal value. 
  
     *obj* must be either a literal or an argument of the caller function, where 
     the argument must be bound to a literal. The literal requirement 
     propagates up the call stack. 
  
     This function is intercepted by the compiler to alter the compilation 
     behavior to wrap the corresponding function parameters as ``Literal``. 
     It has **no effect** outside of nopython-mode (interpreter, and objectmode). 
  
     The current implementation detects literal arguments in two ways: 
  
     1. Scans for uses of ``literally`` via a compiler pass. 
     2. ``literally`` is overloaded to raise ``numba.errors.ForceLiteralArg`` 
        to signal the dispatcher to treat the corresponding parameter 
        differently. This mode is to support indirect use (via a function call). 
  
     The execution semantic of this function is equivalent to an identity 
     function. 
  
     See :ghfile:`numba/tests/test_literal_dispatch.py` for examples. 
     """ 
     return obj 
+        self.check(test_impl, np.empty(sz), sz)
+
+    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.empty(sz), sz)
+
+    def test_tuple_arg_1d_literal(self):
+        def test_impl(x, sz):
+            for i in numba.pndindex(sz):
+                x[i] = 1
+            return x
+        sz = (numba.literally(10),)
+        self.check(test_impl, np.empty(sz), sz)
 
 @skip_parfors_unsupported
 class TestParforsLeaks(MemoryLeakMixin, TestParforsBase):