import sympy as sp from sympy.codegen import ast import itertools as it import networkx as nx from .parser import OPERATORS, sympy_to_dag from .util import DecodeError def isnum(s): try: float(s) return True except ValueError: return False class Implementor: def __init__(self, expr, constants={}, dtype="double"): self.expr = expr self.constants = constants self.cdtype = dtype self.cpf = "lf" if dtype == "double" else "f" self.fdtype = "double precision" if dtype == "double" else "real" def implement(self, impl): if impl == "dag_c": return self.dag_to_c_impl() elif impl == "cse_c": return self.sympy_cse_c_impl() elif impl == "dag_fortran": return self.dag_to_fortran_impl() elif impl == "cse_fortran": return self.sympy_cse_fortran_impl() def op_c_impl(self, f, children): if f == "add": return " + ".join(children); elif f == "mul": return " * ".join(children); elif f == "pow": assert len(children) == 2 if self.cdtype == "double": return f"pow({children[0]}, {children[1]})" else: return f"powf({children[0]}, {children[1]})" elif f == "ln": assert len(children) == 1 if self.cdtype == "double": return f"log({children[0]})" else: return f"logf({children[0]})" else: if f in OPERATORS and OPERATORS[f][1] == 1: assert len(children) == 1 if self.cdtype == "double": return f"{f}({children[0]})" else: return f"{f}f({children[0]})" else: raise DecodeError(f"C impl: operation {f} not handled") def op_f_impl(self, f, children): if f == "add": j = ")+(".join(children) return "(" + j + ")" elif f == "mul": j = ")*(".join(children) return "(" + j + ")" elif f == "pow": assert len(children) == 2 return f"({children[0]})**({children[1]})" elif f == "ln": assert len(children) == 1 return f"log({children[0]})" else: if f in OPERATORS and OPERATORS[f][1] == 1: assert len(children) == 1 return f"{f}({children[0]})" else: raise DecodeError(f"F impl: operation {f} not handled") def full_c_code(self, body): pre = f"#include \n#include \n{self.cdtype} myfunc({self.cdtype} x) {{" post = f"}}\nint main() {{ {self.cdtype} x; scanf(\"%{self.cpf}\", &x); printf(\"%{self.cpf}\", myfunc(x)); }}" return f"{pre}\n{body}\n{post}" def full_f_code(self, body): pre = "function myfunc(x) result(y)\nimplicit none\n" + \ f"{self.fdtype}, intent(in) :: x\n{self.fdtype} :: y, E, pi\n" post = "end function myfunc\nprogram main\nimplicit none\n" + \ f"{self.fdtype} :: x\n{self.fdtype} :: myfunc\n" + \ "read(*, *) x\nprint *, \"y is:\", myfunc(x)\nend program main" return f"{pre}\n{body}\n{post}" def dag_to_c_impl(self): dag = sympy_to_dag(self.expr, csuf="F" if self.cdtype == "float" else "") cstr = "" added_pi, added_E = False, False for c in self.constants: cstr += f"{self.cdtype} {c} = {self.constants[c]};\n" varidx = it.count() for node in reversed(list(nx.topological_sort(dag))): label = dag.nodes[node]["label"] children = [dag.nodes[n]["var"] for n in dag.adj[node]] if len(children) == 0: if label == "pi": if self.cdtype == "float" and not added_pi: cstr += "const float pi = 3.14159265F;\n" added_pi = True else: label = "M_PI" elif label == "E": if self.cdtype == "float" and not added_E: cstr += "const float E = 2.71828183F;\n" added_E = True else: label = "M_E" dag.nodes[node]["var"] = label continue varname = f"t{next(varidx)}" cexpr = self.op_c_impl(label, children) dag.nodes[node]["var"] = varname cstr += f"{self.cdtype} {varname} = {cexpr};\n" retname = varname cstr += f"return {retname};\n" return self.full_c_code(cstr) def dag_to_fortran_impl(self): csuf = "" if self.fdtype == "real" else "d0" dag = sympy_to_dag(self.expr, csuf=csuf) varstr = "" fstr = "parameter E = 2.71828183\nparameter pi = 3.14159265\n" for c in self.constants: varstr += f"{self.fdtype} :: {c}\n" fstr += f"parameter {c} = {self.constants[c]}{csuf}\n" varidx = it.count() allvars = [] for node in reversed(list(nx.topological_sort(dag))): label = dag.nodes[node]["label"] children = [dag.nodes[n]["var"] for n in dag.adj[node]] if len(children) == 0: dag.nodes[node]["var"] = label continue varname = f"t{next(varidx)}" fexpr = self.op_f_impl(label, children) dag.nodes[node]["var"] = varname fstr += f"{varname} = {fexpr}\n" retname = varname varstr += f"{self.fdtype} :: {varname}\n" fstr += f"y = {retname};\n" fstr = varstr + "\n" + fstr return self.full_f_code(fstr) def sympy_cse_c_impl(self): if self.cdtype == "float": extraargs = { "type_aliases": {ast.real: ast.float32}, "math_macros": {}, } else: extraargs = {} cstr = "" for c in self.constants: cstr += f"{self.cdtype} {c} = {self.constants[c]};\n" xvars, xpr = sp.cse(self.expr) for vname, vxpr in xvars: code = sp.ccode(vxpr, assign_to=vname.name, **extraargs) cstr += f"{self.cdtype} {vname.name}; {code};\n" assert len(xpr) == 1 code = sp.ccode(xpr[0], assign_to="y", **extraargs) cstr += f"{self.cdtype} y; {code}; return y;\n" return self.full_c_code(cstr) def sympy_cse_fortran_impl(self): csuf = "" if self.fdtype == "real" else "d0" varstr = "" fstr = "" for c in self.constants: varstr += f"{self.fdtype} :: {c}\n" fstr += f"parameter {c} = {self.constants[c]}{csuf}\n" xvars, xpr = sp.cse(self.expr) for vname, vxpr in xvars: varstr += f"{self.fdtype} :: {vname.name}\n" fstr += sp.fcode(vxpr, assign_to=vname.name, standard=95, source_format="free") + "\n" assert len(xpr) == 1 fstr += sp.fcode(xpr[0], assign_to="y", standard=95, source_format="free") + "\n" fstr = varstr + "\n" + fstr if self.fdtype == "real": # Hack to fix sympy generation fstr = fstr.replace("d0", "") return self.full_f_code(fstr) # For testing only if __name__ == "__main__": from .parser import parse_prefix_to_sympy, sympy_to_dag prefs = "add mul div INT+ 1 INT+ 5 x mul div INT+ 1 INT+ 5 mul x tan pow x INT+ 2".split(" ") exp = parse_prefix_to_sympy(prefs) impl = Implementor(exp, dtype="float") print("DAG C:") print(impl.dag_to_c_impl()) print("DAG Fortran:") print(impl.dag_to_fortran_impl()) print("CSE C:") print(impl.sympy_cse_c_impl()) print("CSE Fortran:") print(impl.sympy_cse_fortran_impl())