Fix slicing for openmp target mappings (#38)

- Pass start of NameSlice to array*data start of slice value.

- Add two slicing tests.

- Fix not equal assertion.

- Change array data struct transfer to use start index and number of elements to match the way the runtime does it.

- Better naming and comments for the tests that do slices of arrays and access either the part transferred or not.

- Remove commented code that isn't needed anymore.  Give better
explanation of why we need lowerer.loadvar for variables created inside
the lowerer.

Author: DrTodd13

numba/openmp.py

@@ -1429,59 +1429,55 @@ def add_struct_tags(self, var_table):
                             the_slice = cur_tag.arg.the_slice[0][0]
                             assert the_slice.step is None
                             if isinstance(the_slice.start, int):
-                                omp_slice_start = the_slice.start
+                                start_index_var = ir.Var(None, f"{cur_tag_var}_start_index_var{target_num}{uniqueness}", array_var.loc)
+                                start_assign = ir.Assign(ir.Const(the_slice.start, array_var.loc), start_index_var, array_var.loc)
+
+                                typemap[start_index_var.name] = types.int64
+                                lowerer.lower_inst(start_assign)
+                                extras_before.append(start_assign)
+                                lowerer._alloca_var(start_index_var.name, typemap[start_index_var.name])
+                                lowerer.loadvar(start_index_var.name)
                             else:
-                                omp_slice_start = the_slice.start
-                                assert isinstance(omp_slice_start, str)
-                                omp_slice_start = ir.Var(None, omp_slice_start, array_var.loc)
+                                start_index_var = the_slice.start
+                                assert isinstance(start_index_var, str)
+                                start_index_var = ir.Var(None, start_index_var, array_var.loc)
                             if isinstance(the_slice.stop, int):
-                                size_var = the_slice.stop
+                                end_index_var = ir.Var(None, f"{cur_tag_var}_end_index_var{target_num}{uniqueness}", array_var.loc)
+                                end_assign = ir.Assign(ir.Const(the_slice.stop, array_var.loc), end_index_var, array_var.loc)
+                                typemap[end_index_var.name] = types.int64
+                                lowerer.lower_inst(end_assign)
+                                extras_before.append(end_assign)
+                                lowerer._alloca_var(end_index_var.name, typemap[end_index_var.name])
+                                lowerer.loadvar(end_index_var.name)
                             else:
-                                size_var = the_slice.stop
-                                assert isinstance(size_var, str)
-                                size_var = ir.Var(None, size_var, array_var.loc)
+                                end_index_var = the_slice.stop
+                                assert isinstance(end_index_var, str)
+                                end_index_var = ir.Var(None, end_index_var, array_var.loc)
+
+                            num_elements_var = ir.Var(None, f"{cur_tag_var}_num_elements_var{target_num}{uniqueness}", array_var.loc)
+                            size_binop = ir.Expr.binop(operator.sub, end_index_var, start_index_var, array_var.loc)
+                            size_assign = ir.Assign(size_binop, num_elements_var, array_var.loc)
+                            calltypes[size_binop] = typing.signature(types.int64, types.int64, types.int64)
                         else:
-                            omp_slice_start = 0
-                            size_var = ir.Var(None, f"{cur_tag_var}_size_var{target_num}{uniqueness}", array_var.loc)
-                            #size_var = array_var.scope.redefine("size_var", array_var.loc)
+                            start_index_var = 0
+                            num_elements_var = ir.Var(None, f"{cur_tag_var}_num_elements_var{target_num}{uniqueness}", array_var.loc)
                             size_getattr = ir.Expr.getattr(array_var, "size", array_var.loc)
-                            size_assign = ir.Assign(size_getattr, size_var, array_var.loc)
-                            typemap[size_var.name] = types.int64
-                            lowerer.lower_inst(size_assign)
-                            extras_before.append(size_assign)
-                            lowerer._alloca_var(size_var.name, typemap[size_var.name])
-
-                            # see core/datamodel/models.py
-                            loaded_size = lowerer.loadvar(size_var.name)
-                            loaded_op = loaded_size.operands[0]
-                            loaded_pointee = loaded_op.type.pointee
-                            loaded_str = str(loaded_pointee) + " * " + loaded_size._get_reference()
-
-                        #--------
-                        """
-                        itemsize_var = ir.Var(None, cur_tag_var + "_itemsize_var", array_var.loc)
-                        itemsize_getattr = ir.Expr.getattr(array_var, "itemsize", array_var.loc)
-                        itemsize_assign = ir.Assign(itemsize_getattr, itemsize_var, array_var.loc)
-                        typemap[itemsize_var.name] = types.int64
-                        lowerer.lower_inst(itemsize_assign)
-                        #--------
-
-                        totalsize_var = ir.Var(None, cur_tag_var + "_totalsize_var", array_var.loc)
-                        totalsize_binop = ir.Expr.binop(operator.mul, size_var, itemsize_var, array_var.loc)
-                        totalsize_assign = ir.Assign(totalsize_binop, totalsize_var, array_var.loc)
-                        calltypes[totalsize_binop] = typing.signature(types.int64, types.int64, types.int64)
-                        typemap[totalsize_var.name] = types.int64
-                        lowerer.lower_inst(totalsize_assign)
-                        #--------
-                        """
+                            size_assign = ir.Assign(size_getattr, num_elements_var, array_var.loc)
+
+                        typemap[num_elements_var.name] = types.int64
+                        lowerer.lower_inst(size_assign)
+                        extras_before.append(size_assign)
+                        lowerer._alloca_var(num_elements_var.name, typemap[num_elements_var.name])
+
+                        # see core/datamodel/models.py
+                        lowerer.loadvar(num_elements_var.name) # alloca the var
 
                         # see core/datamodel/models.py
-                        if isinstance(size_var, ir.Var):
-                            loaded_size = lowerer.loadvar(size_var.name)
-                            loaded_op = loaded_size.operands[0]
-                            loaded_pointee = loaded_op.type.pointee
-                            loaded_str = str(loaded_pointee) + " * " + loaded_size._get_reference()
-                        struct_tags.append(openmp_tag(cur_tag.name + ".STRUCT", cur_tag_var + "*data", non_arg=True, omp_slice=(0, size_var)))
+                        if isinstance(start_index_var, ir.Var):
+                            lowerer.loadvar(start_index_var.name) # alloca the var
+                        if isinstance(num_elements_var, ir.Var):
+                            lowerer.loadvar(num_elements_var.name) # alloca the var
+                        struct_tags.append(openmp_tag(cur_tag.name + ".STRUCT", cur_tag_var + "*data", non_arg=True, omp_slice=(start_index_var, num_elements_var)))
                         struct_tags.append(openmp_tag("QUAL.OMP.MAP.TO.STRUCT", cur_tag_var + "*shape", non_arg=True, omp_slice=(0, 1)))
                         struct_tags.append(openmp_tag("QUAL.OMP.MAP.TO.STRUCT", cur_tag_var + "*strides", non_arg=True, omp_slice=(0, 1)))
                         # Peel off NameSlice, it served its purpose and is not

numba/tests/test_openmp.py

@@ -3575,6 +3575,43 @@ def test_impl(n):
         r = test_impl(n)
         np.testing.assert_array_equal(r, np.full(n, 42))
 
+    def target_map_slice_in_mapping(self, device):
+        target_pragma = f"target device({device}) map(a[50:100]) map(to: b[100:150])"
+        @njit
+        def test_impl(n):
+            a = np.zeros(n)
+            b = np.arange(n)
+            with openmp(target_pragma):
+                for i in range(50):
+                    # These b accesses are within the transferred region.
+                    a[i+50] = b[i+100]
+            return a
+        n = 200
+        r = test_impl(n)
+        np.testing.assert_array_equal(r[0:50], np.zeros(50))
+        np.testing.assert_array_equal(r[50:100], np.arange(n)[100:150])
+        np.testing.assert_array_equal(r[100:200], np.zeros(100))
+
+    def target_map_slice_read_out_mapping(self, device):
+        target_pragma = f"target device({device}) map(a[50:100]) map(to: b[100:150])"
+        @njit
+        def test_impl(n):
+            a = np.zeros(n)
+            b = np.arange(n)
+            with openmp(target_pragma):
+                for i in range(50):
+                    # These b accesses are outside the transferred region.
+                    # Should get whatever happens to be in memory at that point.
+                    # We assume that isn't arange(50:100).
+                    a[i+50] = b[i+50]
+            return a
+        n = 200
+        r = test_impl(n)
+        np.testing.assert_array_equal(r[0:50], np.zeros(50))
+        # Make sure that the range 50-100 was not transferred.
+        assert not np.array_equal(r[50:100], np.arange(n)[50:100])
+        np.testing.assert_array_equal(r[100:200], np.zeros(100))
+
     def target_map_tofrom_array(self, device):
         target_pragma = f"target device({device}) map(tofrom: a)"
         @njit