Add resolve condition pass

dace/transformation/passes/__init__.py

@@ -11,5 +11,6 @@
 from .scalar_to_symbol import ScalarToSymbolPromotion
 from .simplify import SimplifyPass
 from .transient_reuse import TransientReuse
+from .resolve_condition import ResolveCondition
 
 from .util import available_passes

dace/transformation/passes/resolve_condition.py

@@ -0,0 +1,88 @@
+# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved.
+""" Resolve condition transformation """
+
+from typing import Any, Dict
+
+from dace.transformation import pass_pipeline as ppl
+import sympy as sy
+from dace.properties import CodeBlock
+from dace import sdfg as sd, symbolic
+from dace.properties import Property, make_properties, CodeBlock
+from sympy.parsing.sympy_parser import parse_expr
+from dace import symbolic
+from dace.frontend.python.astutils import unparse
+from dace.transformation.pass_pipeline import Pipeline
+import ast
+
+
+def eliminate_branch(sdfg: sd.SDFG, initial_edge: sd.graph.Edge):
+    sdfg.remove_edge(initial_edge)
+    if sdfg.in_degree(initial_edge.dst) > 0:
+        return
+    state_list = [initial_edge.dst]
+    while len(state_list) > 0:
+        new_state_list = []
+        for s in state_list:
+            for e in sdfg.out_edges(s):
+                if len(sdfg.in_edges(e.dst)) == 1:
+                    new_state_list.append(e.dst)
+                    sdfg.remove_edge(e)
+            sdfg.remove_node(s)
+        state_list = new_state_list
+
+
+@make_properties
+class ResolveCondition(ppl.Pass):
+    """
+    Given a condition (e.g. var == 1) assumes the condition always holds true, and removes unreachable states
+    """
+
+    CATEGORY: str = 'Simplification'
+
+    # Properties
+    condition = Property(dtype=str, default="", desc="condition to be parsed")
+
+    def should_reapply(self, modified: ppl.Modifies) -> bool:
+        return modified & (ppl.Modifies.States | ppl.Modifies.InterstateEdges)
+
+    def modifies(self) -> ppl.Modifies:
+        return ppl.Modifies.States | ppl.Modifies.InterstateEdges
+
+    def apply_pass(self, sdfg: sd.SDFG, _: Dict[str, Any]) -> None:
+        # obtain loop information
+        try:
+            parsed_condition = parse_expr(unparse(symbolic.PythonOpToSympyConverter().visit(
+                    ast.parse(self.condition).body[0]))).simplify()
+        except:
+            return
+
+        seen = set()
+        found = True
+        while found:
+            found = False
+            for e in sdfg.edges():
+                if e.data.is_unconditional() or e in seen:
+                    continue
+                # cache seen edges for performance
+                seen.add(e)
+                try:
+                    cond = unparse(symbolic.PythonOpToSympyConverter().visit(
+                            ast.parse(e.data.condition.as_string).body[0]))
+                    cur_parsed_condition = parse_expr(cond)
+                    if sy.And(cur_parsed_condition, parsed_condition).simplify() == False:
+                        found = True
+                        eliminate_branch(sdfg, e)
+                        break
+                except:
+                    pass
+
+        from dace.transformation.passes import DeadStateElimination
+
+        pipeline = Pipeline([DeadStateElimination()])
+        pipeline.apply_pass(sdfg, {})
+
+        # make all lone edges unconditional
+        for n in sdfg.nodes():
+            out_edges = sdfg.out_edges(n)
+            if len(out_edges) == 1 and not out_edges[0].data.is_unconditional():
+                out_edges[0].data.condition = CodeBlock("1")

tests/transformations/conditional_elimination_test.py

@@ -0,0 +1,55 @@
+# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved.
+
+import dace
+from dace.transformation.passes import ResolveCondition
+from dace.transformation.dataflow import MapFusion
+import numpy as np
+import pytest
+
+@pytest.mark.parametrize("condition, after_pass, after_simplify, map_fusion_applications", [
+    ('flag', 3, 1, 2),
+    ('not flag', 3, 1, 2),
+    ('dummy', 4, 4, 0)
+])
+
+
+def test_application_and_fusion(condition: str, after_pass: int, after_simplify: int, map_fusion_applications: int):
+    N = dace.symbol('N', dtype=dace.int32)
+
+    @dace.program
+    def simple_program(flag: dace.bool, arr: dace.float32[N]):
+        tmp1 = np.empty_like(arr)
+        tmp2 = np.empty_like(arr)
+        for i in dace.map[0:N]:
+            tmp1[i] = arr[i] + 1
+        if flag:
+            for i in dace.map[0:N]:
+                tmp2[i] = tmp1[i] * 2
+        else:
+            for i in dace.map[0:N]:
+                tmp2[i] = tmp1[i] * 3
+        for i in dace.map[0:N]:
+            arr[i] = tmp2[i] + 1
+
+
+    sdfg = simple_program.to_sdfg(simplify=True)
+
+    assert len(sdfg.nodes()) == 4
+
+    p = ResolveCondition()
+    p.condition = condition
+    p.apply_pass(sdfg, {})
+
+    assert len(sdfg.nodes()) == after_pass
+
+    sdfg.simplify()
+
+    assert len(sdfg.nodes()) == after_simplify
+
+    assert sdfg.apply_transformations_repeated([MapFusion]) == map_fusion_applications
+
+
+if __name__ == '__main__':
+    test_application_and_fusion('flag', 3, 1, 2)
+    test_application_and_fusion('not flag', 3, 1, 2)
+    test_application_and_fusion('dummy', 4, 4, 0)