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): """Transformer that parses Lark AST mathematical nodes into SymPy expression objects. This recursively traverses the AST for algebraic equations, mapping operations, constants, brackets, and custom derivative definitions directly to standard SymPy nodes. """ def __init__(self, fdict): """Initializes the Lark-to-SymPy transformer. Args: fdict (dict): Dictionary mapping variables to their Field definitions. """ 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 calculate_flops_and_heavy(self, expr): flops = 0 heavy = 0 for node in sympy.preorder_traversal(expr): if isinstance(node, sympy.Add): flops += len(node.args) - 1 elif isinstance(node, sympy.Mul): flops += len(node.args) - 1 elif isinstance(node, sympy.Pow): base, exp = node.args if exp == 0.5 or exp == -0.5: flops += 10 heavy += 1 elif exp == -1: flops += 4 heavy += 1 elif isinstance(exp, sympy.Integer): val = abs(int(exp)) if val > 1: flops += val - 1 else: flops += 10 heavy += 1 elif isinstance(node, (sympy.Derivative, sympy.Function)): name = node.func.__name__ if name == 'sqrt': flops += 10 heavy += 1 elif name in ('exp', 'log', 'sin', 'cos', 'tan', 'rxn_rate'): flops += 10 heavy += 1 elif name == 'Abs': flops += 1 else: flops += 10 heavy += 1 return flops, heavy def count_3d_loads(self, expr, three_d_arrays): count = 0 for node in sympy.preorder_traversal(expr): if isinstance(node, sympy.Symbol) and node.name in three_d_arrays: count += 1 return count def optimize_field(self, name, alloc=None): expr = self.get_sympy_expr(name) three_d_arrays = { k for k, v in self.fdict.items() if hasattr(v, 'dim') and v.dim == ':,:,:' } before_flops, before_heavy = self.calculate_flops_and_heavy(expr) before_loads = self.count_3d_loads(expr, three_d_arrays) 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] after_flops = 0 after_heavy = 0 after_loads = 0 for temp_var, temp_expr in replacements: f_val, h_val = self.calculate_flops_and_heavy(temp_expr) after_flops += f_val after_heavy += h_val after_loads += self.count_3d_loads(temp_expr, three_d_arrays) f_val, h_val = self.calculate_flops_and_heavy(reduced_expr) after_flops += f_val after_heavy += h_val after_loads += self.count_3d_loads(reduced_expr, three_d_arrays) import sys def pct_str(before, after): if before == 0: return "0.0%" if after == 0 else "+inf%" diff = after - before pct = (diff / before) * 100 return f"{pct:+.1f}%" flops_pct = pct_str(before_flops, after_flops) heavy_pct = pct_str(before_heavy, after_heavy) loads_pct = pct_str(before_loads, after_loads) if after_flops < before_flops * 0.5 or after_loads < before_loads * 0.5: est_speedup = "Highly significant" elif after_flops < before_flops or after_loads < before_loads: est_speedup = "Moderate" else: est_speedup = "Minimal / Already optimal" sys.stderr.write(f"\n[SymPy Optimizer Report: {name}]\n") sys.stderr.write(f"- Floating Point Ops : {before_flops} -> {after_flops} ({flops_pct})\n") sys.stderr.write(f"- Heavy Ops (Div/Sqrt): {before_heavy} -> {after_heavy} ({heavy_pct})\n") sys.stderr.write(f"- 3D Array Mem Reads : {before_loads} -> {after_loads} ({loads_pct})\n") sys.stderr.write(f"=> Estimated Speedup in loop: {est_speedup}\n\n") 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): def __init__ (self, primary, derived, averaged): self.primary = primary self.derived = derived self.averaged = averaged def varlist(self, tree): for v in tree.children: self.primary.add(v.value) self.derived[v.value] = PrimaryField(v.value, self.derived) def assign_var (self, tree): if len(tree.children) > 2: lval, lattr, rval = tree.children else: lval, rval = tree.children lattr = None attr_dict = {} if lattr is not None: for t in lattr.children: k, v = t.children attr_dict[k.value] = v.value if lval.value in self.derived: raise ValueError("duplicate definition of " + lval) self.derived[lval.value] = Field(lval.value, attr_dict, rval, self.derived) def assign_avg_var (self, tree): w = tree.children[0] targets = tree.children[1:] if (not w.children) or (w.children[0] is None): self.averaged[None] = set([x.value for x in targets]) else: self.averaged[w.children[0].value] = set([x.value for x in targets]) class ExpInspector(Visitor): def __init__(self): self.fluctuation = False self.dep = set([]) self.deriv = set([]) @classmethod def inspect(cls, tree): self = cls() return self(tree) def __call__(self, tree): self.visit(tree) return self.fluctuation, self.dep, self.deriv def fluc(self, tree): self.fluctuation = True self.dep.add(tree.children[0].value) def var(self, tree): self.dep.add(tree.children[0].value) def dnx (self, tree): op, v = tree.children deriv = "{}_{}".format(op.data, v.value) self.dep.add(deriv) self.deriv.add((op.data, v.value)) @v_args(inline=True) # Affects the signatures of the methods class ExpToLatex(Transformer): def __init__(self, fdict): self.fdict = fdict def arithmatic_rooted(self, name): try: exproot = self.fdict[name].exp.data except AttributeError: exproot = "something_11fasq2afa3rfzsaerqw23" return ((exproot == "add") or (exproot == "sub") or (exproot == "mul") or (exproot == "div")) def parenthise(self, name): try: latex = self.fdict[name].latex latex_given = self.fdict[name].latex_given except KeyError: warnings.warn(name + " is not found") latex = r"\mathrm{{{}}}".format(name) latex_given = None if self.arithmatic_rooted(name) and (latex_given is None): latex = "(" + latex + ")" return latex def number(self, numeral): return numeral def env(self, name): return r"\mathrm{{{}}}".format(name.value) def paren(self, name): return "({})".format(str(name)) def var(self, name): return self.parenthise(name.value) def fluc(self, name): return self.parenthise(name.value)+ "''" def dnx (self, partial, b): fmt = r"\partial_{{{}}}" coord = partial.data[-1] op = fmt.format(coord + coord if len(partial.data) > 3 else coord) signature = "{}_{}".format(partial.data, b.value) try: eq = self.fdict[signature].latex except KeyError: eq = op + self.parenthise(b.value) warnings.warn(signature + " is not found: " + eq) return eq def icall (self, a, b): if a.data == "sqr": fcode = "({0})^2".format(b) elif a.data == "pow3": fcode = "({0})^3".format(b) else: fcode = "({0})".format(b) return fcode def fcall (self, *args): a = args[0] b = ", ".join(args[1:]) if a == 'sqrt': fcode = r"\sqrt{{{}}}".format(b) return fcode elif a == 'abs': fcode = r"\left| {} \right|".format(b) return fcode elif a.startswith("\\"): fcode = r"{}{{({})}}".format(a, b) return fcode else: fcode = r"\mathrm{{{}}}({})".format(a, b) return fcode def neg(self, b): fcode = "(-{})".format(b) return fcode def add(self, a, b): fcode = "{} + {}".format(a, b) return fcode def sub(self, a, b): fcode = "{} - {}".format(a, b) return fcode def mul(self, a, b): fcode = "{} {}".format(a, b) return fcode def div(self, a, b): fcode = "{} / {}".format(a, b) return fcode log = lambda self : "\log" exp = lambda self : "\exp" sqrt = lambda self : "sqrt" abs = lambda self : "abs" rxn_rate = lambda self : "\omega" udf = lambda self, a : a.value @v_args(inline=True) # Affects the signatures of the methods class ExpToCode(Transformer): def __init__(self, fdict): self.fdict = fdict def number(self, numeral): return str(float(numeral)) def env(self, name): return name.value def paren(self, name): return "({})".format(str(name)) def var(self, name): try: arrname = self.fdict[name.value].array except KeyError: arrname = name.value return arrname + "(i,j,k)" def fluc(self, name): try: arrname = self.fdict[name.value].array except KeyError: arrname = name.value fmt = "({0}(i,j,k) - {{0}}avg_{1}(i))" return fmt.format(arrname, name.value) def dnx (self, partial, b): signature = "{}_{}".format(partial.data, b.value) try: arrname = self.fdict[signature].array except KeyError: arrname = signature return arrname + "(i,j,k)" def icall (self, a, b): if a.data == "sqr": fcode = "(({0})*({0}))".format(b) elif a.data == "pow3": fcode = "(({0})*({0})*({0}))".format(b) else: fcode = "({0})".format(b) return fcode def fcall (self, *args): a = args[0] b = ", ".join(args[1:]) fcode = "( {} ( {} ) )".format(a, b) return fcode def neg(self, b): fcode = "( - {} )".format(b) return fcode def add(self, a, b): fcode = "( {} + {} )".format(a, b) return fcode def sub(self, a, b): fcode = "( {} - {} )".format(a, b) return fcode def mul(self, a, b): fcode = "( {} * {} )".format(a, b) return fcode def div(self, a, b): fcode = "( {} / {} )".format(a, b) return fcode log = lambda self : "log" exp = lambda self : "exp" sqrt = lambda self : "sqrt" abs = lambda self : "abs" rxn_rate = lambda self : "rxn_rate" udf = lambda self, a : a.value def make_allocate(name, shape, init_zero=True): alloc_str = f"allocate({name}({shape}), stat=ierr)\n" alloc_str += f"if (ierr /= 0) then\n" alloc_str += f" write(0,*) 'Error: allocation of {name} failed on process', myid\n" alloc_str += f" call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err)\n" alloc_str += f"end if" if init_zero: alloc_str += f"\n{name} = 0." return alloc_str class FieldBase (object): def __init__ (self, name, fdict): self.name = name self.array = name self.dep = set([]) self.fluc = False self.prime = False self.fdict = fdict self.shape = "nxp,nyp,nzp" self.dim = ":,:,:" def depends_on (self, a): return (a in self.dep) def is_primary (self): return self.prime def not_primary (self): return not self.prime def is_fluctuation (self): return self.fluc def export_on (self): return False def __repr__ (self): return self.name def checkFluctuation (self): fset = set([]) for d in map(self.fdict.get, self.dep): fset.update(d.checkFluctuation()) if self.is_fluctuation() or len(fset) > 0: fset.add(self.name) return fset def depClosure (self): fset = set(self.dep) for d in self.dep: fset.update(self.fdict[d].depClosure()) return fset def code_decl (self): real_array_decl = "real(real64), allocatable, dimension({1}) :: {0}" return real_array_decl.format(self.name, self.dim) def code_alloc (self): return make_allocate(self.name, self.shape) def code_free (self): real_array_free = "deallocate({})" return real_array_free.format(self.name) class FieldExporter (object): mpi_io_decl=''' ! field exporter common integer(kind=MPI_OFFSET_KIND) :: offset ''' # Subarray version fmt_decl_subarray=''' ! - file_handles and mpi_infos integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_fh integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_info integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_filetype ''' fmt_init_subarray=''' ! init subarray datatype for {{ field_name }} block integer(4) :: sizes(3), subsizes(3), starts(3) call MPI_INFO_CREATE({{ field_name }}_info, mpi_err) call MPI_FILE_OPEN(MPI_COMM_TASK,'export-{{ field_name }}.dat',MPI_MODE_WRONLY+MPI_MODE_CREATE,{{ field_name }}_info,{{ field_name }}_fh,mpi_err) sizes = (/ nxp, nyp, nzp /) subsizes = (/ {{ len_xpts }}, {{ ye }} - {{ ys }} + 1, {{ ze }} - {{ zs }} + 1 /) starts = (/ {{ xs }} - 1, {{ ys }} - 1, {{ zs }} - 1 /) call MPI_TYPE_CREATE_SUBARRAY(3, sizes, subsizes, starts, MPI_ORDER_FORTRAN, MPI_REAL8, {{ field_name }}_filetype, mpi_err) call MPI_TYPE_COMMIT({{ field_name }}_filetype, mpi_err) end block ''' fmt_final_subarray=''' ! finalize call MPI_FILE_CLOSE({{ field_name }}_fh, mpi_err) call MPI_INFO_FREE({{ field_name }}_info, mpi_err) call MPI_TYPE_FREE({{ field_name }}_filetype, mpi_err) ''' fmt_calc_subarray=''' ! write to file via MPI Subarray count = ({{ len_xpts }}) * ({{ ye }} - {{ ys }} + 1) * ({{ ze }} - {{ zs }} + 1) offset = export_offset(fidx) * count * 8 call MPI_FILE_WRITE_AT({{ field_name }}_fh, offset, {{ work_array }}, 1, {{ field_name }}_filetype, mpi_status, mpi_err) ''' # Legacy copy version (fallback) fmt_decl_legacy=''' ! - file_handles and mpi_infos integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_fh integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_info ! - buffer real(real64), allocatable, dimension(:,:,:) :: {{ field_name }}_export_array integer, allocatable, dimension(:) :: {{ field_name }}_xpts ''' fmt_init_legacy=''' ! init call MPI_INFO_CREATE({{ field_name }}_info, mpi_err) call MPI_FILE_OPEN(MPI_COMM_TASK,'export-{{ field_name }}.dat',MPI_MODE_WRONLY+MPI_MODE_CREATE,{{ field_name }}_info,{{ field_name }}_fh,mpi_err) allocate({{ field_name }}_export_array(1:{{ len_xpts }},{{ ys }}:{{ ye }},{{ zs }}:{{ ze }}), stat=ierr) if (ierr /= 0) then write(0,*) 'Error: allocation of {{ field_name }}_export_array failed on process', myid call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err) end if {{ field_name }}_export_array = 0. allocate({{ field_name }}_xpts(1:{{ len_xpts }}), stat=ierr) if (ierr /= 0) then write(0,*) 'Error: allocation of {{ field_name }}_xpts failed on process', myid call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err) end if {{ xpts_init }} ''' fmt_final_legacy=''' ! finalize call MPI_FILE_CLOSE({{ field_name }}_fh, mpi_err) call MPI_INFO_FREE({{ field_name }}_info, mpi_err) deallocate({{ field_name }}_export_array) deallocate({{ field_name }}_xpts) ''' fmt_calc_legacy=''' ! copy to array for export do k = {{ zs }}, {{ ze }} do j = {{ ys }}, {{ ye }} do i = 1, {{ len_xpts }} {{ field_name }}_export_array(i,j,k) = {{ work_array }}({{ field_name }}_xpts(i),j,k) end do end do end do ! write to file count = ({{ len_xpts }}) * ({{ ye }} - {{ ys }} + 1) * ({{ ze }} - {{ zs }} + 1) offset = export_offset(fidx) * count * 8 call MPI_FILE_WRITE_AT({{ field_name }}_fh, offset, {{ field_name }}_export_array, count, MPI_REAL8, mpi_status, mpi_err) ''' def __init__ (self, name, attr, parent): self.name = name self.attr = attr self.parent = parent self.params = dict(attr) self.params.setdefault("xs", 1) self.params.setdefault("xe", "nxp") self.params.setdefault("ys", 1) self.params.setdefault("ye", "nyp") self.params.setdefault("zs", 1) self.params.setdefault("ze", "nzp") self.params.setdefault("field_name", self.name) self.params.setdefault("len_xpts", f"({self.params['xe']} - {self.params['xs']} + 1)") fmt_xpts_init = ''' do i = {{ xs }}, {{ xe }} {{ field_name }}_xpts(i-{{ xs }}+1) = i end do ''' self.params.setdefault("xpts_init", Template(fmt_xpts_init).render(**self.params)) try: # Sampling at listed x coordinates fmt_xpts_init_list = "{{ field_name }}_xpts = (/ {{ list_xpts }} /)" import sys raw_xpts = self.params["xpts"] int_xpts = list(map(int, raw_xpts.split())) len_xpts = len(int_xpts) self.params["len_xpts"] = len_xpts self.params["list_xpts"] = ",".join(map(str, int_xpts)) self.params["xpts_init"] = Template(fmt_xpts_init_list).render(**self.params) except KeyError: pass self.use_subarray = ("xpts" not in self.params) def code (self): self.params["work_array"] = self.parent.array if self.use_subarray: return Template(FieldExporter.fmt_calc_subarray).render(**self.params) else: return Template(FieldExporter.fmt_calc_legacy).render(**self.params) def code_decl (self): if self.use_subarray: return Template(FieldExporter.fmt_decl_subarray).render(**self.params) else: return Template(FieldExporter.fmt_decl_legacy).render(**self.params) def code_alloc (self): if self.use_subarray: return Template(FieldExporter.fmt_init_subarray).render(**self.params) else: return Template(FieldExporter.fmt_init_legacy).render(**self.params) def code_free (self): if self.use_subarray: return Template(FieldExporter.fmt_final_subarray).render(**self.params) else: return Template(FieldExporter.fmt_final_legacy).render(**self.params) class Field (FieldBase): def __init__ (self, name, attr, exp, fdict): super(Field,self).__init__(name, fdict) self.attr = attr self.exp = exp self.fluc, self.dep, self.derivs = ExpInspector.inspect(exp) self.comment = self.name + " = " + ExpToCode(self.fdict).transform(self.exp) self.latex_given = self.attr.get("latex") if self.latex_given is None: self.latex = ExpToLatex(self.fdict).transform(self.exp) else: self.latex = self.latex_given self.exporter = None try: if self.attr["export"]: self.exporter = FieldExporter(self.name, self.attr, self) except KeyError: pass ''' for a in self.dep: if a not in self.fdict: raise StandardError(a + " is not defined") ''' def export_on (self): return self.exporter is not None def code (self, alloc=None): 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 {% 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 %} """ 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 "") return calculation_code + export_code class FluctuationField (FieldBase): def __init__ (self, w, field, fset, fdict): super(FluctuationField,self).__init__(self.id(w,field), fdict) if w is not None: self.w = w + "_" else: self.w = "" self.field = fdict[field] self.dep = self.field.dep - fset for df in self.field.dep & fset: self.dep.add(self.id(w,df)) self.comment = ExpToCode(self.fdict).transform(self.field.exp) if self.field.is_fluctuation(): self.comment = self.comment.format(self.w) self.comment = self.name + " = " + self.comment def code (self, alloc=None): 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 do i = 1, nxp {{ array }}(i,j,k) = {{ rhs }} end do end do end do """ return Template(real_array_loop).render( comment=self.comment, array=self.array, rhs=rhs ) @classmethod def id (cls, w, field): if w: name = "{}____{}_avg".format(field, w) else: name = "{}____avg".format(field) return name class PrimaryField (FieldBase): def __init__ (self, name, fdict): super(PrimaryField,self).__init__(name, fdict) self.derivs = set([]) self.prime = True self.latex = name self.latex_given = None def code (self, alloc=None): return "! {} is read from file".format(self.name) def code_decl (self): return "! {} is read from file".format(self.name) def code_alloc (self): return "! {} is read from file".format(self.name) def code_free (self): return "! {} is read from file".format(self.name) class DerivedField (FieldBase): def __init__ (self, op, v, fdict): name = "{}_{}".format(op, v) super(DerivedField,self).__init__(name, fdict) self.op = op self.v = v self.dep = set([v]) fmt = r"\partial_{{{}}}" coord = op[-1] partial = fmt.format(coord + coord if len(op) > 3 else coord) self.latex = partial + "(" + fdict[v].latex + ")" def code (self, alloc=None): self.array = alloc[self.name] if alloc else self.name varray = alloc[self.v] if alloc else self.v return "call {0} ( {2}, {1} )".format(self.op, varray, self.array) class AveragedField (FieldBase): @classmethod def id (cls, w, tgt): if w: return "{}_avg_{}".format(w, tgt) else: return "avg_{}".format(tgt) def __init__ (self, w, tgt, fdict): name = self.id(w,tgt) super(AveragedField,self).__init__(name, fdict) self.shape = "nxp" self.dim = ":" self.target = tgt tfield = fdict[tgt] self.fset = tfield.checkFluctuation() self.latex = r"\left\langle {} \right\rangle".format(tfield.latex) if not self.fset: self.tgt = tgt self.dep.add(tgt) else: ftgt = FluctuationField.id(w,tgt) self.tgt = ftgt self.dep.add(ftgt) self.weighted = w if w: self.w = fdict[w] self.dep.add(w) self.latex += ("_{{{}}}".format(w)) def code (self, alloc=None): avg_array_sum = """ do k = 1, nzp do j = 1, nyp do i = 1, nxp {{ name }}(i) = {{ name }}(i) + {{ arrname }} end do end do end do """ arrname = self.fdict[self.tgt].array + "(i,j,k)" if self.weighted is not None: arrname = arrname + " * " + self.w.array + "(i,j,k)" return Template(avg_array_sum).render(name=self.name, arrname=arrname) def code_avg (self): avg_array_divide = """ call MPI_ALLREDUCE(MPI_IN_PLACE, {{ name }}, nxp, MPI_REAL8, MPI_SUM, MPI_COMM_TASK, mpi_err) {{ name }} = {{ name }} {{ dWeight }} / denum """ dWeight = (f"/ avg_{self.weighted}" if self.weighted else "") return Template(avg_array_divide).render(name=self.name, dWeight=dWeight) def isWeighted (self): return self.weighted is not None def pass1 (self): return not self.pass2() def pass2 (self): return len(self.fset) > 0 class Stage1(): """First compilation stage. Performs conversion of Lark AST tree into Python field datasets. This uses a visitor class to collect all algebraic equations, primary variables, derived fields, and average specs from the parsed DSL input. """ def __init__ (self, raw_tree): """Initializes Stage 1 by parsing the raw AST tree. Args: raw_tree (lark.Tree): The parsed AST representation of the DSL input. """ self.primary = set([]) self.derived = {} self.averaged = {} # Construct Field objects CollectDefinitions( self.primary, self.derived, self.averaged ).visit(raw_tree) def __repr__ (self): return "\n".join(map(str, [self.primary, self.derived, self.averaged])) class Stage2(): """Second compilation stage. Expands derivative and fluctuation terms. This stage traverses the collected definitions, resolving and creating matching derived fields for derivatives, fluctuations, and weighted averages. """ def __init__ (self, src): """Initializes Stage 2 by expanding variables from the previous stage. Args: src (Stage1): Completed Stage 1 compilation dataset. """ self.src = src self.primary = src.primary self.derived = src.derived self.derivative = {} self.averaged = {} # Construct Derivative Field objects dset = set([]) for k, v in self.derived.items(): dset.update(v.derivs) for tup in dset: a = DerivedField(tup[0], tup[1], self.derived) self.derived[a.name] = a self.derivative[tup] = a # Construct Averaged Field objects for w, tgts in src.averaged.items(): for t in tgts: a = AveragedField(w, t, self.derived) self.averaged[a.name] = a for ff in a.fset: b = FluctuationField(w, ff, a.fset, self.derived) self.derived[b.name] = b def __repr__ (self): return "\n".join(map(str, [self.derived, self.derivative, self.averaged])) def dependency (self): dgraph = {} for k,v in self.derived.items(): dgraph[k] = v.dep for k,v in self.averaged.items(): dgraph[k] = v.dep return dgraph class Stage3(): """Third compilation stage. Calculates the topological execution order. Resolves dependencies between algebraic equations and calculates the correct order of calculations. It splits calculation passes into two distinct execution blocks (Pre-averaging Pass 1, and Post-averaging Pass 2) using a topological sorter. """ def __init__ (self, src): """Initializes Stage 3 by resolving execution flows. Args: src (Stage2): Completed Stage 2 dataset. """ self.src = src self.primary = src.primary self.derived = src.derived self.averaged = src.averaged self.dependency = src.dependency() assert set(self.derived.keys()).isdisjoint(self.averaged.keys()) self.avg1 = set(filter(AveragedField.pass1, self.averaged.values())) self.avg2 = set(filter(AveragedField.pass2, self.averaged.values())) pass1calc = set(map(repr, self.avg1)) for x in self.avg1: pass1calc.update(x.depClosure()) self.pass1 = self.sort_vars_new(self.dependency, pass1calc - self.primary) # self.sort_vars(self.dependency, pass1calc - self.primary -set(map(repr, self.avg1))) pass2calc = set(map(repr, self.avg2)) for x in self.avg2: pass2calc.update(x.depClosure()) self.pass2 = self.sort_vars_new(self.dependency, pass2calc - self.primary) def __repr__ (self): return "\n".join(map(str, [self.pass1, self.pass2])) def sort_vars (self, dependency, group): order = [] remain = list(group) remain.sort() while len(remain) > 0: for v in remain: if dependency[v].isdisjoint(remain): order.append(v) remain.remove(v) return order def calc_size (self, ordered, remaining): count = 0 dep_union = set([]) g = self.dependency for v in remaining: dep_union |= set(g[v]) for v in ordered: if v in dep_union: count += 1 return count def sort_vars_new (self, dependency, group): order = [] remain = list(group) remain.sort() while len(remain) > 0: candidate = [] for v in remain: if set(dependency[v]).isdisjoint(remain): candidate.append(v) impact = {} size0 = self.calc_size(set(order), set(remain)) for v in candidate: impact[v] = self.calc_size(set(order) | set([v]), set(remain) - set([v])) - size0 candidate.sort(key=impact.get) order.append(candidate[0]) remain.remove(candidate[0]) return order def print_program (self): allvar = dict(self.derived) allvar.update(self.averaged) decl = "\n".join(allvar[v].code_decl() for v in set(self.pass1+self.pass2)) alloc = "\n".join(allvar[v].code_alloc() for v in set(self.pass1+self.pass2)) free = "\n".join(allvar[v].code_free() for v in set(self.pass1+self.pass2)) calc1 = "\n".join(allvar[v].code() for v in self.pass1) calc2 = "\n".join(allvar[v].code() for v in self.pass2) set1 = [a.name for a in filter(AveragedField.pass1, self.averaged.values())] set2 = [a.name for a in filter(AveragedField.pass2, self.averaged.values())] set1.sort() set2.sort() avg1 = "\n".join(allvar[v].code_avg() for v in set1) avg2 = "\n".join(allvar[v].code_avg() for v in set2) hfmt = 'character (len = *), parameter :: output_header="{}"' declh = hfmt.format(" ".join(["x"] + set1 + set2)) avg_array_write = ''' integer :: i open (200, file="qEdge_X.dat") write (200,*) output_header do i=1,nxp write (200,'({0}e20.10)') real(i)*hxp, {1} end do close (200) ''' avgarr = "{}(i)" write_avg = avg_array_write.format( len(self.averaged)+1, ", ".join(map(avgarr.format, set1+set2)) ) md = {} md["module_name"] = "terms" md["module_data"] = "\n".join((declh, decl)) md["module_init"] = alloc md["module_finalize"] = free md["module_pass1"] = calc1 md["module_pass1_avg"] = avg1 md["module_pass2"] = calc2 md["module_pass2_avg"] = avg2 md["module_write_result"] = write_avg return md class Stage4(): """Fourth compilation stage. Performs variable liveness analysis and memory pooling. Analyzes variable lifetimes inside loops to perform cache-friendly array pooling. It leverages SymPy to simplify expressions, count flops, extract Common Subexpressions (CSE), and compile highly optimized Fortran calculations with minimal memory footprint. """ def __init__ (self, src): """Initializes Stage 4. Args: src (Stage3): Completed Stage 3 execution order. """ self.src = src self.primary = src.primary self.derived = src.derived self.averaged = src.averaged self.dependency = src.dependency self.avg1 = src.avg1 self.avg2 = src.avg2 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)) narr2, alloc2 = (self.allocate_arr(self.pass2)) self.narr = max(narr1, narr2) self.alloc1 = alloc1 self.alloc2 = alloc2 def liveness (self, l1, g): import numpy as np img = np.zeros((len(l1), len(l1))) for i, v in enumerate(l1): for j in range(i, len(l1)): img[i,i:j+1] = img[i,i:j+1] + (1 if v in g[l1[j]] else 0) return img > 0 def allocate_arr (self, l): import numpy as np dg = self.dependency mask = self.liveness(l, dg) try: narr = mask.astype(int).sum(axis=0).max() except ValueError: narr = 0 array_pool = set([self.array_name.format(i) for i in range(narr)]) livesets = [set([])] + [set(np.asarray(l)[row]) for row in mask.T] var2arr = { p : p for p in self.primary } for i, (s0, s1) in enumerate(zip(livesets[:-1], livesets[1:])): array_pool.update(map(var2arr.get, s0 - s1)) for new in s1 - s0: var2arr[new] = array_pool.pop() return narr, var2arr def work_array_codes (self): array_names = [self.array_name.format(i) for i in range(self.narr)] real_array_decl = "real(real64), allocatable, dimension(:,:,:) :: {0}" decl = "\n".join([real_array_decl.format(v) for v in array_names]) alloc = "\n".join([make_allocate(v, "nxp,nyp,nzp") for v in array_names]) free = "\n".join(["deallocate({})".format(v) for v in array_names]) return decl, alloc, free def write_avg_codes (self, avglist): avg_array_write = ''' real(real64), dimension(nxp) :: xbuffer integer :: i open (200, file="qEdge_X.dat") write (200,*) output_header do i=1,nxp write (200,'({{ num_args }}e20.10)') real(i)*hxp, {{ formatted_avglist }} end do close (200) open (200, file="d1.dat") {{ deriv1_lines }} close (200) open (200, file="d2.dat") {{ deriv2_lines }} close (200) ''' avgarr = "{}(i)" deriv1_avgarr = """call ddx1d ( xbuffer, {} ) ; write (200,*) xbuffer""" deriv2_avgarr = """call d2dx1d ( xbuffer, {} ) ; write (200,*) xbuffer""" num_args = len(self.averaged) + 1 formatted_avglist = ", ".join(map(avgarr.format, avglist)) deriv1_lines = "\n".join(map(deriv1_avgarr.format, avglist)) deriv2_lines = "\n".join(map(deriv2_avgarr.format, avglist)) write_avg = Template(avg_array_write).render( num_args=num_args, formatted_avglist=formatted_avglist, deriv1_lines=deriv1_lines, deriv2_lines=deriv2_lines ) return write_avg 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) set1 = sorted([a.name for a in filter(AveragedField.pass1, self.averaged.values())]) set2 = sorted([a.name for a in filter(AveragedField.pass2, self.averaged.values())]) set_export_on = list(filter(lambda x: x.export_on(), self.derived.values())) ffmt = 'logical, parameter :: pass2_required={}' declf = ffmt.format('.true.' if len(set2) > 0 else '.false.') hfmt = 'character (len = *), parameter :: output_header="{}"' declh = hfmt.format(" ".join(["x"] + set1 + set2)) declarr, allocarr, freearr = self.work_array_codes() declavg = "\n".join(self.averaged[v].code_decl() for v in sorted(self.averaged)) allocavg = "\n".join(self.averaged[v].code_alloc() for v in sorted(self.averaged)) freeavg = "\n".join(self.averaged[v].code_free() for v in sorted(self.averaged)) decl_export = "\n".join(v.exporter.code_decl() for v in set_export_on) 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 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) sub_write_avg = self.write_avg_codes(set1+set2) md = {} md["module_name"] = "terms" md["module_data"] = "\n".join((declf, declh, declavg, FieldExporter.mpi_io_decl, decl_export, declarr)) md["module_init"] = "\n".join((allocavg, alloc_export, allocarr)) md["module_finalize"] = "\n".join((freeavg, free_export, freearr)) md["module_pass1"] = sub_calc1 md["module_pass1_avg"] = sub_avg1 md["module_pass2"] = sub_calc2 md["module_pass2_avg"] = sub_avg2 md["module_write_result"] = sub_write_avg return md def save_ir (self): import json dg = {k:list(v) for k,v in self.dependency.items()} with open("ir2.py", "w") as irf: print("g = ", json.dumps(dg, indent=4), file=irf) print("l1 = ", json.dumps(self.pass1, indent=4), file=irf) print("l2 = ", json.dumps(self.pass2, indent=4), file=irf) print("avg1 = ", json.dumps(list(map(repr,self.avg1)), indent=4), file=irf) print("avg2 = ", json.dumps(list(map(repr,self.avg2)), indent=4), file=irf) class Stage5(): ''' pass1 and pass2 seperation and calculation ordering ''' class Stage6(): ''' generate fortran code '''