Expand CSEs. Factor make_kernel into multiple parts in a separate file.

MEMO

@@ -57,6 +57,8 @@ To-do
 
 - Scalar insn priority
 
+- What to do about constants in codegen? (...f suffix, complex types)
+
 - If finding a maximum proves troublesome, move parameters into the domain
 
 - : (as in, Matlab full-sclice) in prefetches

loopy/__init__.py

@@ -22,7 +22,9 @@ class LoopyAdvisory(UserWarning):
 
 from loopy.kernel import ScalarArg, ArrayArg, ConstantArrayArg, ImageArg
 
-from loopy.kernel import AutoFitLocalIndexTag, get_dot_dependency_graph, LoopKernel
+from loopy.kernel import (AutoFitLocalIndexTag, get_dot_dependency_graph,
+        LoopKernel, Instruction)
+from loopy.creation import make_kernel
 from loopy.subst import extract_subst, expand_subst
 from loopy.cse import precompute
 from loopy.preprocess import preprocess_kernel, realize_reduction
@@ -32,6 +34,7 @@ from loopy.compiled import CompiledKernel, drive_timing_run, auto_test_vs_ref
 from loopy.check import check_kernels
 
 __all__ = ["ScalarArg", "ArrayArg", "ConstantArrayArg", "ImageArg", "LoopKernel",
+        "Instruction", "make_kernel",
         "get_dot_dependency_graph",
         "preprocess_kernel", "realize_reduction",
         "generate_loop_schedules",
@@ -45,153 +48,6 @@ __all__ = ["ScalarArg", "ArrayArg", "ConstantArrayArg", "ImageArg", "LoopKernel"
 
 # }}}
 
-# {{{ kernel creation
-
-def make_kernel(*args, **kwargs):
-    """Second pass of kernel creation. Think about requests for iname duplication
-    and temporary variable declaration received as part of string instructions.
-    """
-
-    knl = LoopKernel(*args, **kwargs)
-
-    knl = tag_dimensions(
-            knl.copy(iname_to_tag_requests=None),
-            knl.iname_to_tag_requests)
-
-    new_insns = []
-    new_domain = knl.domain
-    new_temp_vars = knl.temporary_variables.copy()
-    new_iname_to_tag = knl.iname_to_tag.copy()
-
-    newly_created_vars = set()
-
-    for insn in knl.instructions:
-
-        # {{{ sanity checking
-
-        if not set(insn.forced_iname_deps) <= knl.all_inames():
-            raise ValueError("In instruction '%s': "
-                    "cannot force dependency on inames '%s'--"
-                    "they don't exist" % (
-                        insn.id,
-                        ",".join(
-                            set(insn.forced_iname_deps)-knl.all_inames())))
-
-        # }}}
-
-        # {{{ iname duplication
-
-        if insn.duplicate_inames_and_tags:
-            insn_dup_iname_to_tag = dict(insn.duplicate_inames_and_tags)
-
-            if not set(insn_dup_iname_to_tag.keys()) <= knl.all_inames():
-                raise ValueError("In instruction '%s': "
-                        "cannot duplicate inames '%s'--"
-                        "they don't exist" % (
-                            insn.id,
-                            ",".join(
-                                set(insn_dup_iname_to_tag.keys())-knl.all_inames())))
-
-            # {{{ duplicate non-reduction inames
-
-            reduction_inames = insn.reduction_inames()
-
-            inames_to_duplicate = [iname
-                    for iname, tag in insn.duplicate_inames_and_tags
-                    if iname not in reduction_inames]
-
-            new_inames = [
-                    knl.make_unique_var_name(
-                        based_on=iname+"_"+insn.id,
-                        extra_used_vars=
-                        newly_created_vars)
-                    for iname in inames_to_duplicate]
-
-            for old_iname, new_iname in zip(inames_to_duplicate, new_inames):
-                new_tag = insn_dup_iname_to_tag[old_iname]
-                if new_tag is None:
-                    new_tag = AutoFitLocalIndexTag()
-                new_iname_to_tag[new_iname] = new_tag
-
-            newly_created_vars.update(new_inames)
-
-            from loopy.isl_helpers import duplicate_axes
-            new_domain = duplicate_axes(new_domain, inames_to_duplicate, new_inames)
-
-            from loopy.symbolic import SubstitutionMapper
-            from pymbolic.mapper.substitutor import make_subst_func
-            from pymbolic import var
-            old_to_new = dict(
-                    (old_iname, var(new_iname))
-                    for old_iname, new_iname in zip(inames_to_duplicate, new_inames))
-            subst_map = SubstitutionMapper(make_subst_func(old_to_new))
-            new_expression = subst_map(insn.expression)
-
-            # }}}
-
-            if len(inames_to_duplicate) < len(insn.duplicate_inames_and_tags):
-                raise RuntimeError("cannot use [|...] syntax to rename reduction "
-                        "inames")
-
-            insn = insn.copy(
-                    assignee=subst_map(insn.assignee),
-                    expression=new_expression,
-                    forced_iname_deps=set(
-                        old_to_new.get(iname, iname) for iname in insn.forced_iname_deps),
-                    duplicate_inames_and_tags=[])
-
-        # }}}
-
-        # {{{ temporary variable creation
-
-        from loopy.kernel import (
-                find_var_base_indices_and_shape_from_inames,
-                TemporaryVariable)
-
-        if insn.temp_var_type is not None:
-            assignee_name = insn.get_assignee_var_name()
-
-            assignee_indices = []
-            from pymbolic.primitives import Variable
-            for index_expr in insn.get_assignee_indices():
-                if (not isinstance(index_expr, Variable)
-                        or not index_expr.name in knl.all_inames()):
-                    raise RuntimeError(
-                            "only plain inames are allowed in "
-                            "the lvalue index when declaring the "
-                            "variable '%s' in an instruction"
-                            % assignee_name)
-
-                assignee_indices.append(index_expr.name)
-
-            base_indices, shape = \
-                    find_var_base_indices_and_shape_from_inames(
-                            new_domain, assignee_indices, knl.cache_manager)
-
-            new_temp_vars[assignee_name] = TemporaryVariable(
-                    name=assignee_name,
-                    dtype=np.dtype(insn.temp_var_type),
-                    is_local=None,
-                    base_indices=base_indices,
-                    shape=shape)
-
-            newly_created_vars.add(assignee_name)
-
-            insn = insn.copy(temp_var_type=None)
-
-        # }}}
-
-        new_insns.append(insn)
-
-    return knl.copy(
-            instructions=new_insns,
-            domain=new_domain,
-            temporary_variables=new_temp_vars,
-            iname_to_tag=new_iname_to_tag,
-            iname_to_tag_requests=[])
-
-# }}}
-
 # {{{ dimension split
 
 def split_dimension(kernel, split_iname, inner_length,

loopy/creation.py

+from __future__ import division
+import numpy as np
+from loopy.symbolic import IdentityMapper
+
+
+# {{{ sanity checking
+
+def check_kernel(knl):
+    for insn in knl.instructions:
+        if not set(insn.forced_iname_deps) <= knl.all_inames():
+            raise ValueError("In instruction '%s': "
+                    "cannot force dependency on inames '%s'--"
+                    "they don't exist" % (
+                        insn.id,
+                        ",".join(
+                            set(insn.forced_iname_deps)-knl.all_inames())))
+
+# }}}
+
+# {{{ expand common subexpressions into assignments
+
+class CSEToAssignmentMapper(IdentityMapper):
+    def __init__(self, add_assignment):
+        self.add_assignment = add_assignment
+        self.expr_to_var = {}
+
+    def map_common_subexpression(self, expr):
+        try:
+            return self.expr_to_var[expr.child]
+        except KeyError:
+            from loopy.symbolic import TypedCSE
+            if isinstance(expr, TypedCSE):
+                dtype = expr.dtype
+            else:
+                dtype = None
+
+            child = self.rec(expr.child)
+
+            var_name = self.add_assignment(expr.prefix, child, dtype)
+            from pymbolic.primitives import Variable
+            var = Variable(var_name)
+            self.expr_to_var[expr.child] = var
+            return var
+
+def expand_cses(knl):
+    def add_assignment(base_name, expr, dtype):
+        kwargs = dict(extra_used_vars=newly_created_vars)
+        if base_name is not None:
+            kwargs["based_on"] = base_name
+        new_var_name = knl.make_unique_var_name(**kwargs)
+        newly_created_vars.add(new_var_name)
+
+        if dtype is None:
+            dtype = tim(expr)
+
+        from loopy.kernel import TemporaryVariable
+        new_temp_vars[new_var_name] = TemporaryVariable(
+                name=new_var_name,
+                dtype=np.dtype(dtype),
+                is_local=None,
+                shape=())
+
+        from pymbolic.primitives import Variable
+        from loopy.kernel import Instruction
+        insn = Instruction(
+                id=knl.make_unique_instruction_id(extra_used_ids=newly_created_insn_ids),
+                assignee=Variable(new_var_name), expression=expr)
+        newly_created_insn_ids.add(insn.id)
+        new_insns.append(insn)
+
+        return new_var_name
+
+    cseam = CSEToAssignmentMapper(add_assignment=add_assignment)
+
+    new_insns = []
+
+    newly_created_vars = set()
+    newly_created_insn_ids = set()
+    new_temp_vars = knl.temporary_variables.copy()
+
+    from loopy.codegen.expression import TypeInferenceMapper
+    tim = TypeInferenceMapper(knl, new_temp_vars)
+
+    for insn in knl.instructions:
+        new_insns.append(insn.copy(expression=cseam(insn.expression)))
+
+    return knl.copy(
+            instructions=new_insns,
+            temporary_variables=new_temp_vars)
+
+# }}}
+
+# {{{ temporary variable creation
+
+def create_temporaries(knl):
+    new_insns = []
+    new_temp_vars = knl.temporary_variables.copy()
+
+    for insn in knl.instructions:
+        from loopy.kernel import (
+                find_var_base_indices_and_shape_from_inames,
+                TemporaryVariable)
+
+        if insn.temp_var_type is not None:
+            assignee_name = insn.get_assignee_var_name()
+
+            assignee_indices = []
+            from pymbolic.primitives import Variable
+            for index_expr in insn.get_assignee_indices():
+                if (not isinstance(index_expr, Variable)
+                        or not index_expr.name in knl.all_inames()):
+                    raise RuntimeError(
+                            "only plain inames are allowed in "
+                            "the lvalue index when declaring the "
+                            "variable '%s' in an instruction"
+                            % assignee_name)
+
+                assignee_indices.append(index_expr.name)
+
+            base_indices, shape = \
+                    find_var_base_indices_and_shape_from_inames(
+                            knl.domain, assignee_indices, knl.cache_manager)
+
+            new_temp_vars[assignee_name] = TemporaryVariable(
+                    name=assignee_name,
+                    dtype=np.dtype(insn.temp_var_type),
+                    is_local=None,
+                    base_indices=base_indices,
+                    shape=shape)
+
+            insn = insn.copy(temp_var_type=None)
+
+        new_insns.append(insn)
+
+    return knl.copy(
+            instructions=new_insns,
+            temporary_variables=new_temp_vars)
+
+# }}}
+
+# {{{ duplicate inames
+
+def duplicate_inames(knl):
+    from loopy.kernel import AutoFitLocalIndexTag
+
+    new_insns = []
+    new_domain = knl.domain
+    new_iname_to_tag = knl.iname_to_tag.copy()
+
+    newly_created_vars = set()
+
+    for insn in knl.instructions:
+        if insn.duplicate_inames_and_tags:
+            insn_dup_iname_to_tag = dict(insn.duplicate_inames_and_tags)
+
+            if not set(insn_dup_iname_to_tag.keys()) <= knl.all_inames():
+                raise ValueError("In instruction '%s': "
+                        "cannot duplicate inames '%s'--"
+                        "they don't exist" % (
+                            insn.id,
+                            ",".join(
+                                set(insn_dup_iname_to_tag.keys())-knl.all_inames())))
+
+            # {{{ duplicate non-reduction inames
+
+            reduction_inames = insn.reduction_inames()
+
+            inames_to_duplicate = [iname
+                    for iname, tag in insn.duplicate_inames_and_tags
+                    if iname not in reduction_inames]
+
+            new_inames = [
+                    knl.make_unique_var_name(
+                        based_on=iname+"_"+insn.id,
+                        extra_used_vars=newly_created_vars)
+                    for iname in inames_to_duplicate]
+
+            for old_iname, new_iname in zip(inames_to_duplicate, new_inames):
+                new_tag = insn_dup_iname_to_tag[old_iname]
+                if new_tag is None:
+                    new_tag = AutoFitLocalIndexTag()
+                new_iname_to_tag[new_iname] = new_tag
+
+            newly_created_vars.update(new_inames)
+
+            from loopy.isl_helpers import duplicate_axes
+            new_domain = duplicate_axes(new_domain, inames_to_duplicate, new_inames)
+
+            from loopy.symbolic import SubstitutionMapper
+            from pymbolic.mapper.substitutor import make_subst_func
+            from pymbolic import var
+            old_to_new = dict(
+                    (old_iname, var(new_iname))
+                    for old_iname, new_iname in zip(inames_to_duplicate, new_inames))
+            subst_map = SubstitutionMapper(make_subst_func(old_to_new))
+            new_expression = subst_map(insn.expression)
+
+            # }}}
+
+            if len(inames_to_duplicate) < len(insn.duplicate_inames_and_tags):
+                raise RuntimeError("cannot use [|...] syntax to rename reduction "
+                        "inames")
+
+            insn = insn.copy(
+                    assignee=subst_map(insn.assignee),
+                    expression=new_expression,
+                    forced_iname_deps=set(
+                        old_to_new.get(iname, iname) for iname in insn.forced_iname_deps),
+                    duplicate_inames_and_tags=[])
+
+        new_insns.append(insn)
+
+    return knl.copy(
+            instructions=new_insns,
+            domain=new_domain,
+            iname_to_tag=new_iname_to_tag)
+# }}}
+
+# {{{ kernel creation top-level
+
+def make_kernel(*args, **kwargs):
+    """Second pass of kernel creation. Think about requests for iname duplication
+    and temporary variable declaration received as part of string instructions.
+    """
+
+    from loopy.kernel import LoopKernel
+    knl = LoopKernel(*args, **kwargs)
+
+    from loopy import tag_dimensions
+    knl = tag_dimensions(
+            knl.copy(iname_to_tag_requests=None),
+            knl.iname_to_tag_requests).copy(
+                    iname_to_tag_requests=[])
+
+    check_kernel(knl)
+
+    knl = create_temporaries(knl)
+    knl = expand_cses(knl)
+    knl = duplicate_inames(knl)
+
+    return knl
+
+# }}}
+
+# vim: fdm=marker