def expressions():
# Case 1: Loop with one state
sdfg = sd.SDFG('_')
sdfg.add_nodes_from([
DetectLoop._loop_guard, DetectLoop._loop_begin,
DetectLoop._exit_state
])
sdfg.add_edge(DetectLoop._loop_guard, DetectLoop._loop_begin,
sd.InterstateEdge())
sdfg.add_edge(DetectLoop._loop_guard, DetectLoop._exit_state,
sd.InterstateEdge())
sdfg.add_edge(DetectLoop._loop_begin, DetectLoop._loop_guard,
sd.InterstateEdge())
# Case 2: Loop with multiple states (no back-edge from state)
msdfg = sd.SDFG('_')
msdfg.add_nodes_from([
DetectLoop._loop_guard, DetectLoop._loop_begin,
DetectLoop._exit_state
])
msdfg.add_edge(DetectLoop._loop_guard, DetectLoop._loop_begin,
sd.InterstateEdge())
msdfg.add_edge(DetectLoop._loop_guard, DetectLoop._exit_state,
sd.InterstateEdge())
return [sdfg, msdfg]
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 apply(self, sdfg):
state = sdfg.nodes()[self.subgraph[FPGATransformState._state]]
# Find source/sink (data) nodes
input_nodes = sdutil.find_source_nodes(state)
output_nodes = sdutil.find_sink_nodes(state)
fpga_data = {}
# Input nodes may also be nodes with WCR memlets
# We have to recur across nested SDFGs to find them
wcr_input_nodes = set()
stack = []
parent_sdfg = {state: sdfg} # Map states to their parent SDFG
for node, graph in state.all_nodes_recursive():
if isinstance(graph, dace.SDFG):
parent_sdfg[node] = graph
if isinstance(node, dace.sdfg.nodes.AccessNode):
for e in graph.all_edges(node):
if e.data.wcr is not None:
trace = dace.sdfg.trace_nested_access(
node, graph, parent_sdfg[graph])
for node_trace, state_trace, sdfg_trace in trace:
# Find the name of the accessed node in our scope
if state_trace == state and sdfg_trace == sdfg:
outer_node = node_trace
break
else:
# This does not trace back to the current state, so
# we don't care
continue
input_nodes.append(outer_node)
wcr_input_nodes.add(outer_node)
if input_nodes:
# create pre_state
pre_state = sd.SDFGState('pre_' + state.label, sdfg)
for node in input_nodes:
if not isinstance(node, dace.sdfg.nodes.AccessNode):
continue
desc = node.desc(sdfg)
if not isinstance(desc, dace.data.Array):
# TODO: handle streams
continue
if node.data in fpga_data:
fpga_array = fpga_data[node.data]
elif node not in wcr_input_nodes:
fpga_array = sdfg.add_array(
'fpga_' + node.data,
desc.shape,
desc.dtype,
materialize_func=desc.materialize_func,
transient=True,
storage=dtypes.StorageType.FPGA_Global,
allow_conflicts=desc.allow_conflicts,
strides=desc.strides,
offset=desc.offset)
fpga_data[node.data] = fpga_array
pre_node = pre_state.add_read(node.data)
pre_fpga_node = pre_state.add_write('fpga_' + node.data)
full_range = subsets.Range([(0, s - 1, 1) for s in desc.shape])
mem = memlet.Memlet(node.data, full_range.num_elements(),
full_range, 1)
pre_state.add_edge(pre_node, None, pre_fpga_node, None, mem)
if node not in wcr_input_nodes:
fpga_node = state.add_read('fpga_' + node.data)
sdutil.change_edge_src(state, node, fpga_node)
state.remove_node(node)
sdfg.add_node(pre_state)
sdutil.change_edge_dest(sdfg, state, pre_state)
sdfg.add_edge(pre_state, state, sd.InterstateEdge())
if output_nodes:
post_state = sd.SDFGState('post_' + state.label, sdfg)
for node in output_nodes:
if not isinstance(node, dace.sdfg.nodes.AccessNode):
continue
desc = node.desc(sdfg)
if not isinstance(desc, dace.data.Array):
# TODO: handle streams
continue
if node.data in fpga_data:
fpga_array = fpga_data[node.data]
else:
fpga_array = sdfg.add_array(
'fpga_' + node.data,
desc.shape,
desc.dtype,
materialize_func=desc.materialize_func,
transient=True,
storage=dtypes.StorageType.FPGA_Global,
allow_conflicts=desc.allow_conflicts,
strides=desc.strides,
offset=desc.offset)
fpga_data[node.data] = fpga_array
# fpga_node = type(node)(fpga_array)
post_node = post_state.add_write(node.data)
post_fpga_node = post_state.add_read('fpga_' + node.data)
full_range = subsets.Range([(0, s - 1, 1) for s in desc.shape])
mem = memlet.Memlet('fpga_' + node.data,
full_range.num_elements(), full_range, 1)
post_state.add_edge(post_fpga_node, None, post_node, None, mem)
fpga_node = state.add_write('fpga_' + node.data)
sdutil.change_edge_dest(state, node, fpga_node)
state.remove_node(node)
sdfg.add_node(post_state)
sdutil.change_edge_src(sdfg, state, post_state)
sdfg.add_edge(state, post_state, sd.InterstateEdge())
veclen_ = 1
# propagate vector info from a nested sdfg
for src, src_conn, dst, dst_conn, mem in state.edges():
# need to go inside the nested SDFG and grab the vector length
if isinstance(dst, dace.sdfg.nodes.NestedSDFG):
# this edge is going to the nested SDFG
for inner_state in dst.sdfg.states():
for n in inner_state.nodes():
if isinstance(n, dace.sdfg.nodes.AccessNode
) and n.data == dst_conn:
# assuming all memlets have the same vector length
veclen_ = inner_state.all_edges(n)[0].data.veclen
if isinstance(src, dace.sdfg.nodes.NestedSDFG):
# this edge is coming from the nested SDFG
for inner_state in src.sdfg.states():
for n in inner_state.nodes():
if isinstance(n, dace.sdfg.nodes.AccessNode
) and n.data == src_conn:
# assuming all memlets have the same vector length
veclen_ = inner_state.all_edges(n)[0].data.veclen
if mem.data is not None and mem.data in fpga_data:
mem.data = 'fpga_' + mem.data
mem.veclen = veclen_
fpga_update(sdfg, state, 0)
def apply(self, sdfg: sd.SDFG):
#######################################################
# Step 0: SDFG metadata
# Find all input and output data descriptors
input_nodes = []
output_nodes = []
global_code_nodes = [[] for _ in sdfg.nodes()]
for i, state in enumerate(sdfg.nodes()):
sdict = state.scope_dict()
for node in state.nodes():
if (isinstance(node, nodes.AccessNode)
and node.desc(sdfg).transient == False):
if (state.out_degree(node) > 0
and node.data not in input_nodes):
# Special case: nodes that lead to top-level dynamic
# map ranges must stay on host
for e in state.out_edges(node):
last_edge = state.memlet_path(e)[-1]
if (isinstance(last_edge.dst, nodes.EntryNode)
and last_edge.dst_conn and
not last_edge.dst_conn.startswith('IN_')
and sdict[last_edge.dst] is None):
break
else:
input_nodes.append((node.data, node.desc(sdfg)))
if (state.in_degree(node) > 0
and node.data not in output_nodes):
output_nodes.append((node.data, node.desc(sdfg)))
elif isinstance(node, nodes.CodeNode) and sdict[node] is None:
if not isinstance(node,
(nodes.LibraryNode, nodes.NestedSDFG)):
global_code_nodes[i].append(node)
# Input nodes may also be nodes with WCR memlets and no identity
for e in state.edges():
if e.data.wcr is not None:
if (e.data.data not in input_nodes
and sdfg.arrays[e.data.data].transient == False):
input_nodes.append(
(e.data.data, sdfg.arrays[e.data.data]))
start_state = sdfg.start_state
end_states = sdfg.sink_nodes()
#######################################################
# Step 1: Create cloned GPU arrays and replace originals
cloned_arrays = {}
for inodename, inode in set(input_nodes):
if isinstance(inode, data.Scalar): # Scalars can remain on host
continue
if inode.storage == dtypes.StorageType.GPU_Global:
continue
newdesc = inode.clone()
newdesc.storage = dtypes.StorageType.GPU_Global
newdesc.transient = True
name = sdfg.add_datadesc('gpu_' + inodename,
newdesc,
find_new_name=True)
cloned_arrays[inodename] = name
for onodename, onode in set(output_nodes):
if onodename in cloned_arrays:
continue
if onode.storage == dtypes.StorageType.GPU_Global:
continue
newdesc = onode.clone()
newdesc.storage = dtypes.StorageType.GPU_Global
newdesc.transient = True
name = sdfg.add_datadesc('gpu_' + onodename,
newdesc,
find_new_name=True)
cloned_arrays[onodename] = name
# Replace nodes
for state in sdfg.nodes():
for node in state.nodes():
if (isinstance(node, nodes.AccessNode)
and node.data in cloned_arrays):
node.data = cloned_arrays[node.data]
# Replace memlets
for state in sdfg.nodes():
for edge in state.edges():
if edge.data.data in cloned_arrays:
edge.data.data = cloned_arrays[edge.data.data]
#######################################################
# Step 2: Create copy-in state
excluded_copyin = self.exclude_copyin.split(',')
copyin_state = sdfg.add_state(sdfg.label + '_copyin')
sdfg.add_edge(copyin_state, start_state, sd.InterstateEdge())
for nname, desc in dtypes.deduplicate(input_nodes):
if nname in excluded_copyin or nname not in cloned_arrays:
continue
src_array = nodes.AccessNode(nname, debuginfo=desc.debuginfo)
dst_array = nodes.AccessNode(cloned_arrays[nname],
debuginfo=desc.debuginfo)
copyin_state.add_node(src_array)
copyin_state.add_node(dst_array)
copyin_state.add_nedge(
src_array, dst_array,
memlet.Memlet.from_array(src_array.data, src_array.desc(sdfg)))
#######################################################
# Step 3: Create copy-out state
excluded_copyout = self.exclude_copyout.split(',')
copyout_state = sdfg.add_state(sdfg.label + '_copyout')
for state in end_states:
sdfg.add_edge(state, copyout_state, sd.InterstateEdge())
for nname, desc in dtypes.deduplicate(output_nodes):
if nname in excluded_copyout or nname not in cloned_arrays:
continue
src_array = nodes.AccessNode(cloned_arrays[nname],
debuginfo=desc.debuginfo)
dst_array = nodes.AccessNode(nname, debuginfo=desc.debuginfo)
copyout_state.add_node(src_array)
copyout_state.add_node(dst_array)
copyout_state.add_nedge(
src_array, dst_array,
memlet.Memlet.from_array(dst_array.data, dst_array.desc(sdfg)))
#######################################################
# Step 4: Modify transient data storage
for state in sdfg.nodes():
sdict = state.scope_dict()
for node in state.nodes():
if isinstance(node,
nodes.AccessNode) and node.desc(sdfg).transient:
nodedesc = node.desc(sdfg)
# Special case: nodes that lead to dynamic map ranges must
# stay on host
if any(
isinstance(
state.memlet_path(e)[-1].dst, nodes.EntryNode)
for e in state.out_edges(node)):
continue
gpu_storage = [
dtypes.StorageType.GPU_Global,
dtypes.StorageType.GPU_Shared,
dtypes.StorageType.CPU_Pinned
]
if sdict[
node] is None and nodedesc.storage not in gpu_storage:
# NOTE: the cloned arrays match too but it's the same
# storage so we don't care
nodedesc.storage = dtypes.StorageType.GPU_Global
# Try to move allocation/deallocation out of loops
if (self.toplevel_trans
and not isinstance(nodedesc, data.Stream)):
nodedesc.lifetime = dtypes.AllocationLifetime.SDFG
elif nodedesc.storage not in gpu_storage:
# Make internal transients registers
if self.register_trans:
nodedesc.storage = dtypes.StorageType.Register
#######################################################
# Step 5: Wrap free tasklets and nested SDFGs with a GPU map
for state, gcodes in zip(sdfg.nodes(), global_code_nodes):
for gcode in gcodes:
if gcode.label in self.exclude_tasklets.split(','):
continue
# Create map and connectors
me, mx = state.add_map(gcode.label + '_gmap',
{gcode.label + '__gmapi': '0:1'},
schedule=dtypes.ScheduleType.GPU_Device)
# Store in/out edges in lists so that they don't get corrupted
# when they are removed from the graph
in_edges = list(state.in_edges(gcode))
out_edges = list(state.out_edges(gcode))
me.in_connectors = {('IN_' + e.dst_conn): None
for e in in_edges}
me.out_connectors = {('OUT_' + e.dst_conn): None
for e in in_edges}
mx.in_connectors = {('IN_' + e.src_conn): None
for e in out_edges}
mx.out_connectors = {('OUT_' + e.src_conn): None
for e in out_edges}
# Create memlets through map
for e in in_edges:
state.remove_edge(e)
state.add_edge(e.src, e.src_conn, me, 'IN_' + e.dst_conn,
e.data)
state.add_edge(me, 'OUT_' + e.dst_conn, e.dst, e.dst_conn,
e.data)
for e in out_edges:
state.remove_edge(e)
state.add_edge(e.src, e.src_conn, mx, 'IN_' + e.src_conn,
e.data)
state.add_edge(mx, 'OUT_' + e.src_conn, e.dst, e.dst_conn,
e.data)
# Map without inputs
if len(in_edges) == 0:
state.add_nedge(me, gcode, memlet.Memlet())
#######################################################
# Step 6: Change all top-level maps and library nodes to GPU schedule
for i, state in enumerate(sdfg.nodes()):
sdict = state.scope_dict()
for node in state.nodes():
if isinstance(node, (nodes.EntryNode, nodes.LibraryNode)):
if sdict[node] is None:
node.schedule = dtypes.ScheduleType.GPU_Device
elif (isinstance(node,
(nodes.EntryNode, nodes.LibraryNode))
and self.sequential_innermaps):
node.schedule = dtypes.ScheduleType.Sequential
#######################################################
# Step 7: Introduce copy-out if data used in outgoing interstate edges
for state in list(sdfg.nodes()):
arrays_used = set()
for e in sdfg.out_edges(state):
# Used arrays = intersection between symbols and cloned arrays
arrays_used.update(
set(e.data.free_symbols)
& set(cloned_arrays.keys()))
# Create a state and copy out used arrays
if len(arrays_used) > 0:
co_state = sdfg.add_state(state.label + '_icopyout')
# Reconnect outgoing edges to after interim copyout state
for e in sdfg.out_edges(state):
sdutil.change_edge_src(sdfg, state, co_state)
# Add unconditional edge to interim state
sdfg.add_edge(state, co_state, sd.InterstateEdge())
# Add copy-out nodes
for nname in arrays_used:
desc = sdfg.arrays[nname]
src_array = nodes.AccessNode(cloned_arrays[nname],
debuginfo=desc.debuginfo)
dst_array = nodes.AccessNode(nname,
debuginfo=desc.debuginfo)
co_state.add_node(src_array)
co_state.add_node(dst_array)
co_state.add_nedge(
src_array, dst_array,
memlet.Memlet.from_array(dst_array.data,
dst_array.desc(sdfg)))
#######################################################
# Step 8: Strict transformations
if not self.strict_transform:
return
# Apply strict state fusions greedily.
sdfg.apply_strict_transformations()
def apply(self, _, sdfg):
state = self.state
# Find source/sink (data) nodes that are relevant outside this FPGA
# kernel
shared_transients = set(sdfg.shared_transients())
input_nodes = [
n for n in sdutil.find_source_nodes(state)
if isinstance(n, nodes.AccessNode) and
(not sdfg.arrays[n.data].transient or n.data in shared_transients)
]
output_nodes = [
n for n in sdutil.find_sink_nodes(state)
if isinstance(n, nodes.AccessNode) and
(not sdfg.arrays[n.data].transient or n.data in shared_transients)
]
fpga_data = {}
# Input nodes may also be nodes with WCR memlets
# We have to recur across nested SDFGs to find them
wcr_input_nodes = set()
stack = []
parent_sdfg = {state: sdfg} # Map states to their parent SDFG
for node, graph in state.all_nodes_recursive():
if isinstance(graph, dace.SDFG):
parent_sdfg[node] = graph
if isinstance(node, dace.sdfg.nodes.AccessNode):
for e in graph.in_edges(node):
if e.data.wcr is not None:
trace = dace.sdfg.trace_nested_access(
node, graph, parent_sdfg[graph])
for node_trace, memlet_trace, state_trace, sdfg_trace in trace:
# Find the name of the accessed node in our scope
if state_trace == state and sdfg_trace == sdfg:
_, outer_node = node_trace
if outer_node is not None:
break
else:
# This does not trace back to the current state, so
# we don't care
continue
input_nodes.append(outer_node)
wcr_input_nodes.add(outer_node)
if input_nodes:
# create pre_state
pre_state = sd.SDFGState('pre_' + state.label, sdfg)
for node in input_nodes:
if not isinstance(node, dace.sdfg.nodes.AccessNode):
continue
desc = node.desc(sdfg)
if not isinstance(desc, dace.data.Array):
# TODO: handle streams
continue
if node.data in fpga_data:
fpga_array = fpga_data[node.data]
elif node not in wcr_input_nodes:
fpga_array = sdfg.add_array(
'fpga_' + node.data,
desc.shape,
desc.dtype,
transient=True,
storage=dtypes.StorageType.FPGA_Global,
allow_conflicts=desc.allow_conflicts,
strides=desc.strides,
offset=desc.offset)
fpga_array[1].location = copy.copy(desc.location)
desc.location.clear()
fpga_data[node.data] = fpga_array
pre_node = pre_state.add_read(node.data)
pre_fpga_node = pre_state.add_write('fpga_' + node.data)
mem = memlet.Memlet(data=node.data,
subset=subsets.Range.from_array(desc))
pre_state.add_edge(pre_node, None, pre_fpga_node, None, mem)
if node not in wcr_input_nodes:
fpga_node = state.add_read('fpga_' + node.data)
sdutil.change_edge_src(state, node, fpga_node)
state.remove_node(node)
sdfg.add_node(pre_state)
sdutil.change_edge_dest(sdfg, state, pre_state)
sdfg.add_edge(pre_state, state, sd.InterstateEdge())
if output_nodes:
post_state = sd.SDFGState('post_' + state.label, sdfg)
for node in output_nodes:
if not isinstance(node, dace.sdfg.nodes.AccessNode):
continue
desc = node.desc(sdfg)
if not isinstance(desc, dace.data.Array):
# TODO: handle streams
continue
if node.data in fpga_data:
fpga_array = fpga_data[node.data]
else:
fpga_array = sdfg.add_array(
'fpga_' + node.data,
desc.shape,
desc.dtype,
transient=True,
storage=dtypes.StorageType.FPGA_Global,
allow_conflicts=desc.allow_conflicts,
strides=desc.strides,
offset=desc.offset)
fpga_array[1].location = copy.copy(desc.location)
desc.location.clear()
fpga_data[node.data] = fpga_array
# fpga_node = type(node)(fpga_array)
post_node = post_state.add_write(node.data)
post_fpga_node = post_state.add_read('fpga_' + node.data)
mem = memlet.Memlet(f"fpga_{node.data}", None,
subsets.Range.from_array(desc))
post_state.add_edge(post_fpga_node, None, post_node, None, mem)
fpga_node = state.add_write('fpga_' + node.data)
sdutil.change_edge_dest(state, node, fpga_node)
state.remove_node(node)
sdfg.add_node(post_state)
sdutil.change_edge_src(sdfg, state, post_state)
sdfg.add_edge(state, post_state, sd.InterstateEdge())
# propagate memlet info from a nested sdfg
for src, src_conn, dst, dst_conn, mem in state.edges():
if mem.data is not None and mem.data in fpga_data:
mem.data = 'fpga_' + mem.data
fpga_update(sdfg, state, 0)
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(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)
class StateFusion(pattern_matching.Transformation):
""" Implements the state-fusion transformation.
State-fusion takes two states that are connected through a single edge,
and fuses them into one state. If strict, only applies if no memory
access hazards are created.
"""
_states_fused = 0
_first_state = sdfg.SDFGState()
_edge = sdfg.InterstateEdge()
_second_state = sdfg.SDFGState()
@staticmethod
def annotates_memlets():
return False
@staticmethod
def expressions():
return [
sdutil.node_path_graph(StateFusion._first_state,
StateFusion._second_state)
]
@staticmethod
def can_be_applied(graph, candidate, expr_index, sdfg, strict=False):
first_state = graph.nodes()[candidate[StateFusion._first_state]]
second_state = graph.nodes()[candidate[StateFusion._second_state]]
out_edges = graph.out_edges(first_state)
in_edges = graph.in_edges(first_state)
# First state must have only one output edge (with dst the second
# state).
if len(out_edges) != 1:
return False
# The interstate edge must not have a condition.
if not out_edges[0].data.is_unconditional():
return False
# The interstate edge may have assignments, as long as there are input
# edges to the first state that can absorb them.
if out_edges[0].data.assignments:
if not in_edges:
return False
# Fail if symbol is set before the state to fuse
# TODO: Also fail if symbol is used in the dataflow of that state
new_assignments = set(out_edges[0].data.assignments.keys())
if any((new_assignments & set(e.data.assignments.keys()))
for e in in_edges):
return False
# There can be no state that have output edges pointing to both the
# first and the second state. Such a case will produce a multi-graph.
for src, _, _ in in_edges:
for _, dst, _ in graph.out_edges(src):
if dst == second_state:
return False
if strict:
# If second state has other input edges, there might be issues
# Exceptions are when none of the states contain dataflow, unless
# the first state is an initial state (in which case the new initial
# state would be ambiguous).
first_in_edges = graph.in_edges(first_state)
second_in_edges = graph.in_edges(second_state)
if ((not second_state.is_empty() or not first_state.is_empty()
or len(first_in_edges) == 0) and len(second_in_edges) != 1):
return False
# Get connected components.
first_cc = [
cc_nodes
for cc_nodes in nx.weakly_connected_components(first_state._nx)
]
second_cc = [
cc_nodes for cc_nodes in nx.weakly_connected_components(
second_state._nx)
]
# Find source/sink (data) nodes
first_input = {
node
for node in sdutil.find_source_nodes(first_state)
if isinstance(node, nodes.AccessNode)
}
first_output = {
node
for node in first_state.nodes() if
isinstance(node, nodes.AccessNode) and node not in first_input
}
second_input = {
node
for node in sdutil.find_source_nodes(second_state)
if isinstance(node, nodes.AccessNode)
}
second_output = {
node
for node in second_state.nodes() if
isinstance(node, nodes.AccessNode) and node not in second_input
}
# Find source/sink (data) nodes by connected component
first_cc_input = [cc.intersection(first_input) for cc in first_cc]
first_cc_output = [
cc.intersection(first_output) for cc in first_cc
]
second_cc_input = [
cc.intersection(second_input) for cc in second_cc
]
second_cc_output = [
cc.intersection(second_output) for cc in second_cc
]
# Apply transformation in case all paths to the second state's
# nodes go through the same access node, which implies sequential
# behavior in SDFG semantics.
check_strict = len(first_cc)
for cc_output in first_cc_output:
out_nodes = [
n for n in first_state.sink_nodes() if n in cc_output
]
# Branching exists, multiple paths may involve same access node
# potentially causing data races
if len(out_nodes) > 1:
continue
# Otherwise, check if any of the second state's connected
# components for matching input
for node in out_nodes:
if (next(
(x for x in second_input if x.label == node.label),
None) is not None):
check_strict -= 1
break
if check_strict > 0:
# Check strict conditions
# RW dependency
for node in first_input:
if (next(
(x for x in second_output if x.label == node.label),
None) is not None):
return False
# WW dependency
for node in first_output:
if (next(
(x for x in second_output if x.label == node.label),
None) is not None):
return False
return True
@staticmethod
def match_to_str(graph, candidate):
first_state = graph.nodes()[candidate[StateFusion._first_state]]
second_state = graph.nodes()[candidate[StateFusion._second_state]]
return " -> ".join(state.label
for state in [first_state, second_state])
def apply(self, sdfg):
first_state = sdfg.nodes()[self.subgraph[StateFusion._first_state]]
second_state = sdfg.nodes()[self.subgraph[StateFusion._second_state]]
# Remove interstate edge(s)
edges = sdfg.edges_between(first_state, second_state)
for edge in edges:
if edge.data.assignments:
for src, dst, other_data in sdfg.in_edges(first_state):
other_data.assignments.update(edge.data.assignments)
sdfg.remove_edge(edge)
# Special case 1: first state is empty
if first_state.is_empty():
sdutil.change_edge_dest(sdfg, first_state, second_state)
sdfg.remove_node(first_state)
return
# Special case 2: second state is empty
if second_state.is_empty():
sdutil.change_edge_src(sdfg, second_state, first_state)
sdutil.change_edge_dest(sdfg, second_state, first_state)
sdfg.remove_node(second_state)
return
# Normal case: both states are not empty
# Find source/sink (data) nodes
first_input = [
node for node in sdutil.find_source_nodes(first_state)
if isinstance(node, nodes.AccessNode)
]
first_output = [
node for node in sdutil.find_sink_nodes(first_state)
if isinstance(node, nodes.AccessNode)
]
second_input = [
node for node in sdutil.find_source_nodes(second_state)
if isinstance(node, nodes.AccessNode)
]
# first input = first input - first output
first_input = [
node for node in first_input
if next((x for x in first_output
if x.label == node.label), None) is None
]
# Merge second state to first state
# First keep a backup of the topological sorted order of the nodes
order = [
x for x in reversed(list(nx.topological_sort(first_state._nx)))
if isinstance(x, nodes.AccessNode)
]
for node in second_state.nodes():
first_state.add_node(node)
for src, src_conn, dst, dst_conn, data in second_state.edges():
first_state.add_edge(src, src_conn, dst, dst_conn, data)
# Merge common (data) nodes
for node in second_input:
if first_state.in_degree(node) == 0:
n = next((x for x in order if x.label == node.label), None)
if n:
sdutil.change_edge_src(first_state, node, n)
first_state.remove_node(node)
n.access = dtypes.AccessType.ReadWrite
# Redirect edges and remove second state
sdutil.change_edge_src(sdfg, second_state, first_state)
sdfg.remove_node(second_state)
if Config.get_bool("debugprint"):
StateFusion._states_fused += 1
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)
