Repository: numba/numba

Issue:
Cannot lower setitem on a record in a parallel loop
## Reporting a bug

When running this code:

```python
import numba
import numpy as np

print(numba.__version__)
# => 0.52.0

state_dtype = np.dtype([
    ('var', np.int32)
 ])

@numba.njit(inline=True)
def inner_function(state):
    state['var'] = 1

@numba.njit(parallel=True)
def outer_function(states):
    for i in numba.prange(1):
        inner_function(states[i])

outer_function(np.ones(shape=1, dtype=state_dtype))
```
numba yields 
```No definition for lowering <built-in function setitem>(Record(var[type=int32;offset=0];4;False), Literal[str](var), Literal[int](1)) -> none```, which seems to be the correct signature and I'm not so sure why there is no lowering available for that situation.

When disabling parallel or inlining the code compiles and produces correct results.

- [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/master/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'.

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->



Discussion:
Many thanks for the issue report - I can reproduce this, with a minor change to the reproducer:

```python
import numba
import numpy as np

print(numba.__version__)
# => 0.52.0

state_dtype = np.dtype([
    ('var', np.int32)
 ])

@numba.njit(inline='always')
def inner_function(state):
    state['var'] = 1

@numba.njit(parallel=True)
def outer_function(states):
    for i in numba.prange(1):
        inner_function(states[i])

outer_function(np.ones(shape=1, dtype=state_dtype))
```

This is because setting `inline=True` resulted in the error:

```
ValueError: Failed in nopython mode pipeline (step: inline inlinable functions)
kwarg 'inline' must be one of the strings 'always' or 'never', or it can be a callable that returns True/False. Found value True
```

so I changed it to `inline=`always`.

As a workaround, you can spell the setitem as a setattr, at least for this reproducer:

```python
import numba
import numpy as np

print(numba.__version__)
# => 0.52.0

state_dtype = np.dtype([
    ('var', np.int32)
 ])

@numba.njit(inline='always')
def inner_function(state):
    state.var = 2

@numba.njit(parallel=True)
def outer_function(states):
    for i in numba.prange(1):
        inner_function(states[i])

data = np.ones(shape=1, dtype=state_dtype)

outer_function(data)

print(data)
```

I edited the setting of data from `1` to `2` so the effect of the inlined function would be visible. This outputs for me:

```
$ python war.py 
0.53.0dev0+739.gf6e735f71
OMP: Info #274: omp_set_nested routine deprecated, please use omp_set_max_active_levels instead.
[(2,)]
```
I don't think this has anything to do with the inliner.
```python
from numba import njit, prange
import numpy as np

state_dtype = np.dtype([
    ('var', np.int32)
 ])

def outer_function(states):
    for i in prange(1):
        states[i]['var'] = 1

print("interpreter", outer_function(np.ones(shape=1, dtype=state_dtype)))
print("njit", njit(outer_function)(np.ones(shape=1, dtype=state_dtype)))
print("parallel accelerator", njit(parallel=True)(outer_function)(np.ones(shape=1, dtype=state_dtype)))
```
it's just the last line that fails, same way as above.
CC @DrTodd13 ^ thanks.
That's weird because on my machine having inline='never' compiles and run fine
> That's weird because on my machine having inline='never' compiles and run fine

The parallel transforms can't "look" through a function call, so if the function remains as a `call` by virtue of it not being inlined the problem would never be hit. Essentially the function would be compiled by "normal" jit (which demonstrably works fine) and the parallel region would just call that function.

As an aside, what's the reasoning behing inlining this function at the Numba IR level? If it's performance, LLVM would probably inline it anyway, there's some notes on what Numba inlining does and its uses here: https://numba.readthedocs.io/en/stable/developer/inlining.html#notes-on-inlining .
Thanks for the explanation @stuartarchibald it makes sense now!

Yeah it is performance. In some instances I've seen up to 30% difference. However I only benchmarked without the `parallel` flag.
@stuartarchibald I guess StaticSetItem is supported for record array writes and not SetItem?

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
@@ -1004,13 +1004,13 @@ def get_or_set(d, k):
     def set_shape_setitem(self, obj, shape):
         """remember shapes of SetItem IR nodes.
         """
-        assert isinstance(obj, ir.StaticSetItem) or isinstance(obj, ir.SetItem)
+        assert isinstance(obj, (ir.StaticSetItem, ir.SetItem))
         self.ext_shapes[obj] = shape
 
     def _get_shape(self, obj):
         """Overload _get_shape to retrieve the shape of SetItem IR nodes.
         """
-        if isinstance(obj, ir.StaticSetItem) or isinstance(obj, ir.SetItem):
+        if isinstance(obj, (ir.StaticSetItem, ir.SetItem)):
             require(obj in self.ext_shapes)
             return self.ext_shapes[obj]
         else:
@@ -1252,12 +1252,6 @@ def _determine_transform(self, cfg, block, label, scope, init_equiv_set):
         pending_transforms = []
         for inst in block.body:
             redefined = set()
-            if isinstance(inst, ir.StaticSetItem):
-                orig_calltype = self.calltypes[inst]
-                inst = ir.SetItem(
-                    inst.target, inst.index_var, inst.value, inst.loc
-                )
-                self.calltypes[inst] = orig_calltype
             pre, post = self._analyze_inst(
                 label, scope, equiv_set, inst, redefined
             )
@@ -1403,7 +1397,7 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
             if not result:
                 return [], []
             if result[0] is not None:
-                assert isinstance(inst, ir.SetItem)
+                assert isinstance(inst, (ir.StaticSetItem, ir.SetItem))
                 inst.index = result[0]
             result = result[1]
             target_shape = result.kwargs['shape']
diff --git a/numba/parfors/parfor.py b/numba/parfors/parfor.py
--- a/numba/parfors/parfor.py
+++ b/numba/parfors/parfor.py
@@ -1689,6 +1689,10 @@ def _type_getitem(self, args):
         return self.pass_states.typingctx.resolve_function_type(fnty, tuple(args), {})
 
 
+def get_index_var(x):
+    return x.index if isinstance(x, ir.SetItem) else x.index_var
+
+
 class ConvertSetItemPass:
     """Parfor subpass to convert setitem on Arrays
     """
@@ -1713,11 +1717,10 @@ def run(self, blocks):
             new_body = []
             equiv_set = pass_states.array_analysis.get_equiv_set(label)
             for instr in block.body:
-                if isinstance(instr, ir.StaticSetItem) or isinstance(instr, ir.SetItem):
+                if isinstance(instr, (ir.StaticSetItem, ir.SetItem)):
                     loc = instr.loc
                     target = instr.target
-                    index = (instr.index if isinstance(instr, ir.SetItem)
-                             else instr.index_var)
+                    index = get_index_var(instr)
                     value = instr.value
                     target_typ = pass_states.typemap[target.name]
                     index_typ = pass_states.typemap[index.name]
@@ -3435,9 +3438,9 @@ def get_parfor_outputs(parfor, parfor_params):
     outputs = []
     for blk in parfor.loop_body.values():
         for stmt in blk.body:
-            if isinstance(stmt, ir.SetItem):
-                if stmt.index.name == parfor.index_var.name:
-                    outputs.append(stmt.target.name)
+            if (isinstance(stmt, (ir.StaticSetItem, ir.SetItem)) and
+                get_index_var(stmt).name == parfor.index_var.name):
+                outputs.append(stmt.target.name)
     # make sure these written arrays are in parfor parameters (live coming in)
     outputs = list(set(outputs) & set(parfor_params))
     return sorted(outputs)
@@ -4135,8 +4138,8 @@ def remove_dead_parfor(parfor, lives, lives_n_aliases, arg_aliases, alias_map, f
             alias_lives = in_lives & alias_set
             for v in alias_lives:
                 in_lives |= alias_map[v]
-            if (isinstance(stmt, ir.SetItem) and
-                stmt.index.name == parfor.index_var.name and
+            if (isinstance(stmt, (ir.StaticSetItem, ir.SetItem)) and
+                get_index_var(stmt).name == parfor.index_var.name and
                 stmt.target.name not in in_lives and
                 stmt.target.name not in arg_aliases):
                 continue
@@ -4173,8 +4176,9 @@ def _update_parfor_get_setitems(block_body, index_var, alias_map,
     replace getitems of a previously set array in a block of parfor loop body
     """
     for stmt in block_body:
-        if (isinstance(stmt, ir.SetItem) and stmt.index.name ==
-                index_var.name and stmt.target.name not in lives):
+        if (isinstance(stmt, (ir.StaticSetItem, ir.SetItem)) and
+            get_index_var(stmt).name == index_var.name and
+            stmt.target.name not in lives):
             # saved values of aliases of SetItem target array are invalid
             for w in alias_map.get(stmt.target.name, []):
                 saved_values.pop(w, None)
diff --git a/numba/parfors/parfor_lowering.py b/numba/parfors/parfor_lowering.py
--- a/numba/parfors/parfor_lowering.py
+++ b/numba/parfors/parfor_lowering.py
@@ -717,7 +717,7 @@ def _hoist_internal(inst, dep_on_param, call_table, hoisted, not_hoisted,
 
 def find_setitems_block(setitems, itemsset, block, typemap):
     for inst in block.body:
-        if isinstance(inst, ir.StaticSetItem) or isinstance(inst, ir.SetItem):
+        if isinstance(inst, (ir.StaticSetItem, ir.SetItem)):
             setitems.add(inst.target.name)
             # If we store a non-mutable object into an array then that is safe to hoist.
             # If the stored object is mutable and you hoist then multiple entries in the


Test patch:
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
@@ -2800,6 +2800,41 @@ def test_impl():
 
         self.prange_tester(test_impl)
 
+    @skip_parfors_unsupported
+    def test_record_array_setitem(self):
+        # issue6704
+        state_dtype = np.dtype([('var', np.int32)])
+
+        def test_impl(states):
+            for i in range(1):
+                states[i]['var'] = 1
+
+        def comparer(a, b):
+            assert(a[0]['var'] == b[0]['var'])
+
+        self.prange_tester(test_impl,
+                           np.zeros(shape=1, dtype=state_dtype),
+                           check_arg_equality=[comparer])
+
+    @skip_parfors_unsupported
+    def test_record_array_setitem_yield_array(self):
+        state_dtype = np.dtype([('x', np.intp)])
+
+        def test_impl(states):
+            n = states.size
+            for i in range(states.size):
+                states["x"][i] = 7 + i
+            return states
+
+        states = np.zeros(10, dtype=state_dtype)
+
+        def comparer(a, b):
+            np.testing.assert_equal(a, b)
+
+        self.prange_tester(test_impl,
+                           states,
+                           check_arg_equality=[comparer])
+
 
 @skip_parfors_unsupported
 @x86_only