Add symbolic sum

.test-conda-env-py3.yml

@@ -12,12 +12,12 @@ dependencies:
 - islpy
 - pyopencl
 - python=3
-- python-symengine=0.6.0
+- python-symengine=0.6.1
 - pyfmmlib
 
 - pip
 - pip:
     - git+https://github.com/inducer/pytools
     - git+https://gitlab.tiker.net/inducer/boxtree
-    - git+https://github.com/inducer/pymbolic
+    - git+https://gitlab.tiker.net/inducer/pymbolic
     - git+https://github.com/inducer/loopy

benchmarks/bench_translations.py

@@ -75,7 +75,7 @@ def track_m2l_op_count(self, param):
         for i, expr in enumerate(result):
             sac.assign_unique("coeff%d" % i, expr)
         sac.run_global_cse()
-        insns = to_loopy_insns(six.iteritems(sac.assignments))
+        insns, _ = to_loopy_insns(six.iteritems(sac.assignments))
         counter = pymbolic.mapper.flop_counter.CSEAwareFlopCounter()
 
         return sum([counter.rec(insn.expression)+1 for insn in insns])

requirements.txt

@@ -1,7 +1,7 @@
 numpy
 sympy==1.1.1
 git+https://github.com/inducer/pytools
-git+https://github.com/inducer/pymbolic
+git+https://gitlab.tiker.net/inducer/pymbolic
 git+https://github.com/inducer/islpy
 git+https://github.com/inducer/pyopencl
 git+https://gitlab.tiker.net/inducer/boxtree

sumpy/assignment_collection.py

@@ -185,16 +185,20 @@ def run_global_cse(self, extra_exprs=[]):
         new_assignments, new_exprs = cse(assign_exprs + extra_exprs,
                 symbols=self.symbol_generator)
 
+        xreplace_dict = {}
+        for name, value in new_assignments:
+            new_value = sym.make_cse(value.xreplace(xreplace_dict))
+            xreplace_dict[name] = new_value
+
+        for i in range(len(new_exprs)):
+            new_exprs[i] = new_exprs[i].xreplace(xreplace_dict)
+
         new_assign_exprs = new_exprs[:len(assign_exprs)]
         new_extra_exprs = new_exprs[len(assign_exprs):]
 
         for name, new_expr in zip(assign_names, new_assign_exprs):
             self.assignments[name] = new_expr
 
-        for name, value in new_assignments:
-            assert isinstance(name, sym.Symbol)
-            self.add_assignment(name.name, value)
-
         logger.info("common subexpression elimination: done after {dur:.2f} s"
                     .format(dur=time.time() - start_time))
         return new_extra_exprs

sumpy/codegen.py

@@ -32,14 +32,16 @@
 import six
 import re
 
-from pymbolic.mapper import IdentityMapper, WalkMapper, CSECachingMapperMixin
+from pymbolic.mapper import CSECachingMapperMixin
 import pymbolic.primitives as prim
 
 from loopy.types import NumpyType
 
 from pytools import memoize_method
 
 from sumpy.symbolic import (SympyToPymbolicMapper as SympyToPymbolicMapperBase)
+from sumpy.symbolic import Series, IdentityMapper, WalkMapper, SubstitutionMapper
+from pymbolic.mapper.substitutor import make_subst_func
 
 import logging
 logger = logging.getLogger(__name__)
@@ -64,13 +66,16 @@
 
 class SympyToPymbolicMapper(SympyToPymbolicMapperBase):
 
+    def map_Sum(self, expr):  # noqa
+        pymbolic_limits = []
+        for name, low, high in expr.limits:
+            pymbolic_limits.append((self.rec(name), self.rec(low), self.rec(high)))
+
+        return Series(self.rec(expr.function), pymbolic_limits)
+
     def not_supported(self, expr):
         if isinstance(expr, int):
             return expr
-        elif getattr(expr, "is_Function", False):
-            func_name = SympyToPymbolicMapperBase.function_name(self, expr)
-            return prim.Variable(func_name)(
-                    *tuple(self.rec(arg) for arg in expr.args))
         else:
             return SympyToPymbolicMapperBase.not_supported(self, expr)
 
@@ -123,8 +128,6 @@ def make_one_step_subst(assignments):
 
     # {{{ make substitution
 
-    from pymbolic import substitute
-
     result = {}
     used_name_to_var = {}
     from pymbolic import evaluate
@@ -133,7 +136,7 @@ def make_one_step_subst(assignments):
 
     for name in toposort:
         value = assignments[name]
-        value = substitute(value, result)
+        value = SubstitutionMapper(make_subst_func(result))(value)
         used_name_to_var.update(
             (used_name, prim.Variable(used_name))
             for used_name in get_dependencies(value)
@@ -177,9 +180,8 @@ def kill_trivial_assignments(assignments, retain_names=set()):
     unsubst_rej = make_one_step_subst(rejected_assignments)
 
     result = []
-    from pymbolic import substitute
     for name, expr in approved_assignments:
-        r = substitute(expr, unsubst_rej)
+        r = SubstitutionMapper(make_subst_func(unsubst_rej))(expr)
         result.append((name, r))
 
     logger.info(
@@ -677,6 +679,35 @@ def map_variable(self, expr):
     map_common_subexpression_uncached = IdentityMapper.map_common_subexpression
 
 
+# {{{ convert pymbolic "Series" class to loopy reduction
+
+class SeriesRewritter(CSECachingMapperMixin, IdentityMapper):
+
+    def __init__(self):
+        self.additional_loop_domain = []
+
+    def map_series(self, expr):
+        function = self.rec(expr.function)
+        inames = []
+
+        for name, low, high in expr.limits:
+            low = self.rec(low)
+            high = self.rec(high)
+            name = str(name)
+
+            inames.append(name)
+            self.additional_loop_domain.append(
+                # +1 is used for converting the closed bound in sympy to open bound
+                (name, low + 1, high + 1)
+            )
+
+        return lp.Reduction("sum", tuple(inames), function)
+
+    map_common_subexpression_uncached = IdentityMapper.map_common_subexpression
+
+# }}}
+
+
 def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[],
                    complex_dtype=None, retain_names=set()):
     logger.info("loopy instruction generation: start")
@@ -703,6 +734,7 @@ def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[],
 
     # do the rest of the conversion
     bessel_sub = BesselSubstitutor(BesselGetter(btog.bessel_j_arg_to_top_order))
+    sr = SeriesRewritter()
     vcr = VectorComponentRewriter(vector_names)
     pwr = PowerRewriter()
     ssg = SumSignGrouper()
@@ -712,6 +744,7 @@ def to_loopy_insns(assignments, vector_names=set(), pymbolic_expr_maps=[],
 
     def convert_expr(name, expr):
         logger.debug("generate expression for: %s" % name)
+        expr = sr(expr)
         expr = bdr(expr)
         expr = bessel_sub(expr)
         expr = vcr(expr)
@@ -735,6 +768,6 @@ def convert_expr(name, expr):
                 for name, expr in assignments]
 
     logger.info("loopy instruction generation: done")
-    return result
+    return result, sr.additional_loop_domain
 
 # vim: fdm=marker

sumpy/e2e.py

@@ -160,13 +160,18 @@ def get_kernel(self):
         #
         # (same for itgt_box, tgt_ibox)
 
+        insns, additional_domain = self.get_translation_loopy_insns()
+
+        from sumpy.tools import get_loopy_domain
+        additional_domain = get_loopy_domain(additional_domain)
+
         from sumpy.tools import gather_loopy_arguments
         loopy_knl = lp.make_kernel(
                 [
                     "{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
                     "{[isrc_box]: isrc_start<=isrc_box<isrc_stop}",
                     "{[idim]: 0<=idim<dim}",
-                    ],
+                ] + additional_domain,
                 ["""
                 for itgt_box
                     <> tgt_ibox = target_boxes[itgt_box]
@@ -190,7 +195,7 @@ def get_kernel(self):
                             {{dep=read_src_ibox}}
                         """.format(coeffidx=i) for i in range(ncoeff_src)] + [
 
-                        ] + self.get_translation_loopy_insns() + ["""
+                        ] + insns + ["""
                     end
 
                     """] + ["""
@@ -276,19 +281,24 @@ def get_kernel(self):
         #
         # (same for itgt_box, tgt_ibox)
 
+        insns, additional_domain = self.get_translation_loopy_insns()
+
+        from sumpy.tools import get_loopy_domain
+        additional_domain = get_loopy_domain(additional_domain)
+
         loopy_insns = [
                 insn.copy(
                     predicates=insn.predicates | frozenset(["is_src_box_valid"]),
                     id=lp.UniqueName("compute_coeff"))
-                for insn in self.get_translation_loopy_insns()]
+                for insn in insns]
 
         from sumpy.tools import gather_loopy_arguments
         loopy_knl = lp.make_kernel(
                 [
                     "{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
                     "{[isrc_box]: 0<=isrc_box<nchildren}",
                     "{[idim]: 0<=idim<dim}",
-                    ],
+                ] + additional_domain,
                 ["""
                 for itgt_box
                     <> tgt_ibox = target_boxes[itgt_box]
@@ -395,12 +405,17 @@ def get_kernel(self):
         #
         # (same for itgt_box, tgt_ibox)
 
+        insns, additional_domain = self.get_translation_loopy_insns()
+
+        from sumpy.tools import get_loopy_domain
+        additional_domain = get_loopy_domain(additional_domain)
+
         from sumpy.tools import gather_loopy_arguments
         loopy_knl = lp.make_kernel(
                 [
                     "{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
-                    "{[idim]: 0<=idim<dim}",
-                    ],
+                    "{[idim]: 0<=idim<dim}"
+                ] + additional_domain,
                 ["""
                 for itgt_box
                     <> tgt_ibox = target_boxes[itgt_box]
@@ -419,7 +434,7 @@ def get_kernel(self):
                         {{id_prefix=read_expn,dep=read_src_ibox}}
                     """.format(i=i) for i in range(ncoeffs)] + [
 
-                    ] + self.get_translation_loopy_insns() + ["""
+                    ] + insns + ["""
 
                     tgt_expansions[tgt_ibox - tgt_base_ibox, {i}] = \
                         tgt_expansions[tgt_ibox - tgt_base_ibox, {i}] + coeff{i} \

sumpy/e2p.py

@@ -102,15 +102,15 @@ def get_loopy_insns_and_result_names(self):
         sac.run_global_cse()
 
         from sumpy.codegen import to_loopy_insns
-        loopy_insns = to_loopy_insns(
+        loopy_insns, additional_domain = to_loopy_insns(
                 six.iteritems(sac.assignments),
                 vector_names=set(["b"]),
                 pymbolic_expr_maps=[self.expansion.get_code_transformer()],
                 retain_names=result_names,
                 complex_dtype=np.complex128  # FIXME
                 )
 
-        return loopy_insns, result_names
+        return loopy_insns, additional_domain, result_names
 
     def get_kernel_scaling_assignment(self):
         from sumpy.symbolic import SympyToPymbolicMapper
@@ -135,13 +135,17 @@ class E2PFromSingleBox(E2PBase):
     def get_kernel(self):
         ncoeffs = len(self.expansion)
 
-        loopy_insns, result_names = self.get_loopy_insns_and_result_names()
+        loopy_insns, additional_domain, result_names = \
+            self.get_loopy_insns_and_result_names()
+
+        from sumpy.tools import get_loopy_domain
+        additional_domain = get_loopy_domain(additional_domain)
 
         loopy_knl = lp.make_kernel(
                 [
                     "{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
                     "{[itgt,idim]: itgt_start<=itgt<itgt_end and 0<=idim<dim}",
-                    ],
+                ] + additional_domain,
                 self.get_kernel_scaling_assignment()
                 + ["""
                 for itgt_box
@@ -232,15 +236,19 @@ class E2PFromCSR(E2PBase):
     def get_kernel(self):
         ncoeffs = len(self.expansion)
 
-        loopy_insns, result_names = self.get_loopy_insns_and_result_names()
+        loopy_insns, additional_domain, result_names = \
+            self.get_loopy_insns_and_result_names()
+
+        from sumpy.tools import get_loopy_domain
+        additional_domain = get_loopy_domain(additional_domain)
 
         loopy_knl = lp.make_kernel(
                 [
                     "{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
                     "{[itgt]: itgt_start<=itgt<itgt_end}",
                     "{[isrc_box]: isrc_box_start<=isrc_box<isrc_box_end }",
                     "{[idim]: 0<=idim<dim}",
-                    ],
+                ] + additional_domain,
                 self.get_kernel_scaling_assignment()
                 + ["""
                 for itgt_box

sumpy/kernel.py

@@ -26,8 +26,8 @@
 
 import loopy as lp
 import numpy as np
-from pymbolic.mapper import IdentityMapper, CSECachingMapperMixin
-from sumpy.symbolic import pymbolic_real_norm_2
+from pymbolic.mapper import CSECachingMapperMixin
+from sumpy.symbolic import pymbolic_real_norm_2, IdentityMapper
 from pymbolic.primitives import make_sym_vector
 from pymbolic import var