Repository: numba/numba

Issue:
Parallel fails when using .real and .imag attributes with numpy functions
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

```
import numpy as np
from numba import njit

@njit(parallel=True)
def fails(z):
    return np.sum(z.real**2 + z.imag**2)

@njit(parallel=False)
def works(z):
    return np.sum(z.real**2 + z.imag**2)

@njit(parallel=True) 
def also_works(z):
    return np.sum(np.real(z)**2 + np.imag(z)**2)

n = 10
z = np.arange(n) + 1j*np.arange(n)
print(works(z))
print(also_works(z))
print(fails(z)) # this one throws an error
```

`parallel=True` fails when sending values derived from `.real` and `.imag` array attribute to a numpy reduction or math function. The example given here uses `np.sum()`, but other functions also show the same error (`np.prod`, `np.min`, `np.max`, `np.argmin`, `np.argmax`, `np.mean`, `np.var`, and `np.std`). Using `np.real()` and `np.imag()` to access the real and imaginary components, instead of `.real` and `.imag` array attributes, works as expected.
<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->



Discussion:
I get the following error. Looking into it.

```python
Traceback (most recent call last):
  File "../debug.py", line 407, in <module>
    test_parallel()
  File "../debug.py", line 405, in test_parallel
    print(fails(z)) # this one throws an error
  File "/Users/czgdp1807/numba_project/numba/numba/core/dispatcher.py", line 487, in _compile_for_args
    raise e
  File "/Users/czgdp1807/numba_project/numba/numba/core/dispatcher.py", line 420, in _compile_for_args
    return_val = self.compile(tuple(argtypes))
  File "/Users/czgdp1807/numba_project/numba/numba/core/dispatcher.py", line 965, in compile
    cres = self._compiler.compile(args, return_type)
  File "/Users/czgdp1807/numba_project/numba/numba/core/dispatcher.py", line 125, in compile
    status, retval = self._compile_cached(args, return_type)
  File "/Users/czgdp1807/numba_project/numba/numba/core/dispatcher.py", line 139, in _compile_cached
    retval = self._compile_core(args, return_type)
  File "/Users/czgdp1807/numba_project/numba/numba/core/dispatcher.py", line 157, in _compile_core
    pipeline_class=self.pipeline_class)
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler.py", line 716, in compile_extra
    return pipeline.compile_extra(func)
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler.py", line 452, in compile_extra
    return self._compile_bytecode()
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler.py", line 520, in _compile_bytecode
    return self._compile_core()
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler.py", line 499, in _compile_core
    raise e
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler.py", line 486, in _compile_core
    pm.run(self.state)
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler_machinery.py", line 353, in run
    raise patched_exception
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler_machinery.py", line 341, in run
    self._runPass(idx, pass_inst, state)
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler_machinery.py", line 296, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/Users/czgdp1807/numba_project/numba/numba/core/compiler_machinery.py", line 269, in check
    mangled = func(compiler_state)
  File "/Users/czgdp1807/numba_project/numba/numba/core/typed_passes.py", line 306, in run_pass
    parfor_pass.run()
  File "/Users/czgdp1807/numba_project/numba/numba/parfors/parfor.py", line 2877, in run
    self.array_analysis.run(self.func_ir.blocks)
  File "/Users/czgdp1807/numba_project/numba/numba/parfors/array_analysis.py", line 1178, in run
    self._run_on_blocks(topo_order, blocks, cfg, init_equiv_set)
  File "/Users/czgdp1807/numba_project/numba/numba/parfors/array_analysis.py", line 1199, in _run_on_blocks
    cfg, block, label, scope, init_equiv_set
  File "/Users/czgdp1807/numba_project/numba/numba/parfors/array_analysis.py", line 1274, in _determine_transform
    label, scope, equiv_set, inst, redefined
  File "/Users/czgdp1807/numba_project/numba/numba/parfors/array_analysis.py", line 1395, in _analyze_inst
    equiv_set, lhs, typ.ndim, shape, post
  File "/Users/czgdp1807/numba_project/numba/numba/parfors/array_analysis.py", line 3192, in _gen_shape_call
    typ = self.typemap[shape[i].name]
KeyError: "Failed in nopython mode pipeline (step: convert to parfors)\n'$8load_attr.3_size0.112'"
```
I noticed that if you deal with only `.real` or `.imag` and don't do any operation with them (like raising to a power) then no error is thrown. I think the problem is with usage of ufuncs on the output of `.real`, etc.

```
def test_parallel():

    @njit(parallel=True)
    def fails(z):
        return np.sum(z.real)

    n = 10
    z = np.arange(n) + 1j*np.arange(n)
    print(fails(z)) # this one throws an error

test_parallel()
```

After some more investigation I noticed when `parallel=True` the same shape/size information is not made available in `typemap`. I think when `parallel=True` then `PreParforPass._replace_parallel_functions` is called which is responsible for filling `typemap` for the `ParforPass`.

https://github.com/numba/numba/blob/e5bbae7de22fa8c0b23dcfc8c9a946e5d37ee509/numba/parfors/parfor.py#L1419

So when I commented out the code below, the code in the OP worked fine. So, I think there is a problem with `_replace_parallel_functions`.

Resolving patch:
diff --git a/numba/parfors/array_analysis.py b/numba/parfors/array_analysis.py
--- a/numba/parfors/array_analysis.py
+++ b/numba/parfors/array_analysis.py
@@ -1579,6 +1579,10 @@ def _analyze_op_getattr(self, scope, equiv_set, expr, lhs):
         elif expr.attr == "shape":
             shape = equiv_set.get_shape(expr.value)
             return ArrayAnalysis.AnalyzeResult(shape=shape)
+        elif expr.attr in ("real", "imag") and self._isarray(expr.value.name):
+            # Shape of real or imag attr is the same as the shape of the array
+            # itself.
+            return ArrayAnalysis.AnalyzeResult(shape=expr.value)
         elif self._isarray(lhs.name):
             canonical_value = get_canonical_alias(
                 expr.value.name, self.alias_map


Test patch:
diff --git a/numba/tests/test_array_analysis.py b/numba/tests/test_array_analysis.py
--- a/numba/tests/test_array_analysis.py
+++ b/numba/tests/test_array_analysis.py
@@ -742,6 +742,17 @@ def test_transpose_3d(m, n, k):
                                        self.with_equiv('e', ('m', 'k', 'n')),
                                        self.with_equiv('et', ('m', 'k', 'n'))])
 
+        def test_real_imag_attr(m, n):
+            a = np.ones((m, n))
+            b = a.real
+            c = a.imag
+
+        self._compile_and_test(test_real_imag_attr, (types.intp, types.intp),
+                               equivs=[self.with_equiv('a', ('m', 'n')),
+                                       self.with_equiv('b', ('m', 'n')),
+                                       self.with_equiv('c', ('m', 'n')),
+                                       ])
+
         def test_random(n):
             a0 = np.random.rand(n)
             a1 = np.random.rand(n, n)
diff --git a/numba/tests/test_parfors.py b/numba/tests/test_parfors.py
--- a/numba/tests/test_parfors.py
+++ b/numba/tests/test_parfors.py
@@ -1246,6 +1246,15 @@ def test_impl(n):
             return np.sum(n) + np.prod(n) + np.min(n) + np.max(n) + np.var(n)
         self.check(test_impl, np.array(7), check_scheduling=False)
 
+    def test_real_imag_attr(self):
+        # See issue 8012
+        def test_impl(z):
+            return np.sum(z.real ** 2 + z.imag ** 2)
+
+        z = np.arange(5) * (1 + 1j)
+        self.check(test_impl, z)
+        self.assertEqual(countParfors(test_impl, (types.complex128[::1],)), 1)
+
 
 class TestParforsUnsupported(TestCase):
     """Tests for unsupported use of parfors"""