Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add resolve condition pass #1645

Draft
wants to merge 5 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions dace/transformation/passes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
88 changes: 88 additions & 0 deletions dace/transformation/passes/resolve_condition.py
Original file line number Diff line number Diff line change
@@ -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):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like yes, it should become a reusable, documented utility function :-)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed! This one is a good candidate for "building block" used by multiple transformations

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")
55 changes: 55 additions & 0 deletions tests/transformations/conditional_elimination_test.py
Original file line number Diff line number Diff line change
@@ -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)
Loading