def apply(self, sdfg: sd.SDFG):
####################################################################
# Obtain loop information
guard: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_guard])
begin: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_begin])
after_state: sd.SDFGState = sdfg.node(
self.subgraph[DetectLoop._exit_state])
# Obtain iteration variable, range, and stride
guard_inedges = sdfg.in_edges(guard)
condition_edge = sdfg.edges_between(guard, begin)[0]
itervar = list(guard_inedges[0].data.assignments.keys())[0]
condition = condition_edge.data.condition_sympy()
rng = self._loop_range(itervar, guard_inedges, condition)
# Find the state prior to the loop
if rng[0] == symbolic.pystr_to_symbolic(
guard_inedges[0].data.assignments[itervar]):
init_edge: sd.InterstateEdge = guard_inedges[0]
before_state: sd.SDFGState = guard_inedges[0].src
last_state: sd.SDFGState = guard_inedges[1].src
else:
init_edge: sd.InterstateEdge = guard_inedges[1]
before_state: sd.SDFGState = guard_inedges[1].src
last_state: sd.SDFGState = guard_inedges[0].src
# Get loop states
loop_states = list(
sdutil.dfs_conditional(sdfg,
sources=[begin],
condition=lambda _, child: child != guard))
first_id = loop_states.index(begin)
last_id = loop_states.index(last_state)
loop_subgraph = gr.SubgraphView(sdfg, loop_states)
####################################################################
# Transform
# If begin, change initialization assignment and prepend states before
# guard
init_edge.data.assignments[itervar] = rng[0] + self.count * rng[2]
append_state = before_state
# Add `count` states, each with instantiated iteration variable
unrolled_states = []
for i in range(self.count):
# Instantiate loop states with iterate value
new_states = self.instantiate_loop(sdfg, loop_states,
loop_subgraph, itervar,
rng[0] + i * rng[2])
# Connect states to before the loop with unconditional edges
sdfg.add_edge(append_state, new_states[first_id],
sd.InterstateEdge())
append_state = new_states[last_id]
# Reconnect edge to guard state from last peeled iteration
if append_state != before_state:
sdfg.remove_edge(init_edge)
sdfg.add_edge(append_state, guard, init_edge.data)
def get_subgraph(self, sdfg: SDFG) -> gr.SubgraphView:
sdfg = sdfg.sdfg_list[self.sdfg_id]
if self.state_id == -1:
return gr.SubgraphView(sdfg, list(map(sdfg.node, self.subgraph)))
state = sdfg.node(self.state_id)
return st.StateSubgraphView(state, list(map(state.node,
self.subgraph)))
def __init__(self,
subgraph: Union[Set[int], gr.SubgraphView],
sdfg_id: int = None,
state_id: int = None):
if (not isinstance(subgraph, (gr.SubgraphView, SDFG, SDFGState))
and (sdfg_id is None or state_id is None)):
raise TypeError(
'Subgraph transformation either expects a SubgraphView or a '
'set of node IDs, SDFG ID and state ID (or -1).')
# An entire graph is given as a subgraph
if isinstance(subgraph, (SDFG, SDFGState)):
subgraph = gr.SubgraphView(subgraph, subgraph.nodes())
if isinstance(subgraph, gr.SubgraphView):
self.subgraph = set(
subgraph.graph.node_id(n) for n in subgraph.nodes())
if isinstance(subgraph.graph, SDFGState):
sdfg = subgraph.graph.parent
self.sdfg_id = sdfg.sdfg_id
self.state_id = sdfg.node_id(subgraph.graph)
elif isinstance(subgraph.graph, SDFG):
self.sdfg_id = subgraph.graph.sdfg_id
self.state_id = -1
else:
raise TypeError('Unrecognized graph type "%s"' %
type(subgraph.graph).__name__)
else:
self.subgraph = subgraph
self.sdfg_id = sdfg_id
self.state_id = state_id
def enumerate_matches(sdfg: SDFG,
pattern: gr.Graph,
node_match=type_or_class_match,
edge_match=None) -> Iterator[gr.SubgraphView]:
"""
Returns a generator of subgraphs that match the given subgraph pattern.
:param sdfg: The SDFG to search in.
:param pattern: A subgraph to look for.
:param node_match: An optional function to use for matching nodes.
:param node_match: An optional function to use for matching edges.
:return: Yields SDFG subgraph view objects.
"""
if len(pattern.nodes()) == 0:
raise ValueError('Subgraph pattern cannot be empty')
# Find if the subgraph is within states or SDFGs
is_interstate = (isinstance(pattern.node(0), SDFGState)
or (isinstance(pattern.node(0), type)
and pattern.node(0) is SDFGState))
# Collapse multigraphs into directed graphs
pattern_digraph = collapse_multigraph_to_nx(pattern)
# Find matches in all SDFGs and nested SDFGs
for graph in sdfg.all_sdfgs_recursive():
if is_interstate:
graph_matcher = iso.DiGraphMatcher(
collapse_multigraph_to_nx(graph),
pattern_digraph,
node_match=node_match,
edge_match=edge_match)
for subgraph in graph_matcher.subgraph_isomorphisms_iter():
yield gr.SubgraphView(graph,
[graph.node(i) for i in subgraph.keys()])
else:
for state in graph.nodes():
graph_matcher = iso.DiGraphMatcher(
collapse_multigraph_to_nx(state),
pattern_digraph,
node_match=node_match,
edge_match=edge_match)
for subgraph in graph_matcher.subgraph_isomorphisms_iter():
yield gr.SubgraphView(
state, [state.node(i) for i in subgraph.keys()])
def _border_arrays(sdfg, parent, subgraph):
""" Returns a set of array names that are local to the fission
subgraph. """
nested = isinstance(parent, sd.SDFGState)
schildren = subgraph.scope_children()
subset = gr.SubgraphView(parent, schildren[None])
if nested:
return set(node.data for node in subset.nodes()
if isinstance(node, nodes.AccessNode)
and sdfg.arrays[node.data].transient)
else:
return set(node.data for node in subset.nodes()
if isinstance(node, nodes.AccessNode))
def setup_match(self,
subgraph: Union[Set[int], gr.SubgraphView],
sdfg_id: int = None,
state_id: int = None):
"""
Sets the transformation to a given subgraph.
:param subgraph: A set of node (or state) IDs or a subgraph view object.
:param sdfg_id: A unique ID of the SDFG.
:param state_id: The node ID of the SDFG state, if applicable. If
transformation does not operate on a single state,
the value should be -1.
"""
if (not isinstance(subgraph, (gr.SubgraphView, SDFG, SDFGState))
and (sdfg_id is None or state_id is None)):
raise TypeError(
'Subgraph transformation either expects a SubgraphView or a '
'set of node IDs, SDFG ID and state ID (or -1).')
self._pipeline_results = None
# An entire graph is given as a subgraph
if isinstance(subgraph, (SDFG, SDFGState)):
subgraph = gr.SubgraphView(subgraph, subgraph.nodes())
if isinstance(subgraph, gr.SubgraphView):
self.subgraph = set(
subgraph.graph.node_id(n) for n in subgraph.nodes())
if isinstance(subgraph.graph, SDFGState):
sdfg = subgraph.graph.parent
self.sdfg_id = sdfg.sdfg_id
self.state_id = sdfg.node_id(subgraph.graph)
elif isinstance(subgraph.graph, SDFG):
self.sdfg_id = subgraph.graph.sdfg_id
self.state_id = -1
else:
raise TypeError('Unrecognized graph type "%s"' %
type(subgraph.graph).__name__)
else:
self.subgraph = subgraph
self.sdfg_id = sdfg_id
self.state_id = state_id
def promote_scalars_to_symbols(sdfg: sd.SDFG,
ignore: Optional[Set[str]] = None,
transients_only: bool = True,
integers_only: bool = True) -> Set[str]:
"""
Promotes all matching transient scalars to SDFG symbols, changing all
tasklets to inter-state assignments. This enables the transformed symbols
to be used within states as part of memlets, and allows further
transformations (such as loop detection) to use the information for
optimization.
:param sdfg: The SDFG to run the pass on.
:param ignore: An optional set of strings of scalars to ignore.
:param transients_only: If False, also considers global data descriptors (e.g., arguments).
:param integers_only: If False, also considers non-integral descriptors for promotion.
:return: Set of promoted scalars.
:note: Operates in-place.
"""
# Process:
# 1. Find scalars to promote
# 2. For every assignment tasklet/access:
# 2.1. Fission state to isolate assignment
# 2.2. Replace assignment with inter-state edge assignment
# 3. For every read of the scalar:
# 3.1. If destination is tasklet, remove node, edges, and connectors
# 3.2. If used in tasklet as subscript or connector, modify tasklet code
# 3.3. If destination is array, change to tasklet that copies symbol data
# 4. Remove newly-isolated access nodes
# 5. Remove data descriptors and add symbols to SDFG
# 6. Replace subscripts in all interstate conditions and assignments
# 7. Make indirections with symbols a single memlet
to_promote = find_promotable_scalars(sdfg,
transients_only=transients_only,
integers_only=integers_only)
if ignore:
to_promote -= ignore
if len(to_promote) == 0:
return to_promote
for state in sdfg.nodes():
scalar_nodes = [
n for n in state.nodes()
if isinstance(n, nodes.AccessNode) and n.data in to_promote
]
# Step 2: Assignment tasklets
for node in scalar_nodes:
if state.in_degree(node) == 0:
continue
in_edge = state.in_edges(node)[0]
input = in_edge.src
# There is only zero or one incoming edges by definition
tasklet_inputs = [e.src for e in state.in_edges(input)]
# Step 2.1
new_state = xfh.state_fission(
sdfg,
gr.SubgraphView(state, set([input, node] + tasklet_inputs)))
new_isedge: sd.InterstateEdge = sdfg.out_edges(new_state)[0]
# Step 2.2
node: nodes.AccessNode = new_state.sink_nodes()[0]
input = new_state.in_edges(node)[0].src
if isinstance(input, nodes.Tasklet):
# Convert tasklet to interstate edge
newcode: str = ''
if input.language is dtypes.Language.Python:
newcode = astutils.unparse(input.code.code[0].value)
elif input.language is dtypes.Language.CPP:
newcode = translate_cpp_tasklet_to_python(
input.code.as_string.strip())
# Replace tasklet inputs with incoming edges
for e in new_state.in_edges(input):
memlet_str: str = e.data.data
if (e.data.subset is not None and not isinstance(
sdfg.arrays[memlet_str], dt.Scalar)):
memlet_str += '[%s]' % e.data.subset
newcode = re.sub(r'\b%s\b' % re.escape(e.dst_conn),
memlet_str, newcode)
# Add interstate edge assignment
new_isedge.data.assignments[node.data] = newcode
elif isinstance(input, nodes.AccessNode):
memlet: mm.Memlet = in_edge.data
if (memlet.src_subset and
not isinstance(sdfg.arrays[memlet.data], dt.Scalar)):
new_isedge.data.assignments[
node.data] = '%s[%s]' % (input.data, memlet.src_subset)
else:
new_isedge.data.assignments[node.data] = input.data
# Clean up all nodes after assignment was transferred
new_state.remove_nodes_from(new_state.nodes())
# Step 3: Scalar reads
remove_scalar_reads(sdfg, {k: k for k in to_promote})
# Step 4: Isolated nodes
for state in sdfg.nodes():
scalar_nodes = [
n for n in state.nodes()
if isinstance(n, nodes.AccessNode) and n.data in to_promote
]
state.remove_nodes_from(
[n for n in scalar_nodes if len(state.all_edges(n)) == 0])
# Step 5: Data descriptor management
for scalar in to_promote:
desc = sdfg.arrays[scalar]
sdfg.remove_data(scalar, validate=False)
# If the scalar is already a symbol (e.g., as part of an array size),
# do not re-add the symbol
if scalar not in sdfg.symbols:
sdfg.add_symbol(scalar, desc.dtype)
# Step 6: Inter-state edge cleanup
cleanup_re = {
s: re.compile(fr'\b{re.escape(s)}\[.*?\]')
for s in to_promote
}
promo = TaskletPromoterDict({k: k for k in to_promote})
for edge in sdfg.edges():
ise: InterstateEdge = edge.data
# Condition
if not edge.data.is_unconditional():
if ise.condition.language is dtypes.Language.Python:
for stmt in ise.condition.code:
promo.visit(stmt)
elif ise.condition.language is dtypes.Language.CPP:
for scalar in to_promote:
ise.condition = cleanup_re[scalar].sub(
scalar, ise.condition.as_string)
# Assignments
for aname, assignment in ise.assignments.items():
for scalar in to_promote:
if scalar in assignment:
ise.assignments[aname] = cleanup_re[scalar].sub(
scalar, assignment.strip())
# Step 7: Indirection
remove_symbol_indirection(sdfg)
return to_promote
def apply(self, sdfg: sd.SDFG):
# Obtain loop information
guard: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_guard])
body: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_begin])
after: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._exit_state])
# Obtain iteration variable, range, and stride
itervar, (start, end, step), (_, body_end) = find_for_loop(
sdfg, guard, body, itervar=self.itervar)
# Find all loop-body states
states = set([body_end])
to_visit = [body]
while to_visit:
state = to_visit.pop(0)
if state is body_end:
continue
for _, dst, _ in sdfg.out_edges(state):
if dst not in states:
to_visit.append(dst)
states.add(state)
# Nest loop-body states
if len(states) > 1:
# Find read/write sets
read_set, write_set = set(), set()
for state in states:
rset, wset = state.read_and_write_sets()
read_set |= rset
write_set |= wset
# Add data from edges
for src in states:
for dst in states:
for edge in sdfg.edges_between(src, dst):
for s in edge.data.free_symbols:
if s in sdfg.arrays:
read_set.add(s)
# Find NestedSDFG's unique data
rw_set = read_set | write_set
unique_set = set()
for name in rw_set:
if not sdfg.arrays[name].transient:
continue
found = False
for state in sdfg.states():
if state in states:
continue
for node in state.nodes():
if (isinstance(node, nodes.AccessNode) and
node.data == name):
found = True
break
if not found:
unique_set.add(name)
# Find NestedSDFG's connectors
read_set = {n for n in read_set if n not in unique_set or not sdfg.arrays[n].transient}
write_set = {n for n in write_set if n not in unique_set or not sdfg.arrays[n].transient}
# Create NestedSDFG and add all loop-body states and edges
# Also, find defined symbols in NestedSDFG
fsymbols = set(sdfg.free_symbols)
new_body = sdfg.add_state('single_state_body')
nsdfg = SDFG("loop_body", constants=sdfg.constants, parent=new_body)
nsdfg.add_node(body, is_start_state=True)
body.parent = nsdfg
exit_state = nsdfg.add_state('exit')
nsymbols = dict()
for state in states:
if state is body:
continue
nsdfg.add_node(state)
state.parent = nsdfg
for state in states:
if state is body:
continue
for src, dst, data in sdfg.in_edges(state):
nsymbols.update({s: sdfg.symbols[s] for s in data.assignments.keys() if s in sdfg.symbols})
nsdfg.add_edge(src, dst, data)
nsdfg.add_edge(body_end, exit_state, InterstateEdge())
# Move guard -> body edge to guard -> new_body
for src, dst, data, in sdfg.edges_between(guard, body):
sdfg.add_edge(src, new_body, data)
# Move body_end -> guard edge to new_body -> guard
for src, dst, data in sdfg.edges_between(body_end, guard):
sdfg.add_edge(new_body, dst, data)
# Delete loop-body states and edges from parent SDFG
for state in states:
for e in sdfg.all_edges(state):
sdfg.remove_edge(e)
sdfg.remove_node(state)
# Add NestedSDFG arrays
for name in read_set | write_set:
nsdfg.arrays[name] = copy.deepcopy(sdfg.arrays[name])
nsdfg.arrays[name].transient = False
for name in unique_set:
nsdfg.arrays[name] = sdfg.arrays[name]
del sdfg.arrays[name]
# Add NestedSDFG node
cnode = new_body.add_nested_sdfg(nsdfg, None, read_set, write_set)
if sdfg.parent:
for s, m in sdfg.parent_nsdfg_node.symbol_mapping.items():
if s not in cnode.symbol_mapping:
cnode.symbol_mapping[s] = m
nsdfg.add_symbol(s, sdfg.symbols[s])
for name in read_set:
r = new_body.add_read(name)
new_body.add_edge(
r, None, cnode, name,
memlet.Memlet.from_array(name, sdfg.arrays[name]))
for name in write_set:
w = new_body.add_write(name)
new_body.add_edge(
cnode, name, w, None,
memlet.Memlet.from_array(name, sdfg.arrays[name]))
# Fix SDFG symbols
for sym in sdfg.free_symbols - fsymbols:
del sdfg.symbols[sym]
for sym, dtype in nsymbols.items():
nsdfg.symbols[sym] = dtype
# Change body state reference
body = new_body
if (step < 0) == True:
# If step is negative, we have to flip start and end to produce a
# correct map with a positive increment
start, end, step = end, start, -step
# If necessary, make a nested SDFG with assignments
isedge = sdfg.edges_between(guard, body)[0]
symbols_to_remove = set()
if len(isedge.data.assignments) > 0:
nsdfg = helpers.nest_state_subgraph(
sdfg, body, gr.SubgraphView(body, body.nodes()))
for sym in isedge.data.free_symbols:
if sym in nsdfg.symbol_mapping or sym in nsdfg.in_connectors:
continue
if sym in sdfg.symbols:
nsdfg.symbol_mapping[sym] = symbolic.pystr_to_symbolic(sym)
nsdfg.sdfg.add_symbol(sym, sdfg.symbols[sym])
elif sym in sdfg.arrays:
if sym in nsdfg.sdfg.arrays:
raise NotImplementedError
rnode = body.add_read(sym)
nsdfg.add_in_connector(sym)
desc = copy.deepcopy(sdfg.arrays[sym])
desc.transient = False
nsdfg.sdfg.add_datadesc(sym, desc)
body.add_edge(rnode, None, nsdfg, sym, memlet.Memlet(sym))
nstate = nsdfg.sdfg.node(0)
init_state = nsdfg.sdfg.add_state_before(nstate)
nisedge = nsdfg.sdfg.edges_between(init_state, nstate)[0]
nisedge.data.assignments = isedge.data.assignments
symbols_to_remove = set(nisedge.data.assignments.keys())
for k in nisedge.data.assignments.keys():
if k in nsdfg.symbol_mapping:
del nsdfg.symbol_mapping[k]
isedge.data.assignments = {}
source_nodes = body.source_nodes()
sink_nodes = body.sink_nodes()
map = nodes.Map(body.label + "_map", [itervar], [(start, end, step)])
entry = nodes.MapEntry(map)
exit = nodes.MapExit(map)
body.add_node(entry)
body.add_node(exit)
# If the map uses symbols from data containers, instantiate reads
containers_to_read = entry.free_symbols & sdfg.arrays.keys()
for rd in containers_to_read:
# We are guaranteed that this is always a scalar, because
# can_be_applied makes sure there are no sympy functions in each of
# the loop expresions
access_node = body.add_read(rd)
body.add_memlet_path(access_node,
entry,
dst_conn=rd,
memlet=memlet.Memlet(rd))
# Reroute all memlets through the entry and exit nodes
for n in source_nodes:
if isinstance(n, nodes.AccessNode):
for e in body.out_edges(n):
body.remove_edge(e)
body.add_edge_pair(entry,
e.dst,
n,
e.data,
internal_connector=e.dst_conn)
else:
body.add_nedge(entry, n, memlet.Memlet())
for n in sink_nodes:
if isinstance(n, nodes.AccessNode):
for e in body.in_edges(n):
body.remove_edge(e)
body.add_edge_pair(exit,
e.src,
n,
e.data,
internal_connector=e.src_conn)
else:
body.add_nedge(n, exit, memlet.Memlet())
# Get rid of the loop exit condition edge
after_edge = sdfg.edges_between(guard, after)[0]
sdfg.remove_edge(after_edge)
# Remove the assignment on the edge to the guard
for e in sdfg.in_edges(guard):
if itervar in e.data.assignments:
del e.data.assignments[itervar]
# Remove the condition on the entry edge
condition_edge = sdfg.edges_between(guard, body)[0]
condition_edge.data.condition = CodeBlock("1")
# Get rid of backedge to guard
sdfg.remove_edge(sdfg.edges_between(body, guard)[0])
# Route body directly to after state, maintaining any other assignments
# it might have had
sdfg.add_edge(
body, after,
sd.InterstateEdge(assignments=after_edge.data.assignments))
# If this had made the iteration variable a free symbol, we can remove
# it from the SDFG symbols
if itervar in sdfg.free_symbols:
sdfg.remove_symbol(itervar)
for sym in symbols_to_remove:
if helpers.is_symbol_unused(sdfg, sym):
sdfg.remove_symbol(sym)
def apply_to(cls,
sdfg: SDFG,
*where: Union[nd.Node, SDFGState, gr.SubgraphView],
verify: bool = True,
**options: Any):
"""
Applies this transformation to a given subgraph, defined by a set of
nodes. Raises an error if arguments are invalid or transformation is
not applicable.
To apply the transformation on a specific subgraph, the `where`
parameter can be used either on a subgraph object (`SubgraphView`), or
on directly on a list of subgraph nodes, given as `Node` or `SDFGState`
objects. Transformation properties can then be given as keyword
arguments. For example, applying `SubgraphFusion` on a subgraph of three
nodes can be called in one of two ways:
```
# Subgraph
SubgraphFusion.apply_to(
sdfg, SubgraphView(state, [node_a, node_b, node_c]))
# Simplified API: list of nodes
SubgraphFusion.apply_to(sdfg, node_a, node_b, node_c)
```
:param sdfg: The SDFG to apply the transformation to.
:param where: A set of nodes in the SDFG/state, or a subgraph thereof.
:param verify: Check that `can_be_applied` returns True before applying.
:param options: A set of parameters to use for applying the
transformation.
"""
subgraph = None
if len(where) == 1:
if isinstance(where[0], (list, tuple)):
where = where[0]
elif isinstance(where[0], gr.SubgraphView):
subgraph = where[0]
if len(where) == 0:
raise ValueError('At least one node is required')
# Check that all keyword arguments are nodes and if interstate or not
if subgraph is None:
sample_node = where[0]
if isinstance(sample_node, SDFGState):
graph = sdfg
state_id = -1
elif isinstance(sample_node, nd.Node):
graph = next(s for s in sdfg.nodes()
if sample_node in s.nodes())
state_id = sdfg.node_id(graph)
else:
raise TypeError('Invalid node type "%s"' %
type(sample_node).__name__)
# Construct subgraph and instantiate transformation
subgraph = gr.SubgraphView(graph, where)
instance = cls(subgraph, sdfg.sdfg_id, state_id)
else:
# Construct instance from subgraph directly
instance = cls(subgraph)
# Construct transformation parameters
for optname, optval in options.items():
if not optname in cls.__properties__:
raise ValueError('Property "%s" not found in transformation' %
optname)
setattr(instance, optname, optval)
if verify:
if not instance.can_be_applied(sdfg, subgraph):
raise ValueError('Transformation cannot be applied on the '
'given subgraph ("can_be_applied" failed)')
# Apply to SDFG
return instance.apply(sdfg)
def subgraph_view(self, sdfg: SDFG) -> gr.SubgraphView:
graph = sdfg.sdfg_list[self.sdfg_id]
if self.state_id != -1:
graph = graph.node(self.state_id)
return gr.SubgraphView(graph,
[graph.node(idx) for idx in self.subgraph])
def apply(self, sdfg):
# Obtain loop information
guard: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_guard])
begin: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_begin])
after_state: sd.SDFGState = sdfg.node(
self.subgraph[DetectLoop._exit_state])
# Obtain iteration variable, range, and stride
guard_inedges = sdfg.in_edges(guard)
condition_edge = sdfg.edges_between(guard, begin)[0]
itervar = list(guard_inedges[0].data.assignments.keys())[0]
condition = condition_edge.data.condition_sympy()
rng = LoopUnroll._loop_range(itervar, guard_inedges, condition)
# Loop must be unrollable
if self.count == 0 and any(
symbolic.issymbolic(r, sdfg.constants) for r in rng):
raise ValueError('Loop cannot be fully unrolled, size is symbolic')
if self.count != 0:
raise NotImplementedError # TODO(later)
# Find the state prior to the loop
if rng[0] == symbolic.pystr_to_symbolic(
guard_inedges[0].data.assignments[itervar]):
before_state: sd.SDFGState = guard_inedges[0].src
last_state: sd.SDFGState = guard_inedges[1].src
else:
before_state: sd.SDFGState = guard_inedges[1].src
last_state: sd.SDFGState = guard_inedges[0].src
# Get loop states
loop_states = list(
sdutil.dfs_conditional(sdfg,
sources=[begin],
condition=lambda _, child: child != guard))
first_id = loop_states.index(begin)
last_id = loop_states.index(last_state)
loop_subgraph = gr.SubgraphView(sdfg, loop_states)
# Evaluate the real values of the loop
start, end, stride = (symbolic.evaluate(r, sdfg.constants)
for r in rng)
# Create states for loop subgraph
unrolled_states = []
for i in range(start, end + 1, stride):
# Instantiate loop states with iterate value
new_states = self.instantiate_loop(sdfg, loop_states,
loop_subgraph, itervar, i)
# Connect iterations with unconditional edges
if len(unrolled_states) > 0:
sdfg.add_edge(unrolled_states[-1][1], new_states[first_id],
sd.InterstateEdge())
unrolled_states.append((new_states[first_id], new_states[last_id]))
# Connect new states to before and after states without conditions
if unrolled_states:
sdfg.add_edge(before_state, unrolled_states[0][0],
sd.InterstateEdge())
sdfg.add_edge(unrolled_states[-1][1], after_state,
sd.InterstateEdge())
# Remove old states from SDFG
sdfg.remove_nodes_from([guard] + loop_states)
def apply(self, sdfg: sd.SDFG):
# Obtain loop information
guard: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_guard])
body: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_begin])
after: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._exit_state])
# Obtain iteration variable, range, and stride
itervar, (start, end, step), _ = find_for_loop(sdfg, guard, body)
if (step < 0) == True:
# If step is negative, we have to flip start and end to produce a
# correct map with a positive increment
start, end, step = end, start, -step
# If necessary, make a nested SDFG with assignments
isedge = sdfg.edges_between(guard, body)[0]
symbols_to_remove = set()
if len(isedge.data.assignments) > 0:
nsdfg = helpers.nest_state_subgraph(
sdfg, body, gr.SubgraphView(body, body.nodes()))
for sym in isedge.data.free_symbols:
if sym in nsdfg.symbol_mapping or sym in nsdfg.in_connectors:
continue
if sym in sdfg.symbols:
nsdfg.symbol_mapping[sym] = symbolic.pystr_to_symbolic(sym)
nsdfg.sdfg.add_symbol(sym, sdfg.symbols[sym])
elif sym in sdfg.arrays:
if sym in nsdfg.sdfg.arrays:
raise NotImplementedError
rnode = body.add_read(sym)
nsdfg.add_in_connector(sym)
desc = copy.deepcopy(sdfg.arrays[sym])
desc.transient = False
nsdfg.sdfg.add_datadesc(sym, desc)
body.add_edge(rnode, None, nsdfg, sym, memlet.Memlet(sym))
nstate = nsdfg.sdfg.node(0)
init_state = nsdfg.sdfg.add_state_before(nstate)
nisedge = nsdfg.sdfg.edges_between(init_state, nstate)[0]
nisedge.data.assignments = isedge.data.assignments
symbols_to_remove = set(nisedge.data.assignments.keys())
for k in nisedge.data.assignments.keys():
if k in nsdfg.symbol_mapping:
del nsdfg.symbol_mapping[k]
isedge.data.assignments = {}
source_nodes = body.source_nodes()
sink_nodes = body.sink_nodes()
map = nodes.Map(body.label + "_map", [itervar], [(start, end, step)])
entry = nodes.MapEntry(map)
exit = nodes.MapExit(map)
body.add_node(entry)
body.add_node(exit)
# If the map uses symbols from data containers, instantiate reads
containers_to_read = entry.free_symbols & sdfg.arrays.keys()
for rd in containers_to_read:
# We are guaranteed that this is always a scalar, because
# can_be_applied makes sure there are no sympy functions in each of
# the loop expresions
access_node = body.add_read(rd)
body.add_memlet_path(access_node,
entry,
dst_conn=rd,
memlet=memlet.Memlet(rd))
# Reroute all memlets through the entry and exit nodes
for n in source_nodes:
if isinstance(n, nodes.AccessNode):
for e in body.out_edges(n):
body.remove_edge(e)
body.add_edge_pair(entry,
e.dst,
n,
e.data,
internal_connector=e.dst_conn)
else:
body.add_nedge(entry, n, memlet.Memlet())
for n in sink_nodes:
if isinstance(n, nodes.AccessNode):
for e in body.in_edges(n):
body.remove_edge(e)
body.add_edge_pair(exit,
e.src,
n,
e.data,
internal_connector=e.src_conn)
else:
body.add_nedge(n, exit, memlet.Memlet())
# Get rid of the loop exit condition edge
after_edge = sdfg.edges_between(guard, after)[0]
sdfg.remove_edge(after_edge)
# Remove the assignment on the edge to the guard
for e in sdfg.in_edges(guard):
if itervar in e.data.assignments:
del e.data.assignments[itervar]
# Remove the condition on the entry edge
condition_edge = sdfg.edges_between(guard, body)[0]
condition_edge.data.condition = CodeBlock("1")
# Get rid of backedge to guard
sdfg.remove_edge(sdfg.edges_between(body, guard)[0])
# Route body directly to after state, maintaining any other assignments
# it might have had
sdfg.add_edge(
body, after,
sd.InterstateEdge(assignments=after_edge.data.assignments))
# If this had made the iteration variable a free symbol, we can remove
# it from the SDFG symbols
if itervar in sdfg.free_symbols:
sdfg.remove_symbol(itervar)
for sym in symbols_to_remove:
if helpers.is_symbol_unused(sdfg, sym):
sdfg.remove_symbol(sym)
def apply(self, sdfg: sd.SDFG):
####################################################################
# Obtain loop information
guard: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_guard])
begin: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_begin])
after_state: sd.SDFGState = sdfg.node(
self.subgraph[DetectLoop._exit_state])
# Obtain iteration variable, range, and stride
guard_inedges = sdfg.in_edges(guard)
condition_edge = sdfg.edges_between(guard, begin)[0]
not_condition_edge = sdfg.edges_between(guard, after_state)[0]
itervar = list(guard_inedges[0].data.assignments.keys())[0]
condition = condition_edge.data.condition_sympy()
rng = self._loop_range(itervar, guard_inedges, condition)
# Find the state prior to the loop
if rng[0] == symbolic.pystr_to_symbolic(
guard_inedges[0].data.assignments[itervar]):
init_edge: sd.InterstateEdge = guard_inedges[0]
before_state: sd.SDFGState = guard_inedges[0].src
last_state: sd.SDFGState = guard_inedges[1].src
else:
init_edge: sd.InterstateEdge = guard_inedges[1]
before_state: sd.SDFGState = guard_inedges[1].src
last_state: sd.SDFGState = guard_inedges[0].src
# Get loop states
loop_states = list(
sdutil.dfs_conditional(sdfg,
sources=[begin],
condition=lambda _, child: child != guard))
first_id = loop_states.index(begin)
last_id = loop_states.index(last_state)
loop_subgraph = gr.SubgraphView(sdfg, loop_states)
####################################################################
# Transform
if self.begin:
# If begin, change initialization assignment and prepend states before
# guard
init_edge.data.assignments[itervar] = str(rng[0] +
self.count * rng[2])
append_state = before_state
# Add `count` states, each with instantiated iteration variable
for i in range(self.count):
# Instantiate loop states with iterate value
state_name: str = 'start_' + itervar + str(i * rng[2])
state_name = state_name.replace('-', 'm').replace(
'+', 'p').replace('*', 'M').replace('/', 'D')
new_states = self.instantiate_loop(
sdfg,
loop_states,
loop_subgraph,
itervar,
rng[0] + i * rng[2],
state_name,
)
# Connect states to before the loop with unconditional edges
sdfg.add_edge(append_state, new_states[first_id],
sd.InterstateEdge())
append_state = new_states[last_id]
# Reconnect edge to guard state from last peeled iteration
if append_state != before_state:
sdfg.remove_edge(init_edge)
sdfg.add_edge(append_state, guard, init_edge.data)
else:
# If begin, change initialization assignment and prepend states before
# guard
itervar_sym = pystr_to_symbolic(itervar)
condition_edge.data.condition = CodeBlock(
self._modify_cond(condition_edge.data.condition, itervar,
rng[2]))
not_condition_edge.data.condition = CodeBlock(
self._modify_cond(not_condition_edge.data.condition, itervar,
rng[2]))
prepend_state = after_state
# Add `count` states, each with instantiated iteration variable
for i in reversed(range(self.count)):
# Instantiate loop states with iterate value
state_name: str = 'end_' + itervar + str(-i * rng[2])
state_name = state_name.replace('-', 'm').replace(
'+', 'p').replace('*', 'M').replace('/', 'D')
new_states = self.instantiate_loop(
sdfg,
loop_states,
loop_subgraph,
itervar,
itervar_sym + i * rng[2],
state_name,
)
# Connect states to before the loop with unconditional edges
sdfg.add_edge(new_states[last_id], prepend_state,
sd.InterstateEdge())
prepend_state = new_states[first_id]
# Reconnect edge to guard state from last peeled iteration
if prepend_state != after_state:
sdfg.remove_edge(not_condition_edge)
sdfg.add_edge(guard, prepend_state, not_condition_edge.data)
def apply(self, _, sdfg: sd.SDFG):
####################################################################
# Obtain loop information
guard: sd.SDFGState = self.loop_guard
begin: sd.SDFGState = self.loop_begin
after_state: sd.SDFGState = self.exit_state
# Obtain iteration variable, range, and stride
condition_edge = sdfg.edges_between(guard, begin)[0]
not_condition_edge = sdfg.edges_between(guard, after_state)[0]
itervar, rng, loop_struct = find_for_loop(sdfg, guard, begin)
# Get loop states
loop_states = list(
sdutil.dfs_conditional(sdfg,
sources=[begin],
condition=lambda _, child: child != guard))
first_id = loop_states.index(begin)
last_state = loop_struct[1]
last_id = loop_states.index(last_state)
loop_subgraph = gr.SubgraphView(sdfg, loop_states)
####################################################################
# Transform
if self.begin:
# If begin, change initialization assignment and prepend states before
# guard
init_edges = []
before_states = loop_struct[0]
for before_state in before_states:
init_edge = sdfg.edges_between(before_state, guard)[0]
init_edge.data.assignments[itervar] = str(rng[0] +
self.count * rng[2])
init_edges.append(init_edge)
append_states = before_states
# Add `count` states, each with instantiated iteration variable
for i in range(self.count):
# Instantiate loop states with iterate value
state_name: str = 'start_' + itervar + str(i * rng[2])
state_name = state_name.replace('-', 'm').replace(
'+', 'p').replace('*', 'M').replace('/', 'D')
new_states = self.instantiate_loop(
sdfg,
loop_states,
loop_subgraph,
itervar,
rng[0] + i * rng[2],
state_name,
)
# Connect states to before the loop with unconditional edges
for append_state in append_states:
sdfg.add_edge(append_state, new_states[first_id],
sd.InterstateEdge())
append_states = [new_states[last_id]]
# Reconnect edge to guard state from last peeled iteration
for append_state in append_states:
if append_state not in before_states:
for init_edge in init_edges:
sdfg.remove_edge(init_edge)
sdfg.add_edge(append_state, guard, init_edges[0].data)
else:
# If begin, change initialization assignment and prepend states before
# guard
itervar_sym = pystr_to_symbolic(itervar)
condition_edge.data.condition = CodeBlock(
self._modify_cond(condition_edge.data.condition, itervar,
rng[2]))
not_condition_edge.data.condition = CodeBlock(
self._modify_cond(not_condition_edge.data.condition, itervar,
rng[2]))
prepend_state = after_state
# Add `count` states, each with instantiated iteration variable
for i in reversed(range(self.count)):
# Instantiate loop states with iterate value
state_name: str = 'end_' + itervar + str(-i * rng[2])
state_name = state_name.replace('-', 'm').replace(
'+', 'p').replace('*', 'M').replace('/', 'D')
new_states = self.instantiate_loop(
sdfg,
loop_states,
loop_subgraph,
itervar,
itervar_sym + i * rng[2],
state_name,
)
# Connect states to before the loop with unconditional edges
sdfg.add_edge(new_states[last_id], prepend_state,
sd.InterstateEdge())
prepend_state = new_states[first_id]
# Reconnect edge to guard state from last peeled iteration
if prepend_state != after_state:
sdfg.remove_edge(not_condition_edge)
sdfg.add_edge(guard, prepend_state, not_condition_edge.data)
def apply(self, sdfg):
# Obtain loop information
guard: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_guard])
begin: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_begin])
after_state: sd.SDFGState = sdfg.node(
self.subgraph[DetectLoop._exit_state])
# Obtain iteration variable, range, and stride
guard_inedges = sdfg.in_edges(guard)
condition_edge = sdfg.edges_between(guard, begin)[0]
itervar = list(guard_inedges[0].data.assignments.keys())[0]
condition = condition_edge.data.condition_sympy()
rng = LoopUnroll._loop_range(itervar, guard_inedges, condition)
# Loop must be unrollable
if self.count == 0 and any(
symbolic.issymbolic(r, sdfg.constants) for r in rng):
raise ValueError('Loop cannot be fully unrolled, size is symbolic')
if self.count != 0:
raise NotImplementedError # TODO(later)
# Find the state prior to the loop
if rng[0] == symbolic.pystr_to_symbolic(
guard_inedges[0].data.assignments[itervar]):
before_state: sd.SDFGState = guard_inedges[0].src
last_state: sd.SDFGState = guard_inedges[1].src
else:
before_state: sd.SDFGState = guard_inedges[1].src
last_state: sd.SDFGState = guard_inedges[0].src
# Get loop states
loop_states = list(
sdutil.dfs_topological_sort(
sdfg,
sources=[begin],
condition=lambda _, child: child != guard))
first_id = loop_states.index(begin)
last_id = loop_states.index(last_state)
loop_subgraph = gr.SubgraphView(sdfg, loop_states)
# Evaluate the real values of the loop
start, end, stride = (symbolic.evaluate(r, sdfg.constants)
for r in rng)
# Create states for loop subgraph
unrolled_states = []
for i in range(start, end + 1, stride):
# Using to/from JSON copies faster than deepcopy (which will also
# copy the parent SDFG)
new_states = [
sd.SDFGState.from_json(s.to_json(), context={'sdfg': sdfg})
for s in loop_states
]
# Replace iterate with value in each state
for state in new_states:
state.set_label(state.label + '_%s_%d' % (itervar, i))
state.replace(itervar, i)
# Add subgraph to original SDFG
for edge in loop_subgraph.edges():
src = new_states[loop_states.index(edge.src)]
dst = new_states[loop_states.index(edge.dst)]
# Replace conditions in subgraph edges
data: sd.InterstateEdge = copy.deepcopy(edge.data)
if data.condition:
ASTFindReplace({itervar: str(i)}).visit(data.condition)
sdfg.add_edge(src, dst, data)
# Connect iterations with unconditional edges
if len(unrolled_states) > 0:
sdfg.add_edge(unrolled_states[-1][1], new_states[first_id],
sd.InterstateEdge())
unrolled_states.append((new_states[first_id], new_states[last_id]))
# Connect new states to before and after states without conditions
if unrolled_states:
sdfg.add_edge(before_state, unrolled_states[0][0],
sd.InterstateEdge())
sdfg.add_edge(unrolled_states[-1][1], after_state,
sd.InterstateEdge())
# Remove old states from SDFG
sdfg.remove_nodes_from([guard] + loop_states)
def apply(self, sdfg):
# Obtain loop information
guard: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_guard])
begin: sd.SDFGState = sdfg.node(self.subgraph[DetectLoop._loop_begin])
after_state: sd.SDFGState = sdfg.node(
self.subgraph[DetectLoop._exit_state])
# Obtain iteration variable, range, and stride, together with the last
# state(s) before the loop and the last loop state.
itervar, rng, loop_struct = find_for_loop(sdfg, guard, begin)
# Loop must be fully unrollable for now.
if self.count != 0:
raise NotImplementedError # TODO(later)
# Get loop states
loop_states = list(
sdutil.dfs_conditional(sdfg,
sources=[begin],
condition=lambda _, child: child != guard))
first_id = loop_states.index(begin)
last_state = loop_struct[1]
last_id = loop_states.index(last_state)
loop_subgraph = gr.SubgraphView(sdfg, loop_states)
try:
start, end, stride = (r for r in rng)
stride = symbolic.evaluate(stride, sdfg.constants)
loop_diff = int(symbolic.evaluate(end - start + 1, sdfg.constants))
is_symbolic = any([symbolic.issymbolic(r) for r in rng[:2]])
except TypeError:
raise TypeError('Loop difference and strides cannot be symbolic.')
# Create states for loop subgraph
unrolled_states = []
for i in range(0, loop_diff, stride):
current_index = start + i
# Instantiate loop states with iterate value
new_states = self.instantiate_loop(sdfg, loop_states,
loop_subgraph, itervar,
current_index,
str(i) if is_symbolic else None)
# Connect iterations with unconditional edges
if len(unrolled_states) > 0:
sdfg.add_edge(unrolled_states[-1][1], new_states[first_id],
sd.InterstateEdge())
unrolled_states.append((new_states[first_id], new_states[last_id]))
# Get any assignments that might be on the edge to the after state
after_assignments = (sdfg.edges_between(
guard, after_state)[0].data.assignments)
# Connect new states to before and after states without conditions
if unrolled_states:
before_states = loop_struct[0]
for before_state in before_states:
sdfg.add_edge(before_state, unrolled_states[0][0],
sd.InterstateEdge())
sdfg.add_edge(unrolled_states[-1][1], after_state,
sd.InterstateEdge(assignments=after_assignments))
# Remove old states from SDFG
sdfg.remove_nodes_from([guard] + loop_states)