Allow conditionals to dependend on parallel axes.

loopy/codegen/bounds.py

@@ -236,11 +236,9 @@ def wrap_in_for_from_constraints(ccm, iname, constraint_bset, stmt,
 
 # {{{ on which variables may a conditional depend?
 
-def get_defined_inames(kernel, sched_index, allow_tag_classes=()):
-    """
-    :param exclude_tags: a tuple of tag classes to exclude
-    """
+def get_defined_inames(kernel, sched_index):
     from loopy.schedule import EnterLoop, LeaveLoop
+    from loopy.kernel import ParallelTag
 
     result = set()
 
@@ -255,10 +253,33 @@ def get_defined_inames(kernel, sched_index, allow_tag_classes=()):
     for iname in kernel.all_inames():
         tag = kernel.iname_to_tag.get(iname)
 
-        if isinstance(tag, allow_tag_classes):
+        # these are always defined
+        if isinstance(tag, ParallelTag):
             result.add(iname)
 
-    return result
+    return frozenset(result)
+
+# }}}
+
+# {{{ get_simple_loop_bounds
+
+def get_simple_loop_bounds(kernel, sched_index, iname, implemented_domain, iname_domain):
+    from loopy.codegen.bounds import get_bounds_constraints, get_defined_inames
+    lower_constraints_orig, upper_constraints_orig, equality_constraints_orig = \
+            get_bounds_constraints(iname_domain, iname,
+                    frozenset([iname])
+                    | get_defined_inames(kernel, sched_index+1),
+                    allow_parameters=True)
+
+    lower_constraints_orig.extend(equality_constraints_orig)
+    upper_constraints_orig.extend(equality_constraints_orig)
+    #assert not equality_constraints_orig
+
+    from loopy.codegen.bounds import pick_simple_constraint
+    lb_cns_orig = pick_simple_constraint(lower_constraints_orig, iname)
+    ub_cns_orig = pick_simple_constraint(upper_constraints_orig, iname)
+
+    return lb_cns_orig, ub_cns_orig
 
 # }}}
 
@@ -283,4 +304,7 @@ def pick_simple_constraint(constraints, iname):
 
 
 
+
+
+
 # vim: foldmethod=marker

loopy/codegen/loop.py

@@ -8,27 +8,6 @@ from loopy.codegen.control import build_loop_nest
 
 
 
-def get_simple_loop_bounds(kernel, sched_index, iname, implemented_domain, iname_domain):
-    from loopy.codegen.bounds import get_bounds_constraints, get_defined_inames
-    lower_constraints_orig, upper_constraints_orig, equality_constraints_orig = \
-            get_bounds_constraints(iname_domain, iname,
-                    frozenset([iname])
-                    | frozenset(get_defined_inames(kernel, sched_index+1)),
-                    allow_parameters=True)
-
-    lower_constraints_orig.extend(equality_constraints_orig)
-    upper_constraints_orig.extend(equality_constraints_orig)
-    #assert not equality_constraints_orig
-
-    from loopy.codegen.bounds import pick_simple_constraint
-    lb_cns_orig = pick_simple_constraint(lower_constraints_orig, iname)
-    ub_cns_orig = pick_simple_constraint(upper_constraints_orig, iname)
-
-    return lb_cns_orig, ub_cns_orig
-
-
-
-
 # {{{ conditional-minimizing slab decomposition
 
 def get_slab_decomposition(kernel, iname, sched_index, codegen_state):
@@ -37,6 +16,7 @@ def get_slab_decomposition(kernel, iname, sched_index, codegen_state):
     if iname_domain.is_empty():
         return ()
 
+    from loopy.codegen.bounds import get_simple_loop_bounds
     lb_cns_orig, ub_cns_orig = get_simple_loop_bounds(kernel, sched_index, iname,
             codegen_state.implemented_domain, iname_domain)