SymbolicAssignmentCollection.run_global_cse: do not modify in-place

Author: inducer

benchmarks/bench_translations.py

@@ -82,7 +82,7 @@ def track_m2l_op_count(self, param):
                                        dvec, tgt_rscale)
         for i, expr in enumerate(result):
             sac.assign_unique(f"coeff{i}", expr)
-        sac.run_global_cse()
+        sac = sac.run_global_cse()
         insns = to_loopy_insns(sac.assignments.items())
         counter = pymbolic.mapper.flop_counter.CSEAwareFlopCounter()
 

sumpy/assignment_collection.py

@@ -25,9 +25,9 @@
 
 import logging
 from collections import defaultdict
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
-from typing_extensions import override
+from typing_extensions import Self, override
 
 import sumpy.symbolic as sym
 
@@ -196,10 +196,18 @@ def assign_temp(self, name_base: str, expr: sym.Basic) -> str:
         new_name = self.symbol_generator(name_base).name
         return self.add_assignment(new_name, expr, retain_name=False)
 
-    def run_global_cse(
-            self,
-            extra_exprs: Sequence[sym.Expr] | None = None
-        ) -> Sequence[sym.Basic]:
+    @overload
+    def run_global_cse(self, extra_exprs: None = None) -> Self: ...
+
+    @overload
+    def run_global_cse(self,
+                extra_exprs: Sequence[sym.Expr]
+            ) -> tuple[Self, Sequence[sym.Basic]]: ...
+
+    def run_global_cse(self,
+                extra_exprs: Sequence[sym.Expr] | None = None
+            ) -> tuple[Self, Sequence[sym.Basic]] | Self:
+        orig_extra_exprs = extra_exprs
         if extra_exprs is None:
             extra_exprs = []
 
@@ -219,33 +227,32 @@ def run_global_cse(
         # from sumpy.symbolic import checked_cse
 
         from sumpy.cse import cse
-        new_assignments, new_exprs = cse(
+        new_cse_assignments, new_exprs = cse(
                 [*assign_exprs, *extra_exprs],
                 symbols=self.symbol_generator)
 
         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, strict=True):
-            self.assignments[name] = new_expr
+        result_assignments: dict[str, sym.Basic] = {}
 
-        for name, value in new_assignments:
+        for name, value in new_cse_assignments:
             assert isinstance(name, sym.Symbol)
-            self.add_assignment(name.name, value)
+            result_assignments[name.name] = value
 
-        for name, new_expr in zip(assign_names, new_assign_exprs, strict=True):
-            # We want the assignment collection to be ordered correctly
-            # to make it easier for loopy to schedule.
-            # Deleting the original assignments and adding them again
-            # makes them occur after the CSE'd expression preserving
-            # the order of operations.
-            del self.assignments[name]
-            self.assignments[name] = new_expr
+        result_assignments = {
+            **result_assignments,
+            **dict(zip(assign_names, new_assign_exprs, strict=True)),
+        }
 
         logger.info("common subexpression elimination: done after %.2f s",
                     time.time() - start_time)
 
-        return new_extra_exprs
+        result = type(self)(result_assignments)
+        if orig_extra_exprs is None:
+            return result
+        else:
+            return result, new_extra_exprs
 
 # }}}
 

sumpy/e2e.py

@@ -121,7 +121,7 @@ def get_translation_loopy_insns(self):
                         self.src_expansion, src_coeff_exprs, src_rscale,
                         dvec=dvec, tgt_rscale=tgt_rscale, sac=sac))]
 
-        sac.run_global_cse()
+        sac = sac.run_global_cse()
 
         from sumpy.codegen import to_loopy_insns
         return to_loopy_insns(
@@ -177,7 +177,7 @@ def get_translation_loopy_insns(self):
                         self.src_expansion, src_coeff_exprs, src_rscale,
                         dvec, tgt_rscale, sac))]
 
-        sac.run_global_cse()
+        sac = sac.run_global_cse()
 
         from sumpy.codegen import to_loopy_insns
         return to_loopy_insns(
@@ -335,7 +335,7 @@ def get_translation_loopy_insns(self, result_dtype):
                         m2l_translation_classes_dependent_data=(
                             m2l_translation_classes_dependent_data)))]
 
-        sac.run_global_cse()
+        sac = sac.run_global_cse()
 
         from sumpy.codegen import to_loopy_insns
         return to_loopy_insns(

sumpy/expansion/loopy.py

@@ -90,7 +90,7 @@ def make_e2p_loopy_kernel(
         for i, knl in enumerate(kernels)
     }
 
-    sac.run_global_cse()
+    sac = sac.run_global_cse()
 
     code_transformers = (
             [expansion.get_code_transformer()]
@@ -194,7 +194,7 @@ def make_p2e_loopy_kernel(
         sac.add_assignment(f"coeffs{i}", coeff) for i, coeff in enumerate(coeffs)
     }
 
-    sac.run_global_cse()
+    sac = sac.run_global_cse()
 
     code_transformers = (
         [expansion.get_code_transformer()]

sumpy/expansion/m2l.py

@@ -1324,7 +1324,7 @@ def loopy_translation_classes_dependent_data(
     tgt_coeff_names = [
             sac.assign_unique(f"m2l_translation_classes_dependent_data{i}", coeff_i)
             for i, coeff_i in enumerate(derivatives)]
-    sac.run_global_cse()
+    sac = sac.run_global_cse()
 
     from sumpy.codegen import to_loopy_insns
     from sumpy.tools import to_complex_dtype

sumpy/p2p.py

@@ -147,7 +147,7 @@ def get_loopy_insns_and_result_names(self):
             for i, expr in enumerate(exprs)
         ]
 
-        sac.run_global_cse()
+        sac = sac.run_global_cse()
 
         from sumpy.codegen import to_loopy_insns
         loopy_insns = to_loopy_insns(sac.assignments.items(),

sumpy/qbx.py

@@ -161,7 +161,7 @@ def get_loopy_insns_and_result_names(self):
 
         logger.info("compute expansion expressions: done")
 
-        sac.run_global_cse()
+        sac = sac.run_global_cse()
 
         pymbolic_expr_maps = [knl.get_code_transformer() for knl in [
             *self.target_kernels, *self.source_kernels]]