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 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(): ''' conversion from tree to python data ''' def __init__ (self, raw_tree): 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(): ''' expand derivatives and averages''' def __init__ (self, src): 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(): ''' calculate execution order ''' def __init__ (self, src): 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(): ''' analyze liveness and allocate array ''' def __init__ (self, src): 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 '''