Repository: numba/numba

Issue:
NumbaIRAssumptionWarning variable not in scope
Thanks to all the Numba devs. This is a great library! 😄 

I've been consistently getting `NumbaIRAssumptionWarning` in my library [galois](https://github.com/mhostetter/galois) since Numba 0.54.0. I previously opened #7623 and I believe this is related.

I've copied code out of my library into this minimal example. I'm calling a JIT function inside a ufunc. While both Numba 0.54.1 and 0.55.1 produce the correct result in *this* example, I've found times in my library where it is *not* producing the correct result. I fear the "variable not in scope" warnings are true and variables are being updated outside of the function scope.

**Q: What does `NumbaIRAssumptionWarning` mean in general?**
I'm trying to debug these warnings from my library and I'm having the hardest time.

**Q: What does `NumbaIRAssumptionWarning` mean in the Numba 0.55.1 case from my example?**
What's the difference between `_power_calculate_v6_b_3` and `_power_calculate_v6_b_4`? How do I know which variable (or one which line) it is referring?

For reference, issues I opened in my repo:
* https://github.com/mhostetter/galois/issues/178
* https://github.com/mhostetter/galois/issues/226

### Minimal example

```python
import numpy as np
import numba
print(np.__version__)
print(numba.__version__)


@numba.extending.register_jitable(inline="always")
def _multiply_calculate(a, b, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY):
    """
    a in GF(2^m), can be represented as a degree m-1 polynomial a(x) in GF(2)[x]
    b in GF(2^m), can be represented as a degree m-1 polynomial b(x) in GF(2)[x]
    p(x) in GF(2)[x] with degree m is the irreducible polynomial of GF(2^m)

    a * b = c
        = (a(x) * b(x)) % p(x) in GF(2)
        = c(x)
        = c
    """
    ORDER = CHARACTERISTIC**DEGREE

    # Re-order operands such that a > b so the while loop has less loops
    if b > a:
        a, b = b, a

    c = 0
    while b > 0:
        if b & 0b1:
            c ^= a  # Add a(x) to c(x)

        b >>= 1  # Divide b(x) by x
        a <<= 1  # Multiply a(x) by x
        if a >= ORDER:
            a ^= IRREDUCIBLE_POLY  # Compute a(x) % p(x)

    return c


MULTIPLY_JIT = numba.jit(["int64(int64, int64, int64, int64, int64)"], nopython=True, cache=True)(_multiply_calculate)


@numba.extending.register_jitable(inline="always")
def _power_calculate(a, b, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY):
    """
    Square and Multiply Algorithm

    a^13 = (1) * (a)^13
            = (a) * (a)^12
            = (a) * (a^2)^6
            = (a) * (a^4)^3
            = (a * a^4) * (a^4)^2
            = (a * a^4) * (a^8)
            = result_m * result_s
    """
    if a == 0 and b < 0:
        raise ZeroDivisionError("Cannot compute the multiplicative inverse of 0 in a Galois field.")

    if b == 0:
        return 1
    # elif b < 0:
    #     a = RECIPROCAL(a, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY)
    #     b = abs(b)

    result_s = a  # The "squaring" part
    result_m = 1  # The "multiplicative" part

    while b > 1:
        if b % 2 == 0:
            result_s = MULTIPLY_JIT(result_s, result_s, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY)
            b //= 2
        else:
            result_m = MULTIPLY_JIT(result_m, result_s, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY)
            b -= 1

    result = MULTIPLY_JIT(result_m, result_s, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY)

    return result


characteristic = 2
degree = 2
irreducible_poly = 7

POWER_UFUNC = numba.vectorize(["int64(int64, int64)"], nopython=True)(lambda aa, bb: _power_calculate(aa, bb, characteristic, degree, irreducible_poly))

x = np.array([0, 1, 2, 3])  # Elements
y = np.array([0, 1, 2, 3])  # Powers

z = POWER_UFUNC.outer(x, y)
print(z)

z_truth = np.array([[1, 0, 0, 0], [1, 1, 1, 1], [1, 2, 3, 1], [1, 3, 2, 1]])
print(z_truth)

print(np.array_equal(z, z_truth))
```

### Numba 0.54.1

```
1.20.3
0.54.1
[[1 0 0 0]
 [1 1 1 1]
 [1 2 3 1]
 [1 3 2 1]]
[[1 0 0 0]
 [1 1 1 1]
 [1 2 3 1]
 [1 3 2 1]]
True
```

### Numba 0.55.1

```
1.20.3
0.55.1
/home/matt/.local/lib/python3.8/site-packages/numba/core/ssa.py:272: NumbaIRAssumptionWarning: variable '_power_calculate_v6_result_s_3' is not in scope.

This warning came from an internal pedantic check. Please report the warning
message and traceback, along with a minimal reproducer at:
https://github.com/numba/numba/issues/new?template=bug_report.md

  warnings.warn(errors.NumbaIRAssumptionWarning(wmsg,
/home/matt/.local/lib/python3.8/site-packages/numba/core/ssa.py:272: NumbaIRAssumptionWarning: variable '_power_calculate_v6_b_4' is not in scope.

This warning came from an internal pedantic check. Please report the warning
message and traceback, along with a minimal reproducer at:
https://github.com/numba/numba/issues/new?template=bug_report.md

  warnings.warn(errors.NumbaIRAssumptionWarning(wmsg,
/home/matt/.local/lib/python3.8/site-packages/numba/core/ssa.py:272: NumbaIRAssumptionWarning: variable '_power_calculate_v6_result_m_3' is not in scope.

This warning came from an internal pedantic check. Please report the warning
message and traceback, along with a minimal reproducer at:
https://github.com/numba/numba/issues/new?template=bug_report.md

  warnings.warn(errors.NumbaIRAssumptionWarning(wmsg,
/home/matt/.local/lib/python3.8/site-packages/numba/core/ssa.py:272: NumbaIRAssumptionWarning: variable '_power_calculate_v6_result_s_2' is not in scope.

This warning came from an internal pedantic check. Please report the warning
message and traceback, along with a minimal reproducer at:
https://github.com/numba/numba/issues/new?template=bug_report.md

  warnings.warn(errors.NumbaIRAssumptionWarning(wmsg,
/home/matt/.local/lib/python3.8/site-packages/numba/core/ssa.py:272: NumbaIRAssumptionWarning: variable '_power_calculate_v6_b_3' is not in scope.

This warning came from an internal pedantic check. Please report the warning
message and traceback, along with a minimal reproducer at:
https://github.com/numba/numba/issues/new?template=bug_report.md

  warnings.warn(errors.NumbaIRAssumptionWarning(wmsg,
/home/matt/.local/lib/python3.8/site-packages/numba/core/ssa.py:272: NumbaIRAssumptionWarning: variable '_power_calculate_v6_result_m_2' is not in scope.

This warning came from an internal pedantic check. Please report the warning
message and traceback, along with a minimal reproducer at:
https://github.com/numba/numba/issues/new?template=bug_report.md

  warnings.warn(errors.NumbaIRAssumptionWarning(wmsg,
[[1 0 0 0]
 [1 1 1 1]
 [1 2 3 1]
 [1 3 2 1]]
[[1 0 0 0]
 [1 1 1 1]
 [1 2 3 1]
 [1 3 2 1]]
True
```


Discussion:
Thanks for the report. MWR:
```python
import numba
from numba import njit
print(numba.__version__)

@numba.extending.register_jitable(inline="always")
def bar():

    b = 3
    while b > 1:
        b //= 2

    return 5

@njit
def foo():
    bar()


foo()

```
I think the issue is triggered by a bug in the implementation of `inline='always'`. I suspect that the inlining algorithms run to inject the IR from (in the MWR) `bar` into `foo` at the call site and then that warning is emitted later during SSA reconstruction, probably because the scope of the variable(s) hasn't been correctly transferred/updated during inlining.

> Q: What does NumbaIRAssumptionWarning mean in general?
> I'm trying to debug these warnings from my library and I'm having the hardest time.

It means that some assumption has been made when constructing or manipulating Numba's intermediate representation (IR). It's reasonably likely that from a user perspective the code will still work and it's just that the IR is for some reason technically invalid, but this it not always the case.

> Q: What does NumbaIRAssumptionWarning mean in the Numba 0.55.1 case from my example?

See the above note RE inlining and scope transfer.

> What's the difference between _power_calculate_v6_b_3 and _power_calculate_v6_b_4? How do I know which variable (or one which line) it is referring?

These are the names in the IR given to variable `b` in the function `_power_calculate`. Numba uses static single assignment form for it's IR, the `v6` and `_3` etc refer to various types of versioning of the variable.
@mhostetter Question: In your code, what is the reason for using `inline='always'`? xref developer docs with notes on uses etc: https://numba.readthedocs.io/en/stable/developer/inlining.html#notes-on-inlining
RE:
> I've found times in my library where it is not producing the correct result.

If you have one of these cases that fail and do so on with Numba 0.55.0+, a code sample would be hugely appreciated, many thanks! In 0.55.0+ this was merged https://github.com/numba/numba/pull/7552 which I think would help with making the code work correctly even if the scopes are not valid.
Upon further examination, the "incorrect result" cases I've found so far have only manifest in my library using Numba 0.54.1 and not Numba 0.55.1. However, both versions throw the `NumbaIRAssumptionWarning`s.

On a positive note, when changing my library from using `@numba.extending.register_jitable(inline="always")` to `@numba.extending.register_jitable`, the `NumbaIRAssumptionWarning`s are suppressed from both 0.54.1 and 0.55.1 *and* the incorrect result in 0.54.1 is resolved.

> @mhostetter Question: In your code, what is the reason for using inline='always'? xref developer docs with notes on uses etc: https://numba.readthedocs.io/en/stable/developer/inlining.html#notes-on-inlining

@stuartarchibald to be honest, I don't know. I thought I needed to "inline" the code so it could be used in a lambda function, like this:

```python
@numba.extending.register_jitable(inline="always")
def _power_calculate(aa, bb, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY):
    ...

POWER_UFUNC = numba.vectorize(["int64(int64, int64)"], nopython=True)(lambda aa, bb: _power_calculate(aa, bb, 2, 2, 7))
```

> If you have one of these cases that fail and do so on with Numba 0.55.0+, a code sample would be hugely appreciated, many thanks! In 0.55.0+ this was merged #7552 which I think would help with making the code work correctly even if the scopes are not valid.

I will keep a lookout for such an example and share if found.
> Upon further examination, the "incorrect result" cases I've found so far have only manifest in my library using Numba 0.54.1 and not Numba 0.55.1. However, both versions throw the `NumbaIRAssumptionWarning`s.

Thanks for checking, what you describe correlates roughly with areas of the code base that targeted fixes for these issues between 0.54 and 0.55, so I think what you are observing makes sense.

> On a positive note, when changing my library from using `@numba.extending.register_jitable(inline="always")` to `@numba.extending.register_jitable`, the `NumbaIRAssumptionWarning`s are suppressed from both 0.54.1 and 0.55.1 _and_ the incorrect result in 0.54.1 is resolved.

Indeed, the offending piece of code is in the inliner part of the compiler, which won't be run unless `inline` is requested.
 
> > @mhostetter Question: In your code, what is the reason for using inline='always'? xref developer docs with notes on uses etc: https://numba.readthedocs.io/en/stable/developer/inlining.html#notes-on-inlining
> 
> @stuartarchibald to be honest, I don't know. I thought I needed to "inline" the code so it could be used in a lambda function, like this:
> 
> ```python
> @numba.extending.register_jitable(inline="always")
> def _power_calculate(aa, bb, CHARACTERISTIC, DEGREE, IRREDUCIBLE_POLY):
>     ...
> 
> POWER_UFUNC = numba.vectorize(["int64(int64, int64)"], nopython=True)(lambda aa, bb: _power_calculate(aa, bb, 2, 2, 7))
> ```

I see, I don't think that's the case, but am also trying to recall if there would be anything similar that might trigger an issue.
 
> > If you have one of these cases that fail and do so on with Numba 0.55.0+, a code sample would be hugely appreciated, many thanks! In 0.55.0+ this was merged #7552 which I think would help with making the code work correctly even if the scopes are not valid.
> 
> I will keep a lookout for such an example and share if found.

Great, thank you. Fixing this scoping problem could well be hard so having an actual failing case would act as good motivation!


Resolving patch:
diff --git a/numba/core/inline_closurecall.py b/numba/core/inline_closurecall.py
--- a/numba/core/inline_closurecall.py
+++ b/numba/core/inline_closurecall.py
@@ -381,7 +381,7 @@ def copy_ir(the_ir):
             self.validator(callee_ir)
 
         # save an unmutated copy of the callee_ir to return
-        callee_ir_original = callee_ir.copy()
+        callee_ir_original = copy_ir(callee_ir)
         scope = block.scope
         instr = block.body[i]
         call_expr = instr.value
@@ -407,11 +407,17 @@ def copy_ir(the_ir):
         assert(len(callee_scopes) == 1)
         callee_scope = callee_scopes[0]
         var_dict = {}
-        for var in callee_scope.localvars._con.values():
+        for var in tuple(callee_scope.localvars._con.values()):
             if not (var.name in callee_freevars):
                 inlined_name = _created_inlined_var_name(
                     callee_ir.func_id.unique_name, var.name)
+                # Update the caller scope with the new names
                 new_var = scope.redefine(inlined_name, loc=var.loc)
+                # Also update the callee scope with the new names. Should the
+                # type and call maps need updating (which requires SSA form) the
+                # transformation to SSA is valid as the IR object is internally
+                # consistent.
+                callee_scope.redefine(inlined_name, loc=var.loc)
                 var_dict[var.name] = new_var
         self.debug_print("var_dict = ", var_dict)
         replace_vars(callee_blocks, var_dict)


Test patch:
diff --git a/numba/tests/test_ir_inlining.py b/numba/tests/test_ir_inlining.py
--- a/numba/tests/test_ir_inlining.py
+++ b/numba/tests/test_ir_inlining.py
@@ -4,6 +4,7 @@
 """
 
 import operator
+import warnings
 from itertools import product
 import numpy as np
 
@@ -29,7 +30,8 @@
 from numba.core.typing import signature
 from numba.tests.support import (TestCase, unittest, skip_py38_or_later,
                                  MemoryLeakMixin, IRPreservingTestPipeline,
-                                 skip_parfors_unsupported)
+                                 skip_parfors_unsupported,
+                                 ignore_internal_warnings)
 
 
 # this global has the same name as the global in inlining_usecases.py, it
@@ -1044,8 +1046,7 @@ def foo(dtype):
         self.check(test_impl, dtype, inline_expect={'foo': True})
 
     def test_inline_always_ssa(self):
-        """make sure IR inlining uses SSA properly. Test for #6721.
-        """
+        # Make sure IR inlining uses SSA properly. Test for #6721.
 
         dummy_true = True
 
@@ -1070,6 +1071,33 @@ def impl():
 
         self.check(impl, block_count='SKIP', inline_expect={'foo': True})
 
+    def test_inline_always_ssa_scope_validity(self):
+        # Make sure IR inlining correctly updates the scope(s). See #7802
+
+        def bar():
+            b = 5
+            while b > 1:
+                b //= 2
+
+            return 10
+
+        @overload(bar, inline="always")
+        def bar_impl():
+            return bar
+
+        @njit
+        def foo():
+            bar()
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always', errors.NumbaIRAssumptionWarning)
+            ignore_internal_warnings()
+            self.assertEqual(foo(), foo.py_func())
+
+        # There should be no warnings as the IR scopes should be consistent with
+        # the IR involved.
+        self.assertEqual(len(w), 0)
+
 
 class TestOverloadMethsAttrsInlining(InliningBase):
     def setUp(self):