checkpoint 0.00016

Author: isuruf

sumpy/expansion/loopy.py

@@ -257,9 +257,14 @@ def l2p_loopy_volume_taylor(expansion, kernels):
         c = max_mi[axis_permutation[0]]
         v = [pymbolic.var(f"x{i}") for i in range(dim)]
         v[axis_permutation[0]], v[0] = v[0], v[axis_permutation[0]]
+        x0 = v[0]
 
         def get_domains(v, iorder):
-            domains = [f"{{ [{v[0]}]: 0<={v[0]}<=order }}"]
+            domains = [f"{{ [{x0}_outer]: 0<={x0}_outer<={order//c} }}"]
+            expr = f"{v[0]}_inner + {c}*{x0}_outer"
+            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 }}"]
             upper_bound = f"{iorder}-{v[0]}"
             for i in range(dim - 1, 1, -1):
@@ -286,25 +291,20 @@ def get_idx(v):
                 continue
             mi_sym = [max_mi_sym[i] + deriv_id.mi[i] for i in range(dim)]
             mi_sym[0] = mi_sym[0] % c
-            expr += (coeffs_copy[(v[0]//c + 1) % 2, wrangler.get_storage_index(mi_sym)]
-                     * (pde_coeff * scale))
+            expr += (coeffs_copy[(v[0]//c + 1) % 2,
+                wrangler.get_storage_index(mi_sym)] * (pde_coeff * scale))
 
-        domains += ["{[dummy]: 0<=dummy<1}"]
-        insns2 = []
         insns.append(lp.Assignment(
-            assignee=coeffs_copy[(v[0]//c)%2, idx],
+            assignee=coeffs_copy[(v[0]//c) % 2, idx],
             expression=expr,
             id="update_coeffs",
             depends_on=frozenset(["copy_coeffs"]),
             depends_on_is_final=True,
-            #within_inames=frozenset(["dummy"]),
             predicates=frozenset([prim.Comparison(v[0], ">=", c)]),
         ))
 
-        x0 = v[0]
         v = [pymbolic.var(f"y{i}") for i in range(dim)]
         v[axis_permutation[0]], v[0] = v[0], v[axis_permutation[0]]
-        v[0] = x0
         iorder = pymbolic.var("iorder2")
         idx = get_idx(v)
         domains += get_domains(v, iorder)[1:]
@@ -318,7 +318,7 @@ def get_idx(v):
             insn = lp.Assignment(
                 assignee=result[ikernel],
                 expression=(result[ikernel]
-                    + coeffs_copy[(v[0]//c)%2, idx] * expr),
+                    + coeffs_copy[(v[0]//c) % 2, idx] * expr),
                 id=f"write_{ikernel}",
                 depends_on=frozenset(["update_monomials", "update_coeffs"]),
                 depends_on_is_final=True,
@@ -328,8 +328,10 @@ def get_idx(v):
             lambda knl: lp.tag_inames(knl, {
                 #"e2p_dummy": "l.0",
                 "e2p_iorder1": "l.0",
-                f"e2p_{v[0]}": "unr",
-                #"e2p_iorder2": "unr",
+                f"e2p_{x0}_outer": "unr",
+                f"e2p_{x0}_inner": "unr",
+                f"e2p_{v[0]}_inner": "unr",
+                "e2p_iorder2": "unr",
             }),
             lambda knl: lp.set_temporary_address_space(knl, "e2p_coeffs_copy",
                 lp.AddressSpace.LOCAL),