Fix collapse loop (#30)

* Updates.

* Redirect TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP.SIMD to TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP.

Author: DrTodd13

numba/openmp.py

@@ -2809,8 +2809,6 @@ def get_dotted_type(x, typemap, lowerer):
 
 
 def is_target_arg(name):
-    #return name in ["QUAL.OMP.FIRSTPRIVATE", "QUAL.OMP.TARGET.IMPLICIT", "QUAL.OMP.THREAD_LIMIT", "QUAL.OMP.NUM_TEAMS"] or name.startswith("QUAL.OMP.MAP")
-    #or name.startswith("QUAL.OMP.NORMALIZED")
     return name in ["QUAL.OMP.FIRSTPRIVATE", "QUAL.OMP.TARGET.IMPLICIT"] or name.startswith("QUAL.OMP.MAP") or name.startswith("QUAL.OMP.REDUCTION")
 
 
@@ -2824,7 +2822,6 @@ def is_pointer_target_arg(name, typ):
     if name.startswith("QUAL.OMP.MAP"):
         if isinstance(typ, types.npytypes.Array):
             return True
-            #return False
         else:
             return True
     if name in ["QUAL.OMP.FIRSTPRIVATE", "QUAL.OMP.PRIVATE"]:
@@ -3408,15 +3405,25 @@ def get_loops_in_region(all_loops):
 
                 # Copy all stmts from the loop entry block up to the ir.Global
                 # for range.
+                call_offset = None
                 for entry_block_index, stmt in enumerate(loop_entry_block.body):
+                    found_range = False
                     if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Global) and stmt.value.name == "range":
+                        found_range = True
                         range_target = stmt.target
-                        call_stmt = loop_entry_block.body[entry_block_index + 1]
-                        assert isinstance(call_stmt, ir.Assign) and isinstance(call_stmt.value, ir.Expr) and call_stmt.value.op == 'call' and call_stmt.value.func == range_target
-                        # Remove stmts that were retained.
-                        loop_entry_block.body = loop_entry_block.body[entry_block_index:]
+                        found_call = False
+                        for call_index in range(entry_block_index + 1, len(loop_entry_block.body)):
+                            call_stmt = loop_entry_block.body[call_index]
+                            if isinstance(call_stmt, ir.Assign) and isinstance(call_stmt.value, ir.Expr) and call_stmt.value.op == 'call' and call_stmt.value.func == range_target:
+                                found_call = True
+                                # Remove stmts that were retained.
+                                loop_entry_block.body = loop_entry_block.body[entry_block_index:]
+                                call_offset = call_index - entry_block_index
+                                break
+                        assert found_call
                         break
                     stmts_to_retain.append(stmt)
+                assert found_range
                 for header_block_index, stmt in enumerate(loop_header_block.body):
                     if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == "iternext":
                         iternext_inst = loop_header_block.body[header_block_index]
@@ -3475,7 +3482,7 @@ def get_loops_in_region(all_loops):
                 new_stmts_for_iterspace.append(ir.Assign(mul_op, new_iterspace_var, self.loc))
             # Change iteration space of innermost loop to the product of all the
             # loops' iteration spaces.
-            last_loop_entry_block.body[1].value.args[0] = new_iterspace_var
+            last_loop_entry_block.body[call_offset].value.args[0] = new_iterspace_var
 
             last_eliminated_loop_header_block.body = new_stmts_for_iterspace + last_eliminated_loop_header_block.body
 
@@ -4324,14 +4331,16 @@ def target_teams_distribute_directive(self, args):
         self.some_target_directive(args, "TARGET.TEAMS.DISTRIBUTE", 3, has_loop=True)
 
     def target_teams_loop_directive(self, args):
-        self.some_target_directive(args, "TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP.SIMD", 3, has_loop=True)
+        self.some_target_directive(args, "TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP", 3, has_loop=True)
+        #self.some_target_directive(args, "TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP.SIMD", 3, has_loop=True)
         #self.some_target_directive(args, "TARGET.TEAMS.LOOP", 3, has_loop=True)
 
     def target_teams_distribute_parallel_for_directive(self, args):
         self.some_target_directive(args, "TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP", 5, has_loop=True)
 
     def target_teams_distribute_parallel_for_simd_directive(self, args):
-        self.some_target_directive(args, "TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP.SIMD", 6, has_loop=True)
+        # Intentionally dropping "SIMD" from string as that typically isn't implemented on GPU.
+        self.some_target_directive(args, "TARGET.TEAMS.DISTRIBUTE.PARALLEL.LOOP", 6, has_loop=True)
 
     def get_clauses_by_name(self, clauses, names, remove_from_orig=False):
         if not isinstance(names, list):

numba/tests/test_openmp.py

@@ -4390,6 +4390,25 @@ def test_impl():
         else:
             raise ValueError(f"Device {device} must be 0 or 1")
 
+    def target_teams_loop_collapse(self, device):
+        target_pragma = f"""target teams loop collapse(2)
+                        device({device})
+                        map(tofrom: a, b, c)"""
+        @njit
+        def test_impl(n):
+            a = np.ones((n,n))
+            b = np.ones((n,n))
+            c = np.zeros((n,n))
+            with openmp(target_pragma):
+                for i in range(n):
+                    for j in range(n):
+                        c[i,j] = a[i,j] + b[i,j]
+            return c
+
+        n = 10
+        c = test_impl(n)
+        np.testing.assert_array_equal(c, np.full((n,n), 2))
+
 
 for memberName in dir(TestOpenmpTarget):
     if memberName.startswith("target"):