Repository: numba/numba

Issue:
Inconsistent behavior of `stencil` decorator
I observed the following strange behavior when using `stencil` as a decorator:
This function works fine (taken from the [stencil examples](https://numba.pydata.org/numba-doc/dev/user/stencil.html): 
```python
@njit(parallel=True)
def stencil_test1(A):
    c = 2
    B = stencil(
        lambda a, c: 0.3 * (a[-c+1] + a[0] + a[c-1]))(A, c)
    return B
stencil_test1(np.arange(10))
# returns array([0. , 0.9, 1.8, 2.7, 3.6, 4.5, 5.4, 6.3, 7.2, 0. ])
```

But this function, which I thought would be equivalent, generates an error at runtime:

```python
@njit(parallel=True)
def stencil_test2(A):
    @stencil
    def fun (a):
        c = 2
        return 0.3 * (a[-c+1] + a[0] + a[c-1])
    B = fun(A)
    return B
```
running `stencil_test2(np.arange(10))` generates a long error with the following message at the end:
```
LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)
stencil kernel index is not constant, 'neighborhood' option required

File "<ipython-input-454-f1261a271023>", line 7:
    def fun (a,c):
        <source elided>
        return 0.3 * (a[-c+1] + a[0] + a[c-1])
    B = fun(A,c)
    ^

[1] During: lowering "$B.9060 = call $fun.9059(A, $c.9058, func=$fun.9059, args=[Var(A, <ipython-input-454-f1261a271023> (3)), Var($c.9058, <ipython-input-454-f1261a271023> (3))], kws=(), vararg=None)" at <ipython-input-454-f1261a271023> (7)
-------------------------------------------------------------------------------
This should not have happened, a problem has occurred in Numba's internals.
You are currently using Numba version 0.45.1.
```
The full traceback can be seen at this gist: https://gist.github.com/d-v-b/630506e0b34a3ee2dd49f6e9c8d2c3b0

In communication via gitter, @stuartarchibald found that removing the `@stencil` decorator and instead calling `stencil(fun)(A,c)` worked fine.


Discussion:
Thanks for the report. This is an exact reproducer:
```python
from numba import njit, stencil
import numpy as np

# OK
@njit(parallel=True)
def stencil_test1(A):
    def fun (a):
        c = 2
        return a[-c+1]
    B = stencil(fun)(A)
    return B

stencil_test1(np.arange(10))

# Broken
@njit(parallel=True)
def stencil_test2(A):
    @stencil # decorator instead of call
    def fun (a):
        c = 2
        return a[-c+1]
    B = fun(A)
    return B

stencil_test2(np.arange(10))
```

@DrTodd13 any ideas? I vaguely recall encountering this before but am struggling to remember what went on.
@stuartarchibald @d-v-b The problem was that in the second case stencilparfor was not recognizing the stencil call because there was an intermediate Var in between.  So:
```
C = stencil(...)
D = A
call D(...)
```
This is due to a problem in get_call_table.  I have a branch that resolves this issue and I will make a PR but I'm not sure the semantics of the solution are what we would always want.  I guess we should discuss in the PR.

Resolving patch:
diff --git a/numba/ir_utils.py b/numba/ir_utils.py
--- a/numba/ir_utils.py
+++ b/numba/ir_utils.py
@@ -1088,6 +1088,13 @@ def get_call_table(blocks, call_table=None, reverse_call_table=None, topological
                     if lhs in reverse_call_table:
                         call_var = reverse_call_table[lhs]
                         call_table[call_var].append(rhs.value)
+                if isinstance(rhs, ir.Var):
+                    if lhs in call_table:
+                        call_table[lhs].append(rhs.name)
+                        reverse_call_table[rhs.name] = lhs
+                    if lhs in reverse_call_table:
+                        call_var = reverse_call_table[lhs]
+                        call_table[call_var].append(rhs.name)
             for T, f in call_table_extensions.items():
                 if isinstance(inst, T):
                     f(inst, call_table, reverse_call_table)
diff --git a/numba/stencilparfor.py b/numba/stencilparfor.py
--- a/numba/stencilparfor.py
+++ b/numba/stencilparfor.py
@@ -48,10 +48,11 @@ def run(self):
         stencil_calls = []
         stencil_dict = {}
         for call_varname, call_list in call_table.items():
-            if len(call_list)==1 and isinstance(call_list[0], StencilFunc):
-                # Remember all calls to StencilFuncs.
-                stencil_calls.append(call_varname)
-                stencil_dict[call_varname] = call_list[0]
+            for one_call in call_list:
+                if isinstance(one_call, StencilFunc):
+                    # Remember all calls to StencilFuncs.
+                    stencil_calls.append(call_varname)
+                    stencil_dict[call_varname] = one_call
         if not stencil_calls:
             return  # return early if no stencil calls found
 


Test patch:
diff --git a/numba/tests/test_stencils.py b/numba/tests/test_stencils.py
--- a/numba/tests/test_stencils.py
+++ b/numba/tests/test_stencils.py
@@ -522,6 +522,28 @@ def test_impl(A):
         cpfunc = self.compile_parallel(test_impl, (numba.float64[:],), stencil=False)
         self.assertNotIn('@do_scheduling', cpfunc.library.get_llvm_str())
 
+    @skip_unsupported
+    def test_stencil_nested1(self):
+        """Tests whether nested stencil decorator works.
+        """
+        @njit(parallel=True)
+        def test_impl(n):
+            @stencil
+            def fun(a):
+                c = 2
+                return a[-c+1]
+            B = fun(n)
+            return B
+
+        def test_impl_seq(n):
+            B = np.zeros(len(n), dtype=int)
+            for i in range(1, len(n)):
+                B[i] = n[i-1]
+            return B
+
+        n = np.arange(10)
+        np.testing.assert_equal(test_impl(n), test_impl_seq(n))
+
     @skip_unsupported
     def test_out_kwarg_w_cval(self):
         """ Issue #3518, out kwarg did not work with cval."""