From 13a6ef823d46debe3942adb41a6badd5677f21b2 Mon Sep 17 00:00:00 2001 From: ignis Date: Sat, 30 May 2026 06:12:49 +0000 Subject: [PATCH] Optimize terms post-processor generation with SymPy, CSE, and selective expansion --- code/code_gen/post.py | 305 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 288 insertions(+), 17 deletions(-) diff --git a/code/code_gen/post.py b/code/code_gen/post.py index ec4fd96..deb2326 100644 --- a/code/code_gen/post.py +++ b/code/code_gen/post.py @@ -1,6 +1,229 @@ from lark import Lark, Visitor, Transformer, v_args, Token import warnings from jinja2 import Template +import sympy +from sympy.printing.fortran import FCodePrinter + +@v_args(inline=True) +class LarkToSympy(Transformer): + def __init__(self, fdict): + self.fdict = fdict + + def number(self, numeral): + return sympy.Float(float(numeral)) + + def env(self, name): + return sympy.Symbol(name.value) + + def paren(self, val): + return val + + def var(self, name): + return sympy.Symbol(name.value) + + def fluc(self, name): + return sympy.Symbol(name.value + "__prime") + + def dnx(self, partial, b): + signature = f"{partial.data}_{b.value}" + return sympy.Symbol(signature) + + def icall(self, op, val): + if op.data == "sqr": + return val**2 + elif op.data == "pow3": + return val**3 + return val + + def fcall(self, *args): + a = args[0] + func_name = a.value if hasattr(a, 'value') else str(a) + if func_name == "sqrt": + return sympy.sqrt(args[1]) + elif func_name == "exp": + return sympy.exp(args[1]) + elif func_name == "log": + return sympy.log(args[1]) + elif func_name == "abs": + return sympy.Abs(args[1]) + elif func_name == "rxn_rate": + return sympy.Function("rxn_rate")(args[1]) + elif func_name == "udf": + return sympy.Function(a.value)(*args[1:]) + return sympy.Function(func_name)(*args[1:]) + + def neg(self, val): + return -val + + def add(self, a, b): + return a + b + + def sub(self, a, b): + return a - b + + def mul(self, a, b): + return a * b + + def div(self, a, b): + return a / b + + def udf(self, a): + return a.value + + log = lambda self: "log" + exp = lambda self: "exp" + sqrt = lambda self: "sqrt" + abs = lambda self: "abs" + rxn_rate = lambda self: "rxn_rate" + + +class ArrayFCodePrinter(FCodePrinter): + def __init__(self, settings=None, array_symbols=None, avg_symbols=None): + settings = settings or {} + settings.setdefault('source_format', 'free') + settings.setdefault('standard', 95) + super().__init__(settings) + self.array_symbols = array_symbols or {} + self.avg_symbols = avg_symbols or {} + + def _print_Float(self, expr): + val = str(expr) + if 'e' in val or 'E' in val: + return val.replace('e', 'd').replace('E', 'd') + if '.' not in val: + return val + ".0d0" + return val + "d0" + + def _print_Symbol(self, expr): + name = expr.name + if name in self.array_symbols: + return f"{self.array_symbols[name]}(i,j,k)" + if name in self.avg_symbols: + return f"{self.avg_symbols[name]}(i)" + if name.startswith("avg_") or "_avg_" in name or name.endswith("_avg"): + return f"{name}(i)" + if name.endswith("__prime"): + base = name[:-7] + arr = self.array_symbols.get(base, base) + return f"({arr}(i,j,k) - {{0}}avg_{base}(i))" + return name + + def _print_Function(self, expr): + try: + return super()._print_Function(expr) + except Exception: + args = ", ".join(self.doprint(arg) for arg in expr.args) + return f"{expr.func.__name__}({args})" + + +class SympyOptimizer: + _instance = None + + @classmethod + def get_instance(cls, fdict): + if cls._instance is None or cls._instance.fdict is not fdict: + cls._instance = cls(fdict) + return cls._instance + + def __init__(self, fdict): + self.fdict = fdict + self.sympy_cache = {} + self.exported_fields = set( + name for name, f in fdict.items() + if hasattr(f, 'attr') and f.attr.get('export') + ) + self.averaged_targets = set() + + def set_averaged(self, averaged_dict): + self.averaged_targets = {a.target for a in averaged_dict.values()} + + def get_sympy_expr(self, name): + if name in self.sympy_cache: + return self.sympy_cache[name] + + field = self.fdict[name] + + if hasattr(field, 'prime') and field.prime: + expr = sympy.Symbol(name) + self.sympy_cache[name] = expr + return expr + + if hasattr(field, 'op'): # DerivedField + expr = sympy.Symbol(name) + self.sympy_cache[name] = expr + return expr + + if hasattr(field, 'weighted'): # AveragedField + expr = sympy.Symbol(name) + self.sympy_cache[name] = expr + return expr + + if hasattr(field, 'field') and hasattr(field, 'w'): # FluctuationField + expr = sympy.Symbol(name) + self.sympy_cache[name] = expr + return expr + + transformer = LarkToSympy(self.fdict) + expr = transformer.transform(field.exp) + + # Recursively substitute derived variables that are not cached + expanded_expr = expr + changed = True + while changed: + changed = False + free_syms = list(expanded_expr.free_symbols) + sub_dict = {} + for sym in free_syms: + sym_name = sym.name + if sym_name in self.fdict: + f = self.fdict[sym_name] + is_derived_field = hasattr(f, 'op') + is_averaged_field = hasattr(f, 'weighted') + is_primary_field = hasattr(f, 'prime') and f.prime + is_exported = sym_name in self.exported_fields + is_averaged_target = sym_name in self.averaged_targets + + if not (is_derived_field or is_averaged_field or is_primary_field or is_exported or is_averaged_target): + sub_dict[sym] = self.get_sympy_expr(sym_name) + changed = True + + if sub_dict: + expanded_expr = expanded_expr.subs(sub_dict) + + self.sympy_cache[name] = expanded_expr + return expanded_expr + + def optimize_field(self, name, alloc=None): + expr = self.get_sympy_expr(name) + + simplified_expr = sympy.simplify(expr) + simplified_expr = sympy.cancel(simplified_expr) + + array_symbols = {} + for k, v in self.fdict.items(): + if hasattr(v, 'array') and v.array: + array_symbols[k] = v.array + elif alloc and k in alloc: + array_symbols[k] = alloc[k] + else: + array_symbols[k] = k + + printer = ArrayFCodePrinter(array_symbols=array_symbols) + + replacements, reduced_exprs = sympy.cse(simplified_expr) + reduced_expr = reduced_exprs[0] + + cse_decls = [] + cse_assigns = [] + + if replacements: + for temp_var, temp_expr in replacements: + cse_decls.append(f"real(real64) :: {temp_var}") + cse_assigns.append(f"{temp_var} = {printer.doprint(temp_expr)}") + + rhs = printer.doprint(reduced_expr) + + return rhs, cse_decls, cse_assigns class CollectDefinitions(Visitor): @@ -545,20 +768,44 @@ class Field (FieldBase): return self.exporter is not None def code (self, alloc=None): - real_array_loop = """ + self.array = alloc[self.name] if alloc else self.name + + # Optimize using SymPy + opt = SympyOptimizer.get_instance(self.fdict) + rhs, cse_decls, cse_assigns = opt.optimize_field(self.name, alloc) + + decls_str = "\n".join(cse_decls) if cse_decls else "" + assigns_str = "\n".join(cse_assigns) if cse_assigns else "" + + real_array_loop = """ ! {{ comment }} +{% if decls_str -%} +block +{{ decls_str | indent(4, True) }} +{%- endif %} do k = 1, nzp do j = 1, nyp do i = 1, nxp -{{ array }}(i,j,k) = {{ rhs }} +{% if assigns_str -%} +{{ assigns_str | indent(4, True) }} + {{ array }}(i,j,k) = {{ rhs }} +{%- else -%} + {{ array }}(i,j,k) = {{ rhs }} +{%- endif %} end do end do end do +{% if decls_str -%} +end block +{%- endif %} """ - self.array = alloc[self.name] if alloc else self.name - - rhs = ExpToCode(self.fdict).transform(self.exp) - calculation_code = Template(real_array_loop).render(comment=self.comment, array=self.array, rhs=rhs) + calculation_code = Template(real_array_loop).render( + comment=self.comment, + decls_str=decls_str, + assigns_str=assigns_str, + array=self.array, + rhs=rhs + ) export_code = ( self.exporter.code() if self.export_on() else "") @@ -589,7 +836,14 @@ class FluctuationField (FieldBase): self.comment = self.name + " = " + self.comment def code (self, alloc=None): - real_array_loop = """ + self.array = alloc[self.name] if alloc else self.name + + rhs = ExpToCode(self.fdict).transform(self.field.exp) + + if self.field.is_fluctuation(): + rhs = rhs.format(self.w) + + real_array_loop = """ ! {{ comment }} do k = 1, nzp do j = 1, nyp @@ -599,14 +853,11 @@ end do end do end do """ - self.array = alloc[self.name] if alloc else self.name - - rhs = ExpToCode(self.fdict).transform(self.field.exp) - - if self.field.is_fluctuation(): - rhs = rhs.format(self.w) - - return Template(real_array_loop).render(comment=self.comment, array=self.array, rhs=rhs) + return Template(real_array_loop).render( + comment=self.comment, + array=self.array, + rhs=rhs + ) @classmethod def id (cls, w, field): @@ -676,6 +927,7 @@ class AveragedField (FieldBase): super(AveragedField,self).__init__(name, fdict) self.shape = "nxp" self.dim = ":" + self.target = tgt tfield = fdict[tgt] self.fset = tfield.checkFluctuation() @@ -957,6 +1209,22 @@ class Stage4(): self.pass1 = src.pass1 self.pass2 = src.pass2 + # Initialize SympyOptimizer and set the averaged fields to collect targets + opt = SympyOptimizer.get_instance(self.derived) + opt.set_averaged(self.averaged) + + # Update dependencies based on SymPy optimized expressions for Field objects + updated_dependency = {} + for name, dep_set in self.dependency.items(): + if name in self.derived and isinstance(self.derived[name], Field): + expr = opt.get_sympy_expr(name) + free_sym_names = {sym.name for sym in expr.free_symbols} + valid_deps = {dep for dep in free_sym_names if dep in self.derived or dep in self.primary} + updated_dependency[name] = valid_deps + else: + updated_dependency[name] = dep_set + self.dependency = updated_dependency + self.array_name = "xyzbuffer{}" narr1, alloc1 = (self.allocate_arr(self.pass1)) @@ -1059,6 +1327,9 @@ close (200) def print_program (self): + # Initialize SympyOptimizer and set the averaged fields to collect targets + opt = SympyOptimizer.get_instance(self.derived) + opt.set_averaged(self.averaged) allvar = dict(self.derived) allvar.update(self.averaged) @@ -1084,8 +1355,8 @@ close (200) alloc_export = "\n".join(v.exporter.code_alloc() for v in set_export_on) free_export = "\n".join(v.exporter.code_free() for v in set_export_on) - sub_calc1 = "\n".join(allvar[v].code(self.alloc1) for v in self.pass1) - sub_calc2 = "\n".join(allvar[v].code(self.alloc2) for v in self.pass2) + sub_calc1 = "\n".join(allvar[v].code(self.alloc1) for v in self.pass1 if v in self.averaged or v in self.alloc1) + sub_calc2 = "\n".join(allvar[v].code(self.alloc2) for v in self.pass2 if v in self.averaged or v in self.alloc2) sub_avg1 = "\n".join(allvar[v].code_avg() for v in set1) sub_avg2 = "\n".join(allvar[v].code_avg() for v in set2)