Skip to content

axis_merge

CartesianAxisMerge

Bases: ScheduleNodeTransformer

Merge a cartesian axis if they are contiguous in code-flow.

Can do
  • merge a given axis with the next maps at the same recursion level
  • can overcompute (eager) to allow for more merging at the cost of an if
It expects
  • All Maps and ForLoop are on a single axis - but doesn't check for it.

Parameters:

Name Type Description Default
axis AxisIterator

AxisIterator to be merged

 required
eager bool

overcompute with a conditional guard

 True

visit_ScheduleTreeRoot(node)

Merge as many maps as possible.

The algorithm works as follows
  • Start merging - move nodes to surface maps as much as possible
  • Try to merge the surfaced maps
  • When done, count the number of actual merges
  • If NO merges - restore the previous children (undo potential changes that didn't lead to map merge) Then exit.
ToDo
  • ForLoop are not merge at the moment, only Maps.
  • Non-cartesian ForLoop should be merged down if the maps below are unique (e.g. if everything has been merged). This is relevant for linear solvers and other iteration-dependent algorithmics
  • The K loops have varied indices name of the form _k_x[_y]. This overcomplicates merging and we don't take care of it at the moment. We could write a pass cleaning those first.

CartesianRefineTransients

Bases: ScheduleNodeTransformer

Refine (reduce dimensionality) of transients based on their true use in the cartesian dimensions.

It expects
  • All Maps and ForLoop are on a single axis - but doesn't check for it.
It can do
  • Looking at usage of a transient in a cartesian axis (e.g. loop over a cartesian axis) it will reduce that axis to 1 if all access are atomic (exactly one element of the array is ever worked on in a single loop)
  • It will refuse to merge if the transient is used in multiple loops of for a given axis - irrigardless of it's access pattern (e.g. even if it could be refine because it's always written first.)
It should but cannot do or will produce bugs if
  • With better dataflow analysis, we can reduce the dimensions to the correct lowest size needed on the axis (e.g. transient[K] and transient[K+1], requires a 2-element buffer), instead of the defensive no refine strategy used now. We have most of the info in the Range
  • Current action when detecting a valid candidate is to reduce the size of the dimension to 1, rather than removing it. This will effectively, if generic compilers do their job, reduce the cache access significantly. This also has been implemented to not deal with offset/slicing downstream impact of removing an axis. Nevertheless the axis should be removed if it's not used.
  • It only knows how to deal with 3D cartesian and 3D cartesian + data dimensions. Anything else will fail _reduce_cartesian_axes_size_to_1 calculation
More tests
  • Test for dataflow with offset
  • Test for I/J refine but not in K
  • Test for J refine but not in I or K
  • Test with dataflow: if/else, while, etc.
  • Test with ForScope (FORWARD/BACKWARD) instead of Map
Coding traps
  • We reduce the "refined" dimensions to 1 which is functionally eliminating it. This is solid. In the case of the one we can't eliminate we don't do anything. We could find the "smallest buffer size" needed and reduce the local dimension to it. BUT if we do this, we have to take into account the offset into memory (e.g. halo) for the RebuildMemletsFromContainers!

CleanUpScheduleTree

Bases: ScheduleNodeTransformer

Remove StateBoundary nodes from children of ScheduleTreeScopes.