Fix fetch_idx

Author: isuruf

sumpy/expansion/loopy.py

@@ -276,17 +276,17 @@ def get_domains(v, iorder, with_sync):
             domains = [f"{{ [{x0}_outer]: 0<={x0}_outer<={order//c} }}"]
             if with_sync:
                 expr = f"{c//sync_split}*{x0}_sync_outer + {c}*{x0}_outer"
-                domains += [f"{{ [{x0}_sync_outer]: 0<={expr}<=order "
+                domains += [f"{{ [{x0}_sync_outer]: 0<={expr}<={order} "
                     f"and 0<={x0}_sync_outer<{c//sync_split} }}"]
                 expr += f" + {v[0]}_inner"
-                domains += [f"{{ [{v[0]}_inner]: 0<={expr}<=order "
+                domains += [f"{{ [{v[0]}_inner]: 0<={expr}<={order} "
                     f"and 0<={v[0]}_inner<{sync_split} }}"]
             else:
                 expr = f"{v[0]}_inner + {c}*{x0}_outer"
-                domains += [f"{{ [{v[0]}_inner]: 0<={expr}<=order "
+                domains += [f"{{ [{v[0]}_inner]: 0<={expr}<={order} "
                     f"and 0<={v[0]}_inner<{c} }}"]
             domains += [f"{{ [{v[0]}]: {expr}<={v[0]}<={expr} }}"]
-            domains += [f"{{ [{iorder}]: {v[0]}<={iorder}<=order }}"]
+            domains += [f"{{ [{iorder}]: {v[0]}<={iorder}<={order} }}"]
             upper_bound = f"{iorder}-{v[0]}"
             for i in range(dim - 1, 1, -1):
                 domains += [f"{{ [{v[i]}]: 0<={v[i]}<={upper_bound} }}"]
@@ -313,7 +313,6 @@ def get_idx(v):
             # a synchronization step.
             prev_copy_idx = (v[0]//c - 1) % 2
             curr_copy_idx = (v[0]//c) % 2
-            fetch_idx = (v[0]//c) % 2
         else:
             # We need to sync within the c rows.
             # Using the biharmonic 2D example:
@@ -330,7 +329,6 @@ def get_idx(v):
             # - Read the rows 4, 5, 6, 7 from coeffs_copy[0, :]
             prev_copy_idx = 0
             curr_copy_idx = 1
-            fetch_idx = 0
 
         max_mi_sym = [v[i] - max_mi[i] for i in range(dim)]
         scale = -1/deriv_id_to_coeff[max_deriv_id]
@@ -377,6 +375,15 @@ def get_idx(v):
         idx = get_idx(v)
         domains += get_domains(v, iorder, with_sync=False)[1:]
 
+        if c == sync_split:
+            # We did not have to sync within the c rows.
+            # We last wrote to coeffs_copy[v[0]//c % 2, :] and we read from it.
+            fetch_idx = (v[0]//c) % 2
+        else:
+            # We need to sync within the c rows.
+            # We last wrote to coeffs_copy[0, :] and we read from it.
+            fetch_idx = 0
+
         for ikernel, expr_dict in enumerate(sym_expr_dicts):
             expr = sum(coeff * prod(powers[i,
                 v[i] + max_deriv_order - mi[i]] for i in range(dim))