scal2sym: Try to symbolically simplify expressions on promotion

spcl · Aug 7, 2023 · 38e623f · 38e623f
1 parent 245b028
commit 38e623f
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/dace/transformation/passes/scalar_to_symbol.py b/dace/transformation/passes/scalar_to_symbol.py
@@ -5,6 +5,7 @@
 import collections
 import re
 from dataclasses import dataclass
+import sympy as sp
 from typing import Any, DefaultDict, Dict, Set, Tuple
 
 import dace
@@ -661,6 +662,13 @@ def apply_pass(self, sdfg: SDFG, _: Dict[Any, Any]) -> Set[str]:
                     elif input.language is dtypes.Language.CPP:
                         newcode = translate_cpp_tasklet_to_python(input.code.as_string.strip())
 
+                    # Try to symbolically simplify expression
+                    try:
+                        symcode = symbolic.pystr_to_symbolic(newcode)
+                        newcode = symbolic.symstr(symcode)
+                    except (TypeError, sp.SympifyError):
+                        pass
+
                     # Replace tasklet inputs with incoming edges
                     for e in new_state.in_edges(input):
                         memlet_str: str = e.data.data