b37f869cc5
All checks were successful
Deploy Documentation / build-and-deploy (push) Successful in 1m10s
1543 lines
46 KiB
Python
1543 lines
46 KiB
Python
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 '''
|