Summarize changes to support prediction

loopy/__init__.py

@@ -120,6 +120,7 @@
 from loopy.target.execution import ExecutorBase
 from loopy.target.ispc import ISPCTarget
 from loopy.target.opencl import OpenCLTarget
+from loopy.target.pycuda import PyCudaTarget, PyCudaWithPackedArgsTarget
 from loopy.target.pyopencl import PyOpenCLTarget
 from loopy.tools import Optional, clear_in_mem_caches, memoize_on_disk, t_unit_to_python
 from loopy.transform.add_barrier import add_barrier
@@ -151,6 +152,7 @@
     tag_array_axes,
     tag_data_axes,
 )
+from loopy.transform.domain import decouple_domain
 from loopy.transform.fusion import fuse_kernels
 from loopy.transform.iname import (
     add_inames_for_unused_hw_axes,
@@ -206,6 +208,10 @@
     unprivatize_temporaries_with_inames,
 )
 from loopy.transform.realize_reduction import realize_reduction
+from loopy.transform.reduction import (
+    hoist_invariant_multiplicative_terms_in_sum_reduction,
+    extract_multiplicative_terms_in_sum_reduction_as_subst)
+from loopy.transform.reindex import reindex_temporary_using_seghir_loechner_scheme
 from loopy.transform.save import save_and_reload_temporaries
 from loopy.transform.subst import (
     assignment_to_subst,
@@ -277,6 +283,8 @@
     "Options",
     "OrderedAtomic",
     "PreambleInfo",
+    "PyCudaTarget",
+    "PyCudaWithPackedArgsTarget",
     "PyOpenCLTarget",
     "Reduction",
     "ScalarCallable",
@@ -317,8 +325,10 @@
     "clear_in_mem_caches",
     "collect_common_factors_on_increment",
     "concatenate_arrays",
+    "decouple_domain",
     "duplicate_inames",
     "expand_subst",
+    "extract_multiplicative_terms_in_sum_reduction_as_subst",
     "extract_subst",
     "find_instructions",
     "find_most_recent_global_barrier",
@@ -349,6 +359,7 @@
     "get_subkernels",
     "get_synchronization_map",
     "has_schedulable_iname_nesting",
+    "hoist_invariant_multiplicative_terms_in_sum_reduction",
     "infer_arg_descr",
     "infer_unknown_types",
     "inline_callable_kernel",
@@ -378,6 +389,7 @@
     "register_preamble_generators",
     "register_reduction_parser",
     "register_symbol_manglers",
+    "reindex_temporary_using_seghir_loechner_scheme",
     "remove_inames_from_insn",
     "remove_instructions",
     "remove_predicates_from_insn",
@@ -416,6 +428,14 @@
     "untag_inames",
     ]
 
+try:
+    import loopy.relations as relations
+except ImportError:
+    # catching ImportErrors to avoid making minikanren a hard-dep
+    pass
+else:
+    __all__ += ["relations"]
+
 # }}}
 
 

loopy/codegen/loop.py

@@ -29,10 +29,12 @@
 import islpy as isl
 from islpy import dim_type
 from pymbolic.mapper.stringifier import PREC_NONE
+from typing import FrozenSet
 
 from loopy.codegen.control import build_loop_nest
 from loopy.codegen.result import CodeGenerationResult, merge_codegen_results
 from loopy.diagnostic import LoopyError, warn
+from loopy.kernel import LoopKernel
 from loopy.symbolic import flatten
 
 
@@ -369,6 +371,16 @@
 
 # {{{ sequential loop
 
+def _get_intersecting_inames(kernel: LoopKernel, iname: str) -> FrozenSet[str]:
+    from functools import reduce
+    return reduce(frozenset.union,
+                  ((kernel.id_to_insn[insn].within_inames
+                    | kernel.id_to_insn[insn].reduction_inames()
+                    | kernel.id_to_insn[insn].sub_array_ref_inames())
+                   for insn in kernel.iname_to_insns()[iname]),
+                  frozenset())
+
+
 def generate_sequential_loop_dim_code(
             codegen_state: CodeGenerationState,
             sched_index: int,
@@ -386,8 +398,18 @@
     from loopy.codegen.bounds import get_usable_inames_for_conditional
 
     # Note: this does not include loop_iname itself!
+
+    # usable_inames = get_usable_inames_for_conditional(
+    #    kernel, sched_index, codegen_state.codegen_cachemanager)
+
+    # # get rid of disjoint loop nests, see
+    # # <www.github.com/inducer/loopy/issues/724>
+    # usable_inames = usable_inames & _get_intersecting_inames(kernel,
+    #                                                        loop_iname)
+    # =======
     usable_inames = get_usable_inames_for_conditional(kernel, sched_index,
             codegen_state.codegen_cache_manager)
+    # >>>>>>> main
 
     domain = kernel.get_inames_domain(loop_iname)
 

loopy/kernel/tools.py

@@ -2218,6 +2218,65 @@
 
 # }}}
 
+# {{{ get access map from an instruction
+
+
+class _IndexCollector(CombineMapper):
+    def __init__(self, var):
+        self.var = var
+        super().__init__()
+
+    def combine(self, values):
+        import operator
+        return reduce(operator.or_, values, frozenset())
+
+    def map_subscript(self, expr):
+        if expr.aggregate.name == self.var:
+            return (super().map_subscript(expr) | frozenset([expr.index_tuple]))
+        else:
+            return super().map_subscript(expr)
+
+    def map_algebraic_leaf(self, expr):
+        return frozenset()
+
+    map_constant = map_algebraic_leaf
+
+
+def _project_out_inames_from_maps(amaps, inames_to_project_out):
+    new_amaps = []
+    for amap in amaps:
+        for iname in inames_to_project_out:
+            dt, pos = amap.get_var_dict()[iname]
+            amap = amap.project_out(dt, pos, 1)
+
+        new_amaps.append(amap)
+
+    return new_amaps
+
+
+def _union_amaps(amaps):
+    import islpy as isl
+    return reduce(isl.Map.union, amaps[1:], amaps[0])
+
+
+def get_insn_access_map(kernel, insn_id, var):
+    from loopy.transform.subst import expand_subst
+    from loopy.symbolic import get_access_map
+
+    insn = kernel.id_to_insn[insn_id]
+
+    kernel = expand_subst(kernel)
+    indices = list(_IndexCollector(var)((insn.expression,
+                                         insn.assignees,
+                                         tuple(insn.predicates))))
+
+    amaps = [get_access_map(kernel.get_inames_domain(insn.within_inames),
+                            idx, kernel.assumptions) for idx in indices]
+
+    return _union_amaps(amaps)
+
+# }}}
+
 
 def get_hw_axis_base_for_codegen(kernel: LoopKernel, iname: str) -> isl.Aff:
     """

loopy/relations.py

@@ -0,0 +1,122 @@
+from kanren import Relation, facts
+
+
+def get_inameo(kernel):
+    inameo = Relation()
+    for iname in kernel.all_inames():
+        facts(inameo, (iname,))
+    return inameo
+
+
+def get_argo(kernel):
+    argo = Relation()
+    for arg in kernel.args:
+        facts(argo, (arg.name,))
+
+    return argo
+
+
+def get_tempo(kernel):
+    tempo = Relation()
+    for tv in kernel.temporary_variables:
+        facts(tempo, (tv,))
+
+    return tempo
+
+
+def get_insno(kernel):
+    insno = Relation()
+    for insn in kernel.instructions:
+        facts(insno, (insn.id,))
+
+    return insno
+
+
+def get_taggedo(kernel):
+    taggedo = Relation()
+
+    for arg_name, arg in kernel.arg_dict.items():
+        for tag in arg.tags:
+            facts(taggedo, (arg_name, tag))
+
+    for iname_name, iname in kernel.inames.items():
+        for tag in iname.tags:
+            facts(taggedo, (iname_name, tag))
+
+    for insn in kernel.instructions:
+        for tag in insn.tags:
+            facts(taggedo, (insn.id, tag))
+
+    return taggedo
+
+
+def get_taggedo_of_type(kernel, tag_type):
+    taggedo = Relation()
+
+    for arg_name, arg in kernel.arg_dict.items():
+        for tag in arg.tags_of_type(tag_type):
+            facts(taggedo, (arg_name, tag))
+
+    for iname_name, iname in kernel.inames.items():
+        for tag in iname.tags_of_type(tag_type):
+            facts(taggedo, (iname_name, tag))
+
+    for insn in kernel.instructions:
+        for tag in insn.tags_of_type(tag_type):
+            facts(taggedo, (insn.id, tag))
+
+    return taggedo
+
+
+def get_producero(kernel):
+    producero = Relation()
+
+    for insn in kernel.instructions:
+        for var in insn.assignee_var_names():
+            facts(producero, (insn.id, var))
+
+    return producero
+
+
+def get_consumero(kernel):
+    consumero = Relation()
+
+    for insn in kernel.instructions:
+        for var in insn.read_dependency_names():
+            facts(consumero, (insn.id, var))
+
+    return consumero
+
+
+def get_withino(kernel):
+    withino = Relation()
+
+    for insn in kernel.instructions:
+        facts(withino, (insn.id, insn.within_inames))
+
+    return withino
+
+
+def get_reduce_insno(kernel):
+    reduce_insno = Relation()
+
+    for insn in kernel.instructions:
+        if insn.reduction_inames():
+            facts(reduce_insno, (insn.id,))
+
+    return reduce_insno
+
+
+def get_reduce_inameo(kernel):
+    from functools import reduce
+    reduce_inameo = Relation()
+
+    for iname in reduce(frozenset.union,
+                        (insn.reduction_inames()
+                         for insn in kernel.instructions),
+                        frozenset()):
+        facts(reduce_inameo, (iname,))
+
+    return reduce_inameo
+
+# vim: fdm=marker