Experimental feature: Schedule tree representation#2262
Open
Experimental feature: Schedule tree representation#2262
Conversation
… undefined variables
Backport of PR spcl#1917. Co-authored-by: Roman Cattaneo <1116746+romanc@users.noreply.github.com>
romanc
changed the title
romanc
added a commit
to GridTools/gt4py
that referenced
this pull request
Feb 17, 2026
…aCe (#2458) ## Description In this PR we pull a preliminary version of the `ScheduleTree` concept based on top of current mainline DaCe, e.g., what will be known as DaCe v2.x. The work in this branch is based on PR spcl/dace#2262 in DaCe, which is functionally complete for our use-cases, but not yet fully ready to merge down to mainline DaCe. Once spcl/dace#2262 is merged into mainline DaCe, we'll create a follow-up PR (which might simplify the DaCe integration in GT4Py since now both, cartesian and next, are on a compatible version of DaCe). ## Requirements - [x] All fixes and/or new features come with corresponding tests. - [ ] Important design decisions have been documented in the appropriate ADR inside the [docs/development/ADRs/](docs/development/ADRs/README.md) folder. N/A
tbennun
requested changes
Mar 2, 2026
Collaborator
tbennun
left a comment
tbennun
left a comment
There was a problem hiding this comment.
Looking great! I would like to see it more feature complete (as it is very close to being!) and I left some comments on the code. Maybe more preprocessing passes can make the implementations simpler
| raise RuntimeError("Expected schedule tree root.") | ||
|
|
||
| for child in stree.children: | ||
| if id(child.parent) != id(stree): |
Collaborator
There was a problem hiding this comment.
should use child.parent is not stree - better and has similar effect
Comment on lines
+26
to
+66
| @dataclass | ||
| class Context: | ||
| root: 'ScheduleTreeRoot' | ||
| current_scope: Optional['ScheduleTreeScope'] | ||
|
|
||
| access_cache: Dict[Tuple[SDFGState, str], Dict[str, nodes.AccessNode]] | ||
| """Per scope (hashed by id(scope_node) access_cache.""" | ||
|
|
||
|
|
||
| class ContextPushPop: | ||
| """Append the given node to the scope, then push/pop the scope.""" | ||
|
|
||
| def __init__(self, ctx: Context, state: SDFGState, node: 'ScheduleTreeScope') -> None: | ||
| if ctx.current_scope is None and not isinstance(node, ScheduleTreeRoot): | ||
| raise ValueError("ctx.current_scope is only allowed to be 'None' when node it tree root.") | ||
|
|
||
| self._ctx = ctx | ||
| self._parent_scope = ctx.current_scope | ||
| self._node = node | ||
| self._state = state | ||
|
|
||
| cache_key = (state, id(node)) | ||
| assert cache_key not in self._ctx.access_cache | ||
| self._ctx.access_cache[cache_key] = {} | ||
|
|
||
| def __enter__(self) -> None: | ||
| assert not self._ctx.access_cache[(self._state, id( | ||
| self._node))], "Expecting an empty access_cache when entering the context." | ||
|
|
||
| self._ctx.current_scope = self._node | ||
|
|
||
| def __exit__( | ||
| self, | ||
| exc_type: Optional[type[BaseException]], | ||
| exc_val: Optional[BaseException], | ||
| exc_tb: Optional[TracebackType], | ||
| ) -> None: | ||
| cache_key = (self._state, id(self._node)) | ||
| assert cache_key in self._ctx.access_cache | ||
|
|
||
| self._ctx.current_scope = self._parent_scope |
Collaborator
There was a problem hiding this comment.
what are these classes? They are missing docstrings and their purpose
| return [tn.StateBoundaryNode(), visited] | ||
| return visited | ||
|
|
||
| stree = NestedSDFGStateBoundaryInserter().visit(stree) |
Collaborator
There was a problem hiding this comment.
oh dear, needs a RefSet state boundary inserter
| :return: The newly created state. | ||
| """ | ||
| if behavior != StateBoundaryBehavior.STATE_TRANSITION: | ||
| raise NotImplementedError("Only STATE_TRANSITION is supported as StateBoundaryBehavior in this prototype.") |
Collaborator
There was a problem hiding this comment.
Suggested change
| raise NotImplementedError("Only STATE_TRANSITION is supported as StateBoundaryBehavior in this prototype.") | |
| raise NotImplementedError("Only STATE_TRANSITION is currently supported as StateBoundaryBehavior.") |
| if behavior != StateBoundaryBehavior.STATE_TRANSITION: | ||
| raise NotImplementedError("Only STATE_TRANSITION is supported as StateBoundaryBehavior in this prototype.") | ||
|
|
||
| # TODO: Some boundaries (control flow, state labels with goto) could not be fulfilled with every |
Collaborator
There was a problem hiding this comment.
Suggested change
| # TODO: Some boundaries (control flow, state labels with goto) could not be fulfilled with every | |
| # TODO: Some boundaries (control flow, state labels with goto, pending assignments) could not be fulfilled with every |
| after_state: Optional[ControlFlowBlock] = None, | ||
| *, | ||
| label: Optional[str] = None, | ||
| assignments: Optional[Dict] = None, |
Collaborator
There was a problem hiding this comment.
type hint incomplete (key, value?)
| insertion order since python 3.7 (which we rely on in this function | ||
| too). Depending on code generation it could(TM) be that we can | ||
| weaken (best case remove) the corresponding check from the sdfg | ||
| validator. |
Collaborator
There was a problem hiding this comment.
The check on the validator is on purpose (and we have seen cases where transformations trigger race conditions, e.g., in special cases where constant propagation modifies something). The semantics of the SDFG should be clear in that you should not make assignments that can cause a cycle or assume any order in the application of inter-state edge assignments. The only order is that conditions precede assignments.
Avoid costly calls to `inspect` by adding a dummy `DebugInfo` object. The call to `inspect` would anyway just return the line in the stree -> sdfg translation layer. If we really wanted to have real `DebugInfo`, we'd need a way to annotate the reads/writes in the stree and then propagate this.
romanc
force-pushed
the
romanc/stree-v2
branch
5 times, most recently
from
948e333 to
4b22dc5
Compare
Much has changed on the pyFV3 / NDSL side that needs updating in this workflow configuration. We're using a temporary branch for pyFV3 since there are changes needed and that side too 🙈.
romanc
force-pushed
the
romanc/stree-v2
branch
8 times, most recently
from
37255cd to
a20292b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Schedule trees are an alternative representation of code in DaCe. It is a tree of nodes that represent the execution order of the SDFG. The tree can be used to perform schedule transformations on the SDFG, i.e., erasing an empty
ifbranch, or merging two consecutivefor-loops.A first version of the schedule tree representation was merged with PR #1145, which adds support for creating schedule trees from SDFGs. This PR builds on #1466 and brings a first version of the transformation from schedule tree to SDFG.