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.

Sign up for GitHub

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

Updated the add_state_{after, before}() function. #1556

Merged
merged 1 commit into from
Apr 10, 2024
Merged

Updated the add_state_{after, before}() function. #1556

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
38 changes: 28 additions & 10 deletions dace/sdfg/state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2477,38 +2477,56 @@ def add_state(self, label=None, is_start_block=False, *, is_start_state: bool=No
self.add_node(state, is_start_block=start_block)
return state

def add_state_before(self, state: SDFGState, label=None, is_start_state=False) -> SDFGState:
def add_state_before(self,
state: SDFGState,
label=None,
is_start_block=False,
condition: CodeBlock = None,
assignments=None,
*,
is_start_state: bool=None) -> SDFGState:
""" Adds a new SDFG state before an existing state, reconnecting predecessors to it instead.

:param state: The state to prepend the new state before.
:param label: State label.
:param is_start_state: If True, resets scope block starting state to this state.
:param is_start_block: If True, resets scope block starting state to this state.
:param condition: Transition condition of the newly created edge between state and the new state.
:param assignments: Assignments to perform upon transition.
:return: A new SDFGState object.
"""
new_state = self.add_state(label, is_start_state)
new_state = self.add_state(label, is_start_block=is_start_block, is_start_state=is_start_state)
# Reconnect
for e in self.in_edges(state):
self.remove_edge(e)
self.add_edge(e.src, new_state, e.data)
# Add unconditional connection between the new state and the current
self.add_edge(new_state, state, dace.sdfg.InterstateEdge())
# Add the new edge
self.add_edge(new_state, state, dace.sdfg.InterstateEdge(condition=condition, assignments=assignments))
return new_state

def add_state_after(self, state: SDFGState, label=None, is_start_state=False) -> SDFGState:
def add_state_after(self,
state: SDFGState,
label=None,
is_start_block=False,
condition: CodeBlock = None,
assignments=None,
*,
is_start_state: bool=None) -> SDFGState:
""" Adds a new SDFG state after an existing state, reconnecting it to the successors instead.

:param state: The state to append the new state after.
:param label: State label.
:param is_start_state: If True, resets SDFG starting state to this state.
:param is_start_block: If True, resets scope block starting state to this state.
:param condition: Transition condition of the newly created edge between state and the new state.
:param assignments: Assignments to perform upon transition.
:return: A new SDFGState object.
"""
new_state = self.add_state(label, is_start_state)
new_state = self.add_state(label, is_start_block=is_start_block, is_start_state=is_start_state)
# Reconnect
for e in self.out_edges(state):
self.remove_edge(e)
self.add_edge(new_state, e.dst, e.data)
# Add unconditional connection between the current and the new state
self.add_edge(state, new_state, dace.sdfg.InterstateEdge())
# Add the new edge
self.add_edge(state, new_state, dace.sdfg.InterstateEdge(condition=condition, assignments=assignments))
return new_state

@abc.abstractmethod
Expand Down
Loading