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